1
0

acceptencoder.go 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231
  1. // Copyright 2015 beego Author. All Rights Reserved.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package context
  15. import (
  16. "bytes"
  17. "compress/flate"
  18. "compress/gzip"
  19. "compress/zlib"
  20. "io"
  21. "net/http"
  22. "os"
  23. "strconv"
  24. "strings"
  25. "sync"
  26. )
  27. var (
  28. //Default size==20B same as nginx
  29. defaultGzipMinLength = 20
  30. //Content will only be compressed if content length is either unknown or greater than gzipMinLength.
  31. gzipMinLength = defaultGzipMinLength
  32. //The compression level used for deflate compression. (0-9).
  33. gzipCompressLevel int
  34. //List of HTTP methods to compress. If not set, only GET requests are compressed.
  35. includedMethods map[string]bool
  36. getMethodOnly bool
  37. )
  38. func InitGzip(minLength, compressLevel int, methods []string) {
  39. if minLength >= 0 {
  40. gzipMinLength = minLength
  41. }
  42. gzipCompressLevel = compressLevel
  43. if gzipCompressLevel < flate.NoCompression || gzipCompressLevel > flate.BestCompression {
  44. gzipCompressLevel = flate.BestSpeed
  45. }
  46. getMethodOnly = (len(methods) == 0) || (len(methods) == 1 && strings.ToUpper(methods[0]) == "GET")
  47. includedMethods = make(map[string]bool, len(methods))
  48. for _, v := range methods {
  49. includedMethods[strings.ToUpper(v)] = true
  50. }
  51. }
  52. type resetWriter interface {
  53. io.Writer
  54. Reset(w io.Writer)
  55. }
  56. type nopResetWriter struct {
  57. io.Writer
  58. }
  59. func (n nopResetWriter) Reset(w io.Writer) {
  60. //do nothing
  61. }
  62. type acceptEncoder struct {
  63. name string
  64. levelEncode func(int) resetWriter
  65. customCompressLevelPool *sync.Pool
  66. bestCompressionPool *sync.Pool
  67. }
  68. func (ac acceptEncoder) encode(wr io.Writer, level int) resetWriter {
  69. if ac.customCompressLevelPool == nil || ac.bestCompressionPool == nil {
  70. return nopResetWriter{wr}
  71. }
  72. var rwr resetWriter
  73. switch level {
  74. case flate.BestSpeed:
  75. rwr = ac.customCompressLevelPool.Get().(resetWriter)
  76. case flate.BestCompression:
  77. rwr = ac.bestCompressionPool.Get().(resetWriter)
  78. default:
  79. rwr = ac.levelEncode(level)
  80. }
  81. rwr.Reset(wr)
  82. return rwr
  83. }
  84. func (ac acceptEncoder) put(wr resetWriter, level int) {
  85. if ac.customCompressLevelPool == nil || ac.bestCompressionPool == nil {
  86. return
  87. }
  88. wr.Reset(nil)
  89. //notice
  90. //compressionLevel==BestCompression DOES NOT MATTER
  91. //sync.Pool will not memory leak
  92. switch level {
  93. case gzipCompressLevel:
  94. ac.customCompressLevelPool.Put(wr)
  95. case flate.BestCompression:
  96. ac.bestCompressionPool.Put(wr)
  97. }
  98. }
  99. var (
  100. noneCompressEncoder = acceptEncoder{"", nil, nil, nil}
  101. gzipCompressEncoder = acceptEncoder{
  102. name: "gzip",
  103. levelEncode: func(level int) resetWriter { wr, _ := gzip.NewWriterLevel(nil, level); return wr },
  104. customCompressLevelPool: &sync.Pool{New: func() interface{} { wr, _ := gzip.NewWriterLevel(nil, gzipCompressLevel); return wr }},
  105. bestCompressionPool: &sync.Pool{New: func() interface{} { wr, _ := gzip.NewWriterLevel(nil, flate.BestCompression); return wr }},
  106. }
  107. //according to the sec :http://tools.ietf.org/html/rfc2616#section-3.5 ,the deflate compress in http is zlib indeed
  108. //deflate
  109. //The "zlib" format defined in RFC 1950 [31] in combination with
  110. //the "deflate" compression mechanism described in RFC 1951 [29].
  111. deflateCompressEncoder = acceptEncoder{
  112. name: "deflate",
  113. levelEncode: func(level int) resetWriter { wr, _ := zlib.NewWriterLevel(nil, level); return wr },
  114. customCompressLevelPool: &sync.Pool{New: func() interface{} { wr, _ := zlib.NewWriterLevel(nil, gzipCompressLevel); return wr }},
  115. bestCompressionPool: &sync.Pool{New: func() interface{} { wr, _ := zlib.NewWriterLevel(nil, flate.BestCompression); return wr }},
  116. }
  117. )
  118. var (
  119. encoderMap = map[string]acceptEncoder{ // all the other compress methods will ignore
  120. "gzip": gzipCompressEncoder,
  121. "deflate": deflateCompressEncoder,
  122. "*": gzipCompressEncoder, // * means any compress will accept,we prefer gzip
  123. "identity": noneCompressEncoder, // identity means none-compress
  124. }
  125. )
  126. // WriteFile reads from file and writes to writer by the specific encoding(gzip/deflate)
  127. func WriteFile(encoding string, writer io.Writer, file *os.File) (bool, string, error) {
  128. return writeLevel(encoding, writer, file, flate.BestCompression)
  129. }
  130. // WriteBody reads writes content to writer by the specific encoding(gzip/deflate)
  131. func WriteBody(encoding string, writer io.Writer, content []byte) (bool, string, error) {
  132. if encoding == "" || len(content) < gzipMinLength {
  133. _, err := writer.Write(content)
  134. return false, "", err
  135. }
  136. return writeLevel(encoding, writer, bytes.NewReader(content), gzipCompressLevel)
  137. }
  138. // writeLevel reads from reader,writes to writer by specific encoding and compress level
  139. // the compress level is defined by deflate package
  140. func writeLevel(encoding string, writer io.Writer, reader io.Reader, level int) (bool, string, error) {
  141. var outputWriter resetWriter
  142. var err error
  143. var ce = noneCompressEncoder
  144. if cf, ok := encoderMap[encoding]; ok {
  145. ce = cf
  146. }
  147. encoding = ce.name
  148. outputWriter = ce.encode(writer, level)
  149. defer ce.put(outputWriter, level)
  150. _, err = io.Copy(outputWriter, reader)
  151. if err != nil {
  152. return false, "", err
  153. }
  154. switch outputWriter.(type) {
  155. case io.WriteCloser:
  156. outputWriter.(io.WriteCloser).Close()
  157. }
  158. return encoding != "", encoding, nil
  159. }
  160. // ParseEncoding will extract the right encoding for response
  161. // the Accept-Encoding's sec is here:
  162. // http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.3
  163. func ParseEncoding(r *http.Request) string {
  164. if r == nil {
  165. return ""
  166. }
  167. if (getMethodOnly && r.Method == "GET") || includedMethods[r.Method] {
  168. return parseEncoding(r)
  169. }
  170. return ""
  171. }
  172. type q struct {
  173. name string
  174. value float64
  175. }
  176. func parseEncoding(r *http.Request) string {
  177. acceptEncoding := r.Header.Get("Accept-Encoding")
  178. if acceptEncoding == "" {
  179. return ""
  180. }
  181. var lastQ q
  182. for _, v := range strings.Split(acceptEncoding, ",") {
  183. v = strings.TrimSpace(v)
  184. if v == "" {
  185. continue
  186. }
  187. vs := strings.Split(v, ";")
  188. var cf acceptEncoder
  189. var ok bool
  190. if cf, ok = encoderMap[vs[0]]; !ok {
  191. continue
  192. }
  193. if len(vs) == 1 {
  194. return cf.name
  195. }
  196. if len(vs) == 2 {
  197. f, _ := strconv.ParseFloat(strings.Replace(vs[1], "q=", "", -1), 64)
  198. if f == 0 {
  199. continue
  200. }
  201. if f > lastQ.value {
  202. lastQ = q{cf.name, f}
  203. }
  204. }
  205. }
  206. return lastQ.name
  207. }