gzip.go 6.6 KB

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