...
Run Format

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

View as plain text