From e187eb2bb5f0c3f899fe643e95d1af8c57e89a73 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 9624b2de74..a817b6e3b9 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -276,7 +276,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 6d7756f9b6..44e193d9fa 100644
--- a/pkg/archive/archive.go
+++ b/pkg/archive/archive.go
@@ -22,6 +22,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"
 )
@@ -83,6 +84,8 @@ const (
 	Gzip
 	// Xz is xz compression algorithm.
 	Xz
+	// Zstd is zstd compression algorithm.
+	Zstd
 )
 
 const (
@@ -127,6 +130,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
@@ -215,6 +219,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())
 	}
@@ -341,6 +352,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 da13d38912..0331c825ba 100644
--- a/pkg/archive/archive_test.go
+++ b/pkg/archive/archive_test.go
@@ -135,6 +135,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 {
@@ -210,6 +217,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")