compress.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148
  1. // Copyright 2013 The Gorilla 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 handlers
  5. import (
  6. "compress/flate"
  7. "compress/gzip"
  8. "io"
  9. "net/http"
  10. "strings"
  11. )
  12. type compressResponseWriter struct {
  13. io.Writer
  14. http.ResponseWriter
  15. http.Hijacker
  16. http.Flusher
  17. http.CloseNotifier
  18. }
  19. func (w *compressResponseWriter) WriteHeader(c int) {
  20. w.ResponseWriter.Header().Del("Content-Length")
  21. w.ResponseWriter.WriteHeader(c)
  22. }
  23. func (w *compressResponseWriter) Header() http.Header {
  24. return w.ResponseWriter.Header()
  25. }
  26. func (w *compressResponseWriter) Write(b []byte) (int, error) {
  27. h := w.ResponseWriter.Header()
  28. if h.Get("Content-Type") == "" {
  29. h.Set("Content-Type", http.DetectContentType(b))
  30. }
  31. h.Del("Content-Length")
  32. return w.Writer.Write(b)
  33. }
  34. type flusher interface {
  35. Flush() error
  36. }
  37. func (w *compressResponseWriter) Flush() {
  38. // Flush compressed data if compressor supports it.
  39. if f, ok := w.Writer.(flusher); ok {
  40. f.Flush()
  41. }
  42. // Flush HTTP response.
  43. if w.Flusher != nil {
  44. w.Flusher.Flush()
  45. }
  46. }
  47. // CompressHandler gzip compresses HTTP responses for clients that support it
  48. // via the 'Accept-Encoding' header.
  49. //
  50. // Compressing TLS traffic may leak the page contents to an attacker if the
  51. // page contains user input: http://security.stackexchange.com/a/102015/12208
  52. func CompressHandler(h http.Handler) http.Handler {
  53. return CompressHandlerLevel(h, gzip.DefaultCompression)
  54. }
  55. // CompressHandlerLevel gzip compresses HTTP responses with specified compression level
  56. // for clients that support it via the 'Accept-Encoding' header.
  57. //
  58. // The compression level should be gzip.DefaultCompression, gzip.NoCompression,
  59. // or any integer value between gzip.BestSpeed and gzip.BestCompression inclusive.
  60. // gzip.DefaultCompression is used in case of invalid compression level.
  61. func CompressHandlerLevel(h http.Handler, level int) http.Handler {
  62. if level < gzip.DefaultCompression || level > gzip.BestCompression {
  63. level = gzip.DefaultCompression
  64. }
  65. return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  66. L:
  67. for _, enc := range strings.Split(r.Header.Get("Accept-Encoding"), ",") {
  68. switch strings.TrimSpace(enc) {
  69. case "gzip":
  70. w.Header().Set("Content-Encoding", "gzip")
  71. w.Header().Add("Vary", "Accept-Encoding")
  72. gw, _ := gzip.NewWriterLevel(w, level)
  73. defer gw.Close()
  74. h, hok := w.(http.Hijacker)
  75. if !hok { /* w is not Hijacker... oh well... */
  76. h = nil
  77. }
  78. f, fok := w.(http.Flusher)
  79. if !fok {
  80. f = nil
  81. }
  82. cn, cnok := w.(http.CloseNotifier)
  83. if !cnok {
  84. cn = nil
  85. }
  86. w = &compressResponseWriter{
  87. Writer: gw,
  88. ResponseWriter: w,
  89. Hijacker: h,
  90. Flusher: f,
  91. CloseNotifier: cn,
  92. }
  93. break L
  94. case "deflate":
  95. w.Header().Set("Content-Encoding", "deflate")
  96. w.Header().Add("Vary", "Accept-Encoding")
  97. fw, _ := flate.NewWriter(w, level)
  98. defer fw.Close()
  99. h, hok := w.(http.Hijacker)
  100. if !hok { /* w is not Hijacker... oh well... */
  101. h = nil
  102. }
  103. f, fok := w.(http.Flusher)
  104. if !fok {
  105. f = nil
  106. }
  107. cn, cnok := w.(http.CloseNotifier)
  108. if !cnok {
  109. cn = nil
  110. }
  111. w = &compressResponseWriter{
  112. Writer: fw,
  113. ResponseWriter: w,
  114. Hijacker: h,
  115. Flusher: f,
  116. CloseNotifier: cn,
  117. }
  118. break L
  119. }
  120. }
  121. h.ServeHTTP(w, r)
  122. })
  123. }