-
Notifications
You must be signed in to change notification settings - Fork 4.9k
/
brotli.go
374 lines (314 loc) · 10.4 KB
/
brotli.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
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
package brotli
import (
"bufio"
"errors"
"fmt"
"io"
"mime"
"net"
"net/http"
"github.com/andybalholm/brotli"
)
const (
vary = "Vary"
acceptEncoding = "Accept-Encoding"
contentEncoding = "Content-Encoding"
contentLength = "Content-Length"
contentType = "Content-Type"
)
// Config is the Brotli handler configuration.
type Config struct {
// ExcludedContentTypes is the list of content types for which we should not compress.
// Mutually exclusive with the IncludedContentTypes option.
ExcludedContentTypes []string
// IncludedContentTypes is the list of content types for which compression should be exclusively enabled.
// Mutually exclusive with the ExcludedContentTypes option.
IncludedContentTypes []string
// MinSize is the minimum size (in bytes) required to enable compression.
MinSize int
}
// NewWrapper returns a new Brotli compressing wrapper.
func NewWrapper(cfg Config) (func(http.Handler) http.HandlerFunc, error) {
if cfg.MinSize < 0 {
return nil, errors.New("minimum size must be greater than or equal to zero")
}
if len(cfg.ExcludedContentTypes) > 0 && len(cfg.IncludedContentTypes) > 0 {
return nil, errors.New("excludedContentTypes and includedContentTypes options are mutually exclusive")
}
var excludedContentTypes []parsedContentType
for _, v := range cfg.ExcludedContentTypes {
mediaType, params, err := mime.ParseMediaType(v)
if err != nil {
return nil, fmt.Errorf("parsing excluded media type: %w", err)
}
excludedContentTypes = append(excludedContentTypes, parsedContentType{mediaType, params})
}
var includedContentTypes []parsedContentType
for _, v := range cfg.IncludedContentTypes {
mediaType, params, err := mime.ParseMediaType(v)
if err != nil {
return nil, fmt.Errorf("parsing included media type: %w", err)
}
includedContentTypes = append(includedContentTypes, parsedContentType{mediaType, params})
}
return func(h http.Handler) http.HandlerFunc {
return func(rw http.ResponseWriter, r *http.Request) {
rw.Header().Add(vary, acceptEncoding)
brw := &responseWriter{
rw: rw,
bw: brotli.NewWriter(rw),
minSize: cfg.MinSize,
statusCode: http.StatusOK,
excludedContentTypes: excludedContentTypes,
includedContentTypes: includedContentTypes,
}
defer brw.close()
h.ServeHTTP(brw, r)
}
}, nil
}
// TODO: check whether we want to implement content-type sniffing (as gzip does)
// TODO: check whether we should support Accept-Ranges (as gzip does, see https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Ranges)
type responseWriter struct {
rw http.ResponseWriter
bw *brotli.Writer
minSize int
excludedContentTypes []parsedContentType
includedContentTypes []parsedContentType
buf []byte
hijacked bool
compressionStarted bool
compressionDisabled bool
headersSent bool
// Mostly needed to avoid calling bw.Flush/bw.Close when no data was
// written in bw.
seenData bool
statusCodeSet bool
statusCode int
}
func (r *responseWriter) Header() http.Header {
return r.rw.Header()
}
func (r *responseWriter) WriteHeader(statusCode int) {
if r.statusCodeSet {
return
}
r.statusCode = statusCode
r.statusCodeSet = true
}
func (r *responseWriter) Write(p []byte) (int, error) {
// i.e. has write ever been called at least once with non nil data.
if !r.seenData && len(p) > 0 {
r.seenData = true
}
// We do not compress, either for contentEncoding or contentType reasons.
if r.compressionDisabled {
return r.rw.Write(p)
}
// We have already buffered more than minSize,
// We are now in compression cruise mode until the end of times.
if r.compressionStarted {
// If compressionStarted we assume we have sent headers already
return r.bw.Write(p)
}
// If we detect a contentEncoding, we know we are never going to compress.
if r.rw.Header().Get(contentEncoding) != "" {
r.compressionDisabled = true
r.rw.WriteHeader(r.statusCode)
return r.rw.Write(p)
}
// Disable compression according to user wishes in excludedContentTypes or includedContentTypes.
if ct := r.rw.Header().Get(contentType); ct != "" {
mediaType, params, err := mime.ParseMediaType(ct)
if err != nil {
return 0, fmt.Errorf("parsing content-type media type: %w", err)
}
if len(r.includedContentTypes) > 0 {
var found bool
for _, includedContentType := range r.includedContentTypes {
if includedContentType.equals(mediaType, params) {
found = true
break
}
}
if !found {
r.compressionDisabled = true
return r.rw.Write(p)
}
}
for _, excludedContentType := range r.excludedContentTypes {
if excludedContentType.equals(mediaType, params) {
r.compressionDisabled = true
return r.rw.Write(p)
}
}
}
// We buffer until we know whether to compress (i.e. when we reach minSize received).
if len(r.buf)+len(p) < r.minSize {
r.buf = append(r.buf, p...)
return len(p), nil
}
// If we ever make it here, we have received at least minSize, which means we want to compress,
// and we are going to send headers right away.
r.compressionStarted = true
// Since we know we are going to compress we will never be able to know the actual length.
r.rw.Header().Del(contentLength)
r.rw.Header().Set(contentEncoding, "br")
r.rw.WriteHeader(r.statusCode)
r.headersSent = true
// Start with sending what we have previously buffered, before actually writing
// the bytes in argument.
n, err := r.bw.Write(r.buf)
if err != nil {
r.buf = r.buf[n:]
// Return zero because we haven't taken care of the bytes in argument yet.
return 0, err
}
// If we wrote less than what we wanted, we need to reclaim the leftovers + the bytes in argument,
// and keep them for a subsequent Write.
if n < len(r.buf) {
r.buf = r.buf[n:]
r.buf = append(r.buf, p...)
return len(p), nil
}
// Otherwise just reset the buffer.
r.buf = r.buf[:0]
// Now that we emptied the buffer, we can actually write the given bytes.
return r.bw.Write(p)
}
// Flush flushes data to the appropriate underlying writer(s), although it does
// not guarantee that all buffered data will be sent.
// If not enough bytes have been written to determine whether to enable compression,
// no flushing will take place.
func (r *responseWriter) Flush() {
if !r.seenData {
// we should not flush if there never was any data, because flushing the bw
// (just like closing) would send some extra end of compressionStarted stream bytes.
return
}
// It was already established by Write that compression is disabled, we only
// have to flush the uncompressed writer.
if r.compressionDisabled {
if rw, ok := r.rw.(http.Flusher); ok {
rw.Flush()
}
return
}
// Here, nothing was ever written either to rw or to bw (since we're still
// waiting to decide whether to compress), so we do not need to flush anything.
// Note that we diverge with klauspost's gzip behavior, where they instead
// force compression and flush whatever was in the buffer in this case.
if !r.compressionStarted {
return
}
// Conversely, we here know that something was already written to bw (or is
// going to be written right after anyway), so bw will have to be flushed.
// Also, since we know that bw writes to rw, but (apparently) never flushes it,
// we have to do it ourselves.
defer func() {
// because we also ignore the error returned by Write anyway
_ = r.bw.Flush()
if rw, ok := r.rw.(http.Flusher); ok {
rw.Flush()
}
}()
// We empty whatever is left of the buffer that Write never took care of.
n, err := r.bw.Write(r.buf)
if err != nil {
return
}
// And just like in Write we also handle "short writes".
if n < len(r.buf) {
r.buf = r.buf[n:]
return
}
r.buf = r.buf[:0]
}
func (r *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
if hijacker, ok := r.rw.(http.Hijacker); ok {
// We only make use of r.hijacked in close (and not in Write/WriteHeader)
// because we want to let the stdlib catch the error on writes, as
// they already do a good job of logging it.
r.hijacked = true
return hijacker.Hijack()
}
return nil, nil, fmt.Errorf("%T is not a http.Hijacker", r.rw)
}
// close closes the underlying writers if/when appropriate.
// Note that the compressed writer should not be closed if we never used it,
// as it would otherwise send some extra "end of compression" bytes.
// Close also makes sure to flush whatever was left to write from the buffer.
func (r *responseWriter) close() error {
if r.hijacked {
return nil
}
// We have to take care of statusCode ourselves (in case there was never any
// call to Write or WriteHeader before us) as it's the only header we buffer.
if !r.headersSent {
r.rw.WriteHeader(r.statusCode)
r.headersSent = true
}
// Nothing was ever written anywhere, nothing to flush.
if !r.seenData {
return nil
}
// If compression was disabled, there never was anything in the buffer to flush,
// and nothing was ever written to bw.
if r.compressionDisabled {
return nil
}
if len(r.buf) == 0 {
// If we got here we know compression has started, so we can safely flush on bw.
return r.bw.Close()
}
// There is still data in the buffer, because we never reached minSize (to
// determine whether to compress). We therefore flush it uncompressed.
if !r.compressionStarted {
n, err := r.rw.Write(r.buf)
if err != nil {
return err
}
if n < len(r.buf) {
return io.ErrShortWrite
}
return nil
}
// There is still data in the buffer, simply because Write did not take care of it all.
// We flush it to the compressed writer.
n, err := r.bw.Write(r.buf)
if err != nil {
r.bw.Close()
return err
}
if n < len(r.buf) {
r.bw.Close()
return io.ErrShortWrite
}
return r.bw.Close()
}
// parsedContentType is the parsed representation of one of the inputs to ContentTypes.
// From https://github.com/klauspost/compress/blob/master/gzhttp/compress.go#L401.
type parsedContentType struct {
mediaType string
params map[string]string
}
// equals returns whether this content type matches another content type.
func (p parsedContentType) equals(mediaType string, params map[string]string) bool {
if p.mediaType != mediaType {
return false
}
// if p has no params, don't care about other's params
if len(p.params) == 0 {
return true
}
// if p has any params, they must be identical to other's.
if len(p.params) != len(params) {
return false
}
for k, v := range p.params {
if w, ok := params[k]; !ok || v != w {
return false
}
}
return true
}