zip.go 3.4 KB

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