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