combined tests from rebase

Author: CamrynCarter

pkg/store/store_test.go

@@ -990,3 +990,206 @@ func TestAddImage_OriginalRefAnnotation(t *testing.T) {
 		}
 	})
 }
+
+// newTestRegistry starts an in-process registry and returns its host and the
+// remote.Option needed to talk to it over plain HTTP.
+func newTestRegistry(t *testing.T) (string, []remote.Option) {
+	t.Helper()
+	srv := httptest.NewServer(registry.New())
+	t.Cleanup(srv.Close)
+	host := strings.TrimPrefix(srv.URL, "http://")
+	return host, []remote.Option{remote.WithTransport(srv.Client().Transport)}
+}
+
+// newTestStore creates a fresh OCI layout store rooted in a temp directory.
+func newTestStore(t *testing.T) *store.Layout {
+	t.Helper()
+	s, err := store.NewLayout(t.TempDir())
+	if err != nil {
+		t.Fatalf("new layout: %v", err)
+	}
+	return s
+}
+
+// seedImage pushes a random image to host/repo:tag and returns it, so a test
+// can later assert on the exact bytes/digest that were pushed.
+func seedImage(t *testing.T, host, repo, tag string, opts ...remote.Option) v1.Image {
+	t.Helper()
+	img, err := random.Image(1024, 3)
+	if err != nil {
+		t.Fatalf("random.Image: %v", err)
+	}
+	ref, err := gname.NewTag(host+"/"+repo+":"+tag, gname.Insecure)
+	if err != nil {
+		t.Fatalf("new tag: %v", err)
+	}
+	if err := remote.Write(ref, img, opts...); err != nil {
+		t.Fatalf("remote.Write: %v", err)
+	}
+	return img
+}
+
+// TestAddImagePinnedDigestIgnoresMovedTag proves the TOCTOU fix: once a caller
+// pins the digest it verified, a tag that moves to different content between
+// verification and the fetch cannot substitute its bytes into the store.
+func TestAddImagePinnedDigestIgnoresMovedTag(t *testing.T) {
+	host, remoteOpts := newTestRegistry(t)
+	original := seedImage(t, host, "test/pinned", "v1", remoteOpts...)
+	originalDigest, err := original.Digest()
+	if err != nil {
+		t.Fatalf("original digest: %v", err)
+	}
+
+	// Move the tag to different content, exactly as a mutable tag could be
+	// re-pushed between verification and the pull.
+	replacement, err := random.Image(1024, 3)
+	if err != nil {
+		t.Fatalf("random.Image: %v", err)
+	}
+	ref, err := gname.ParseReference(host + "/test/pinned:v1")
+	if err != nil {
+		t.Fatalf("parse: %v", err)
+	}
+	if err := remote.Write(ref, replacement, remoteOpts...); err != nil {
+		t.Fatalf("remote.Write replacement: %v", err)
+	}
+
+	s := newTestStore(t)
+	got, err := s.AddImage(context.Background(), host+"/test/pinned:v1", "", true,
+		originalDigest.String(), false, "", remoteOpts...)
+	if err != nil {
+		t.Fatalf("AddImage: %v", err)
+	}
+	if got != originalDigest.String() {
+		t.Fatalf("stored digest = %s, want the pinned %s (the moved tag won)", got, originalDigest.String())
+	}
+}
+
+// TestAddImageEmptyPinResolvesTag confirms the unpinned path is untouched:
+// an empty pinnedDigest still resolves the tag normally.
+func TestAddImageEmptyPinResolvesTag(t *testing.T) {
+	host, remoteOpts := newTestRegistry(t)
+	img := seedImage(t, host, "test/unpinned", "v1", remoteOpts...)
+	want, err := img.Digest()
+	if err != nil {
+		t.Fatalf("digest: %v", err)
+	}
+
+	s := newTestStore(t)
+	got, err := s.AddImage(context.Background(), host+"/test/unpinned:v1", "", true, "", false, "", remoteOpts...)
+	if err != nil {
+		t.Fatalf("AddImage: %v", err)
+	}
+	if got != want.String() {
+		t.Fatalf("stored digest = %s, want %s", got, want.String())
+	}
+}
+
+// TestAddImage_OriginalRefAnnotation verifies that AddImage captures the original,
+// fully pullable containerd-style reference (registry/repo:tag) under
+// consts.OriginalRefAnnotation, for both single-platform images (writeImage) and
+// multi-platform indices (writeIndex), so that provenance survives even if the
+// ref/containerd-name annotations are later overwritten by a rewrite.
+func TestAddImage_OriginalRefAnnotation(t *testing.T) {
+	srv := httptest.NewServer(registry.New())
+	t.Cleanup(srv.Close)
+	host := strings.TrimPrefix(srv.URL, "http://")
+
+	remoteOpts := []remote.Option{
+		remote.WithTransport(srv.Client().Transport),
+	}
+
+	t.Run("single-platform image", func(t *testing.T) {
+		tag, err := gname.NewTag(host+"/test/image:v1", gname.Insecure)
+		if err != nil {
+			t.Fatalf("new tag: %v", err)
+		}
+		img, err := random.Image(512, 2)
+		if err != nil {
+			t.Fatalf("random image: %v", err)
+		}
+		if err := remote.Write(tag, img, remoteOpts...); err != nil {
+			t.Fatalf("push image: %v", err)
+		}
+
+		s, err := store.NewLayout(t.TempDir())
+		if err != nil {
+			t.Fatalf("new layout: %v", err)
+		}
+		if _, err := s.AddImage(context.Background(), tag.Name(), "", false, remoteOpts...); err != nil {
+			t.Fatalf("AddImage: %v", err)
+		}
+
+		wantOriginalRef := tag.Name()
+		found := false
+		if err := s.Walk(func(_ string, desc ocispec.Descriptor) error {
+			if desc.Annotations[consts.OriginalRefAnnotation] == wantOriginalRef {
+				found = true
+			}
+			return nil
+		}); err != nil {
+			t.Fatalf("Walk: %v", err)
+		}
+		if !found {
+			t.Errorf("expected an artifact with OriginalRefAnnotation=%q, none found", wantOriginalRef)
+		}
+	})
+
+	t.Run("multi-platform index", func(t *testing.T) {
+		amd64Img, err := random.Image(512, 2)
+		if err != nil {
+			t.Fatalf("random image amd64: %v", err)
+		}
+		arm64Img, err := random.Image(512, 2)
+		if err != nil {
+			t.Fatalf("random image arm64: %v", err)
+		}
+		idx := mutate.AppendManifests(
+			empty.Index,
+			mutate.IndexAddendum{
+				Add: amd64Img,
+				Descriptor: v1.Descriptor{
+					MediaType: types.OCIManifestSchema1,
+					Platform:  &v1.Platform{OS: "linux", Architecture: "amd64"},
+				},
+			},
+			mutate.IndexAddendum{
+				Add: arm64Img,
+				Descriptor: v1.Descriptor{
+					MediaType: types.OCIManifestSchema1,
+					Platform:  &v1.Platform{OS: "linux", Architecture: "arm64"},
+				},
+			},
+		)
+		tag, err := gname.NewTag(host+"/test/multiarch:v1", gname.Insecure)
+		if err != nil {
+			t.Fatalf("new tag: %v", err)
+		}
+		if err := remote.WriteIndex(tag, idx, remoteOpts...); err != nil {
+			t.Fatalf("push index: %v", err)
+		}
+
+		s, err := store.NewLayout(t.TempDir())
+		if err != nil {
+			t.Fatalf("new layout: %v", err)
+		}
+		if _, err := s.AddImage(context.Background(), tag.Name(), "", false, remoteOpts...); err != nil {
+			t.Fatalf("AddImage: %v", err)
+		}
+
+		wantOriginalRef := tag.Name()
+		found := false
+		if err := s.Walk(func(_ string, desc ocispec.Descriptor) error {
+			if desc.Annotations[consts.KindAnnotationName] == consts.KindAnnotationIndex &&
+				desc.Annotations[consts.OriginalRefAnnotation] == wantOriginalRef {
+				found = true
+			}
+			return nil
+		}); err != nil {
+			t.Fatalf("Walk: %v", err)
+		}
+		if !found {
+			t.Errorf("expected the index artifact to have OriginalRefAnnotation=%q, none found", wantOriginalRef)
+		}
+	})
+}