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
5package gzip
6
7import (
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".
18const (
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.
27type 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.
48func 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.
59func 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
68func (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.
94func (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).
99func put2(p []byte, v uint16) {
100	p[0] = uint8(v >> 0)
101	p[1] = uint8(v >> 8)
102}
103
104func 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.
112func (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).
127func (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.
158func (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.
228func (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.
247func (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