github.com/miolini/go@v0.0.0-20160405192216-fca68c8cb408/src/compress/gzip/gzip.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 gzip
     6  
     7  import (
     8  	"compress/flate"
     9  	"errors"
    10  	"fmt"
    11  	"hash/crc32"
    12  	"io"
    13  )
    14  
    15  // These constants are copied from the flate package, so that code that imports
    16  // "compress/gzip" does not also have to import "compress/flate".
    17  const (
    18  	NoCompression      = flate.NoCompression
    19  	BestSpeed          = flate.BestSpeed
    20  	BestCompression    = flate.BestCompression
    21  	DefaultCompression = flate.DefaultCompression
    22  )
    23  
    24  // A Writer is an io.WriteCloser.
    25  // Writes to a Writer are compressed and written to w.
    26  type Writer struct {
    27  	Header      // written at first call to Write, Flush, or Close
    28  	w           io.Writer
    29  	level       int
    30  	wroteHeader bool
    31  	compressor  *flate.Writer
    32  	digest      uint32 // CRC-32, IEEE polynomial (section 8)
    33  	size        uint32 // Uncompressed size (section 2.3.1)
    34  	closed      bool
    35  	buf         [10]byte
    36  	err         error
    37  }
    38  
    39  // NewWriter returns a new Writer.
    40  // Writes to the returned writer are compressed and written to w.
    41  //
    42  // It is the caller's responsibility to call Close on the WriteCloser when done.
    43  // Writes may be buffered and not flushed until Close.
    44  //
    45  // Callers that wish to set the fields in Writer.Header must do so before
    46  // the first call to Write, Flush, or Close.
    47  func NewWriter(w io.Writer) *Writer {
    48  	z, _ := NewWriterLevel(w, DefaultCompression)
    49  	return z
    50  }
    51  
    52  // NewWriterLevel is like NewWriter but specifies the compression level instead
    53  // of assuming DefaultCompression.
    54  //
    55  // The compression level can be DefaultCompression, NoCompression, or any
    56  // integer value between BestSpeed and BestCompression inclusive. The error
    57  // returned will be nil if the level is valid.
    58  func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
    59  	if level < DefaultCompression || level > BestCompression {
    60  		return nil, fmt.Errorf("gzip: invalid compression level: %d", level)
    61  	}
    62  	z := new(Writer)
    63  	z.init(w, level)
    64  	return z, nil
    65  }
    66  
    67  func (z *Writer) init(w io.Writer, level int) {
    68  	compressor := z.compressor
    69  	if compressor != nil {
    70  		compressor.Reset(w)
    71  	}
    72  	*z = Writer{
    73  		Header: Header{
    74  			OS: 255, // unknown
    75  		},
    76  		w:          w,
    77  		level:      level,
    78  		compressor: compressor,
    79  	}
    80  }
    81  
    82  // Reset discards the Writer z's state and makes it equivalent to the
    83  // result of its original state from NewWriter or NewWriterLevel, but
    84  // writing to w instead. This permits reusing a Writer rather than
    85  // allocating a new one.
    86  func (z *Writer) Reset(w io.Writer) {
    87  	z.init(w, z.level)
    88  }
    89  
    90  // GZIP (RFC 1952) is little-endian, unlike ZLIB (RFC 1950).
    91  func put2(p []byte, v uint16) {
    92  	p[0] = uint8(v >> 0)
    93  	p[1] = uint8(v >> 8)
    94  }
    95  
    96  func put4(p []byte, v uint32) {
    97  	p[0] = uint8(v >> 0)
    98  	p[1] = uint8(v >> 8)
    99  	p[2] = uint8(v >> 16)
   100  	p[3] = uint8(v >> 24)
   101  }
   102  
   103  // writeBytes writes a length-prefixed byte slice to z.w.
   104  func (z *Writer) writeBytes(b []byte) error {
   105  	if len(b) > 0xffff {
   106  		return errors.New("gzip.Write: Extra data is too large")
   107  	}
   108  	put2(z.buf[:2], uint16(len(b)))
   109  	_, err := z.w.Write(z.buf[:2])
   110  	if err != nil {
   111  		return err
   112  	}
   113  	_, err = z.w.Write(b)
   114  	return err
   115  }
   116  
   117  // writeString writes a UTF-8 string s in GZIP's format to z.w.
   118  // GZIP (RFC 1952) specifies that strings are NUL-terminated ISO 8859-1 (Latin-1).
   119  func (z *Writer) writeString(s string) (err error) {
   120  	// GZIP stores Latin-1 strings; error if non-Latin-1; convert if non-ASCII.
   121  	needconv := false
   122  	for _, v := range s {
   123  		if v == 0 || v > 0xff {
   124  			return errors.New("gzip.Write: non-Latin-1 header string")
   125  		}
   126  		if v > 0x7f {
   127  			needconv = true
   128  		}
   129  	}
   130  	if needconv {
   131  		b := make([]byte, 0, len(s))
   132  		for _, v := range s {
   133  			b = append(b, byte(v))
   134  		}
   135  		_, err = z.w.Write(b)
   136  	} else {
   137  		_, err = io.WriteString(z.w, s)
   138  	}
   139  	if err != nil {
   140  		return err
   141  	}
   142  	// GZIP strings are NUL-terminated.
   143  	z.buf[0] = 0
   144  	_, err = z.w.Write(z.buf[:1])
   145  	return err
   146  }
   147  
   148  // Write writes a compressed form of p to the underlying io.Writer. The
   149  // compressed bytes are not necessarily flushed until the Writer is closed.
   150  func (z *Writer) Write(p []byte) (int, error) {
   151  	if z.err != nil {
   152  		return 0, z.err
   153  	}
   154  	var n int
   155  	// Write the GZIP header lazily.
   156  	if !z.wroteHeader {
   157  		z.wroteHeader = true
   158  		z.buf[0] = gzipID1
   159  		z.buf[1] = gzipID2
   160  		z.buf[2] = gzipDeflate
   161  		z.buf[3] = 0
   162  		if z.Extra != nil {
   163  			z.buf[3] |= 0x04
   164  		}
   165  		if z.Name != "" {
   166  			z.buf[3] |= 0x08
   167  		}
   168  		if z.Comment != "" {
   169  			z.buf[3] |= 0x10
   170  		}
   171  		put4(z.buf[4:8], uint32(z.ModTime.Unix()))
   172  		if z.level == BestCompression {
   173  			z.buf[8] = 2
   174  		} else if z.level == BestSpeed {
   175  			z.buf[8] = 4
   176  		} else {
   177  			z.buf[8] = 0
   178  		}
   179  		z.buf[9] = z.OS
   180  		n, z.err = z.w.Write(z.buf[:10])
   181  		if z.err != nil {
   182  			return n, z.err
   183  		}
   184  		if z.Extra != nil {
   185  			z.err = z.writeBytes(z.Extra)
   186  			if z.err != nil {
   187  				return n, z.err
   188  			}
   189  		}
   190  		if z.Name != "" {
   191  			z.err = z.writeString(z.Name)
   192  			if z.err != nil {
   193  				return n, z.err
   194  			}
   195  		}
   196  		if z.Comment != "" {
   197  			z.err = z.writeString(z.Comment)
   198  			if z.err != nil {
   199  				return n, z.err
   200  			}
   201  		}
   202  		if z.compressor == nil {
   203  			z.compressor, _ = flate.NewWriter(z.w, z.level)
   204  		}
   205  	}
   206  	z.size += uint32(len(p))
   207  	z.digest = crc32.Update(z.digest, crc32.IEEETable, p)
   208  	n, z.err = z.compressor.Write(p)
   209  	return n, z.err
   210  }
   211  
   212  // Flush flushes any pending compressed data to the underlying writer.
   213  //
   214  // It is useful mainly in compressed network protocols, to ensure that
   215  // a remote reader has enough data to reconstruct a packet. Flush does
   216  // not return until the data has been written. If the underlying
   217  // writer returns an error, Flush returns that error.
   218  //
   219  // In the terminology of the zlib library, Flush is equivalent to Z_SYNC_FLUSH.
   220  func (z *Writer) Flush() error {
   221  	if z.err != nil {
   222  		return z.err
   223  	}
   224  	if z.closed {
   225  		return nil
   226  	}
   227  	if !z.wroteHeader {
   228  		z.Write(nil)
   229  		if z.err != nil {
   230  			return z.err
   231  		}
   232  	}
   233  	z.err = z.compressor.Flush()
   234  	return z.err
   235  }
   236  
   237  // Close closes the Writer, flushing any unwritten data to the underlying
   238  // io.Writer, but does not close the underlying io.Writer.
   239  func (z *Writer) Close() error {
   240  	if z.err != nil {
   241  		return z.err
   242  	}
   243  	if z.closed {
   244  		return nil
   245  	}
   246  	z.closed = true
   247  	if !z.wroteHeader {
   248  		z.Write(nil)
   249  		if z.err != nil {
   250  			return z.err
   251  		}
   252  	}
   253  	z.err = z.compressor.Close()
   254  	if z.err != nil {
   255  		return z.err
   256  	}
   257  	put4(z.buf[:4], z.digest)
   258  	put4(z.buf[4:8], z.size)
   259  	_, z.err = z.w.Write(z.buf[:8])
   260  	return z.err
   261  }