-
Notifications
You must be signed in to change notification settings - Fork 5
/
compress.go
217 lines (194 loc) · 6.22 KB
/
compress.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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
// MIT License
// Copyright (c) 2020 Tree Xie
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
package middleware
import (
"bytes"
"errors"
"regexp"
"strings"
"github.com/vicanso/elton"
)
var (
// DefaultCompressRegexp compress text, javascript, json and wasm
DefaultCompressRegexp = regexp.MustCompile("text|javascript|json|wasm|font")
)
const (
// DefaultCompressMinLength min compress length(1KB)
DefaultCompressMinLength = 1024
)
const IgnoreCompression = -128
type (
// Compressor compressor interface
Compressor interface {
// Accept accept check function
Accept(c *elton.Context, bodySize int) (acceptable bool, encoding string)
// Compress compress function
Compress([]byte, ...int) (*bytes.Buffer, error)
// Pipe pipe function
Pipe(*elton.Context) error
}
// Config compress config
CompressConfig struct {
// Checker check the data is compressable
Checker *regexp.Regexp
// Compressors compressor list
Compressors []Compressor
// Skipper skipper function
Skipper elton.Skipper
// DynamicLevel return dynamic level
DynamicLevel func(c *elton.Context, bodySize int, encoding string) int
// OnBeforeCompress before compress event
OnBeforeCompress func(c *elton.Context) error
}
)
// AcceptEncoding check request accept encoding
func AcceptEncoding(c *elton.Context, encoding string) (bool, string) {
acceptEncoding := c.GetRequestHeader(elton.HeaderAcceptEncoding)
if strings.Contains(acceptEncoding, encoding) {
return true, encoding
}
return false, ""
}
// AddCompressor to the compress config
func (conf *CompressConfig) AddCompressor(compressor Compressor) {
if conf.Compressors == nil {
conf.Compressors = make([]Compressor, 0)
}
conf.Compressors = append(conf.Compressors, compressor)
}
// NewCompressConfig returns a compress config with multi-compressor
func NewCompressConfig(compressors ...Compressor) CompressConfig {
cfg := CompressConfig{}
for _, compressor := range compressors {
cfg.AddCompressor(compressor)
}
return cfg
}
// NewDefaultCompress return a new compress middleware, it include gzip compress
func NewDefaultCompress() elton.Handler {
cfg := NewCompressConfig(new(GzipCompressor))
return NewCompress(cfg)
}
// NewCompress return a new compress middleware.
// It will use 'text|javascript|json|wasm|font' as default content type checker for compress.
// It will throw a panic if the compressors is empty.
func NewCompress(config CompressConfig) elton.Handler {
skipper := config.Skipper
if skipper == nil {
skipper = elton.DefaultSkipper
}
checker := config.Checker
if checker == nil {
checker = DefaultCompressRegexp
}
compressorList := config.Compressors
if len(compressorList) == 0 {
panic(errors.New("compressor can't be empty"))
}
dynamicLevel := config.DynamicLevel
return func(c *elton.Context) error {
if skipper(c) {
return c.Next()
}
err := c.Next()
if err != nil {
return err
}
isReaderBody := c.IsReaderBody()
// 如果数据为空,而且body不是reader,直接跳过
if c.BodyBuffer == nil && !isReaderBody {
return nil
}
// encoding 不为空,已做处理,无需要压缩
if c.GetHeader(elton.HeaderContentEncoding) != "" {
return nil
}
contentType := c.GetHeader(elton.HeaderContentType)
// 数据类型为非可压缩,则返回
if !checker.MatchString(contentType) {
return nil
}
if config.OnBeforeCompress != nil {
err = config.OnBeforeCompress(c)
if err != nil {
return err
}
}
var body []byte
if c.BodyBuffer != nil {
body = c.BodyBuffer.Bytes()
}
// 对于reader类,无法判断长度,认为长度为-1
bodySize := -1
if !isReaderBody {
// 如果数据长度少于最小压缩长度
bodySize = len(body)
}
fillHeader := func(encoding string) {
c.SetHeader(elton.HeaderContentEncoding, encoding)
c.AddHeader("Vary", "Accept-Encoding")
etagValue := c.GetHeader(elton.HeaderETag)
// after compress, etag should be weak etag
if etagValue != "" && !strings.HasPrefix(etagValue, "W/") {
c.SetHeader(elton.HeaderETag, "W/"+etagValue)
}
}
for _, compressor := range compressorList {
acceptable, encoding := compressor.Accept(c, bodySize)
if !acceptable {
continue
}
if isReaderBody {
// 压缩时清除content length
c.Header().Del(elton.HeaderContentLength)
// 执行pipe之前先设置http响应头
fillHeader(encoding)
err = compressor.Pipe(c)
// 如果出错直接返回,此时也有可能已经开始写入数据,导致http后续无法再写入status code
if err != nil {
return err
}
// 成功跳出循环
// pipe 将数据直接转至原有的Response,因此设置committed为true
c.Committed = true
// 清除 response body
c.Body = nil
break
}
var levels []int
// 如果获取压缩级别函数有设置
if dynamicLevel != nil {
levels = []int{
dynamicLevel(c, len(body), encoding),
}
}
newBuf, e := compressor.Compress(body, levels...)
// 如果压缩成功,则使用压缩数据
// 失败则忽略(不修改原数据,仅触发error)
if e != nil {
c.Elton().EmitError(c, e)
} else {
fillHeader(encoding)
c.BodyBuffer = newBuf
}
// 已有符合的压缩,处理完成跳出当前循环
break
}
return nil
}
}