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