chunked: refactor code in new functions

Signed-off-by: Giuseppe Scrivano <gscrivan@redhat.com>

storage/pkg/chunked/compression_linux.go

@@ -2,8 +2,6 @@ package chunked
 
 import (
 	archivetar "archive/tar"
-	"bytes"
-	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
@@ -35,13 +33,6 @@ func typeToTarType(t string) (byte, error) {
 	return r, nil
 }
 
-func isZstdChunkedFrameMagic(data []byte) bool {
-	if len(data) < 8 {
-		return false
-	}
-	return bytes.Equal(internal.ZstdChunkedFrameMagic, data[:8])
-}
-
 func readEstargzChunkedManifest(blobStream ImageSourceSeekable, blobSize int64, annotations map[string]string) ([]byte, int64, error) {
 	// information on the format here https://github.com/containerd/stargz-snapshotter/blob/main/docs/stargz-estargz.md
 	footerSize := int64(51)
@@ -155,27 +146,14 @@ func readZstdChunkedManifest(blobStream ImageSourceSeekable, blobSize int64, ann
 		return nil, nil, 0, errors.New("blob too small")
 	}
 
-	manifestChecksumAnnotation := annotations[internal.ManifestChecksumKey]
+	var footerData internal.ZstdChunkedFooterData
-	if manifestChecksumAnnotation == "" {
-		return nil, nil, 0, fmt.Errorf("manifest checksum annotation %q not found", internal.ManifestChecksumKey)
-	}
-
-	var offset, length, lengthUncompressed, manifestType uint64
-
-	var offsetTarSplit, lengthTarSplit, lengthUncompressedTarSplit uint64
-	tarSplitChecksumAnnotation := ""
 
 	if offsetMetadata := annotations[internal.ManifestInfoKey]; offsetMetadata != "" {
-		if _, err := fmt.Sscanf(offsetMetadata, "%d:%d:%d:%d", &offset, &length, &lengthUncompressed, &manifestType); err != nil {
+		var err error
+		footerData, err = internal.ReadFooterDataFromAnnotations(annotations)
+		if err != nil {
 			return nil, nil, 0, err
 		}
-
-		if tarSplitInfoKeyAnnotation, found := annotations[internal.TarSplitInfoKey]; found {
-			if _, err := fmt.Sscanf(tarSplitInfoKeyAnnotation, "%d:%d:%d", &offsetTarSplit, &lengthTarSplit, &lengthUncompressedTarSplit); err != nil {
-				return nil, nil, 0, err
-			}
-			tarSplitChecksumAnnotation = annotations[internal.TarSplitChecksumKey]
-		}
 	} else {
 		chunk := ImageSourceChunk{
 			Offset: uint64(blobSize - footerSize),
@@ -197,38 +175,35 @@ func readZstdChunkedManifest(blobStream ImageSourceSeekable, blobSize int64, ann
 			return nil, nil, 0, err
 		}
 
-		offset = binary.LittleEndian.Uint64(footer[0:8])
+		footerData, err = internal.ReadFooterDataFromBlob(footer)
-		length = binary.LittleEndian.Uint64(footer[8:16])
+		if err != nil {
-		lengthUncompressed = binary.LittleEndian.Uint64(footer[16:24])
+			return nil, nil, 0, err
-		manifestType = binary.LittleEndian.Uint64(footer[24:32])
-		if !isZstdChunkedFrameMagic(footer[48:56]) {
-			return nil, nil, 0, errors.New("invalid magic number")
 		}
 	}
 
-	if manifestType != internal.ManifestTypeCRFS {
+	if footerData.ManifestType != internal.ManifestTypeCRFS {
 		return nil, nil, 0, errors.New("invalid manifest type")
 	}
 
 	// set a reasonable limit
-	if length > (1<<20)*50 {
+	if footerData.LengthCompressed > (1<<20)*50 {
 		return nil, nil, 0, errors.New("manifest too big")
 	}
-	if lengthUncompressed > (1<<20)*50 {
+	if footerData.LengthUncompressed > (1<<20)*50 {
 		return nil, nil, 0, errors.New("manifest too big")
 	}
 
 	chunk := ImageSourceChunk{
-		Offset: offset,
+		Offset: footerData.Offset,
-		Length: length,
+		Length: footerData.LengthCompressed,
 	}
 
 	chunks := []ImageSourceChunk{chunk}
 
-	if offsetTarSplit > 0 {
+	if footerData.OffsetTarSplit > 0 {
 		chunkTarSplit := ImageSourceChunk{
-			Offset: offsetTarSplit,
+			Offset: footerData.OffsetTarSplit,
-			Length: lengthTarSplit,
+			Length: footerData.LengthCompressedTarSplit,
 		}
 		chunks = append(chunks, chunkTarSplit)
 	}
@@ -258,28 +233,28 @@ func readZstdChunkedManifest(blobStream ImageSourceSeekable, blobSize int64, ann
 		return blob, nil
 	}
 
-	manifest, err := readBlob(length)
+	manifest, err := readBlob(footerData.LengthCompressed)
 	if err != nil {
 		return nil, nil, 0, err
 	}
 
-	decodedBlob, err := decodeAndValidateBlob(manifest, lengthUncompressed, manifestChecksumAnnotation)
+	decodedBlob, err := decodeAndValidateBlob(manifest, footerData.LengthUncompressed, footerData.ChecksumAnnotation)
 	if err != nil {
 		return nil, nil, 0, err
 	}
 	decodedTarSplit := []byte{}
-	if offsetTarSplit > 0 {
+	if footerData.OffsetTarSplit > 0 {
-		tarSplit, err := readBlob(lengthTarSplit)
+		tarSplit, err := readBlob(footerData.LengthCompressedTarSplit)
 		if err != nil {
 			return nil, nil, 0, err
 		}
 
-		decodedTarSplit, err = decodeAndValidateBlob(tarSplit, lengthUncompressedTarSplit, tarSplitChecksumAnnotation)
+		decodedTarSplit, err = decodeAndValidateBlob(tarSplit, footerData.LengthUncompressedTarSplit, footerData.ChecksumAnnotationTarSplit)
 		if err != nil {
 			return nil, nil, 0, err
 		}
 	}
-	return decodedBlob, decodedTarSplit, int64(offset), err
+	return decodedBlob, decodedTarSplit, int64(footerData.Offset), err
 }
 
 func decodeAndValidateBlob(blob []byte, lengthUncompressed uint64, expectedUncompressedChecksum string) ([]byte, error) {

storage/pkg/chunked/internal/compression.go

@@ -8,6 +8,7 @@ import (
 	"archive/tar"
 	"bytes"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"io"
 	"time"
@@ -198,3 +199,67 @@ func ZstdWriterWithLevel(dest io.Writer, level int) (*zstd.Encoder, error) {
 	el := zstd.EncoderLevelFromZstd(level)
 	return zstd.NewWriter(dest, zstd.WithEncoderLevel(el))
 }
+
+// ZstdChunkedFooterData contains all the data stored in the zstd:chunked footer.
+type ZstdChunkedFooterData struct {
+	ManifestType uint64
+
+	Offset             uint64
+	LengthCompressed   uint64
+	LengthUncompressed uint64
+	ChecksumAnnotation string // Only used when reading a layer, not when creating it
+
+	OffsetTarSplit             uint64
+	LengthCompressedTarSplit   uint64
+	LengthUncompressedTarSplit uint64
+	ChecksumAnnotationTarSplit string // Only used when reading a layer, not when creating it
+}
+
+// ReadFooterDataFromAnnotations reads the zstd:chunked footer data from the given annotations.
+func ReadFooterDataFromAnnotations(annotations map[string]string) (ZstdChunkedFooterData, error) {
+	var footerData ZstdChunkedFooterData
+
+	footerData.ChecksumAnnotation = annotations[ManifestChecksumKey]
+	if footerData.ChecksumAnnotation == "" {
+		return footerData, fmt.Errorf("manifest checksum annotation %q not found", ManifestChecksumKey)
+	}
+
+	offsetMetadata := annotations[ManifestInfoKey]
+
+	if _, err := fmt.Sscanf(offsetMetadata, "%d:%d:%d:%d", &footerData.Offset, &footerData.LengthCompressed, &footerData.LengthUncompressed, &footerData.ManifestType); err != nil {
+		return footerData, err
+	}
+
+	if tarSplitInfoKeyAnnotation, found := annotations[TarSplitInfoKey]; found {
+		if _, err := fmt.Sscanf(tarSplitInfoKeyAnnotation, "%d:%d:%d", &footerData.OffsetTarSplit, &footerData.LengthCompressedTarSplit, &footerData.LengthUncompressedTarSplit); err != nil {
+			return footerData, err
+		}
+		footerData.ChecksumAnnotationTarSplit = annotations[TarSplitChecksumKey]
+	}
+	return footerData, nil
+}
+
+// IsZstdChunkedFrameMagic returns true if the given data starts with the zstd:chunked magic.
+func IsZstdChunkedFrameMagic(data []byte) bool {
+	if len(data) < 8 {
+		return false
+	}
+	return bytes.Equal(ZstdChunkedFrameMagic, data[:8])
+}
+
+// ReadFooterDataFromBlob reads the zstd:chunked footer from the binary buffer.
+func ReadFooterDataFromBlob(footer []byte) (ZstdChunkedFooterData, error) {
+	var footerData ZstdChunkedFooterData
+
+	if len(footer) < FooterSizeSupported {
+		return footerData, errors.New("blob too small")
+	}
+	footerData.Offset = binary.LittleEndian.Uint64(footer[0:8])
+	footerData.LengthCompressed = binary.LittleEndian.Uint64(footer[8:16])
+	footerData.LengthUncompressed = binary.LittleEndian.Uint64(footer[16:24])
+	footerData.ManifestType = binary.LittleEndian.Uint64(footer[24:32])
+	if !IsZstdChunkedFrameMagic(footer[48:56]) {
+		return footerData, errors.New("invalid magic number")
+	}
+	return footerData, nil
+}

storage/pkg/chunked/internal/compression_test.go

@@ -0,0 +1,20 @@
+//go:build linux
+// +build linux
+
+package internal
+
+import (
+	"testing"
+)
+
+func TestIsZstdChunkedFrameMagic(t *testing.T) {
+	b := append(ZstdChunkedFrameMagic[:], make([]byte, 200)...)
+	if !IsZstdChunkedFrameMagic(b) {
+		t.Fatal("Chunked frame magic not found")
+	}
+	// change a byte
+	b[0] = -b[0]
+	if IsZstdChunkedFrameMagic(b) {
+		t.Fatal("Invalid chunked frame magic found")
+	}
+}

storage/pkg/chunked/zstdchunked_test.go

@@ -16,18 +16,6 @@ import (
 	"github.com/opencontainers/go-digest"
 )
 
-func TestIsZstdChunkedFrameMagic(t *testing.T) {
-	b := append(internal.ZstdChunkedFrameMagic[:], make([]byte, 200)...)
-	if !isZstdChunkedFrameMagic(b) {
-		t.Fatal("Chunked frame magic not found")
-	}
-	// change a byte
-	b[0] = -b[0]
-	if isZstdChunkedFrameMagic(b) {
-		t.Fatal("Invalid chunked frame magic found")
-	}
-}
-
 type seekable struct {
 	data         []byte
 	tarSplitData []byte