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