From 3d6895fac5a57106eb5aac687d1da6de61442a55 Mon Sep 17 00:00:00 2001
From: Giuseppe Scrivano <gscrivan@redhat.com>
Date: Tue, 8 Dec 2020 10:56:32 +0100
Subject: [PATCH] compression: add support for the zstd algorithm

zstd is a compression algorithm that has a very fast decoder, while
providing also good compression ratios.  The fast decoder makes it
suitable for container images, as decompressing the tarballs is a very
expensive operation.

https://github.com/opencontainers/image-spec/pull/788 added support
for zstd to the OCI image specs.

Signed-off-by: Giuseppe Scrivano <gscrivan@redhat.com>
---
 Dockerfile                  |  3 ++-
 pkg/archive/archive.go      | 13 +++++++++++++
 pkg/archive/archive_test.go | 14 ++++++++++++++
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 20dea15b82243..b3021d02de171 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -292,7 +292,8 @@ RUN --mount=type=cache,sharing=locked,id=moby-dev-aptlib,target=/var/lib/apt \
             vim-common \
             xfsprogs \
             xz-utils \
-            zip
+            zip \
+            zstd
 
 
 # Switch to use iptables instead of nftables (to match the CI hosts)
diff --git a/pkg/archive/archive.go b/pkg/archive/archive.go
index ebc60ca0f75e1..2f4a83146569e 100644
--- a/pkg/archive/archive.go
+++ b/pkg/archive/archive.go
@@ -23,6 +23,7 @@ import (
 	"github.com/docker/docker/pkg/ioutils"
 	"github.com/docker/docker/pkg/pools"
 	"github.com/docker/docker/pkg/system"
+	"github.com/klauspost/compress/zstd"
 	"github.com/sirupsen/logrus"
 	exec "golang.org/x/sys/execabs"
 )
@@ -84,6 +85,8 @@ const (
 	Gzip
 	// Xz is xz compression algorithm.
 	Xz
+	// Zstd is zstd compression algorithm.
+	Zstd
 )
 
 const (
@@ -128,6 +131,7 @@ func DetectCompression(source []byte) Compression {
 		Bzip2: {0x42, 0x5A, 0x68},
 		Gzip:  {0x1F, 0x8B, 0x08},
 		Xz:    {0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00},
+		Zstd:  {0x28, 0xb5, 0x2f, 0xfd},
 	} {
 		if bytes.HasPrefix(source, m) {
 			return compression
@@ -221,6 +225,13 @@ func DecompressStream(archive io.Reader) (io.ReadCloser, error) {
 		}
 		readBufWrapper := p.NewReadCloserWrapper(buf, xzReader)
 		return wrapReadCloser(readBufWrapper, cancel), nil
+	case Zstd:
+		zstdReader, err := zstd.NewReader(buf)
+		if err != nil {
+			return nil, err
+		}
+		readBufWrapper := p.NewReadCloserWrapper(buf, zstdReader)
+		return readBufWrapper, nil
 	default:
 		return nil, fmt.Errorf("Unsupported compression format %s", (&compression).Extension())
 	}
@@ -345,6 +356,8 @@ func (compression *Compression) Extension() string {
 		return "tar.gz"
 	case Xz:
 		return "tar.xz"
+	case Zstd:
+		return "tar.zst"
 	}
 	return ""
 }
diff --git a/pkg/archive/archive_test.go b/pkg/archive/archive_test.go
index d7632e1f06d3e..c3490cf5d32c0 100644
--- a/pkg/archive/archive_test.go
+++ b/pkg/archive/archive_test.go
@@ -136,6 +136,13 @@ func TestDecompressStreamXz(t *testing.T) {
 	testDecompressStream(t, "xz", "xz -f")
 }
 
+func TestDecompressStreamZstd(t *testing.T) {
+	if _, err := exec.LookPath("zstd"); err != nil {
+		t.Skip("zstd not installed")
+	}
+	testDecompressStream(t, "zst", "zstd -f")
+}
+
 func TestCompressStreamXzUnsupported(t *testing.T) {
 	dest, err := os.Create(tmp + "dest")
 	if err != nil {
@@ -211,6 +218,13 @@ func TestExtensionXz(t *testing.T) {
 		t.Fatalf("The extension of a xz archive should be 'tar.xz'")
 	}
 }
+func TestExtensionZstd(t *testing.T) {
+	compression := Zstd
+	output := compression.Extension()
+	if output != "tar.zst" {
+		t.Fatalf("The extension of a zstd archive should be 'tar.zst'")
+	}
+}
 
 func TestCmdStreamLargeStderr(t *testing.T) {
 	cmd := exec.Command("sh", "-c", "dd if=/dev/zero bs=1k count=1000 of=/dev/stderr; echo hello")