Merge pull request #292 from yaroot/zstd

add zstd support

Author: schmichael

decompress.go

@@ -26,6 +26,7 @@ func init() {
 	tbzDecompressor := new(TarBzip2Decompressor)
 	tgzDecompressor := new(TarGzipDecompressor)
 	txzDecompressor := new(TarXzDecompressor)
+	tzstDecompressor := new(TarZstdDecompressor)
 
 	Decompressors = map[string]Decompressor{
 		"bz2":     new(Bzip2Decompressor),
@@ -34,10 +35,13 @@ func init() {
 		"tar.bz2": tbzDecompressor,
 		"tar.gz":  tgzDecompressor,
 		"tar.xz":  txzDecompressor,
+		"tar.zst": tzstDecompressor,
 		"tbz2":    tbzDecompressor,
 		"tgz":     tgzDecompressor,
 		"txz":     txzDecompressor,
+		"tzst":    tzstDecompressor,
 		"zip":     new(ZipDecompressor),
+		"zst":     new(ZstdDecompressor),
 	}
 }
 

decompress_tzst.go

@@ -0,0 +1,39 @@
+package getter
+
+import (
+	"fmt"
+	"github.com/klauspost/compress/zstd"
+	"os"
+	"path/filepath"
+)
+
+// TarZstdDecompressor is an implementation of Decompressor that can
+// decompress tar.zstd files.
+type TarZstdDecompressor struct{}
+
+func (d *TarZstdDecompressor) Decompress(dst, src string, dir bool, umask os.FileMode) error {
+	// If we're going into a directory we should make that first
+	mkdir := dst
+	if !dir {
+		mkdir = filepath.Dir(dst)
+	}
+	if err := os.MkdirAll(mkdir, mode(0755, umask)); err != nil {
+		return err
+	}
+
+	// File first
+	f, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	// Zstd compression is second
+	zstdR, err := zstd.NewReader(f)
+	if err != nil {
+		return fmt.Errorf("Error opening a zstd reader for %s: %s", src, err)
+	}
+	defer zstdR.Close()
+
+	return untar(zstdR, dst, src, dir, umask)
+}

decompress_tzst_test.go

@@ -0,0 +1,95 @@
+package getter
+
+import (
+	"path/filepath"
+	"testing"
+)
+
+func TestTarZstdDecompressor(t *testing.T) {
+
+	multiplePaths := []string{"dir/", "dir/test2", "test1"}
+	orderingPaths := []string{"workers/", "workers/mq/", "workers/mq/__init__.py"}
+
+	cases := []TestDecompressCase{
+		{
+			"empty.tar.zst",
+			false,
+			true,
+			nil,
+			"",
+			nil,
+		},
+
+		{
+			"single.tar.zst",
+			false,
+			false,
+			nil,
+			"d3b07384d113edec49eaa6238ad5ff00",
+			nil,
+		},
+
+		{
+			"single.tar.zst",
+			true,
+			false,
+			[]string{"file"},
+			"",
+			nil,
+		},
+
+		{
+			"multiple.tar.zst",
+			true,
+			false,
+			[]string{"file1", "file2"},
+			"",
+			nil,
+		},
+
+		{
+			"multiple.tar.zst",
+			false,
+			true,
+			nil,
+			"",
+			nil,
+		},
+
+		{
+			"multiple_dir.tar.zst",
+			true,
+			false,
+			multiplePaths,
+			"",
+			nil,
+		},
+
+		// Tests when the file is listed before the parent folder
+		{
+			"ordering.tar.zst",
+			true,
+			false,
+			orderingPaths,
+			"",
+			nil,
+		},
+
+		// Tests that a tar.zst can't contain references with "..".
+		// GNU `tar` also disallows this.
+		{
+			"outside_parent.tar.zst",
+			true,
+			true,
+			nil,
+			"",
+			nil,
+		},
+	}
+
+	for i, tc := range cases {
+		cases[i].Input = filepath.Join("./testdata", "decompress-tzst", tc.Input)
+	}
+
+	TestDecompressor(t, new(TarZstdDecompressor), cases)
+}

decompress_zstd.go

@@ -0,0 +1,40 @@
+package getter
+
+import (
+	"fmt"
+	"github.com/klauspost/compress/zstd"
+	"os"
+	"path/filepath"
+)
+
+// ZstdDecompressor is an implementation of Decompressor that
+// can decompress .zst files.
+type ZstdDecompressor struct{}
+
+func (d *ZstdDecompressor) Decompress(dst, src string, dir bool, umask os.FileMode) error {
+	if dir {
+		return fmt.Errorf("zstd-compressed files can only unarchive to a single file")
+	}
+
+	// If we're going into a directory we should make that first
+	if err := os.MkdirAll(filepath.Dir(dst), mode(0755, umask)); err != nil {
+		return err
+	}
+
+	// File first
+	f, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	// zstd compression is second
+	zstdR, err := zstd.NewReader(f)
+	if err != nil {
+		return err
+	}
+	defer zstdR.Close()
+
+	// Copy it out
+	return copyReader(dst, zstdR, 0622, umask)
+}

decompress_zstd_test.go

@@ -0,0 +1,34 @@
+package getter
+
+import (
+	"path/filepath"
+	"testing"
+)
+
+func TestZstdDecompressor(t *testing.T) {
+	cases := []TestDecompressCase{
+		{
+			"single.zst",
+			false,
+			false,
+			nil,
+			"d3b07384d113edec49eaa6238ad5ff00",
+			nil,
+		},
+
+		{
+			"single.zst",
+			true,
+			true,
+			nil,
+			"",
+			nil,
+		},
+	}
+
+	for i, tc := range cases {
+		cases[i].Input = filepath.Join("./testdata", "decompress-zst", tc.Input)
+	}
+
+	TestDecompressor(t, new(ZstdDecompressor), cases)
+}

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/hashicorp/go-cleanhttp v0.5.0
 	github.com/hashicorp/go-safetemp v1.0.0
 	github.com/hashicorp/go-version v1.1.0
+	github.com/klauspost/compress v1.11.2
 	github.com/mattn/go-colorable v0.0.9 // indirect
 	github.com/mattn/go-isatty v0.0.4 // indirect
 	github.com/mattn/go-runewidth v0.0.4 // indirect

go.sum

@@ -54,6 +54,8 @@ github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ
 github.com/jmespath/go-jmespath v0.0.0-20160202185014-0b12d6b521d8 h1:12VvqtR6Aowv3l/EQUlocDHW2Cp4G9WJVH7uyH8QFJE=
 github.com/jmespath/go-jmespath v0.0.0-20160202185014-0b12d6b521d8/go.mod h1:Nht3zPeWKUH0NzdCt2Blrr5ys8VGpn0CEB0cQHVjt7k=
 github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU=
+github.com/klauspost/compress v1.11.2 h1:MiK62aErc3gIiVEtyzKfeOHgW7atJb5g/KNX5m3c2nQ=
+github.com/klauspost/compress v1.11.2/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
 github.com/mattn/go-colorable v0.0.9 h1:UVL0vNpWh04HeJXV0KLcaT7r06gOH2l4OW6ddYRUIY4=
 github.com/mattn/go-colorable v0.0.9/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU=
 github.com/mattn/go-isatty v0.0.4 h1:bnP0vzxcAdeI1zdubAl5PjU6zsERjGZb7raWodagDYs=