Unmarshal the TOC already in readZstdChunkedManifest

Other TOC formats don't fill the data in. For now, this only increases memory usage, but we will need the data soon. Signed-off-by: Miloslav Trmač <[email protected]>
containers · Apr 23, 2024 · 471b895 · 471b895
1 parent 8b0c90d
commit 471b895
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 19 deletions.
diff --git a/pkg/chunked/compression_linux.go b/pkg/chunked/compression_linux.go
@@ -133,37 +133,38 @@ func readEstargzChunkedManifest(blobStream ImageSourceSeekable, blobSize int64,
 }
 
 // readZstdChunkedManifest reads the zstd:chunked manifest from the seekable stream blobStream.
-func readZstdChunkedManifest(blobStream ImageSourceSeekable, tocDigest digest.Digest, annotations map[string]string) ([]byte, []byte, int64, error) {
+// Returns (manifest blob, parsed manifest, tar-split blob, manifest offset).
+func readZstdChunkedManifest(blobStream ImageSourceSeekable, tocDigest digest.Digest, annotations map[string]string) ([]byte, *internal.TOC, []byte, int64, error) {
 	offsetMetadata := annotations[internal.ManifestInfoKey]
 	if offsetMetadata == "" {
-		return nil, nil, 0, fmt.Errorf("%q annotation missing", internal.ManifestInfoKey)
+		return nil, nil, nil, 0, fmt.Errorf("%q annotation missing", internal.ManifestInfoKey)
 	}
 	var manifestChunk ImageSourceChunk
 	var manifestLengthUncompressed, manifestType uint64
 	if _, err := fmt.Sscanf(offsetMetadata, "%d:%d:%d:%d", &manifestChunk.Offset, &manifestChunk.Length, &manifestLengthUncompressed, &manifestType); err != nil {
-		return nil, nil, 0, err
+		return nil, nil, nil, 0, err
 	}
 	// The tarSplit… values are valid if tarSplitChunk.Offset > 0
 	var tarSplitChunk ImageSourceChunk
 	var tarSplitLengthUncompressed uint64
 	var tarSplitChecksum string
 	if tarSplitInfoKeyAnnotation, found := annotations[internal.TarSplitInfoKey]; found {
 		if _, err := fmt.Sscanf(tarSplitInfoKeyAnnotation, "%d:%d:%d", &tarSplitChunk.Offset, &tarSplitChunk.Length, &tarSplitLengthUncompressed); err != nil {
-			return nil, nil, 0, err
+			return nil, nil, nil, 0, err
 		}
 		tarSplitChecksum = annotations[internal.TarSplitChecksumKey]
 	}
 
 	if manifestType != internal.ManifestTypeCRFS {
-		return nil, nil, 0, errors.New("invalid manifest type")
+		return nil, nil, nil, 0, errors.New("invalid manifest type")
 	}
 
 	// set a reasonable limit
 	if manifestChunk.Length > (1<<20)*50 {
-		return nil, nil, 0, errors.New("manifest too big")
+		return nil, nil, nil, 0, errors.New("manifest too big")
 	}
 	if manifestLengthUncompressed > (1<<20)*50 {
-		return nil, nil, 0, errors.New("manifest too big")
+		return nil, nil, nil, 0, errors.New("manifest too big")
 	}
 
 	chunks := []ImageSourceChunk{manifestChunk}
@@ -172,7 +173,7 @@ func readZstdChunkedManifest(blobStream ImageSourceSeekable, tocDigest digest.Di
 	}
 	parts, errs, err := blobStream.GetBlobAt(chunks)
 	if err != nil {
-		return nil, nil, 0, err
+		return nil, nil, nil, 0, err
 	}
 
 	readBlob := func(len uint64) ([]byte, error) {
@@ -197,26 +198,31 @@ func readZstdChunkedManifest(blobStream ImageSourceSeekable, tocDigest digest.Di
 
 	manifest, err := readBlob(manifestChunk.Length)
 	if err != nil {
-		return nil, nil, 0, err
+		return nil, nil, nil, 0, err
 	}
 
 	decodedBlob, err := decodeAndValidateBlob(manifest, manifestLengthUncompressed, tocDigest.String())
 	if err != nil {
-		return nil, nil, 0, err
+		return nil, nil, nil, 0, err
 	}
+	toc, err := unmarshalToc(decodedBlob)
+	if err != nil {
+		return nil, nil, nil, 0, err
+	}
+
 	decodedTarSplit := []byte{}
 	if tarSplitChunk.Offset > 0 {
 		tarSplit, err := readBlob(tarSplitChunk.Length)
 		if err != nil {
-			return nil, nil, 0, err
+			return nil, nil, nil, 0, err
 		}
 
 		decodedTarSplit, err = decodeAndValidateBlob(tarSplit, tarSplitLengthUncompressed, tarSplitChecksum)
 		if err != nil {
-			return nil, nil, 0, err
+			return nil, nil, nil, 0, err
 		}
 	}
-	return decodedBlob, decodedTarSplit, int64(manifestChunk.Offset), err
+	return decodedBlob, toc, decodedTarSplit, int64(manifestChunk.Offset), err
 }
 
 func decodeAndValidateBlob(blob []byte, lengthUncompressed uint64, expectedCompressedChecksum string) ([]byte, error) {

diff --git a/pkg/chunked/storage_linux.go b/pkg/chunked/storage_linux.go
@@ -79,6 +79,7 @@ type compressedFileType int
 type chunkedDiffer struct {
 	stream      ImageSourceSeekable
 	manifest    []byte
+	toc         *internal.TOC // The parsed contents of manifest, or nil if not yet available
 	tarSplit    []byte
 	layersCache *layersCache
 	tocOffset   int64
@@ -314,7 +315,7 @@ func makeConvertFromRawDiffer(ctx context.Context, store storage.Store, blobDige
 }
 
 func makeZstdChunkedDiffer(ctx context.Context, store storage.Store, blobSize int64, tocDigest digest.Digest, annotations map[string]string, iss ImageSourceSeekable, storeOpts *types.StoreOptions) (*chunkedDiffer, error) {
-	manifest, tarSplit, tocOffset, err := readZstdChunkedManifest(iss, tocDigest, annotations)
+	manifest, toc, tarSplit, tocOffset, err := readZstdChunkedManifest(iss, tocDigest, annotations)
 	if err != nil {
 		return nil, fmt.Errorf("read zstd:chunked manifest: %w", err)
 	}
@@ -331,6 +332,7 @@ func makeZstdChunkedDiffer(ctx context.Context, store storage.Store, blobSize in
 		fileType:        fileTypeZstdChunked,
 		layersCache:     layersCache,
 		manifest:        manifest,
+		toc:             toc,
 		storeOpts:       storeOpts,
 		stream:          iss,
 		tarSplit:        tarSplit,
@@ -1701,7 +1703,7 @@ func (c *chunkedDiffer) ApplyDiff(dest string, options *archive.TarOptions, diff
 		if tocDigest == nil {
 			return graphdriver.DriverWithDifferOutput{}, fmt.Errorf("internal error: just-created zstd:chunked missing TOC digest")
 		}
-		manifest, tarSplit, tocOffset, err := readZstdChunkedManifest(fileSource, *tocDigest, annotations)
+		manifest, toc, tarSplit, tocOffset, err := readZstdChunkedManifest(fileSource, *tocDigest, annotations)
 		if err != nil {
 			return graphdriver.DriverWithDifferOutput{}, fmt.Errorf("read zstd:chunked manifest: %w", err)
 		}
@@ -1712,6 +1714,7 @@ func (c *chunkedDiffer) ApplyDiff(dest string, options *archive.TarOptions, diff
 		// fill the chunkedDiffer with the data we just read.
 		c.fileType = fileTypeZstdChunked
 		c.manifest = manifest
+		c.toc = toc
 		c.tarSplit = tarSplit
 		c.tocOffset = tocOffset
 
@@ -1732,9 +1735,13 @@ func (c *chunkedDiffer) ApplyDiff(dest string, options *archive.TarOptions, diff
 	}
 
 	// Generate the manifest
-	toc, err := unmarshalToc(c.manifest)
-	if err != nil {
-		return graphdriver.DriverWithDifferOutput{}, err
+	toc := c.toc
+	if toc == nil {
+		toc_, err := unmarshalToc(c.manifest)
+		if err != nil {
+			return graphdriver.DriverWithDifferOutput{}, err
+		}
+		toc = toc_
 	}
 
 	output := graphdriver.DriverWithDifferOutput{

diff --git a/pkg/chunked/zstdchunked_test.go b/pkg/chunked/zstdchunked_test.go
@@ -15,6 +15,7 @@ import (
 	"github.com/containers/storage/pkg/chunked/toc"
 	"github.com/klauspost/compress/zstd"
 	"github.com/opencontainers/go-digest"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -153,7 +154,7 @@ func TestGenerateAndParseManifest(t *testing.T) {
 	tocDigest, err := toc.GetTOCDigest(annotations)
 	require.NoError(t, err)
 	require.NotNil(t, tocDigest)
-	manifest, _, _, err := readZstdChunkedManifest(s, *tocDigest, annotations)
+	manifest, decodedTOC, _, _, err := readZstdChunkedManifest(s, *tocDigest, annotations)
 	if err != nil {
 		t.Error(err)
 	}
@@ -169,6 +170,7 @@ func TestGenerateAndParseManifest(t *testing.T) {
 	if len(toc.Entries) != len(someFiles) {
 		t.Fatal("Manifest mismatch")
 	}
+	assert.Equal(t, toc, decodedTOC)
 }
 
 func TestGetTarType(t *testing.T) {