github.com/bgentry/go@v0.0.0-20150121062915-6cf5a733d54d/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 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  }