1 // Copyright 2009 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 zlib
6 7 import (
8 "compress/flate"
9 "encoding/binary"
10 "fmt"
11 "hash"
12 "hash/adler32"
13 "io"
14 )
15 16 // These constants are copied from the [flate] package, so that code that imports
17 // [compress/zlib] 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 HuffmanOnly = flate.HuffmanOnly
24 )
25 26 // A Writer takes data written to it and writes the compressed
27 // form of that data to an underlying writer (see [NewWriter]).
28 type Writer struct {
29 w io.Writer
30 level int
31 dict []byte
32 compressor *flate.Writer
33 digest hash.Hash32
34 err error
35 scratch [4]byte
36 wroteHeader bool
37 }
38 39 // NewWriter creates a new [Writer].
40 // Writes to the returned Writer are compressed and written to w.
41 //
42 // It is the caller's responsibility to call Close on the Writer when done.
43 // Writes may be buffered and not flushed until Close.
44 func NewWriter(w io.Writer) *Writer {
45 z, _ := NewWriterLevelDict(w, DefaultCompression, nil)
46 return z
47 }
48 49 // NewWriterLevel is like [NewWriter] but specifies the compression level instead
50 // of assuming [DefaultCompression].
51 //
52 // The compression level can be [DefaultCompression], [NoCompression], [HuffmanOnly]
53 // or any integer value between [BestSpeed] and [BestCompression] inclusive.
54 // The error returned will be nil if the level is valid.
55 func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
56 return NewWriterLevelDict(w, level, nil)
57 }
58 59 // NewWriterLevelDict is like [NewWriterLevel] but specifies a dictionary to
60 // compress with.
61 //
62 // The dictionary may be nil. If not, its contents should not be modified until
63 // the Writer is closed.
64 func NewWriterLevelDict(w io.Writer, level int, dict []byte) (*Writer, error) {
65 if level < HuffmanOnly || level > BestCompression {
66 return nil, fmt.Errorf("zlib: invalid compression level: %d", level)
67 }
68 return &Writer{
69 w: w,
70 level: level,
71 dict: dict,
72 }, nil
73 }
74 75 // Reset clears the state of the [Writer] z such that it is equivalent to its
76 // initial state from [NewWriterLevel] or [NewWriterLevelDict], but instead writing
77 // to w.
78 func (z *Writer) Reset(w io.Writer) {
79 z.w = w
80 // z.level and z.dict left unchanged.
81 if z.compressor != nil {
82 z.compressor.Reset(w)
83 }
84 if z.digest != nil {
85 z.digest.Reset()
86 }
87 z.err = nil
88 z.scratch = [4]byte{}
89 z.wroteHeader = false
90 }
91 92 // writeHeader writes the ZLIB header.
93 func (z *Writer) writeHeader() (err error) {
94 z.wroteHeader = true
95 // ZLIB has a two-byte header (as documented in RFC 1950).
96 // The first four bits is the CINFO (compression info), which is 7 for the default deflate window size.
97 // The next four bits is the CM (compression method), which is 8 for deflate.
98 z.scratch[0] = 0x78
99 // The next two bits is the FLEVEL (compression level). The four values are:
100 // 0=fastest, 1=fast, 2=default, 3=best.
101 // The next bit, FDICT, is set if a dictionary is given.
102 // The final five FCHECK bits form a mod-31 checksum.
103 switch z.level {
104 case -2, 0, 1:
105 z.scratch[1] = 0 << 6
106 case 2, 3, 4, 5:
107 z.scratch[1] = 1 << 6
108 case 6, -1:
109 z.scratch[1] = 2 << 6
110 case 7, 8, 9:
111 z.scratch[1] = 3 << 6
112 default:
113 panic("unreachable")
114 }
115 if z.dict != nil {
116 z.scratch[1] |= 1 << 5
117 }
118 z.scratch[1] += uint8(31 - binary.BigEndian.Uint16(z.scratch[:2])%31)
119 if _, err = z.w.Write(z.scratch[0:2]); err != nil {
120 return err
121 }
122 if z.dict != nil {
123 // The next four bytes are the Adler-32 checksum of the dictionary.
124 binary.BigEndian.PutUint32(z.scratch[:], adler32.Checksum(z.dict))
125 if _, err = z.w.Write(z.scratch[0:4]); err != nil {
126 return err
127 }
128 }
129 if z.compressor == nil {
130 // Initialize deflater unless the Writer is being reused
131 // after a Reset call.
132 z.compressor, err = flate.NewWriterDict(z.w, z.level, z.dict)
133 if err != nil {
134 return err
135 }
136 z.digest = adler32.New()
137 }
138 return nil
139 }
140 141 // Write writes a compressed form of p to the underlying [io.Writer]. The
142 // compressed bytes are not necessarily flushed until the [Writer] is closed or
143 // explicitly flushed.
144 func (z *Writer) Write(p []byte) (n int, err error) {
145 if !z.wroteHeader {
146 z.err = z.writeHeader()
147 }
148 if z.err != nil {
149 return 0, z.err
150 }
151 if len(p) == 0 {
152 return 0, nil
153 }
154 n, err = z.compressor.Write(p)
155 if err != nil {
156 z.err = err
157 return
158 }
159 z.digest.Write(p)
160 return
161 }
162 163 // Flush flushes the Writer to its underlying [io.Writer].
164 func (z *Writer) Flush() error {
165 if !z.wroteHeader {
166 z.err = z.writeHeader()
167 }
168 if z.err != nil {
169 return z.err
170 }
171 z.err = z.compressor.Flush()
172 return z.err
173 }
174 175 // Close closes the Writer, flushing any unwritten data to the underlying
176 // [io.Writer], but does not close the underlying io.Writer.
177 func (z *Writer) Close() error {
178 if !z.wroteHeader {
179 z.err = z.writeHeader()
180 }
181 if z.err != nil {
182 return z.err
183 }
184 z.err = z.compressor.Close()
185 if z.err != nil {
186 return z.err
187 }
188 checksum := z.digest.Sum32()
189 // ZLIB (RFC 1950) is big-endian, unlike GZIP (RFC 1952).
190 binary.BigEndian.PutUint32(z.scratch[:], checksum)
191 _, z.err = z.w.Write(z.scratch[0:4])
192 return z.err
193 }
194