compression: support zstd with skippable frame

As a matter of fact, there are two frame formats defined by Zstandard: Zstandard frames and Skippable frames.
So we should probably support zstd algorithms with skippable frames.
See https://tools.ietf.org/id/draft-kucherawy-dispatch-zstd-00.html#rfc.section.2 for more details.

Signed-off-by: Da McGrady <dabkb@aol.com>

pkg/archive/archive.go

@@ -7,6 +7,7 @@ import (
 	"compress/bzip2"
 	"compress/gzip"
 	"context"
+	"encoding/binary"
 	"fmt"
 	"io"
 	"os"
@@ -124,15 +125,66 @@ func IsArchivePath(path string) bool {
 	return err == nil
 }
 
+const (
+	ZstdMagicSkippableStart = 0x184D2A50
+	ZstdMagicSkippableMask  = 0xFFFFFFF0
+)
+
+var (
+	Bzip2Magic = []byte{0x42, 0x5A, 0x68}
+	GzipMagic  = []byte{0x1F, 0x8B, 0x08}
+	XzMagic    = []byte{0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00}
+	ZstdMagic  = []byte{0x28, 0xb5, 0x2f, 0xfd}
+)
+
+type matcher = func([]byte) bool
+
+func magicNumberMatcher(m []byte) matcher {
+	return func(source []byte) bool {
+		return bytes.HasPrefix(source, m)
+	}
+}
+
+// ZstDetector detects zstd compression algorithm.
+// Zstandard compressed data is made of one or more frames.
+// There are two frame formats defined by Zstandard: Zstandard frames and Skippable frames.
+// See https://tools.ietf.org/id/draft-kucherawy-dispatch-zstd-00.html#rfc.section.2 for more details.
+func ZstDetector(buf []byte) bool {
+	if bytes.HasPrefix(buf, ZstdMagic) {
+		// Zstandard frame
+		return true
+	} else {
+		// skippable frame
+		if len(buf) < 8 {
+			return false
+		}
+		if binary.LittleEndian.Uint32(buf[:4])&ZstdMagicSkippableMask == ZstdMagicSkippableStart {
+			userDataLength := binary.LittleEndian.Uint32(buf[4:8])
+			if len(buf) < 8+int(userDataLength) {
+				return false
+			}
+			nextFrame := buf[8+userDataLength:]
+			return ZstDetector(nextFrame)
+		}
+		return false
+	}
+}
+
+func zstdMatcher() matcher {
+	return func(source []byte) bool {
+		return ZstDetector(source)
+	}
+}
+
 // DetectCompression detects the compression algorithm of the source.
 func DetectCompression(source []byte) Compression {
-	for compression, m := range map[Compression][]byte{
-		Bzip2: {0x42, 0x5A, 0x68},
-		Gzip:  {0x1F, 0x8B, 0x08},
-		Xz:    {0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00},
-		Zstd:  {0x28, 0xb5, 0x2f, 0xfd},
+	for compression, fn := range map[Compression]matcher{
+		Bzip2: magicNumberMatcher(Bzip2Magic),
+		Gzip:  magicNumberMatcher(GzipMagic),
+		Xz:    magicNumberMatcher(XzMagic),
+		Zstd:  zstdMatcher(),
 	} {
-		if bytes.HasPrefix(source, m) {
+		if fn(source) {
 			return compression
 		}
 	}

pkg/archive/archive_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/containerd/containerd/pkg/userns"
 	"github.com/docker/docker/pkg/idtools"
 	"github.com/docker/docker/pkg/ioutils"
+	"github.com/klauspost/compress/zstd"
 	"gotest.tools/v3/assert"
 	is "gotest.tools/v3/assert/cmp"
 	"gotest.tools/v3/skip"
@@ -700,6 +701,56 @@ func tarUntar(t *testing.T, origin string, options *TarOptions) ([]Change, error
 	return ChangesDirs(origin, tmp)
 }
 
+func DecompressWithZstd(in []byte) []byte {
+	input := bytes.NewBuffer(in)
+	output := &bytes.Buffer{}
+	d, err := zstd.NewReader(input)
+	if err != nil {
+		return output.Bytes()
+	}
+	defer d.Close()
+	io.Copy(output, d)
+	return output.Bytes()
+}
+
+func CompressWithZstd(in []byte) []byte {
+	input := bytes.NewBuffer(in)
+	output := &bytes.Buffer{}
+	enc, _ := zstd.NewWriter(output)
+	io.Copy(enc, input)
+	enc.Close()
+	return output.Bytes()
+}
+
+func TestDetectCompression(t *testing.T) {
+	// test zstd compression without skippable frames.
+	compressedData := CompressWithZstd([]byte("hello world"))
+	compression := DetectCompression(compressedData)
+	if compression != Zstd {
+		t.Fatal("Unexpected compression")
+	}
+	// test zstd compression with skippable frames.
+	hex := []byte{
+		0x50, 0x2a, 0x4d, 0x18, // magic number of skippable frame: 0x184D2A50 to 0x184D2A5F
+		0x04, 0x00, 0x00, 0x00, // frame size
+		0x5d, 0x00, 0x00, 0x00, // user data
+		0x28, 0xb5, 0x2f, 0xfd, // magic number of Zstandard frame: 0xFD2FB528
+		0x04, 0x00, 0x31, 0x00, 0x00, // frame header
+		0x64, 0x6f, 0x63, 0x6b, 0x65, 0x72, // data block "docker"
+		0x16, 0x0e, 0x21, 0xc3, // content checksum
+	}
+	compression = DetectCompression(hex)
+	if compression != Zstd {
+		t.Fatal("Unexpected compression")
+	}
+
+	result := DecompressWithZstd(hex)
+	if string(result) != "docker" {
+		fmt.Println(result)
+		t.Fatal("Unexpected decompression result: " + string(result))
+	}
+}
+
 func TestTarUntar(t *testing.T) {
 	origin, err := os.MkdirTemp("", "docker-test-untar-origin")
 	if err != nil {