本站源代码
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

151 lines
3.5KB

  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. r.Header.Del("Accept-Encoding")
  72. w.Header().Add("Vary", "Accept-Encoding")
  73. gw, _ := gzip.NewWriterLevel(w, level)
  74. defer gw.Close()
  75. h, hok := w.(http.Hijacker)
  76. if !hok { /* w is not Hijacker... oh well... */
  77. h = nil
  78. }
  79. f, fok := w.(http.Flusher)
  80. if !fok {
  81. f = nil
  82. }
  83. cn, cnok := w.(http.CloseNotifier)
  84. if !cnok {
  85. cn = nil
  86. }
  87. w = &compressResponseWriter{
  88. Writer: gw,
  89. ResponseWriter: w,
  90. Hijacker: h,
  91. Flusher: f,
  92. CloseNotifier: cn,
  93. }
  94. break L
  95. case "deflate":
  96. w.Header().Set("Content-Encoding", "deflate")
  97. r.Header.Del("Accept-Encoding")
  98. w.Header().Add("Vary", "Accept-Encoding")
  99. fw, _ := flate.NewWriter(w, level)
  100. defer fw.Close()
  101. h, hok := w.(http.Hijacker)
  102. if !hok { /* w is not Hijacker... oh well... */
  103. h = nil
  104. }
  105. f, fok := w.(http.Flusher)
  106. if !fok {
  107. f = nil
  108. }
  109. cn, cnok := w.(http.CloseNotifier)
  110. if !cnok {
  111. cn = nil
  112. }
  113. w = &compressResponseWriter{
  114. Writer: fw,
  115. ResponseWriter: w,
  116. Hijacker: h,
  117. Flusher: f,
  118. CloseNotifier: cn,
  119. }
  120. break L
  121. }
  122. }
  123. h.ServeHTTP(w, r)
  124. })
  125. }
上海开阖软件有限公司 沪ICP备12045867号-1