gzip.go 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252
  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. "io"
  9. "github.com/klauspost/compress/flate"
  10. "github.com/klauspost/crc32"
  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. )
  22. // A Writer is an io.WriteCloser.
  23. // Writes to a Writer are compressed and written to w.
  24. type Writer struct {
  25. Header // written at first call to Write, Flush, or Close
  26. w io.Writer
  27. level int
  28. wroteHeader bool
  29. compressor *flate.Writer
  30. digest uint32 // CRC-32, IEEE polynomial (section 8)
  31. size uint32 // Uncompressed size (section 2.3.1)
  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, Flush, or Close.
  44. func NewWriter(w io.Writer) *Writer {
  45. z, _ := NewWriterLevel(w, DefaultCompression)
  46. return z
  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, or any
  52. // integer value between BestSpeed and BestCompression inclusive. The error
  53. // returned will be nil if the level is valid.
  54. func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
  55. if level < HuffmanOnly || level > BestCompression {
  56. return nil, fmt.Errorf("gzip: invalid compression level: %d", level)
  57. }
  58. z := new(Writer)
  59. z.init(w, level)
  60. return z, nil
  61. }
  62. func (z *Writer) init(w io.Writer, level int) {
  63. compressor := z.compressor
  64. if compressor != nil {
  65. compressor.Reset(w)
  66. }
  67. *z = Writer{
  68. Header: Header{
  69. OS: 255, // unknown
  70. },
  71. w: w,
  72. level: level,
  73. compressor: compressor,
  74. }
  75. }
  76. // Reset discards the Writer z's state and makes it equivalent to the
  77. // result of its original state from NewWriter or NewWriterLevel, but
  78. // writing to w instead. This permits reusing a Writer rather than
  79. // allocating a new one.
  80. func (z *Writer) Reset(w io.Writer) {
  81. z.init(w, z.level)
  82. }
  83. // writeBytes writes a length-prefixed byte slice to z.w.
  84. func (z *Writer) writeBytes(b []byte) error {
  85. if len(b) > 0xffff {
  86. return errors.New("gzip.Write: Extra data is too large")
  87. }
  88. le.PutUint16(z.buf[:2], uint16(len(b)))
  89. _, err := z.w.Write(z.buf[:2])
  90. if err != nil {
  91. return err
  92. }
  93. _, err = z.w.Write(b)
  94. return err
  95. }
  96. // writeString writes a UTF-8 string s in GZIP's format to z.w.
  97. // GZIP (RFC 1952) specifies that strings are NUL-terminated ISO 8859-1 (Latin-1).
  98. func (z *Writer) writeString(s string) (err error) {
  99. // GZIP stores Latin-1 strings; error if non-Latin-1; convert if non-ASCII.
  100. needconv := false
  101. for _, v := range s {
  102. if v == 0 || v > 0xff {
  103. return errors.New("gzip.Write: non-Latin-1 header string")
  104. }
  105. if v > 0x7f {
  106. needconv = true
  107. }
  108. }
  109. if needconv {
  110. b := make([]byte, 0, len(s))
  111. for _, v := range s {
  112. b = append(b, byte(v))
  113. }
  114. _, err = z.w.Write(b)
  115. } else {
  116. _, err = io.WriteString(z.w, s)
  117. }
  118. if err != nil {
  119. return err
  120. }
  121. // GZIP strings are NUL-terminated.
  122. z.buf[0] = 0
  123. _, err = z.w.Write(z.buf[:1])
  124. return err
  125. }
  126. // Write writes a compressed form of p to the underlying io.Writer. The
  127. // compressed bytes are not necessarily flushed until the Writer is closed.
  128. func (z *Writer) Write(p []byte) (int, error) {
  129. if z.err != nil {
  130. return 0, z.err
  131. }
  132. var n int
  133. // Write the GZIP header lazily.
  134. if !z.wroteHeader {
  135. z.wroteHeader = true
  136. z.buf[0] = gzipID1
  137. z.buf[1] = gzipID2
  138. z.buf[2] = gzipDeflate
  139. z.buf[3] = 0
  140. if z.Extra != nil {
  141. z.buf[3] |= 0x04
  142. }
  143. if z.Name != "" {
  144. z.buf[3] |= 0x08
  145. }
  146. if z.Comment != "" {
  147. z.buf[3] |= 0x10
  148. }
  149. le.PutUint32(z.buf[4:8], uint32(z.ModTime.Unix()))
  150. if z.level == BestCompression {
  151. z.buf[8] = 2
  152. } else if z.level == BestSpeed {
  153. z.buf[8] = 4
  154. } else {
  155. z.buf[8] = 0
  156. }
  157. z.buf[9] = z.OS
  158. n, z.err = z.w.Write(z.buf[:10])
  159. if z.err != nil {
  160. return n, z.err
  161. }
  162. if z.Extra != nil {
  163. z.err = z.writeBytes(z.Extra)
  164. if z.err != nil {
  165. return n, z.err
  166. }
  167. }
  168. if z.Name != "" {
  169. z.err = z.writeString(z.Name)
  170. if z.err != nil {
  171. return n, z.err
  172. }
  173. }
  174. if z.Comment != "" {
  175. z.err = z.writeString(z.Comment)
  176. if z.err != nil {
  177. return n, z.err
  178. }
  179. }
  180. if z.compressor == nil {
  181. z.compressor, _ = flate.NewWriter(z.w, z.level)
  182. }
  183. }
  184. z.size += uint32(len(p))
  185. z.digest = crc32.Update(z.digest, crc32.IEEETable, p)
  186. n, z.err = z.compressor.Write(p)
  187. return n, z.err
  188. }
  189. // Flush flushes any pending compressed data to the underlying writer.
  190. //
  191. // It is useful mainly in compressed network protocols, to ensure that
  192. // a remote reader has enough data to reconstruct a packet. Flush does
  193. // not return until the data has been written. If the underlying
  194. // writer returns an error, Flush returns that error.
  195. //
  196. // In the terminology of the zlib library, Flush is equivalent to Z_SYNC_FLUSH.
  197. func (z *Writer) Flush() error {
  198. if z.err != nil {
  199. return z.err
  200. }
  201. if z.closed {
  202. return nil
  203. }
  204. if !z.wroteHeader {
  205. z.Write(nil)
  206. if z.err != nil {
  207. return z.err
  208. }
  209. }
  210. z.err = z.compressor.Flush()
  211. return z.err
  212. }
  213. // Close closes the Writer, flushing any unwritten data to the underlying
  214. // io.Writer, but does not close the underlying io.Writer.
  215. func (z *Writer) Close() error {
  216. if z.err != nil {
  217. return z.err
  218. }
  219. if z.closed {
  220. return nil
  221. }
  222. z.closed = true
  223. if !z.wroteHeader {
  224. z.Write(nil)
  225. if z.err != nil {
  226. return z.err
  227. }
  228. }
  229. z.err = z.compressor.Close()
  230. if z.err != nil {
  231. return z.err
  232. }
  233. le.PutUint32(z.buf[:4], z.digest)
  234. le.PutUint32(z.buf[4:8], z.size)
  235. _, z.err = z.w.Write(z.buf[:8])
  236. return z.err
  237. }