From 1534cc984afeb5337d19952a3f544b7f8a459582 Mon Sep 17 00:00:00 2001
From: Jon Johnson <jon.johnson@chainguard.dev>
Date: Tue, 10 Oct 2023 19:19:57 -0700
Subject: [PATCH] Fix ReplaceSignatures

This was correctly replacing the base of the signatures but was also
appending those signatures to the Get() response. So the v1.Image
representation was getting out of sync with the oci.Signatures
representation.

Normally you won't see this because cosign doesn't do multiple
attestations in the same command for the same image, so you'd roundtrip
the signatures through the v1.Image being pushed and pulled, which
corrects the discrepancy.

If you (specifically, me) attempt to attach an attestation multiple
times, the Get() will get very out of sync with the v1.Image and Replace
no longer does what you'd expect because it's operating on incorrect
Get() results instead of directly on the v1.Image representation.

Signed-off-by: Jon Johnson <jon.johnson@chainguard.dev>
---
 pkg/oci/mutate/mutate_test.go | 58 ++++++++++++++++++++++++++++++++++-
 pkg/oci/mutate/signatures.go  |  2 +-
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/pkg/oci/mutate/mutate_test.go b/pkg/oci/mutate/mutate_test.go
index 4a3152683b5..b0e85be2a43 100644
--- a/pkg/oci/mutate/mutate_test.go
+++ b/pkg/oci/mutate/mutate_test.go
@@ -351,7 +351,7 @@ func TestSignEntity(t *testing.T) {
 				if al, err := atts.Get(); err != nil {
 					t.Fatalf("Get() = %v", err)
 				} else if len(al) != 1 {
-					t.Errorf("len(Get()) = %d, wanted %d", len(al), i)
+					t.Errorf("len(Get()) = %d, wanted %d", len(al), 1)
 				}
 			}
 		}
@@ -424,6 +424,43 @@ func TestSignEntity(t *testing.T) {
 			}
 		}
 	})
+
+	t.Run("with replace op (attestation)", func(t *testing.T) {
+		for _, se := range []oci.SignedEntity{si, sii, sunk} {
+			orig, err := static.NewAttestation([]byte("blah"))
+			if err != nil {
+				t.Fatalf("static.NewAttestation() = %v", err)
+			}
+			se, err = AttachAttestationToEntity(se, orig)
+			if err != nil {
+				t.Fatalf("AttachAttestationToEntity() = %v", err)
+			}
+
+			ro := &replaceAll{}
+
+			for i := 2; i < 10; i++ {
+				sig, err := static.NewAttestation([]byte(fmt.Sprintf("%d", i)))
+				if err != nil {
+					t.Fatalf("static.NewAttestation() = %v", err)
+				}
+
+				se, err = AttachAttestationToEntity(se, sig, WithReplaceOp(ro))
+				if err != nil {
+					t.Fatalf("AttachAttestationToEntity() = %v", err)
+				}
+
+				atts, err := se.Attestations()
+				if err != nil {
+					t.Fatalf("Attestations() = %v", err)
+				}
+				if al, err := atts.Get(); err != nil {
+					t.Fatalf("Get() = %v", err)
+				} else if len(al) != 1 {
+					t.Errorf("len(Get()) = %d, wanted %d", len(al), 1)
+				}
+			}
+		}
+	})
 }
 
 type dupe struct {
@@ -437,3 +474,22 @@ var _ DupeDetector = (*dupe)(nil)
 func (d *dupe) Find(oci.Signatures, oci.Signature) (oci.Signature, error) {
 	return d.sig, d.err
 }
+
+type replaceAll struct {
+}
+
+func (r *replaceAll) Replace(signatures oci.Signatures, o oci.Signature) (oci.Signatures, error) {
+	return &replaceOCISignatures{
+		Signatures:   signatures,
+		attestations: []oci.Signature{o},
+	}, nil
+}
+
+type replaceOCISignatures struct {
+	oci.Signatures
+	attestations []oci.Signature
+}
+
+func (r *replaceOCISignatures) Get() ([]oci.Signature, error) {
+	return r.attestations, nil
+}
diff --git a/pkg/oci/mutate/signatures.go b/pkg/oci/mutate/signatures.go
index 4ac356fe7ad..648cab04148 100644
--- a/pkg/oci/mutate/signatures.go
+++ b/pkg/oci/mutate/signatures.go
@@ -74,7 +74,7 @@ func ReplaceSignatures(base oci.Signatures) (oci.Signatures, error) {
 	return &sigAppender{
 		Image: img,
 		base:  base,
-		sigs:  sigs,
+		sigs:  []oci.Signature{},
 	}, nil
 }