diff --git a/tar.go b/tar.go index 9ef79b0..15ef541 100644 --- a/tar.go +++ b/tar.go @@ -4,13 +4,13 @@ import ( "archive/tar" "bytes" "compress/bzip2" - "compress/gzip" "errors" "fmt" "io" "io/fs" "strings" + "github.com/klauspost/compress/gzip" "github.com/klauspost/compress/zstd" "github.com/ulikunitz/xz" ) diff --git a/tar_read_test.go b/tar_read_test.go index 51a414f..ebd54ea 100644 --- a/tar_read_test.go +++ b/tar_read_test.go @@ -178,6 +178,35 @@ func TestTarTruncatedPayload(t *testing.T) { } } +func TestTarGzipMultistream(t *testing.T) { + // Alpine .apk packages are concatenated gzip members with the tar + // stream spanning the boundary; the reader must transparently continue + // into the next member. + raw := tarWithPayloads(t, []byte("aaa"), []byte("bbbb")) + split := 1024 // after first header + 512-byte padded payload + data := append(gzipTar(t, raw[:split]), gzipTar(t, raw[split:])...) + r, err := OpenBytes("alpine.apk", data) + if err != nil { + t.Fatal(err) + } + defer func() { _ = r.Close() }() + files, err := r.List() + if err != nil || len(files) != 2 { + t.Fatalf("List = %d entries, %v", len(files), err) + } + assertTarPayload(t, r, "package/lib/file0", []byte("aaa")) + assertTarPayload(t, r, "package/lib/file1", []byte("bbbb")) +} + +func TestTarGzipHeaderError(t *testing.T) { + data := gzipTar(t, tarWithPayloads(t, []byte("x"))) + data[0] = 0 // clobber gzip magic + _, err := OpenBytes("test.tar.gz", data) + if !errors.Is(err, gzip.ErrHeader) { + t.Fatalf("got %v, want gzip header error", err) + } +} + func TestTarPayloadLimitBoundary(t *testing.T) { oldMax := maxDecompressedSize maxDecompressedSize = 1024 diff --git a/zip.go b/zip.go index 838fff9..ac00e48 100644 --- a/zip.go +++ b/zip.go @@ -2,11 +2,15 @@ package archives import ( "archive/zip" + "bufio" "bytes" "encoding/binary" "fmt" "io" "strings" + "sync" + + "github.com/klauspost/compress/flate" ) // zipUnixModeShift is the bit offset of the Unix st_mode field within a @@ -48,6 +52,7 @@ func openZip(raw []byte) (*zipReader, error) { if err != nil { return nil, fmt.Errorf("opening zip: %w", err) } + reader.RegisterDecompressor(zip.Deflate, newFlateReader) if err := checkArchiveEntryCount(len(reader.File)); err != nil { return nil, err } @@ -243,6 +248,49 @@ func zipHasUnixMode(h *zip.FileHeader) bool { } } +// RegisterDecompressor bypasses archive/zip's internal flate reader pool, +// so pool klauspost readers here. archive/zip hands the decompressor an +// io.SectionReader with no ReadByte, which klauspost's Reset would wrap in +// a fresh bufio.Reader on every call; pooling one alongside the decoder and +// passing it to Reset avoids that and hits klauspost's *bufio.Reader fast +// path. +var flateReaderPool sync.Pool + +type flateState struct { + br *bufio.Reader + fr io.ReadCloser +} + +type pooledFlateReader struct { + s *flateState +} + +func newFlateReader(r io.Reader) io.ReadCloser { + s, _ := flateReaderPool.Get().(*flateState) + if s == nil { + br := bufio.NewReader(r) + return &pooledFlateReader{s: &flateState{br: br, fr: flate.NewReader(br)}} + } + s.br.Reset(r) + _ = s.fr.(flate.Resetter).Reset(s.br, nil) + return &pooledFlateReader{s: s} +} + +func (p *pooledFlateReader) Read(b []byte) (int, error) { + return p.s.fr.Read(b) +} + +func (p *pooledFlateReader) Close() error { + if p.s == nil { + return nil + } + err := p.s.fr.Close() + p.s.br.Reset(nil) + flateReaderPool.Put(p.s) + p.s = nil + return err +} + func extractName(path string) string { path = strings.TrimSuffix(path, "/") if idx := strings.LastIndex(path, "/"); idx >= 0 { diff --git a/zip_bench_test.go b/zip_bench_test.go new file mode 100644 index 0000000..37d6f5b --- /dev/null +++ b/zip_bench_test.go @@ -0,0 +1,73 @@ +package archives + +import ( + "archive/zip" + "bytes" + "fmt" + "io" + "math/rand" + "testing" +) + +// createBenchZip builds a zip of fileCount deflated members. Payload bytes +// are drawn from a 4-symbol alphabet giving a ~34% deflate ratio, so the +// decoder does real Huffman and match-copy work; fully random data would +// be emitted as stored blocks and skip the decoder. +func createBenchZip(fileCount, fileSize int) []byte { + rnd := rand.New(rand.NewSource(1)) //nolint:gosec + payload := make([]byte, fileSize) + buf := new(bytes.Buffer) + w := zip.NewWriter(buf) + for i := range fileCount { + for j := range payload { + payload[j] = byte(rnd.Intn(4)) + } + f, _ := w.CreateHeader(&zip.FileHeader{ + Name: fmt.Sprintf("lib/file%04d.dat", i), + Method: zip.Deflate, + }) + _, _ = f.Write(payload) + } + _ = w.Close() + return buf.Bytes() +} + +const ( + benchZipFileCount = 64 + benchZipFileSize = 16 * 1024 +) + +var benchZipArchive = createBenchZip(benchZipFileCount, benchZipFileSize) + +// BenchmarkZipExtract measures Open + draining every deflated member. +// zip decompresses lazily on Extract, so this is where the flate decoder +// cost lands; ExtractAll adds filesystem noise on top of the same loop. +func BenchmarkZipExtract(b *testing.B) { + b.SetBytes(int64(benchZipFileCount * benchZipFileSize)) + b.ReportAllocs() + for b.Loop() { + r, err := OpenBytes("bench.zip", benchZipArchive) + if err != nil { + b.Fatal(err) + } + files, err := r.List() + if err != nil { + b.Fatal(err) + } + for _, f := range files { + src, err := r.Extract(f.Path) + if err != nil { + b.Fatal(err) + } + if _, err := io.Copy(io.Discard, src); err != nil { + b.Fatal(err) + } + if err := src.Close(); err != nil { + b.Fatal(err) + } + } + if err := r.Close(); err != nil { + b.Fatal(err) + } + } +}