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