1 // Copyright 2010 The Go Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 package zip 6 7 import ( 8 "compress/flate" 9 "errors" 10 "io" 11 "io/ioutil" 12 "sync" 13 ) 14 15 // A Compressor returns a new compressing writer, writing to w. 16 // The WriteCloser's Close method must be used to flush pending data to w. 17 // The Compressor itself must be safe to invoke from multiple goroutines 18 // simultaneously, but each returned writer will be used only by 19 // one goroutine at a time. 20 type Compressor func(w io.Writer) (io.WriteCloser, error) 21 22 // A Decompressor returns a new decompressing reader, reading from r. 23 // The ReadCloser's Close method must be used to release associated resources. 24 // The Decompressor itself must be safe to invoke from multiple goroutines 25 // simultaneously, but each returned reader will be used only by 26 // one goroutine at a time. 27 type Decompressor func(r io.Reader) io.ReadCloser 28 29 var flateWriterPool sync.Pool 30 31 func newFlateWriter(w io.Writer) io.WriteCloser { 32 fw, ok := flateWriterPool.Get().(*flate.Writer) 33 if ok { 34 fw.Reset(w) 35 } else { 36 fw, _ = flate.NewWriter(w, 5) 37 } 38 return &pooledFlateWriter{fw: fw} 39 } 40 41 type pooledFlateWriter struct { 42 mu sync.Mutex // guards Close and Write 43 fw *flate.Writer 44 } 45 46 func (w *pooledFlateWriter) Write(p []byte) (n int, err error) { 47 w.mu.Lock() 48 defer w.mu.Unlock() 49 if w.fw == nil { 50 return 0, errors.New("Write after Close") 51 } 52 return w.fw.Write(p) 53 } 54 55 func (w *pooledFlateWriter) Close() error { 56 w.mu.Lock() 57 defer w.mu.Unlock() 58 var err error 59 if w.fw != nil { 60 err = w.fw.Close() 61 flateWriterPool.Put(w.fw) 62 w.fw = nil 63 } 64 return err 65 } 66 67 var flateReaderPool sync.Pool 68 69 func newFlateReader(r io.Reader) io.ReadCloser { 70 fr, ok := flateReaderPool.Get().(io.ReadCloser) 71 if ok { 72 fr.(flate.Resetter).Reset(r, nil) 73 } else { 74 fr = flate.NewReader(r) 75 } 76 return &pooledFlateReader{fr: fr} 77 } 78 79 type pooledFlateReader struct { 80 mu sync.Mutex // guards Close and Read 81 fr io.ReadCloser 82 } 83 84 func (r *pooledFlateReader) Read(p []byte) (n int, err error) { 85 r.mu.Lock() 86 defer r.mu.Unlock() 87 if r.fr == nil { 88 return 0, errors.New("Read after Close") 89 } 90 return r.fr.Read(p) 91 } 92 93 func (r *pooledFlateReader) Close() error { 94 r.mu.Lock() 95 defer r.mu.Unlock() 96 var err error 97 if r.fr != nil { 98 err = r.fr.Close() 99 flateReaderPool.Put(r.fr) 100 r.fr = nil 101 } 102 return err 103 } 104 105 var ( 106 mu sync.RWMutex // guards compressor and decompressor maps 107 108 compressors = map[uint16]Compressor{ 109 Store: func(w io.Writer) (io.WriteCloser, error) { return &nopCloser{w}, nil }, 110 Deflate: func(w io.Writer) (io.WriteCloser, error) { return newFlateWriter(w), nil }, 111 } 112 113 decompressors = map[uint16]Decompressor{ 114 Store: ioutil.NopCloser, 115 Deflate: newFlateReader, 116 } 117 ) 118 119 // RegisterDecompressor allows custom decompressors for a specified method ID. 120 // The common methods Store and Deflate are built in. 121 func RegisterDecompressor(method uint16, dcomp Decompressor) { 122 mu.Lock() 123 defer mu.Unlock() 124 125 if _, ok := decompressors[method]; ok { 126 panic("decompressor already registered") 127 } 128 decompressors[method] = dcomp 129 } 130 131 // RegisterCompressor registers custom compressors for a specified method ID. 132 // The common methods Store and Deflate are built in. 133 func RegisterCompressor(method uint16, comp Compressor) { 134 mu.Lock() 135 defer mu.Unlock() 136 137 if _, ok := compressors[method]; ok { 138 panic("compressor already registered") 139 } 140 compressors[method] = comp 141 } 142 143 func compressor(method uint16) Compressor { 144 mu.RLock() 145 defer mu.RUnlock() 146 return compressors[method] 147 } 148 149 func decompressor(method uint16) Decompressor { 150 mu.RLock() 151 defer mu.RUnlock() 152 return decompressors[method] 153 } 154
