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