gzip.go 6.5 KB

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