forked from google/go-containerregistry
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
(cherry picked from commit c645201)
- Loading branch information
Showing
2 changed files
with
226 additions
and
0 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,128 @@ | ||
package zstd | ||
|
||
import ( | ||
"bufio" | ||
"bytes" | ||
"github.com/google/go-containerregistry/internal/and" | ||
"github.com/klauspost/compress/zstd" | ||
"io" | ||
) | ||
|
||
var zstdMagicHeader = []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 gzip.BestSpeed for the compression level. | ||
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, gzip.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 gzip writer to Flush it and write gzip 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, zstdMagicHeader), nil | ||
} | ||
|
||
// PeekReader is an io.Reader that also implements Peek a la bufio.Reader. | ||
type PeekReader interface { | ||
io.Reader | ||
Peek(n int) ([]byte, error) | ||
} | ||
|
||
// Peek detects whether the input stream is gzip compressed. | ||
// | ||
// 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 Peek(r io.Reader) (bool, PeekReader, error) { | ||
var pr PeekReader | ||
if p, ok := r.(PeekReader); ok { | ||
pr = p | ||
} else { | ||
pr = bufio.NewReader(r) | ||
} | ||
header, err := pr.Peek(2) | ||
if err != nil { | ||
// https://github.com/google/go-containerregistry/issues/367 | ||
if err == io.EOF { | ||
return false, pr, nil | ||
} | ||
return false, pr, err | ||
} | ||
return bytes.Equal(header, zstdMagicHeader), 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,98 @@ | ||
// 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/ioutil" | ||
"strings" | ||
"testing" | ||
) | ||
|
||
func TestReader(t *testing.T) { | ||
want := "This is the input string." | ||
buf := bytes.NewBufferString(want) | ||
zipped := ReadCloser(ioutil.NopCloser(buf)) | ||
unzipped, err := UnzipReadCloser(zipped) | ||
if err != nil { | ||
t.Error("UnzipReadCloser() =", err) | ||
} | ||
|
||
b, err := ioutil.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 := ioutil.NopCloser(fr) | ||
if _, err := UnzipReadCloser(frc); err != errRead { | ||
t.Error("UnzipReadCloser: expected errRead, got", err) | ||
} | ||
|
||
zr := ReadCloser(ioutil.NopCloser(fr)) | ||
if _, err := zr.Read(nil); err != errRead { | ||
t.Error("ReadCloser: expected errRead, got", err) | ||
} | ||
|
||
zr = ReadCloserLevel(ioutil.NopCloser(strings.NewReader("zip me")), -10) | ||
if _, err := zr.Read(nil); err == nil { | ||
t.Error("Expected invalid level error, got:", err) | ||
} | ||
} |