gzip.go

Documentation: github.com/klauspost/compress/gzip

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