summaryrefslogtreecommitdiff
path: root/vendor/github.com/klauspost/compress/zstd/zip.go
blob: 29c15c8c4efe59449cd2028aacd636c377731c74 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
// Copyright 2019+ Klaus Post. All rights reserved.
// License information can be found in the LICENSE file.

package zstd

import (
	"errors"
	"io"
	"sync"
)

// ZipMethodWinZip is the method for Zstandard compressed data inside Zip files for WinZip.
// See https://www.winzip.com/win/en/comp_info.html
const ZipMethodWinZip = 93

// ZipMethodPKWare is the original method number used by PKWARE to indicate Zstandard compression.
// Deprecated: This has been deprecated by PKWARE, use ZipMethodWinZip instead for compression.
// See https://pkware.cachefly.net/webdocs/APPNOTE/APPNOTE-6.3.9.TXT
const ZipMethodPKWare = 20

// zipReaderPool is the default reader pool.
var zipReaderPool = sync.Pool{New: func() interface{} {
	z, err := NewReader(nil, WithDecoderLowmem(true), WithDecoderMaxWindow(128<<20), WithDecoderConcurrency(1))
	if err != nil {
		panic(err)
	}
	return z
}}

// newZipReader creates a pooled zip decompressor.
func newZipReader(opts ...DOption) func(r io.Reader) io.ReadCloser {
	pool := &zipReaderPool
	if len(opts) > 0 {
		opts = append([]DOption{WithDecoderLowmem(true), WithDecoderMaxWindow(128 << 20)}, opts...)
		// Force concurrency 1
		opts = append(opts, WithDecoderConcurrency(1))
		// Create our own pool
		pool = &sync.Pool{}
	}
	return func(r io.Reader) io.ReadCloser {
		dec, ok := pool.Get().(*Decoder)
		if ok {
			dec.Reset(r)
		} else {
			d, err := NewReader(r, opts...)
			if err != nil {
				panic(err)
			}
			dec = d
		}
		return &pooledZipReader{dec: dec, pool: pool}
	}
}

type pooledZipReader struct {
	mu   sync.Mutex // guards Close and Read
	pool *sync.Pool
	dec  *Decoder
}

func (r *pooledZipReader) Read(p []byte) (n int, err error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.dec == nil {
		return 0, errors.New("read after close or EOF")
	}
	dec, err := r.dec.Read(p)
	if err == io.EOF {
		r.dec.Reset(nil)
		r.pool.Put(r.dec)
		r.dec = nil
	}
	return dec, err
}

func (r *pooledZipReader) Close() error {
	r.mu.Lock()
	defer r.mu.Unlock()
	var err error
	if r.dec != nil {
		err = r.dec.Reset(nil)
		r.pool.Put(r.dec)
		r.dec = nil
	}
	return err
}

type pooledZipWriter struct {
	mu   sync.Mutex // guards Close and Read
	enc  *Encoder
	pool *sync.Pool
}

func (w *pooledZipWriter) Write(p []byte) (n int, err error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.enc == nil {
		return 0, errors.New("Write after Close")
	}
	return w.enc.Write(p)
}

func (w *pooledZipWriter) Close() error {
	w.mu.Lock()
	defer w.mu.Unlock()
	var err error
	if w.enc != nil {
		err = w.enc.Close()
		w.pool.Put(w.enc)
		w.enc = nil
	}
	return err
}

// ZipCompressor returns a compressor that can be registered with zip libraries.
// The provided encoder options will be used on all encodes.
func ZipCompressor(opts ...EOption) func(w io.Writer) (io.WriteCloser, error) {
	var pool sync.Pool
	return func(w io.Writer) (io.WriteCloser, error) {
		enc, ok := pool.Get().(*Encoder)
		if ok {
			enc.Reset(w)
		} else {
			var err error
			enc, err = NewWriter(w, opts...)
			if err != nil {
				return nil, err
			}
		}
		return &pooledZipWriter{enc: enc, pool: &pool}, nil
	}
}

// ZipDecompressor returns a decompressor that can be registered with zip libraries.
// See ZipCompressor for example.
// Options can be specified. WithDecoderConcurrency(1) is forced,
// and by default a 128MB maximum decompression window is specified.
// The window size can be overridden if required.
func ZipDecompressor(opts ...DOption) func(r io.Reader) io.ReadCloser {
	return newZipReader(opts...)
}