gzip.go 7.6 KB

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