-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathbrotli.go
100 lines (81 loc) · 2.63 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
// Copyright 2021 Flamego. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
package brotli
import (
"bufio"
"net"
"net/http"
"strings"
"github.com/andybalholm/brotli"
"github.com/pkg/errors"
"github.com/flamego/flamego"
)
const (
headerAcceptEncoding = "Accept-Encoding"
headerContentEncoding = "Content-Encoding"
headerVary = "Vary"
)
// Options represents a struct for specifying configuration options for the
// brotli middleware.
type Options struct {
// CompressionLevel indicates the compression level, possible values are between
// 0 to 11 (from best speed to best compression). Default is 5.
CompressionLevel int
}
func isCompressionLevelValid(level int) bool {
return level == brotli.DefaultCompression ||
(level >= brotli.BestSpeed && level <= brotli.BestCompression)
}
func prepareOptions(options []Options) Options {
var opt Options
if len(options) > 0 {
opt = options[0]
}
if !isCompressionLevelValid(opt.CompressionLevel) {
// For web content, level 5 seems to be a sweet spot.
opt.CompressionLevel = 5
}
return opt
}
// Brotli returns a Handler that adds brotli compression to all requests.
// Make sure to include the brotli middleware above other middleware
// that alter the response body (like the render middleware).
func Brotli(options ...Options) flamego.Handler {
opt := prepareOptions(options)
return flamego.ContextInvoker(func(ctx flamego.Context) {
if !strings.Contains(ctx.Request().Header.Get(headerAcceptEncoding), "br") {
return
}
headers := ctx.ResponseWriter().Header()
headers.Set(headerContentEncoding, "br")
headers.Set(headerVary, headerAcceptEncoding)
// We've made sure compression level is valid in prepareOptions, no need to
// check same error again.
br := brotli.NewWriterLevel(ctx.ResponseWriter(), opt.CompressionLevel)
defer func() { _ = br.Close() }()
w := &responseWriter{
writer: br,
ResponseWriter: ctx.ResponseWriter(),
}
ctx.MapTo(w, (*http.ResponseWriter)(nil))
ctx.Next()
// Delete content length after we know we have been written to
ctx.ResponseWriter().Header().Del("Content-Length")
})
}
type responseWriter struct {
writer *brotli.Writer
flamego.ResponseWriter
}
func (w *responseWriter) Write(p []byte) (int, error) {
return w.writer.Write(p)
}
var _ http.Hijacker = (*responseWriter)(nil)
func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hijacker, ok := w.ResponseWriter.(http.Hijacker)
if !ok {
return nil, nil, errors.New("the ResponseWriter doesn't support the Hijacker interface")
}
return hijacker.Hijack()
}