-
Notifications
You must be signed in to change notification settings - Fork 3
/
middleware_gzip2.go
90 lines (72 loc) · 2 KB
/
middleware_gzip2.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
package ginKit
import (
"bytes"
"compress/gzip"
"github.com/gin-gonic/gin"
"github.com/richelieu-yang/chimera/v3/src/compress/gzipKit"
"strings"
)
// NewGzipMiddleware2
/*
@param minLength 等效于 Nginx 的 gzip_min_length 配置项
*/
func NewGzipMiddleware2(level, minLength int) (middleware gin.HandlerFunc, err error) {
if err = gzipKit.AssertValidLevel(level); err != nil {
return
}
middleware = func(ctx *gin.Context) {
req := ctx.Request
if !strings.Contains(req.Header.Get("Accept-Encoding"), "gzip") ||
strings.Contains(req.Header.Get("Connection"), "Upgrade") ||
strings.Contains(req.Header.Get("Accept"), "text/event-stream") {
ctx.Next()
return
}
ctx.Header("Content-Encoding", "gzip")
ctx.Header("Vary", "Accept-Encoding")
// 创建一个 ResponseWriter 接口的包装器
grw := &gzipResponseWriter{
ResponseWriter: ctx.Writer,
level: level,
minLength: minLength,
buffer: bytes.NewBuffer(nil),
}
// 替换原始的 ResponseWriter 接口
ctx.Writer = grw
// 继续处理请求
ctx.Next()
grw.WriteBody(ctx)
}
return
}
// 自定义的 ResponseWriter 接口的包装器
type gzipResponseWriter struct {
gin.ResponseWriter
level int
minLength int
buffer *bytes.Buffer
}
// 重写 Write 方法,将响应体写入缓冲区
func (grw *gzipResponseWriter) Write(data []byte) (int, error) {
return grw.buffer.Write(data)
}
func (grw *gzipResponseWriter) WriteString(s string) (int, error) {
return grw.buffer.WriteString(s)
}
// WriteBody 在请求结束时进行 gzip 压缩
func (grw *gzipResponseWriter) WriteBody(ctx *gin.Context) {
length := grw.buffer.Len()
// (1) 不进行 gzip 压缩
if length < grw.minLength {
ctx.Writer.Header().Del("Content-Encoding")
if length > 0 {
_, _ = grw.buffer.WriteTo(grw.ResponseWriter)
}
return
}
// (2)进行 gzip 压缩
gzipWriter, _ := gzip.NewWriterLevel(grw.ResponseWriter, grw.level)
defer gzipWriter.Close()
_, _ = grw.buffer.WriteTo(gzipWriter)
//_ = gzipWriter.Flush()
}