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

View as plain text