gzip.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121
  1. // Copyright 2013 Martini Authors
  2. // Copyright 2015 The Macaron Authors
  3. //
  4. // Licensed under the Apache License, Version 2.0 (the "License"): you may
  5. // not use this file except in compliance with the License. You may obtain
  6. // a copy of the License at
  7. //
  8. // http://www.apache.org/licenses/LICENSE-2.0
  9. //
  10. // Unless required by applicable law or agreed to in writing, software
  11. // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  12. // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  13. // License for the specific language governing permissions and limitations
  14. // under the License.
  15. package gzip
  16. import (
  17. "bufio"
  18. "fmt"
  19. "net"
  20. "net/http"
  21. "strings"
  22. "github.com/klauspost/compress/gzip"
  23. "gopkg.in/macaron.v1"
  24. )
  25. const (
  26. _HEADER_ACCEPT_ENCODING = "Accept-Encoding"
  27. _HEADER_CONTENT_ENCODING = "Content-Encoding"
  28. _HEADER_CONTENT_LENGTH = "Content-Length"
  29. _HEADER_CONTENT_TYPE = "Content-Type"
  30. _HEADER_VARY = "Vary"
  31. )
  32. // Options represents a struct for specifying configuration options for the GZip middleware.
  33. type Options struct {
  34. // Compression level. Can be DefaultCompression(-1), ConstantCompression(-2)
  35. // or any integer value between BestSpeed(1) and BestCompression(9) inclusive.
  36. CompressionLevel int
  37. }
  38. func isCompressionLevelValid(level int) bool {
  39. return level == gzip.DefaultCompression ||
  40. level == gzip.ConstantCompression ||
  41. (level >= gzip.BestSpeed && level <= gzip.BestCompression)
  42. }
  43. func prepareOptions(options []Options) Options {
  44. var opt Options
  45. if len(options) > 0 {
  46. opt = options[0]
  47. }
  48. if !isCompressionLevelValid(opt.CompressionLevel) {
  49. // For web content, level 4 seems to be a sweet spot.
  50. opt.CompressionLevel = 4
  51. }
  52. return opt
  53. }
  54. // Gziper returns a Handler that adds gzip compression to all requests.
  55. // Make sure to include the Gzip middleware above other middleware
  56. // that alter the response body (like the render middleware).
  57. func Gziper(options ...Options) macaron.Handler {
  58. opt := prepareOptions(options)
  59. return func(ctx *macaron.Context) {
  60. if !strings.Contains(ctx.Req.Header.Get(_HEADER_ACCEPT_ENCODING), "gzip") {
  61. return
  62. }
  63. headers := ctx.Resp.Header()
  64. headers.Set(_HEADER_CONTENT_ENCODING, "gzip")
  65. headers.Set(_HEADER_VARY, _HEADER_ACCEPT_ENCODING)
  66. // We've made sure compression level is valid in prepareGzipOptions,
  67. // no need to check same error again.
  68. gz, err := gzip.NewWriterLevel(ctx.Resp, opt.CompressionLevel)
  69. if err != nil {
  70. panic(err.Error())
  71. }
  72. defer gz.Close()
  73. gzw := gzipResponseWriter{gz, ctx.Resp}
  74. ctx.Resp = gzw
  75. ctx.MapTo(gzw, (*http.ResponseWriter)(nil))
  76. // Check if render middleware has been registered,
  77. // if yes, we need to modify ResponseWriter for it as well.
  78. if _, ok := ctx.Render.(*macaron.DummyRender); !ok {
  79. ctx.Render.SetResponseWriter(gzw)
  80. }
  81. ctx.Next()
  82. // delete content length after we know we have been written to
  83. gzw.Header().Del("Content-Length")
  84. }
  85. }
  86. type gzipResponseWriter struct {
  87. w *gzip.Writer
  88. macaron.ResponseWriter
  89. }
  90. func (grw gzipResponseWriter) Write(p []byte) (int, error) {
  91. if len(grw.Header().Get(_HEADER_CONTENT_TYPE)) == 0 {
  92. grw.Header().Set(_HEADER_CONTENT_TYPE, http.DetectContentType(p))
  93. }
  94. return grw.w.Write(p)
  95. }
  96. func (grw gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
  97. hijacker, ok := grw.ResponseWriter.(http.Hijacker)
  98. if !ok {
  99. return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
  100. }
  101. return hijacker.Hijack()
  102. }