Check first header in each layer

In awslabs#1147, we cherry-picked a commit from stargz where we don't make
registry calls if layers are fully pulled. This works fine for stargz,
but for SOCI, we skip reading the first header for each layer, as we
don't need it. This caused a bug where our fetched size never matched
our expected size, so this condition was never met.

This commit fixes this by reading the initially skipped header. This
commit also checks this header to ensure that it is not malformed for
any reason.

Signed-off-by: David Son <[email protected]>

Author: sondavidb

fs/span-manager/span_manager.go

@@ -27,6 +27,7 @@ import (
 	"github.com/awslabs/soci-snapshotter/cache"
 	"github.com/awslabs/soci-snapshotter/ztoc"
 	"github.com/awslabs/soci-snapshotter/ztoc/compression"
+	"github.com/containerd/log"
 	"github.com/opencontainers/go-digest"
 	"golang.org/x/sync/errgroup"
 )
@@ -88,12 +89,25 @@ 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 {
 	index, err := ztoc.Zinfo()
 	if err != nil {
 		return nil
 	}
 	ztoc.Checkpoints = nil
+
+	// 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)
+	}
+	if err := index.VerifyHeader(bytes.NewReader(b)); err != nil {
+		log.L.Errorf("unable to verify %v header: %v", ztoc.CompressionAlgorithm, err)
+	}
+
 	spans := make([]*span, ztoc.MaxSpanID+1)
 	m := &SpanManager{
 		cache:                             cache,

fs/span-manager/span_manager_test.go

@@ -60,6 +60,7 @@ func TestSpanManager(t *testing.T) {
 	}
 
 	for _, tc := range testCases {
+		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			var err error
 			defer func() {

integration/pull_test.go

@@ -36,6 +36,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"math"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -872,3 +873,31 @@ min_layer_size=` + strconv.FormatInt(middleSize, 10) + `
 
 	checkFuseMounts(t, sh, layerCount-middleIndex)
 }
+
+// This checks if the initial header is properly recorded in the image descriptor on image pull
+func TestFullLayerRead(t *testing.T) {
+	regConfig := newRegistryConfig()
+	sh, done := newShellWithRegistry(t, regConfig)
+	defer done()
+	rebootContainerd(t, sh, getContainerdConfigToml(t, false), "")
+
+	containerImage := alpineImage
+	copyImage(sh, dockerhub(containerImage), regConfig.mirror(containerImage))
+	indexDigest := buildIndex(sh, regConfig.mirror(containerImage), withMinLayerSize(0), withSpanSize(math.MaxInt64))
+	// Max span size is used to ensure that the entire image will always be fetched at once.
+
+	sh.X(append(imagePullCmd, "--soci-index-digest", indexDigest, regConfig.mirror(containerImage).ref)...)
+	sh.X(append(runSociCmd, "-d", "--name", t.Name(), regConfig.mirror(containerImage).ref, "sleep", "infinity")...)
+	jsonFile := sh.O("nerdctl", "exec", t.Name(), "ls", "/.soci-snapshotter")
+	rawJSON := sh.O("nerdctl", "exec", t.Name(), "cat", "/.soci-snapshotter/"+strings.TrimSpace(string(jsonFile)))
+
+	var layers struct {
+		FetchedPercent float64
+	}
+	if err := json.Unmarshal(rawJSON, &layers); err != nil {
+		t.Fatalf("cannot unmarshal image layer JSON: %v", err)
+	}
+	if layers.FetchedPercent != 100 {
+		t.Fatalf("Expected 100%% fetchedPercent, found %v", layers.FetchedPercent)
+	}
+}

ztoc/compression/gzip_zinfo.go

@@ -24,7 +24,9 @@ package compression
 import "C"
 
 import (
+	"compress/gzip"
 	"fmt"
+	"io"
 	"unsafe"
 )
 
@@ -181,6 +183,15 @@ func (i *GzipZinfo) EndUncompressedOffset(spanID SpanID, fileSize Offset) Offset
 	return i.getUncompressedOffset(spanID + 1)
 }
 
+// VerifyHeader checks if the given zinfo has a proper header
+func (i *GzipZinfo) VerifyHeader(r io.Reader) error {
+	gz, err := gzip.NewReader(r)
+	if gz != nil {
+		gz.Close()
+	}
+	return err
+}
+
 // getCompressedOffset wraps `C.get_comp_off` and returns the offset for the span in the compressed stream.
 func (i *GzipZinfo) getCompressedOffset(spanID SpanID) Offset {
 	return Offset(C.get_comp_off(i.cZinfo, C.int(spanID)))

ztoc/compression/tar_zinfo.go

@@ -18,6 +18,7 @@ package compression
 
 import (
 	"fmt"
+	"io"
 	"os"
 
 	zinfo_flatbuffers "github.com/awslabs/soci-snapshotter/ztoc/compression/fbs/zinfo"
@@ -186,6 +187,14 @@ func (i *TarZinfo) EndUncompressedOffset(spanID SpanID, fileSize Offset) Offset
 	return i.spanIDToOffset(spanID + 1)
 }
 
+// VerifyHeader checks if the given zinfo has a proper header
+func (i *TarZinfo) VerifyHeader(r io.Reader) error {
+	// As this is a catch-all for all compression algorithms,
+	// there's not really a way to verify the header,
+	// so blindly assume it's correct.
+	return nil
+}
+
 func (i *TarZinfo) spanIDToOffset(spanID SpanID) Offset {
 	return Offset(i.spanSize * int64(spanID))
 }

ztoc/compression/zinfo.go

@@ -18,6 +18,7 @@ package compression
 
 import (
 	"fmt"
+	"io"
 )
 
 // Zinfo is the interface for dealing with compressed data efficiently. It chunks
@@ -75,6 +76,8 @@ type Zinfo interface {
 	// EndUncompressedOffset returns the offset (in uncompressed stream)
 	// of the last byte belonging to `spanID`. If it's the last span, `fileSize` is returned.
 	EndUncompressedOffset(spanID SpanID, fileSize Offset) Offset
+	// VerifyHeader checks if the given zinfo has a proper header
+	VerifyHeader(r io.Reader) error
 }
 
 // NewZinfo deseralizes given zinfo bytes into a zinfo struct.