-
Notifications
You must be signed in to change notification settings - Fork 538
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Squashed commits: [6563fc8] Run ./hack/update-codegen.sh [c86a013] Fix linting issues [353dd69] Remove redundant break statements [a06b4cc] Fix some comments [fb0678f] Expose compression algorithms enum [3321d0c] Update mediaType according to compression for tarball layers (cherry picked from commit 56d4df5) [5124d2c] Fix unit tests (cherry picked from commit 7ce08d1) [2958f30] Optimize compression detection (cherry picked from commit 0a3437f) [c211a61] Add support for zstd-compressedd layers (cherry picked from commit d8fb0f3) [0b9a996] Add zstd compression utilities (cherry picked from commit c645201) [7dd71ab] Add layer types with zstd compression (cherry picked from commit 328a14b) Signed-off-by: Lavrenti Frobeen <lavrenti@northflank.com>
- Loading branch information
Showing
10 changed files
with
382 additions
and
67 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,84 @@ | ||
| package compression | ||
|
|
||
| import ( | ||
| "bufio" | ||
| "bytes" | ||
| "github.com/google/go-containerregistry/internal/gzip" | ||
| "github.com/google/go-containerregistry/internal/zstd" | ||
| "io" | ||
| ) | ||
|
|
||
| type Compression string | ||
|
|
||
| // The collection of known MediaType values. | ||
| const ( | ||
| None Compression = "none" | ||
| GZip Compression = "gzip" | ||
| ZStd Compression = "zstd" | ||
| ) | ||
|
|
||
| type Opener = func() (io.ReadCloser, error) | ||
|
|
||
| func GetCompression(opener Opener) (Compression, error) { | ||
| rc, err := opener() | ||
| if err != nil { | ||
| return None, err | ||
| } | ||
| defer rc.Close() | ||
|
|
||
| compression, _, err := PeekCompression(rc) | ||
| if err != nil { | ||
| return None, err | ||
| } | ||
|
|
||
| return compression, nil | ||
| } | ||
|
|
||
| // PeekReader is an io.Reader that also implements Peek a la bufio.Reader. | ||
| type PeekReader interface { | ||
| io.Reader | ||
| Peek(n int) ([]byte, error) | ||
| } | ||
|
|
||
| // PeekCompression detects whether the input stream is compressed and which algorithm is used. | ||
| // | ||
| // If r implements Peek, we will use that directly, otherwise a small number | ||
| // of bytes are buffered to Peek at the gzip header, and the returned | ||
| // PeekReader can be used as a replacement for the consumed input io.Reader. | ||
| func PeekCompression(r io.Reader) (Compression, PeekReader, error) { | ||
| var pr PeekReader | ||
| if p, ok := r.(PeekReader); ok { | ||
| pr = p | ||
| } else { | ||
| pr = bufio.NewReader(r) | ||
| } | ||
|
|
||
| var header []byte | ||
| var err error | ||
|
|
||
| if header, err = pr.Peek(2); err != nil { | ||
| // https://github.com/google/go-containerregistry/issues/367 | ||
| if err == io.EOF { | ||
| return None, pr, nil | ||
| } | ||
| return None, pr, err | ||
| } | ||
|
|
||
| if bytes.Equal(header, gzip.MagicHeader) { | ||
| return GZip, pr, nil | ||
| } | ||
|
|
||
| if header, err = pr.Peek(4); err != nil { | ||
| // https://github.com/google/go-containerregistry/issues/367 | ||
| if err == io.EOF { | ||
| return None, pr, nil | ||
| } | ||
| return None, pr, err | ||
| } | ||
|
|
||
| if bytes.Equal(header, zstd.MagicHeader) { | ||
| return ZStd, pr, nil | ||
| } | ||
|
|
||
| return None, pr, nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,99 @@ | ||
| package zstd | ||
|
|
||
| import ( | ||
| "bufio" | ||
| "bytes" | ||
| "github.com/google/go-containerregistry/internal/and" | ||
| "github.com/klauspost/compress/zstd" | ||
| "io" | ||
| ) | ||
|
|
||
| var MagicHeader = []byte{'\x28', '\xb5', '\x2f', '\xfd'} | ||
|
|
||
| // ReadCloser reads uncompressed input data from the io.ReadCloser and | ||
| // returns an io.ReadCloser from which compressed data may be read. | ||
| // This uses zstd level 1 for the compression. | ||
| func ReadCloser(r io.ReadCloser) io.ReadCloser { | ||
| return ReadCloserLevel(r, 1) | ||
| } | ||
|
|
||
| // ReadCloserLevel reads uncompressed input data from the io.ReadCloser and | ||
| // returns an io.ReadCloser from which compressed data may be read. | ||
| func ReadCloserLevel(r io.ReadCloser, level int) io.ReadCloser { | ||
| pr, pw := io.Pipe() | ||
|
|
||
| // For highly compressible layers, zstd.Writer will output a very small | ||
| // number of bytes per Write(). This is normally fine, but when pushing | ||
| // to a registry, we want to ensure that we're taking full advantage of | ||
| // the available bandwidth instead of sending tons of tiny writes over | ||
| // the wire. | ||
| // 64K ought to be small enough for anybody. | ||
| bw := bufio.NewWriterSize(pw, 2<<16) | ||
|
|
||
| // Returns err so we can pw.CloseWithError(err) | ||
| go func() error { | ||
| // TODO(go1.14): Just defer {pw,gw,r}.Close like you'd expect. | ||
| // Context: https://golang.org/issue/24283 | ||
| gw, err := zstd.NewWriter(bw, zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(level))) | ||
| if err != nil { | ||
| return pw.CloseWithError(err) | ||
| } | ||
|
|
||
| if _, err := io.Copy(gw, r); err != nil { | ||
| defer r.Close() | ||
| defer gw.Close() | ||
| return pw.CloseWithError(err) | ||
| } | ||
|
|
||
| // Close zstd writer to Flush it and write zstd trailers. | ||
| if err := gw.Close(); err != nil { | ||
| return pw.CloseWithError(err) | ||
| } | ||
|
|
||
| // Flush bufio writer to ensure we write out everything. | ||
| if err := bw.Flush(); err != nil { | ||
| return pw.CloseWithError(err) | ||
| } | ||
|
|
||
| // We don't really care if these fail. | ||
| defer pw.Close() | ||
| defer r.Close() | ||
|
|
||
| return nil | ||
| }() | ||
|
|
||
| return pr | ||
| } | ||
|
|
||
| // UnzipReadCloser reads compressed input data from the io.ReadCloser and | ||
| // returns an io.ReadCloser from which uncompessed data may be read. | ||
| func UnzipReadCloser(r io.ReadCloser) (io.ReadCloser, error) { | ||
| gr, err := zstd.NewReader(r) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| return &and.ReadCloser{ | ||
| Reader: gr, | ||
| CloseFunc: func() error { | ||
| // If the unzip fails, then this seems to return the same | ||
| // error as the read. We don't want this to interfere with | ||
| // us closing the main ReadCloser, since this could leave | ||
| // an open file descriptor (fails on Windows). | ||
| gr.Close() | ||
| return r.Close() | ||
| }, | ||
| }, nil | ||
| } | ||
|
|
||
| // Is detects whether the input stream is compressed. | ||
| func Is(r io.Reader) (bool, error) { | ||
| magicHeader := make([]byte, 4) | ||
| n, err := r.Read(magicHeader) | ||
| if n == 0 && err == io.EOF { | ||
| return false, nil | ||
| } | ||
| if err != nil { | ||
| return false, err | ||
| } | ||
| return bytes.Equal(magicHeader, MagicHeader), nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,101 @@ | ||
| // Copyright 2020 Google LLC All Rights Reserved. | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package zstd | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "fmt" | ||
| "io" | ||
| "testing" | ||
| ) | ||
|
|
||
| func TestReader(t *testing.T) { | ||
| want := "This is the input string." | ||
| buf := bytes.NewBufferString(want) | ||
| zipped := ReadCloser(io.NopCloser(buf)) | ||
| unzipped, err := UnzipReadCloser(zipped) | ||
| if err != nil { | ||
| t.Error("UnzipReadCloser() =", err) | ||
| } | ||
|
|
||
| b, err := io.ReadAll(unzipped) | ||
| if err != nil { | ||
| t.Error("ReadAll() =", err) | ||
| } | ||
| if got := string(b); got != want { | ||
| t.Errorf("ReadAll(); got %q, want %q", got, want) | ||
| } | ||
| if err := unzipped.Close(); err != nil { | ||
| t.Error("Close() =", err) | ||
| } | ||
| } | ||
|
|
||
| func TestIs(t *testing.T) { | ||
| tests := []struct { | ||
| in []byte | ||
| out bool | ||
| err error | ||
| }{ | ||
| {[]byte{}, false, nil}, | ||
| {[]byte{'\x00', '\x00', '\x00', '\x00', '\x00'}, false, nil}, | ||
| {[]byte{'\x28', '\xb5', '\x2f', '\xfd', '\x1b'}, true, nil}, | ||
| } | ||
| for _, test := range tests { | ||
| reader := bytes.NewReader(test.in) | ||
| got, err := Is(reader) | ||
| if got != test.out { | ||
| t.Errorf("Is; n: got %v, wanted %v\n", got, test.out) | ||
| } | ||
| if err != test.err { | ||
| t.Errorf("Is; err: got %v, wanted %v\n", err, test.err) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| var ( | ||
| errRead = fmt.Errorf("read failed") | ||
| ) | ||
|
|
||
| type failReader struct{} | ||
|
|
||
| func (f failReader) Read(_ []byte) (int, error) { | ||
| return 0, errRead | ||
| } | ||
|
|
||
| func TestReadErrors(t *testing.T) { | ||
| fr := failReader{} | ||
| if _, err := Is(fr); err != errRead { | ||
| t.Error("Is: expected errRead, got", err) | ||
| } | ||
|
|
||
| frc := io.NopCloser(fr) | ||
| if r, err := UnzipReadCloser(frc); err != errRead { | ||
| data := make([]byte, 100) | ||
| _, err := r.Read(data) | ||
| if err != errRead { | ||
| t.Error("UnzipReadCloser: expected errRead, got", err) | ||
| } | ||
| } | ||
|
|
||
| zr := ReadCloser(io.NopCloser(fr)) | ||
| if _, err := zr.Read(nil); err != errRead { | ||
| t.Error("ReadCloser: expected errRead, got", err) | ||
| } | ||
|
|
||
| //zr = ReadCloserLevel(io.NopCloser(strings.NewReader("zip me")), -10) | ||
| //if _, err := zr.Read(nil); err == nil { | ||
| // t.Error("Expected invalid level error, got:", err) | ||
| //} | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.