Fast-fail on incorrect MaxSpanID

If a MaxSpanID is incorrect on a malformed zTOC, we will not find any
error until the spans start getting pulled and verified from
verifySpanContents. This will usually happen after the FUSE setup is
finished. We should fail much faster since we already have the info we
need to make this determination beforehand.

This should cover any failures in accessing the data in
verifySpanContents, but figured it can't hurt to add an extra sanity
check in the func to be extra sure.

This change required reworking the SpanManager's New function to handle
fast fails, so the commit also covers all the changes needed for that.

Signed-off-by: David Son <davbson@amazon.com>

Author: sondavidb

fs/backgroundfetcher/background_fetcher_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/awslabs/soci-snapshotter/util/testutil"
 	"github.com/awslabs/soci-snapshotter/ztoc"
 	"github.com/opencontainers/go-digest"
+	"github.com/stretchr/testify/assert"
 )
 
 func withPauser(p pauser) Option {
@@ -117,7 +118,8 @@ func TestBackgroundFetcherRun(t *testing.T) {
 					t.Fatalf("error building span manager and section reader: %v", err)
 				}
 				cache := &countingCache{}
-				sm := spanmanager.New(ztoc, sr, cache, 0)
+				sm, err := spanmanager.New(ztoc, sr, cache, 0)
+				assert.Nil(t, err)
 				infos = append(infos, testInfo{sm, cache, ztoc})
 			}
 

