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 compressing writer, writing to the 16 // provided writer. On Close, any pending data should be flushed. 17 type Compressor func(io.Writer) (io.WriteCloser, error) 18 19 // Decompressor is a function that wraps a Reader with a decompressing Reader. 20 // The decompressed ReadCloser is returned to callers who open files from 21 // within the archive. These callers are responsible for closing this reader 22 // when they're finished reading. 23 type Decompressor func(io.Reader) io.ReadCloser 24 25 var flateWriterPool sync.Pool 26 27 func newFlateWriter(w io.Writer) io.WriteCloser { 28 fw, ok := flateWriterPool.Get().(*flate.Writer) 29 if ok { 30 fw.Reset(w) 31 } else { 32 fw, _ = flate.NewWriter(w, 5) 33 } 34 return &pooledFlateWriter{fw: fw} 35 } 36 37 type pooledFlateWriter struct { 38 mu sync.Mutex // guards Close and Write 39 fw *flate.Writer 40 } 41 42 func (w *pooledFlateWriter) Write(p []byte) (n int, err error) { 43 w.mu.Lock() 44 defer w.mu.Unlock() 45 if w.fw == nil { 46 return 0, errors.New("Write after Close") 47 } 48 return w.fw.Write(p) 49 } 50 51 func (w *pooledFlateWriter) Close() error { 52 w.mu.Lock() 53 defer w.mu.Unlock() 54 var err error 55 if w.fw != nil { 56 err = w.fw.Close() 57 flateWriterPool.Put(w.fw) 58 w.fw = nil 59 } 60 return err 61 } 62 63 var ( 64 mu sync.RWMutex // guards compressor and decompressor maps 65 66 compressors = map[uint16]Compressor{ 67 Store: func(w io.Writer) (io.WriteCloser, error) { return &nopCloser{w}, nil }, 68 Deflate: func(w io.Writer) (io.WriteCloser, error) { return newFlateWriter(w), nil }, 69 } 70 71 decompressors = map[uint16]Decompressor{ 72 Store: ioutil.NopCloser, 73 Deflate: flate.NewReader, 74 } 75 ) 76 77 // RegisterDecompressor allows custom decompressors for a specified method ID. 78 func RegisterDecompressor(method uint16, d Decompressor) { 79 mu.Lock() 80 defer mu.Unlock() 81 82 if _, ok := decompressors[method]; ok { 83 panic("decompressor already registered") 84 } 85 decompressors[method] = d 86 } 87 88 // RegisterCompressor registers custom compressors for a specified method ID. 89 // The common methods Store and Deflate are built in. 90 func RegisterCompressor(method uint16, comp Compressor) { 91 mu.Lock() 92 defer mu.Unlock() 93 94 if _, ok := compressors[method]; ok { 95 panic("compressor already registered") 96 } 97 compressors[method] = comp 98 } 99 100 func compressor(method uint16) Compressor { 101 mu.RLock() 102 defer mu.RUnlock() 103 return compressors[method] 104 } 105 106 func decompressor(method uint16) Decompressor { 107 mu.RLock() 108 defer mu.RUnlock() 109 return decompressors[method] 110 } 111
