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 that satisfies writes by compressing data written
    26	// to its wrapped io.Writer.
    27	type Writer struct {
    28		Header
    29		w          io.Writer
    30		level      int
    31		compressor io.WriteCloser
    32		digest     hash.Hash32
    33		size       uint32
    34		closed     bool
    35		buf        [10]byte
    36		err        error
    37	}
    38	
    39	// NewWriter creates a new Writer that satisfies writes by compressing data
    40	// 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 or Close. The Comment and Name header fields are
    47	// UTF-8 strings in Go, but the underlying format requires NUL-terminated ISO
    48	// 8859-1 (Latin-1). NUL or non-Latin-1 runes in those strings will lead to an
    49	// error on Write.
    50	func NewWriter(w io.Writer) *Writer {
    51		z, _ := NewWriterLevel(w, DefaultCompression)
    52		return z
    53	}
    54	
    55	// NewWriterLevel is like NewWriter but specifies the compression level instead
    56	// of assuming DefaultCompression.
    57	//
    58	// The compression level can be DefaultCompression, NoCompression, or any
    59	// integer value between BestSpeed and BestCompression inclusive. The error
    60	// returned will be nil if the level is valid.
    61	func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
    62		if level < DefaultCompression || level > BestCompression {
    63			return nil, fmt.Errorf("gzip: invalid compression level: %d", level)
    64		}
    65		return &Writer{
    66			Header: Header{
    67				OS: 255, // unknown
    68			},
    69			w:      w,
    70			level:  level,
    71			digest: crc32.NewIEEE(),
    72		}, nil
    73	}
    74	
    75	// GZIP (RFC 1952) is little-endian, unlike ZLIB (RFC 1950).
    76	func put2(p []byte, v uint16) {
    77		p[0] = uint8(v >> 0)
    78		p[1] = uint8(v >> 8)
    79	}
    80	
    81	func put4(p []byte, v uint32) {
    82		p[0] = uint8(v >> 0)
    83		p[1] = uint8(v >> 8)
    84		p[2] = uint8(v >> 16)
    85		p[3] = uint8(v >> 24)
    86	}
    87	
    88	// writeBytes writes a length-prefixed byte slice to z.w.
    89	func (z *Writer) writeBytes(b []byte) error {
    90		if len(b) > 0xffff {
    91			return errors.New("gzip.Write: Extra data is too large")
    92		}
    93		put2(z.buf[0:2], uint16(len(b)))
    94		_, err := z.w.Write(z.buf[0:2])
    95		if err != nil {
    96			return err
    97		}
    98		_, err = z.w.Write(b)
    99		return err
   100	}
   101	
   102	// writeString writes a UTF-8 string s in GZIP's format to z.w.
   103	// GZIP (RFC 1952) specifies that strings are NUL-terminated ISO 8859-1 (Latin-1).
   104	func (z *Writer) writeString(s string) (err error) {
   105		// GZIP stores Latin-1 strings; error if non-Latin-1; convert if non-ASCII.
   106		needconv := false
   107		for _, v := range s {
   108			if v == 0 || v > 0xff {
   109				return errors.New("gzip.Write: non-Latin-1 header string")
   110			}
   111			if v > 0x7f {
   112				needconv = true
   113			}
   114		}
   115		if needconv {
   116			b := make([]byte, 0, len(s))
   117			for _, v := range s {
   118				b = append(b, byte(v))
   119			}
   120			_, err = z.w.Write(b)
   121		} else {
   122			_, err = io.WriteString(z.w, s)
   123		}
   124		if err != nil {
   125			return err
   126		}
   127		// GZIP strings are NUL-terminated.
   128		z.buf[0] = 0
   129		_, err = z.w.Write(z.buf[0:1])
   130		return err
   131	}
   132	
   133	// Write writes a compressed form of p to the underlying io.Writer. The
   134	// compressed bytes are not necessarily flushed until the Writer is closed.
   135	func (z *Writer) Write(p []byte) (int, error) {
   136		if z.err != nil {
   137			return 0, z.err
   138		}
   139		var n int
   140		// Write the GZIP header lazily.
   141		if z.compressor == nil {
   142			z.buf[0] = gzipID1
   143			z.buf[1] = gzipID2
   144			z.buf[2] = gzipDeflate
   145			z.buf[3] = 0
   146			if z.Extra != nil {
   147				z.buf[3] |= 0x04
   148			}
   149			if z.Name != "" {
   150				z.buf[3] |= 0x08
   151			}
   152			if z.Comment != "" {
   153				z.buf[3] |= 0x10
   154			}
   155			put4(z.buf[4:8], uint32(z.ModTime.Unix()))
   156			if z.level == BestCompression {
   157				z.buf[8] = 2
   158			} else if z.level == BestSpeed {
   159				z.buf[8] = 4
   160			} else {
   161				z.buf[8] = 0
   162			}
   163			z.buf[9] = z.OS
   164			n, z.err = z.w.Write(z.buf[0:10])
   165			if z.err != nil {
   166				return n, z.err
   167			}
   168			if z.Extra != nil {
   169				z.err = z.writeBytes(z.Extra)
   170				if z.err != nil {
   171					return n, z.err
   172				}
   173			}
   174			if z.Name != "" {
   175				z.err = z.writeString(z.Name)
   176				if z.err != nil {
   177					return n, z.err
   178				}
   179			}
   180			if z.Comment != "" {
   181				z.err = z.writeString(z.Comment)
   182				if z.err != nil {
   183					return n, z.err
   184				}
   185			}
   186			z.compressor, _ = flate.NewWriter(z.w, z.level)
   187		}
   188		z.size += uint32(len(p))
   189		z.digest.Write(p)
   190		n, z.err = z.compressor.Write(p)
   191		return n, z.err
   192	}
   193	
   194	// Close closes the Writer. It does not close the underlying io.Writer.
   195	func (z *Writer) Close() error {
   196		if z.err != nil {
   197			return z.err
   198		}
   199		if z.closed {
   200			return nil
   201		}
   202		z.closed = true
   203		if z.compressor == nil {
   204			z.Write(nil)
   205			if z.err != nil {
   206				return z.err
   207			}
   208		}
   209		z.err = z.compressor.Close()
   210		if z.err != nil {
   211			return z.err
   212		}
   213		put4(z.buf[0:4], z.digest.Sum32())
   214		put4(z.buf[4:8], z.size)
   215		_, z.err = z.w.Write(z.buf[0:8])
   216		return z.err
   217	}