1 // Copyright 2022 Google LLC All Rights Reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 // Package zstd provides helper functions for interacting with zstd streams. 16 package zstd 17 18 import ( 19 "bufio" 20 "bytes" 21 "io" 22 23 "github.com/google/go-containerregistry/internal/and" 24 "github.com/klauspost/compress/zstd" 25 ) 26 27 // MagicHeader is the start of zstd files. 28 var MagicHeader = []byte{'\x28', '\xb5', '\x2f', '\xfd'} 29 30 // ReadCloser reads uncompressed input data from the io.ReadCloser and 31 // returns an io.ReadCloser from which compressed data may be read. 32 // This uses zstd level 1 for the compression. 33 func ReadCloser(r io.ReadCloser) io.ReadCloser { 34 return ReadCloserLevel(r, 1) 35 } 36 37 // ReadCloserLevel reads uncompressed input data from the io.ReadCloser and 38 // returns an io.ReadCloser from which compressed data may be read. 39 func ReadCloserLevel(r io.ReadCloser, level int) io.ReadCloser { 40 pr, pw := io.Pipe() 41 42 // For highly compressible layers, zstd.Writer will output a very small 43 // number of bytes per Write(). This is normally fine, but when pushing 44 // to a registry, we want to ensure that we're taking full advantage of 45 // the available bandwidth instead of sending tons of tiny writes over 46 // the wire. 47 // 64K ought to be small enough for anybody. 48 bw := bufio.NewWriterSize(pw, 2<<16) 49 50 // Returns err so we can pw.CloseWithError(err) 51 go func() error { 52 // TODO(go1.14): Just defer {pw,zw,r}.Close like you'd expect. 53 // Context: https://golang.org/issue/24283 54 zw, err := zstd.NewWriter(bw, zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(level))) 55 if err != nil { 56 return pw.CloseWithError(err) 57 } 58 59 if _, err := io.Copy(zw, r); err != nil { 60 defer r.Close() 61 defer zw.Close() 62 return pw.CloseWithError(err) 63 } 64 65 // Close zstd writer to Flush it and write zstd trailers. 66 if err := zw.Close(); err != nil { 67 return pw.CloseWithError(err) 68 } 69 70 // Flush bufio writer to ensure we write out everything. 71 if err := bw.Flush(); err != nil { 72 return pw.CloseWithError(err) 73 } 74 75 // We don't really care if these fail. 76 defer pw.Close() 77 defer r.Close() 78 79 return nil 80 }() 81 82 return pr 83 } 84 85 // UnzipReadCloser reads compressed input data from the io.ReadCloser and 86 // returns an io.ReadCloser from which uncompressed data may be read. 87 func UnzipReadCloser(r io.ReadCloser) (io.ReadCloser, error) { 88 gr, err := zstd.NewReader(r) 89 if err != nil { 90 return nil, err 91 } 92 return &and.ReadCloser{ 93 Reader: gr, 94 CloseFunc: func() error { 95 // If the unzip fails, then this seems to return the same 96 // error as the read. We don't want this to interfere with 97 // us closing the main ReadCloser, since this could leave 98 // an open file descriptor (fails on Windows). 99 gr.Close() 100 return r.Close() 101 }, 102 }, nil 103 } 104 105 // Is detects whether the input stream is compressed. 106 func Is(r io.Reader) (bool, error) { 107 magicHeader := make([]byte, 4) 108 n, err := r.Read(magicHeader) 109 if n == 0 && err == io.EOF { 110 return false, nil 111 } 112 if err != nil { 113 return false, err 114 } 115 return bytes.Equal(magicHeader, MagicHeader), nil 116 } 117