-
Notifications
You must be signed in to change notification settings - Fork 0
/
gzip.go
116 lines (103 loc) · 2.52 KB
/
gzip.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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
package rider
import (
"compress/gzip"
"io/ioutil"
"net/http"
"path/filepath"
"strings"
"sync"
)
const (
BestCompression = gzip.BestCompression
BestSpeed = gzip.BestSpeed
DefaultCompression = gzip.DefaultCompression
NoCompression = gzip.NoCompression
)
type gzipWriter struct {
http.ResponseWriter
writer *gzip.Writer
}
func Gzip(level int) HandlerFunc {
var gzPool sync.Pool
gzPool.New = func() interface{} {
gz, err := gzip.NewWriterLevel(ioutil.Discard, level)
if err != nil {
panic(err)
}
return gz
}
return func(c Context) {
if !shouldCompress(c.Request().request) {
c.Next()
return
}
res := c.Response()
originW := res.writer
gz := gzPool.Get().(*gzip.Writer)
defer gzPool.Put(gz)
gz.Reset(res.writer)
c.SetHeader(HeaderContentEncoding, "gzip")
c.AddHeader(HeaderVary, HeaderAcceptEncoding)
gw := &gzipWriter{}
gw.ResponseWriter = res.writer
gw.writer = gz
res.writer = gw
defer func() {
if res.Size == 0 {
//当发生panic的时候会走到这一步;将response的writer恢复为原来的writer,因为gzip的writer无法正确处理错误
if res.Header().Get(HeaderContentEncoding) == "gzip" {
res.Header().Del(HeaderContentEncoding)
}
res.writer = originW
}
gz.Close()
}()
c.Next()
}
}
func (g *gzipWriter) WriteString(s string) (int, error) {
return g.writer.Write([]byte(s))
}
func (g *gzipWriter) WriteHeader(code int) {
if code == http.StatusNoContent {
g.Header().Del(HeaderContentEncoding)
g.Header().Del(HeaderContentLength)
}
g.ResponseWriter.WriteHeader(code)
}
func (g *gzipWriter) Write(data []byte) (int, error) {
return g.writer.Write(data)
}
func shouldCompress(req *http.Request) bool {
if !strings.Contains(req.Header.Get(HeaderAcceptEncoding), "gzip") {
return false
}
extension := filepath.Ext(req.URL.Path)
if len(extension) < 4 { // fast path
return true
}
switch extension {
case ".png", ".gif", ".jpeg", ".jpg":
return false
default:
return true
}
}
/*
func (g *gzipWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
rw := &bufio.ReadWriter{}
writer := bufio.NewWriter(g)
rw.Writer = writer
n, _ := net.Dial("tcp", "mv.51mzzk.com:5000")
return n, rw, nil
}
*/
//在客户端关闭连接但还未发送响应体时,关闭连接
func (g *gzipWriter) CloseNotify() <-chan bool {
/*notify := r.writer.(http.CloseNotifier).CloseNotify()
go func() {
<-notify
r.server.logger.WARNING("HTTP connection just closed.")
}()*/
return g.ResponseWriter.(http.CloseNotifier).CloseNotify()
}