Add FetchSociArtifacts test

This adds tests to verify that we correctly reject SOCI indexes if
they are somehow modified by the registry. We have always done this, but
this verifies it via a test that's run continuously.

Signed-off-by: Kern Walster <walster@amazon.com>

Author: Kern--

fs/artifact_fetcher_test.go

@@ -19,11 +19,14 @@ package fs
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"testing"
 
+	"github.com/awslabs/soci-snapshotter/soci"
+	"github.com/awslabs/soci-snapshotter/soci/store"
 	"github.com/containerd/containerd/reference"
 	"github.com/google/go-cmp/cmp"
 	"github.com/opencontainers/go-digest"
@@ -214,6 +217,49 @@ func TestArtifactFetcherFetchOnlyOnce(t *testing.T) {
 	}
 }
 
+func TestArtifactFetcherStore(t *testing.T) {
+	testCases := []struct {
+		name          string
+		contents      []byte
+		digest        digest.Digest
+		expectedError error
+	}{
+		{
+			name:          "correct digest succeeds on store",
+			contents:      []byte("test"),
+			digest:        digest.FromBytes([]byte("test")),
+			expectedError: nil,
+		},
+		{
+			name:          "incorrect digest fails on store",
+			contents:      []byte("test"),
+			digest:        digest.FromBytes([]byte("different data")),
+			expectedError: content.ErrMismatchedDigest,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			fetcher, err := newFakeArtifactFetcher(imageRef, tc.contents)
+			if err != nil {
+				t.Fatalf("could not create artifact fetcher: %v", err)
+			}
+			size := len(tc.contents)
+			desc := ocispec.Descriptor{
+				Digest: tc.digest,
+				Size:   int64(size),
+			}
+			ctx := context.Background()
+
+			err = fetcher.Store(ctx, desc, bytes.NewReader(tc.contents))
+			if !errors.Is(err, tc.expectedError) {
+				t.Fatalf("unexpected error, expected = %v, got = %v", tc.expectedError, err)
+			}
+		})
+	}
+
+}
+
 func TestNewRemoteStore(t *testing.T) {
 	client := http.Client{}
 	testCases := []struct {
@@ -255,6 +301,80 @@ func TestNewRemoteStore(t *testing.T) {
 	}
 }
 
+func TestFetchSociArtifacts(t *testing.T) {
+	fakeZtoc := []byte("test data")
+	fakeZtocDesc := ocispec.Descriptor{
+		Size:   int64(len(fakeZtoc)),
+		Digest: digest.FromBytes(fakeZtoc),
+	}
+
+	blobs := []ocispec.Descriptor{
+		{
+			MediaType: soci.SociLayerMediaType,
+			Digest:    fakeZtocDesc.Digest,
+			Size:      fakeZtocDesc.Size,
+		},
+	}
+	sociIndex := soci.NewIndex(soci.V2, blobs, nil, nil)
+	sociBytes, err := soci.MarshalIndex(sociIndex)
+	if err != nil {
+		t.Fatalf("failed to serialize soci index: %v", err)
+	}
+	sociIndexDesc := ocispec.Descriptor{
+		MediaType: ocispec.MediaTypeImageManifest,
+		Size:      int64(len(sociBytes)),
+		Digest:    digest.FromBytes(sociBytes),
+	}
+
+	modifiedSociIndex := soci.NewIndex(soci.V2, blobs, nil, map[string]string{"a": "b"})
+	modifiedSociBytes, err := soci.MarshalIndex(modifiedSociIndex)
+	if err != nil {
+		t.Fatalf("failed to serialize modified soci index: %v", err)
+	}
+	modifiedZtocBytes := []byte("modified test data")
+
+	tests := []struct {
+		name           string
+		remoteContents map[digest.Digest][]byte
+		expectedError  error
+	}{
+		{
+			name: "correct data succeeds",
+			remoteContents: map[digest.Digest][]byte{
+				sociIndexDesc.Digest: sociBytes,
+				fakeZtocDesc.Digest:  fakeZtoc,
+			},
+			expectedError: nil,
+		},
+		{
+			name: "modified index data fails",
+			remoteContents: map[digest.Digest][]byte{
+				sociIndexDesc.Digest: modifiedSociBytes,
+				fakeZtocDesc.Digest:  fakeZtoc,
+			},
+			expectedError: content.ErrMismatchedDigest,
+		},
+		{
+			name: "modified ztoc data fails",
+			remoteContents: map[digest.Digest][]byte{
+				sociIndexDesc.Digest: sociBytes,
+				fakeZtocDesc.Digest:  modifiedZtocBytes,
+			},
+			expectedError: content.ErrTrailingData,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			ctx := context.Background()
+			_, err = FetchSociArtifacts(ctx, reference.Spec{}, sociIndexDesc, newFakeLocalStore(), newFakeRemoteStoreWithContents(test.remoteContents))
+			if !errors.Is(err, test.expectedError) {
+				t.Fatalf("unexpected error, got: %v. expected: %v", err, test.expectedError)
+			}
+		})
+	}
+}
+
 func newFakeArtifactFetcher(ref string, contents []byte) (*artifactFetcher, error) {
 	refspec, err := reference.Parse(ref)
 	if err != nil {
@@ -263,20 +383,55 @@ func newFakeArtifactFetcher(ref string, contents []byte) (*artifactFetcher, erro
 	return newArtifactFetcher(refspec, memory.New(), newFakeRemoteStore(contents))
 }
 
+func newFakeLocalStore() store.Store {
+	return &fakeLocalStore{
+		Store: memory.New(),
+	}
+}
+
+type fakeLocalStore struct {
+	*memory.Store
+}
+
+func (s *fakeLocalStore) BatchOpen(ctx context.Context) (context.Context, store.CleanupFunc, error) {
+	return ctx, func(ctx context.Context) error { return nil }, nil
+}
+
+func (s *fakeLocalStore) Delete(_ context.Context, _ digest.Digest) error {
+	return nil
+}
+
+func (s *fakeLocalStore) Label(_ context.Context, _ ocispec.Descriptor, _, _ string) error {
+	return nil
+}
+
 func newFakeRemoteStore(contents []byte) resolverStorage {
 	return &fakeRemoteStore{
-		contents: contents,
+		defaultContents: contents,
+		contents:        make(map[digest.Digest][]byte),
+	}
+}
+
+func newFakeRemoteStoreWithContents(contents map[digest.Digest][]byte) resolverStorage {
+	return &fakeRemoteStore{
+		defaultContents: []byte{},
+		contents:        contents,
 	}
 }
 
 type fakeRemoteStore struct {
-	contents []byte
+	defaultContents []byte
+	contents        map[digest.Digest][]byte
 }
 
 var _ content.Storage = &fakeRemoteStore{}
 
 func (f *fakeRemoteStore) Fetch(_ context.Context, desc ocispec.Descriptor) (io.ReadCloser, error) {
-	return io.NopCloser(bytes.NewReader(f.contents)), nil
+	if data, ok := f.contents[desc.Digest]; ok {
+		return io.NopCloser(bytes.NewReader(data)), nil
+	}
+
+	return io.NopCloser(bytes.NewReader(f.defaultContents)), nil
 }
 
 func (f *fakeRemoteStore) Push(_ context.Context, desc ocispec.Descriptor, ra io.Reader) error {
@@ -289,6 +444,6 @@ func (f *fakeRemoteStore) Exists(_ context.Context, desc ocispec.Descriptor) (bo
 
 func (f *fakeRemoteStore) Resolve(_ context.Context, ref string) (ocispec.Descriptor, error) {
 	return ocispec.Descriptor{
-		Size: int64(len(f.contents)),
+		Size: int64(len(f.defaultContents)),
 	}, nil
 }