fs/backgroundfetcher/resolver_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/awslabs/soci-snapshotter/util/testutil"
 	"github.com/awslabs/soci-snapshotter/ztoc"
 	"github.com/opencontainers/go-digest"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestSequentialResolver(t *testing.T) {
@@ -50,7 +51,8 @@ func TestSequentialResolver(t *testing.T) {
 			if err != nil {
 				t.Fatalf("error build ztoc and section reader: %v", err)
 			}
-			sm := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+			sm, err := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+			assert.Nil(t, err)
 			sequentialResolver := NewSequentialResolver(digest.FromString("test"), sm)
 
 			var resolvedSpans []int

fs/layer/layer.go

@@ -348,7 +348,10 @@ func (r *Resolver) Resolve(ctx context.Context, hosts []docker.RegistryHost, ref
 	ztoc.TOC.FileMetadata = nil
 	log.G(ctx).Debugf("[Resolver.Resolve]Initialized metadata store for layer sha=%v", desc.Digest)
 
-	spanManager := spanmanager.New(ztoc, sr, spanCache, r.config.BlobConfig.MaxSpanVerificationRetries, cache.Direct())
+	spanManager, err := spanmanager.New(ztoc, sr, spanCache, r.config.BlobConfig.MaxSpanVerificationRetries, cache.Direct())
+	if err != nil {
+		return nil, fmt.Errorf("error creating span manager: %v", err)
+	}
 	var bgLayerResolver backgroundfetcher.Resolver
 	if r.bgFetcher != nil {
 		bgLayerResolver = backgroundfetcher.NewSequentialResolver(desc.Digest, spanManager)

fs/layer/util_test.go

@@ -67,6 +67,7 @@ import (
 	"github.com/hanwen/go-fuse/v2/fuse"
 	digest "github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/stretchr/testify/assert"
 	"golang.org/x/sys/unix"
 )
 
@@ -172,7 +173,8 @@ func makeNodeReader(t *testing.T, contents []byte, spanSize int64, factory metad
 	if err != nil {
 		t.Fatalf("failed to create reader: %v", err)
 	}
-	spanManager := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+	spanManager, err := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+	assert.Nil(t, err)
 	r, err := reader.NewReader(mr, digest.FromString(""), spanManager, false)
 	if err != nil {
 		mr.Close()
@@ -333,7 +335,8 @@ func testExistenceWithOpaque(t *testing.T, factory metadata.Store, opaque Overla
 					t.Fatalf("failed to create reader: %v", err)
 				}
 				defer mr.Close()
-				spanManager := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+				spanManager, err := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+				assert.Nil(t, err)
 				r, err := reader.NewReader(mr, digest.FromString(""), spanManager, false)
 				if err != nil {
 					t.Fatalf("failed to make new reader: %v", err)
@@ -595,7 +598,7 @@ func hasEntry(t *testing.T, name string, ents fusefs.DirStream) (fuse.DirEntry,
 	return fuse.DirEntry{}, false
 }
 
-func hasStateFile(t *testing.T, id string) check {
+func hasStateFile(_ *testing.T, id string) check {
 	return func(t *testing.T, root *node) {
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
@@ -774,7 +777,8 @@ func testStatfs(t *testing.T, factory metadata.Store) {
 	}
 	defer mr.Close()
 
-	spanManager := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+	spanManager, err := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+	assert.Nil(t, err)
 	r, err := reader.NewReader(mr, digest.FromString(""), spanManager, false)
 	if err != nil {
 		t.Fatalf("failed to create reader: %v", err)

fs/reader/reader_test.go

@@ -52,6 +52,7 @@ import (
 	"github.com/awslabs/soci-snapshotter/util/testutil"
 	"github.com/awslabs/soci-snapshotter/ztoc"
 	digest "github.com/opencontainers/go-digest"
+	"github.com/stretchr/testify/assert"
 )
 
 const (
@@ -161,7 +162,8 @@ func makeFile(t *testing.T, contents []byte, prefix string, factory metadata.Sto
 	if err != nil {
 		t.Fatalf("failed to create reader: %v", err)
 	}
-	spanManager := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+	spanManager, err := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+	assert.Nil(t, err)
 	r, err := NewReader(mr, digest.FromString(""), spanManager, false)
 	if err != nil {
 		mr.Close()
@@ -216,7 +218,8 @@ func testFailReader(t *testing.T, factory metadata.Store) {
 			if !found {
 				t.Fatalf("free ID not found")
 			}
-			spanManager := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+			spanManager, err := spanmanager.New(ztoc, sr, cache.NewMemoryCache(), 0)
+			assert.Nil(t, err)
 			r, err := NewReader(mr, digest.FromString(""), spanManager, false)
 			if err != nil {
 				mr.Close()

fs/span-manager/span_manager.go

@@ -39,6 +39,7 @@ var (
 	ErrSpanNotAvailable    = errors.New("span not available in cache")
 	ErrIncorrectSpanDigest = errors.New("span digests do not match")
 	ErrExceedMaxSpan       = errors.New("span id larger than max span id")
+	ErrIncorrectMaxSpanID  = errors.New("given max span ID differs from calculated max span ID")
 )
 
 // SpanManager fetches and caches spans of a given layer.
@@ -69,22 +70,30 @@ type spanInfo struct {
 // New creates a SpanManager with given ztoc and content reader, and builds all
 // spans based on the ztoc.
 
-// TODO: return errors/nil objects on failure
-func New(ztoc *ztoc.Ztoc, r *io.SectionReader, cache cache.BlobCache, retries int, cacheOpt ...cache.Option) *SpanManager {
+func New(ztoc *ztoc.Ztoc, r *io.SectionReader, cache cache.BlobCache, retries int, cacheOpt ...cache.Option) (*SpanManager, error) {
 	index, err := ztoc.Zinfo()
 	if err != nil {
-		return nil
+		return nil, err
 	}
 	ztoc.Checkpoints = nil
 
+	// Check zTOC contents against generated Zinfo
+	maxSpanID := index.MaxSpanID()
+	if ztoc.MaxSpanID != maxSpanID {
+		return nil, fmt.Errorf("%w (zTOC: %d, checkpoints: %d)", ErrIncorrectMaxSpanID, maxSpanID, ztoc.MaxSpanID)
+	}
+	if int32(len(ztoc.SpanDigests)) != int32(maxSpanID+1) {
+		return nil, fmt.Errorf("%w (len span digests: %d, MaxSpanID+1: %d", ErrIncorrectMaxSpanID, len(ztoc.SpanDigests), maxSpanID+1)
+	}
+
 	// We don't need the header anywhere else, but we want to ensure we read every byte from the layer.
 	// This also serves as a sanity check to make sure the file isn't corrupted.
 	b := make([]byte, index.StartCompressedOffset(0))
 	if _, err := r.Read(b); err != nil {
-		log.L.Errorf("unable to read ztoc header: %v", err)
+		return nil, fmt.Errorf("unable to read ztoc header: %w", err)
 	}
 	if err := index.VerifyHeader(bytes.NewReader(b)); err != nil {
-		log.L.Errorf("unable to verify %v header: %v", ztoc.CompressionAlgorithm, err)
+		return nil, fmt.Errorf("unable to verify %v header: %w", ztoc.CompressionAlgorithm, err)
 	}
 
 	spans := make([]*span, ztoc.MaxSpanID+1)
@@ -105,7 +114,7 @@ func New(ztoc *ztoc.Ztoc, r *io.SectionReader, cache cache.BlobCache, retries in
 		m.Close()
 	})
 
-	return m
+	return m, nil
 }
 
 func (m *SpanManager) buildAllSpans() {
@@ -452,6 +461,10 @@ func (m *SpanManager) getSpanFromCache(spanID compression.SpanID, offset, size c
 // verifySpanContents calculates span digest from its compressed bytes, and compare
 // with the digest stored in ztoc.
 func (m *SpanManager) verifySpanContents(compressedData []byte, spanID compression.SpanID) error {
+	if int32(spanID) >= int32(len(m.ztoc.SpanDigests)) {
+		return ErrExceedMaxSpan
+	}
+
 	actual := digest.FromBytes(compressedData)
 	expected := m.ztoc.SpanDigests[spanID]
 	if actual != expected {

fs/span-manager/span_manager_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/awslabs/soci-snapshotter/util/testutil"
 	"github.com/awslabs/soci-snapshotter/ztoc"
 	"github.com/awslabs/soci-snapshotter/ztoc/compression"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestSpanManager(t *testing.T) {
@@ -48,14 +49,38 @@ func TestSpanManager(t *testing.T) {
 			maxSpans: 100,
 		},
 		{
-			name:     "span digest verification fails",
+			name:          "bad MaxSpanID",
+			maxSpans:      5,
+			expectedError: ErrIncorrectMaxSpanID,
+		},
+		{
+			name:     "header verification fails",
 			maxSpans: 100,
 			sectionReader: io.NewSectionReader(readerFn(func(b []byte, _ int64) (int, error) {
 				sz := compression.Offset(len(b))
 				r := testutil.NewTestRand(t)
 				copy(b, r.RandomByteData(int64(sz))) // populate with garbage data
 				return len(b), nil
 			}), 0, 10000000),
+			expectedError: gzip.ErrHeader,
+		},
+		{
+			name:     "span digest verification fails",
+			maxSpans: 10,
+			sectionReader: io.NewSectionReader(readerFn(func(b []byte, _ int64) (int, error) {
+				var r bytes.Buffer
+				w := gzip.NewWriter(&r)
+				w.Write([]byte("failing digest verification"))
+				w.Close()
+
+				gz, err := io.ReadAll(&r)
+				if err != nil {
+					t.Fatalf("error creating SectionReader: %v", err)
+				}
+
+				copy(b, gz)
+				return len(b), nil
+			}), 0, int64(spanSize)*10),
 			expectedError: ErrIncorrectSpanDigest,
 		},
 	}
@@ -84,13 +109,20 @@ func TestSpanManager(t *testing.T) {
 				return
 			}
 
+			if tc.expectedError == ErrIncorrectMaxSpanID {
+				toc.MaxSpanID += 1
+			}
+
 			if tc.sectionReader != nil {
 				r = tc.sectionReader
 			}
 
 			cache := cache.NewMemoryCache()
 			defer cache.Close()
-			m := New(toc, r, cache, 0)
+			m, err := New(toc, r, cache, 0)
+			if err != nil {
+				return
+			}
 
 			// Test GetContent
 			fileContentFromSpans, err := getFileContentFromSpans(m, toc, fileName)
@@ -133,7 +165,8 @@ func TestSpanManagerCache(t *testing.T) {
 	}
 	cache := cache.NewMemoryCache()
 	defer cache.Close()
-	m := New(toc, r, cache, 0)
+	m, err := New(toc, r, cache, 0)
+	assert.Nil(t, err)
 	spanID := 0
 	err = m.resolveSpan(compression.SpanID(spanID))
 	if err != nil {
@@ -188,7 +221,8 @@ func TestStateTransition(t *testing.T) {
 	}
 	cache := cache.NewMemoryCache()
 	defer cache.Close()
-	m := New(toc, r, cache, 0)
+	m, err := New(toc, r, cache, 0)
+	assert.Nil(t, err)
 
 	// check initial span states
 	for i := uint32(0); i <= uint32(toc.MaxSpanID); i++ {
@@ -347,16 +381,19 @@ func TestSpanManagerRetries(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			r := testutil.NewTestRand(t)
+			randStr := string(r.RandomByteData(10000000))
+
 			entries := []testutil.TarEntry{
-				testutil.File("test", string(r.RandomByteData(10000000))),
+				testutil.File("test", randStr),
 			}
 			ztoc, sr, err := ztoc.BuildZtocReader(t, entries, gzip.DefaultCompression, 100000)
 			if err != nil {
 				t.Fatal(err)
 			}
-			rdr := &retryableReaderAt{inner: sr, maxErrors: tc.readerErrors}
+			rdr := newRetryableReaderAt(sr, tc.readerErrors)
 			sr = io.NewSectionReader(rdr, 0, 10000000)
-			sm := New(ztoc, sr, cache.NewMemoryCache(), tc.spanManagerRetries)
+			sm, err := New(ztoc, sr, cache.NewMemoryCache(), tc.spanManagerRetries)
+			assert.Nil(t, err)
 
 			for i := 0; i < int(ztoc.MaxSpanID); i++ {
 				rdr.errCount = 0
@@ -381,14 +418,24 @@ type retryableReaderAt struct {
 	maxErrors int
 }
 
+func newRetryableReaderAt(inner *io.SectionReader, maxErrors int) *retryableReaderAt {
+	return &retryableReaderAt{
+		inner:     inner,
+		maxErrors: maxErrors,
+		errCount:  -1, // First read needs to succeed to create spanmanager
+	}
+}
+
 func (r *retryableReaderAt) ReadAt(buf []byte, off int64) (int, error) {
 	n, err := r.inner.ReadAt(buf, off)
 	if (err != nil && err != io.EOF) || n != len(buf) {
 		return n, err
 	}
 	if r.errCount < r.maxErrors {
 		r.errCount++
-		buf[0] = buf[0] ^ 0xff
+		if r.errCount > 0 {
+			buf[0] = buf[0] ^ 0xff
+		}
 	}
 	return n, err
 }