zip.go raw

   1  // Copyright 2019+ Klaus Post. All rights reserved.
   2  // License information can be found in the LICENSE file.
   3  
   4  package zstd
   5  
   6  import (
   7  	"errors"
   8  	"io"
   9  	"sync"
  10  )
  11  
  12  // ZipMethodWinZip is the method for Zstandard compressed data inside Zip files for WinZip.
  13  // See https://www.winzip.com/win/en/comp_info.html
  14  const ZipMethodWinZip = 93
  15  
  16  // ZipMethodPKWare is the original method number used by PKWARE to indicate Zstandard compression.
  17  // Deprecated: This has been deprecated by PKWARE, use ZipMethodWinZip instead for compression.
  18  // See https://pkware.cachefly.net/webdocs/APPNOTE/APPNOTE-6.3.9.TXT
  19  const ZipMethodPKWare = 20
  20  
  21  // zipReaderPool is the default reader pool.
  22  var zipReaderPool = sync.Pool{New: func() any {
  23  	z, err := NewReader(nil, WithDecoderLowmem(true), WithDecoderMaxWindow(128<<20), WithDecoderConcurrency(1))
  24  	if err != nil {
  25  		panic(err)
  26  	}
  27  	return z
  28  }}
  29  
  30  // newZipReader creates a pooled zip decompressor.
  31  func newZipReader(opts ...DOption) func(r io.Reader) io.ReadCloser {
  32  	pool := &zipReaderPool
  33  	if len(opts) > 0 {
  34  		opts = append([]DOption{WithDecoderLowmem(true), WithDecoderMaxWindow(128 << 20)}, opts...)
  35  		// Force concurrency 1
  36  		opts = append(opts, WithDecoderConcurrency(1))
  37  		// Create our own pool
  38  		pool = &sync.Pool{}
  39  	}
  40  	return func(r io.Reader) io.ReadCloser {
  41  		dec, ok := pool.Get().(*Decoder)
  42  		if ok {
  43  			dec.Reset(r)
  44  		} else {
  45  			d, err := NewReader(r, opts...)
  46  			if err != nil {
  47  				panic(err)
  48  			}
  49  			dec = d
  50  		}
  51  		return &pooledZipReader{dec: dec, pool: pool}
  52  	}
  53  }
  54  
  55  type pooledZipReader struct {
  56  	mu   sync.Mutex // guards Close and Read
  57  	pool *sync.Pool
  58  	dec  *Decoder
  59  }
  60  
  61  func (r *pooledZipReader) Read(p []byte) (n int, err error) {
  62  	r.mu.Lock()
  63  	defer r.mu.Unlock()
  64  	if r.dec == nil {
  65  		return 0, errors.New("read after close or EOF")
  66  	}
  67  	dec, err := r.dec.Read(p)
  68  	if err == io.EOF {
  69  		r.dec.Reset(nil)
  70  		r.pool.Put(r.dec)
  71  		r.dec = nil
  72  	}
  73  	return dec, err
  74  }
  75  
  76  func (r *pooledZipReader) Close() error {
  77  	r.mu.Lock()
  78  	defer r.mu.Unlock()
  79  	var err error
  80  	if r.dec != nil {
  81  		err = r.dec.Reset(nil)
  82  		r.pool.Put(r.dec)
  83  		r.dec = nil
  84  	}
  85  	return err
  86  }
  87  
  88  type pooledZipWriter struct {
  89  	mu   sync.Mutex // guards Close and Read
  90  	enc  *Encoder
  91  	pool *sync.Pool
  92  }
  93  
  94  func (w *pooledZipWriter) Write(p []byte) (n int, err error) {
  95  	w.mu.Lock()
  96  	defer w.mu.Unlock()
  97  	if w.enc == nil {
  98  		return 0, errors.New("Write after Close")
  99  	}
 100  	return w.enc.Write(p)
 101  }
 102  
 103  func (w *pooledZipWriter) Close() error {
 104  	w.mu.Lock()
 105  	defer w.mu.Unlock()
 106  	var err error
 107  	if w.enc != nil {
 108  		err = w.enc.Close()
 109  		w.pool.Put(w.enc)
 110  		w.enc = nil
 111  	}
 112  	return err
 113  }
 114  
 115  // ZipCompressor returns a compressor that can be registered with zip libraries.
 116  // The provided encoder options will be used on all encodes.
 117  func ZipCompressor(opts ...EOption) func(w io.Writer) (io.WriteCloser, error) {
 118  	var pool sync.Pool
 119  	return func(w io.Writer) (io.WriteCloser, error) {
 120  		enc, ok := pool.Get().(*Encoder)
 121  		if ok {
 122  			enc.Reset(w)
 123  		} else {
 124  			var err error
 125  			enc, err = NewWriter(w, opts...)
 126  			if err != nil {
 127  				return nil, err
 128  			}
 129  		}
 130  		return &pooledZipWriter{enc: enc, pool: &pool}, nil
 131  	}
 132  }
 133  
 134  // ZipDecompressor returns a decompressor that can be registered with zip libraries.
 135  // See ZipCompressor for example.
 136  // Options can be specified. WithDecoderConcurrency(1) is forced,
 137  // and by default a 128MB maximum decompression window is specified.
 138  // The window size can be overridden if required.
 139  func ZipDecompressor(opts ...DOption) func(r io.Reader) io.ReadCloser {
 140  	return newZipReader(opts...)
 141  }
 142