From 000cae8822676c7e7020f1f3447c9876751267a1 Mon Sep 17 00:00:00 2001 From: Yan Su Date: Thu, 5 Nov 2020 10:06:35 +0900 Subject: [PATCH] add zstd support --- decompress.go | 4 + decompress_tzst.go | 39 +++++++ decompress_tzst_test.go | 95 ++++++++++++++++++ decompress_zstd.go | 40 ++++++++ decompress_zstd_test.go | 34 +++++++ go.mod | 1 + go.sum | 2 + testdata/decompress-tzst/empty.tar.zst | Bin 0 -> 22 bytes testdata/decompress-tzst/multiple.tar.zst | Bin 0 -> 121 bytes testdata/decompress-tzst/multiple_dir.tar.zst | Bin 0 -> 153 bytes testdata/decompress-tzst/ordering.tar.zst | Bin 0 -> 164 bytes .../decompress-tzst/outside_parent.tar.zst | Bin 0 -> 151 bytes testdata/decompress-tzst/single.tar.zst | Bin 0 -> 104 bytes testdata/decompress-zst/single.zst | Bin 0 -> 17 bytes 14 files changed, 215 insertions(+) create mode 100644 decompress_tzst.go create mode 100644 decompress_tzst_test.go create mode 100644 decompress_zstd.go create mode 100644 decompress_zstd_test.go create mode 100644 testdata/decompress-tzst/empty.tar.zst create mode 100644 testdata/decompress-tzst/multiple.tar.zst create mode 100644 testdata/decompress-tzst/multiple_dir.tar.zst create mode 100644 testdata/decompress-tzst/ordering.tar.zst create mode 100644 testdata/decompress-tzst/outside_parent.tar.zst create mode 100644 testdata/decompress-tzst/single.tar.zst create mode 100644 testdata/decompress-zst/single.zst diff --git a/decompress.go b/decompress.go index 62ca1060f..d0ad43900 100644 --- a/decompress.go +++ b/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), } } diff --git a/decompress_tzst.go b/decompress_tzst.go new file mode 100644 index 000000000..a9f3da51e --- /dev/null +++ b/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) +} diff --git a/decompress_tzst_test.go b/decompress_tzst_test.go new file mode 100644 index 000000000..9787c7267 --- /dev/null +++ b/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) +} diff --git a/decompress_zstd.go b/decompress_zstd.go new file mode 100644 index 000000000..6ff6c86a9 --- /dev/null +++ b/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) +} diff --git a/decompress_zstd_test.go b/decompress_zstd_test.go new file mode 100644 index 000000000..6a0683589 --- /dev/null +++ b/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) +} diff --git a/go.mod b/go.mod index 14af6b718..9c0ae09b7 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index c16d1b4c2..99c5ea66b 100644 --- a/go.sum +++ b/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= diff --git a/testdata/decompress-tzst/empty.tar.zst b/testdata/decompress-tzst/empty.tar.zst new file mode 100644 index 0000000000000000000000000000000000000000..d123f19dfd728faef8209ec9b3a285107ab4ce20 GIT binary patch literal 22 dcmdPcs{fZI!IyzSfPsOL;deWG#Gbu%*8oX~2YLVi literal 0 HcmV?d00001 diff --git a/testdata/decompress-tzst/multiple.tar.zst b/testdata/decompress-tzst/multiple.tar.zst new file mode 100644 index 0000000000000000000000000000000000000000..5eb8008bffdf2368481b046314657c87f25c6aec GIT binary patch literal 121 zcmV-<0EYi4wJ-eyU}XaUV#N>?prH*UZqb2u;#buVN=izs`L$R1r1;JvwLlPBx2j2n zhut#(5(@klm{bh>@-Q3r${&s+m}iA}Ct3cx5c9w7EYc7luv6Os{oljbdNA+8y{#*klBug9tVu{5BYy4ld<9LwjYSK{Tpom3OUD%zqb#WKa+R79b!4 zqqxuk{oRAHT%7~R1jz^gw`KNJH<)U=iy28%2QUN)r!XwAl!10p3Be?sg$Hx@GCDaU HjbMqH8T3J+ literal 0 HcmV?d00001 diff --git a/testdata/decompress-tzst/ordering.tar.zst b/testdata/decompress-tzst/ordering.tar.zst new file mode 100644 index 0000000000000000000000000000000000000000..d9cbbeeefc3ac468495999a5dec089e17fa593d9 GIT binary patch literal 164 zcmV;V09*ekwJ-eyV7&wYf=3q_fT;~Y*zQIxL)?oKGNB}WZjz5G95$Pg0_9rJmzzKP zTG6N9T=e>1+h%~J&EE5g^cO?froVV=CDW2N*+6oN(1iC=?j3T__P++=N||N&-`5 FVvE3mK#c$Z literal 0 HcmV?d00001 diff --git a/testdata/decompress-tzst/single.tar.zst b/testdata/decompress-tzst/single.tar.zst new file mode 100644 index 0000000000000000000000000000000000000000..551371c387edb9e965aee440d95469b7fe2b434e GIT binary patch literal 104 zcmV-u0GIzLwJ-eyVBG=$(u5BbkO2nJ Kr=SH^v;eulpD`f- literal 0 HcmV?d00001 diff --git a/testdata/decompress-zst/single.zst b/testdata/decompress-zst/single.zst new file mode 100644 index 0000000000000000000000000000000000000000..5a56316bf434503831e56eb69600d6f0392eb521 GIT binary patch literal 17 YcmdPcs{dDoMUjCaEkB=2H&jIe04*m3UjP6A literal 0 HcmV?d00001