forked from influxdata/influxdb
/
http_compression.go
64 lines (54 loc) · 1.55 KB
/
http_compression.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
package http
import (
"compress/gzip"
"compress/zlib"
"io"
libhttp "net/http"
"strings"
)
type CompressedResponseWriter struct {
responseWriter libhttp.ResponseWriter
writer io.Writer
}
func NewCompressionResponseWriter(useCompression bool, rw libhttp.ResponseWriter, req *libhttp.Request) *CompressedResponseWriter {
var writer io.Writer = rw
if req.Header.Get("Accept-Encoding") != "" {
encodings := strings.Split(req.Header.Get("Accept-Encoding"), ",")
for _, val := range encodings {
if val == "gzip" {
rw.Header().Set("Content-Encoding", "gzip")
writer, _ = gzip.NewWriterLevel(writer, gzip.BestSpeed)
break
} else if val == "deflate" {
rw.Header().Set("Content-Encoding", "deflate")
writer, _ = zlib.NewWriterLevel(writer, zlib.BestSpeed)
break
}
}
}
return &CompressedResponseWriter{rw, writer}
}
func (self *CompressedResponseWriter) Header() libhttp.Header {
return self.responseWriter.Header()
}
func (self *CompressedResponseWriter) Write(bs []byte) (int, error) {
return self.writer.Write(bs)
}
func (self *CompressedResponseWriter) WriteHeader(responseCode int) {
self.responseWriter.WriteHeader(responseCode)
}
func CompressionHandler(enableCompression bool, handler libhttp.HandlerFunc) libhttp.HandlerFunc {
if !enableCompression {
return handler
}
return func(rw libhttp.ResponseWriter, req *libhttp.Request) {
crw := NewCompressionResponseWriter(true, rw, req)
handler(crw, req)
switch x := crw.writer.(type) {
case *gzip.Writer:
x.Close()
case *zlib.Writer:
x.Close()
}
}
}