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