...
Run Format

Source file src/compress/gzip/gzip.go

     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	// writeBytes writes a length-prefixed byte slice to z.w.
    91	func (z *Writer) writeBytes(b []byte) error {
    92		if len(b) > 0xffff {
    93			return errors.New("gzip.Write: Extra data is too large")
    94		}
    95		le.PutUint16(z.buf[:2], uint16(len(b)))
    96		_, err := z.w.Write(z.buf[:2])
    97		if err != nil {
    98			return err
    99		}
   100		_, err = z.w.Write(b)
   101		return err
   102	}
   103	
   104	// writeString writes a UTF-8 string s in GZIP's format to z.w.
   105	// GZIP (RFC 1952) specifies that strings are NUL-terminated ISO 8859-1 (Latin-1).
   106	func (z *Writer) writeString(s string) (err error) {
   107		// GZIP stores Latin-1 strings; error if non-Latin-1; convert if non-ASCII.
   108		needconv := false
   109		for _, v := range s {
   110			if v == 0 || v > 0xff {
   111				return errors.New("gzip.Write: non-Latin-1 header string")
   112			}
   113			if v > 0x7f {
   114				needconv = true
   115			}
   116		}
   117		if needconv {
   118			b := make([]byte, 0, len(s))
   119			for _, v := range s {
   120				b = append(b, byte(v))
   121			}
   122			_, err = z.w.Write(b)
   123		} else {
   124			_, err = io.WriteString(z.w, s)
   125		}
   126		if err != nil {
   127			return err
   128		}
   129		// GZIP strings are NUL-terminated.
   130		z.buf[0] = 0
   131		_, err = z.w.Write(z.buf[:1])
   132		return err
   133	}
   134	
   135	// Write writes a compressed form of p to the underlying io.Writer. The
   136	// compressed bytes are not necessarily flushed until the Writer is closed.
   137	func (z *Writer) Write(p []byte) (int, error) {
   138		if z.err != nil {
   139			return 0, z.err
   140		}
   141		var n int
   142		// Write the GZIP header lazily.
   143		if !z.wroteHeader {
   144			z.wroteHeader = true
   145			z.buf[0] = gzipID1
   146			z.buf[1] = gzipID2
   147			z.buf[2] = gzipDeflate
   148			z.buf[3] = 0
   149			if z.Extra != nil {
   150				z.buf[3] |= 0x04
   151			}
   152			if z.Name != "" {
   153				z.buf[3] |= 0x08
   154			}
   155			if z.Comment != "" {
   156				z.buf[3] |= 0x10
   157			}
   158			le.PutUint32(z.buf[4:8], uint32(z.ModTime.Unix()))
   159			if z.level == BestCompression {
   160				z.buf[8] = 2
   161			} else if z.level == BestSpeed {
   162				z.buf[8] = 4
   163			} else {
   164				z.buf[8] = 0
   165			}
   166			z.buf[9] = z.OS
   167			n, z.err = z.w.Write(z.buf[:10])
   168			if z.err != nil {
   169				return n, z.err
   170			}
   171			if z.Extra != nil {
   172				z.err = z.writeBytes(z.Extra)
   173				if z.err != nil {
   174					return n, z.err
   175				}
   176			}
   177			if z.Name != "" {
   178				z.err = z.writeString(z.Name)
   179				if z.err != nil {
   180					return n, z.err
   181				}
   182			}
   183			if z.Comment != "" {
   184				z.err = z.writeString(z.Comment)
   185				if z.err != nil {
   186					return n, z.err
   187				}
   188			}
   189			if z.compressor == nil {
   190				z.compressor, _ = flate.NewWriter(z.w, z.level)
   191			}
   192		}
   193		z.size += uint32(len(p))
   194		z.digest = crc32.Update(z.digest, crc32.IEEETable, p)
   195		n, z.err = z.compressor.Write(p)
   196		return n, z.err
   197	}
   198	
   199	// Flush flushes any pending compressed data to the underlying writer.
   200	//
   201	// It is useful mainly in compressed network protocols, to ensure that
   202	// a remote reader has enough data to reconstruct a packet. Flush does
   203	// not return until the data has been written. If the underlying
   204	// writer returns an error, Flush returns that error.
   205	//
   206	// In the terminology of the zlib library, Flush is equivalent to Z_SYNC_FLUSH.
   207	func (z *Writer) Flush() error {
   208		if z.err != nil {
   209			return z.err
   210		}
   211		if z.closed {
   212			return nil
   213		}
   214		if !z.wroteHeader {
   215			z.Write(nil)
   216			if z.err != nil {
   217				return z.err
   218			}
   219		}
   220		z.err = z.compressor.Flush()
   221		return z.err
   222	}
   223	
   224	// Close closes the Writer, flushing any unwritten data to the underlying
   225	// io.Writer, but does not close the underlying io.Writer.
   226	func (z *Writer) Close() error {
   227		if z.err != nil {
   228			return z.err
   229		}
   230		if z.closed {
   231			return nil
   232		}
   233		z.closed = true
   234		if !z.wroteHeader {
   235			z.Write(nil)
   236			if z.err != nil {
   237				return z.err
   238			}
   239		}
   240		z.err = z.compressor.Close()
   241		if z.err != nil {
   242			return z.err
   243		}
   244		le.PutUint32(z.buf[:4], z.digest)
   245		le.PutUint32(z.buf[4:8], z.size)
   246		_, z.err = z.w.Write(z.buf[:8])
   247		return z.err
   248	}
   249	

View as plain text