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