github.com/MangoDowner/go-gm@v0.0.0-20180818020936-8baa2bd4408c/src/archive/zip/register.go (about) 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 compressors sync.Map // map[uint16]Compressor 107 decompressors sync.Map // map[uint16]Decompressor 108 ) 109 110 func init() { 111 compressors.Store(Store, Compressor(func(w io.Writer) (io.WriteCloser, error) { return &nopCloser{w}, nil })) 112 compressors.Store(Deflate, Compressor(func(w io.Writer) (io.WriteCloser, error) { return newFlateWriter(w), nil })) 113 114 decompressors.Store(Store, Decompressor(ioutil.NopCloser)) 115 decompressors.Store(Deflate, Decompressor(newFlateReader)) 116 } 117 118 // RegisterDecompressor allows custom decompressors for a specified method ID. 119 // The common methods Store and Deflate are built in. 120 func RegisterDecompressor(method uint16, dcomp Decompressor) { 121 if _, dup := decompressors.LoadOrStore(method, dcomp); dup { 122 panic("decompressor already registered") 123 } 124 } 125 126 // RegisterCompressor registers custom compressors for a specified method ID. 127 // The common methods Store and Deflate are built in. 128 func RegisterCompressor(method uint16, comp Compressor) { 129 if _, dup := compressors.LoadOrStore(method, comp); dup { 130 panic("compressor already registered") 131 } 132 } 133 134 func compressor(method uint16) Compressor { 135 ci, ok := compressors.Load(method) 136 if !ok { 137 return nil 138 } 139 return ci.(Compressor) 140 } 141 142 func decompressor(method uint16) Decompressor { 143 di, ok := decompressors.Load(method) 144 if !ok { 145 return nil 146 } 147 return di.(Decompressor) 148 }