gzip.go 9.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333
  1. package gziphandler
  2. import (
  3. "bufio"
  4. "compress/gzip"
  5. "fmt"
  6. "io"
  7. "net"
  8. "net/http"
  9. "strconv"
  10. "strings"
  11. "sync"
  12. )
  13. const (
  14. vary = "Vary"
  15. acceptEncoding = "Accept-Encoding"
  16. contentEncoding = "Content-Encoding"
  17. contentType = "Content-Type"
  18. contentLength = "Content-Length"
  19. )
  20. type codings map[string]float64
  21. const (
  22. // DefaultQValue is the default qvalue to assign to an encoding if no explicit qvalue is set.
  23. // This is actually kind of ambiguous in RFC 2616, so hopefully it's correct.
  24. // The examples seem to indicate that it is.
  25. DefaultQValue = 1.0
  26. // DefaultMinSize defines the minimum size to reach to enable compression.
  27. // It's 512 bytes.
  28. DefaultMinSize = 512
  29. )
  30. // gzipWriterPools stores a sync.Pool for each compression level for reuse of
  31. // gzip.Writers. Use poolIndex to covert a compression level to an index into
  32. // gzipWriterPools.
  33. var gzipWriterPools [gzip.BestCompression - gzip.BestSpeed + 2]*sync.Pool
  34. func init() {
  35. for i := gzip.BestSpeed; i <= gzip.BestCompression; i++ {
  36. addLevelPool(i)
  37. }
  38. addLevelPool(gzip.DefaultCompression)
  39. }
  40. // poolIndex maps a compression level to its index into gzipWriterPools. It
  41. // assumes that level is a valid gzip compression level.
  42. func poolIndex(level int) int {
  43. // gzip.DefaultCompression == -1, so we need to treat it special.
  44. if level == gzip.DefaultCompression {
  45. return gzip.BestCompression - gzip.BestSpeed + 1
  46. }
  47. return level - gzip.BestSpeed
  48. }
  49. func addLevelPool(level int) {
  50. gzipWriterPools[poolIndex(level)] = &sync.Pool{
  51. New: func() interface{} {
  52. // NewWriterLevel only returns error on a bad level, we are guaranteeing
  53. // that this will be a valid level so it is okay to ignore the returned
  54. // error.
  55. w, _ := gzip.NewWriterLevel(nil, level)
  56. return w
  57. },
  58. }
  59. }
  60. // GzipResponseWriter provides an http.ResponseWriter interface, which gzips
  61. // bytes before writing them to the underlying response. This doesn't close the
  62. // writers, so don't forget to do that.
  63. // It can be configured to skip response smaller than minSize.
  64. type GzipResponseWriter struct {
  65. http.ResponseWriter
  66. index int // Index for gzipWriterPools.
  67. gw *gzip.Writer
  68. code int // Saves the WriteHeader value.
  69. minSize int // Specifed the minimum response size to gzip. If the response length is bigger than this value, it is compressed.
  70. buf []byte // Holds the first part of the write before reaching the minSize or the end of the write.
  71. }
  72. // Write appends data to the gzip writer.
  73. func (w *GzipResponseWriter) Write(b []byte) (int, error) {
  74. // If content type is not set.
  75. if _, ok := w.Header()[contentType]; !ok {
  76. // It infer it from the uncompressed body.
  77. w.Header().Set(contentType, http.DetectContentType(b))
  78. }
  79. // GZIP responseWriter is initialized. Use the GZIP responseWriter.
  80. if w.gw != nil {
  81. n, err := w.gw.Write(b)
  82. return n, err
  83. }
  84. // Save the write into a buffer for later use in GZIP responseWriter (if content is long enough) or at close with regular responseWriter.
  85. // On the first write, w.buf changes from nil to a valid slice
  86. w.buf = append(w.buf, b...)
  87. // If the global writes are bigger than the minSize, compression is enable.
  88. if len(w.buf) >= w.minSize {
  89. err := w.startGzip()
  90. if err != nil {
  91. return 0, err
  92. }
  93. }
  94. return len(b), nil
  95. }
  96. // startGzip initialize any GZIP specific informations.
  97. func (w *GzipResponseWriter) startGzip() error {
  98. // Set the GZIP header.
  99. w.Header().Set(contentEncoding, "gzip")
  100. // if the Content-Length is already set, then calls to Write on gzip
  101. // will fail to set the Content-Length header since its already set
  102. // See: https://github.com/golang/go/issues/14975.
  103. w.Header().Del(contentLength)
  104. // Write the header to gzip response.
  105. if w.code != 0 {
  106. w.ResponseWriter.WriteHeader(w.code)
  107. }
  108. // Initialize the GZIP response.
  109. w.init()
  110. // Flush the buffer into the gzip reponse.
  111. n, err := w.gw.Write(w.buf)
  112. // This should never happen (per io.Writer docs), but if the write didn't
  113. // accept the entire buffer but returned no specific error, we have no clue
  114. // what's going on, so abort just to be safe.
  115. if err == nil && n < len(w.buf) {
  116. return io.ErrShortWrite
  117. }
  118. w.buf = nil
  119. return err
  120. }
  121. // WriteHeader just saves the response code until close or GZIP effective writes.
  122. func (w *GzipResponseWriter) WriteHeader(code int) {
  123. w.code = code
  124. }
  125. // init graps a new gzip writer from the gzipWriterPool and writes the correct
  126. // content encoding header.
  127. func (w *GzipResponseWriter) init() {
  128. // Bytes written during ServeHTTP are redirected to this gzip writer
  129. // before being written to the underlying response.
  130. gzw := gzipWriterPools[w.index].Get().(*gzip.Writer)
  131. gzw.Reset(w.ResponseWriter)
  132. w.gw = gzw
  133. }
  134. // Close will close the gzip.Writer and will put it back in the gzipWriterPool.
  135. func (w *GzipResponseWriter) Close() error {
  136. if w.gw == nil {
  137. // Gzip not trigged yet, write out regular response.
  138. if w.code != 0 {
  139. w.ResponseWriter.WriteHeader(w.code)
  140. }
  141. if w.buf != nil {
  142. _, writeErr := w.ResponseWriter.Write(w.buf)
  143. // Returns the error if any at write.
  144. if writeErr != nil {
  145. return fmt.Errorf("gziphandler: write to regular responseWriter at close gets error: %q", writeErr.Error())
  146. }
  147. }
  148. return nil
  149. }
  150. err := w.gw.Close()
  151. gzipWriterPools[w.index].Put(w.gw)
  152. w.gw = nil
  153. return err
  154. }
  155. // Flush flushes the underlying *gzip.Writer and then the underlying
  156. // http.ResponseWriter if it is an http.Flusher. This makes GzipResponseWriter
  157. // an http.Flusher.
  158. func (w *GzipResponseWriter) Flush() {
  159. if w.gw != nil {
  160. w.gw.Flush()
  161. }
  162. if fw, ok := w.ResponseWriter.(http.Flusher); ok {
  163. fw.Flush()
  164. }
  165. }
  166. // Hijack implements http.Hijacker. If the underlying ResponseWriter is a
  167. // Hijacker, its Hijack method is returned. Otherwise an error is returned.
  168. func (w *GzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
  169. if hj, ok := w.ResponseWriter.(http.Hijacker); ok {
  170. return hj.Hijack()
  171. }
  172. return nil, nil, fmt.Errorf("http.Hijacker interface is not supported")
  173. }
  174. // verify Hijacker interface implementation
  175. var _ http.Hijacker = &GzipResponseWriter{}
  176. // MustNewGzipLevelHandler behaves just like NewGzipLevelHandler except that in
  177. // an error case it panics rather than returning an error.
  178. func MustNewGzipLevelHandler(level int) func(http.Handler) http.Handler {
  179. wrap, err := NewGzipLevelHandler(level)
  180. if err != nil {
  181. panic(err)
  182. }
  183. return wrap
  184. }
  185. // NewGzipLevelHandler returns a wrapper function (often known as middleware)
  186. // which can be used to wrap an HTTP handler to transparently gzip the response
  187. // body if the client supports it (via the Accept-Encoding header). Responses will
  188. // be encoded at the given gzip compression level. An error will be returned only
  189. // if an invalid gzip compression level is given, so if one can ensure the level
  190. // is valid, the returned error can be safely ignored.
  191. func NewGzipLevelHandler(level int) (func(http.Handler) http.Handler, error) {
  192. return NewGzipLevelAndMinSize(level, DefaultMinSize)
  193. }
  194. // NewGzipLevelAndMinSize behave as NewGzipLevelHandler except it let the caller
  195. // specify the minimum size before compression.
  196. func NewGzipLevelAndMinSize(level, minSize int) (func(http.Handler) http.Handler, error) {
  197. if level != gzip.DefaultCompression && (level < gzip.BestSpeed || level > gzip.BestCompression) {
  198. return nil, fmt.Errorf("invalid compression level requested: %d", level)
  199. }
  200. if minSize < 0 {
  201. return nil, fmt.Errorf("minimum size must be more than zero")
  202. }
  203. return func(h http.Handler) http.Handler {
  204. index := poolIndex(level)
  205. return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  206. w.Header().Add(vary, acceptEncoding)
  207. if acceptsGzip(r) {
  208. gw := &GzipResponseWriter{
  209. ResponseWriter: w,
  210. index: index,
  211. minSize: minSize,
  212. }
  213. defer gw.Close()
  214. h.ServeHTTP(gw, r)
  215. } else {
  216. h.ServeHTTP(w, r)
  217. }
  218. })
  219. }, nil
  220. }
  221. // GzipHandler wraps an HTTP handler, to transparently gzip the response body if
  222. // the client supports it (via the Accept-Encoding header). This will compress at
  223. // the default compression level.
  224. func GzipHandler(h http.Handler) http.Handler {
  225. wrapper, _ := NewGzipLevelHandler(gzip.DefaultCompression)
  226. return wrapper(h)
  227. }
  228. // acceptsGzip returns true if the given HTTP request indicates that it will
  229. // accept a gzipped response.
  230. func acceptsGzip(r *http.Request) bool {
  231. acceptedEncodings, _ := parseEncodings(r.Header.Get(acceptEncoding))
  232. return acceptedEncodings["gzip"] > 0.0
  233. }
  234. // parseEncodings attempts to parse a list of codings, per RFC 2616, as might
  235. // appear in an Accept-Encoding header. It returns a map of content-codings to
  236. // quality values, and an error containing the errors encountered. It's probably
  237. // safe to ignore those, because silently ignoring errors is how the internet
  238. // works.
  239. //
  240. // See: http://tools.ietf.org/html/rfc2616#section-14.3.
  241. func parseEncodings(s string) (codings, error) {
  242. c := make(codings)
  243. var e []string
  244. for _, ss := range strings.Split(s, ",") {
  245. coding, qvalue, err := parseCoding(ss)
  246. if err != nil {
  247. e = append(e, err.Error())
  248. } else {
  249. c[coding] = qvalue
  250. }
  251. }
  252. // TODO (adammck): Use a proper multi-error struct, so the individual errors
  253. // can be extracted if anyone cares.
  254. if len(e) > 0 {
  255. return c, fmt.Errorf("errors while parsing encodings: %s", strings.Join(e, ", "))
  256. }
  257. return c, nil
  258. }
  259. // parseCoding parses a single conding (content-coding with an optional qvalue),
  260. // as might appear in an Accept-Encoding header. It attempts to forgive minor
  261. // formatting errors.
  262. func parseCoding(s string) (coding string, qvalue float64, err error) {
  263. for n, part := range strings.Split(s, ";") {
  264. part = strings.TrimSpace(part)
  265. qvalue = DefaultQValue
  266. if n == 0 {
  267. coding = strings.ToLower(part)
  268. } else if strings.HasPrefix(part, "q=") {
  269. qvalue, err = strconv.ParseFloat(strings.TrimPrefix(part, "q="), 64)
  270. if qvalue < 0.0 {
  271. qvalue = 0.0
  272. } else if qvalue > 1.0 {
  273. qvalue = 1.0
  274. }
  275. }
  276. }
  277. if coding == "" {
  278. err = fmt.Errorf("empty content-coding")
  279. }
  280. return
  281. }