-
Notifications
You must be signed in to change notification settings - Fork 18
/
gzip.go
147 lines (124 loc) · 4.17 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
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
// Package gzip implements a gzip compression handler middleware for Negroni.
package gzip
import (
"compress/gzip"
"io/ioutil"
"net/http"
"strings"
"sync"
"github.com/urfave/negroni"
)
// These compression constants are copied from the compress/gzip package.
const (
encodingGzip = "gzip"
headerAcceptEncoding = "Accept-Encoding"
headerContentEncoding = "Content-Encoding"
headerContentLength = "Content-Length"
headerContentType = "Content-Type"
headerVary = "Vary"
headerSecWebSocketKey = "Sec-WebSocket-Key"
BestCompression = gzip.BestCompression
BestSpeed = gzip.BestSpeed
DefaultCompression = gzip.DefaultCompression
NoCompression = gzip.NoCompression
)
// gzipResponseWriter is the ResponseWriter that negroni.ResponseWriter is
// wrapped in.
type gzipResponseWriter struct {
w *gzip.Writer
negroni.ResponseWriter
wroteHeader bool
}
// Check whether underlying response is already pre-encoded and disable
// gzipWriter before the body gets written, otherwise encoding headers
func (grw *gzipResponseWriter) WriteHeader(code int) {
headers := grw.ResponseWriter.Header()
if headers.Get(headerContentEncoding) == "" {
headers.Set(headerContentEncoding, encodingGzip)
headers.Add(headerVary, headerAcceptEncoding)
} else {
grw.w.Reset(ioutil.Discard)
grw.w = nil
}
// Avoid sending Content-Length header before compression. The length would
// be invalid, and some browsers like Safari will report
// "The network connection was lost." errors
grw.Header().Del(headerContentLength)
grw.ResponseWriter.WriteHeader(code)
grw.wroteHeader = true
}
// Write writes bytes to the gzip.Writer. It will also set the Content-Type
// header using the net/http library content type detection if the Content-Type
// header was not set yet.
func (grw *gzipResponseWriter) Write(b []byte) (int, error) {
if !grw.wroteHeader {
grw.WriteHeader(http.StatusOK)
}
if grw.w == nil {
return grw.ResponseWriter.Write(b)
}
if len(grw.Header().Get(headerContentType)) == 0 {
grw.Header().Set(headerContentType, http.DetectContentType(b))
}
return grw.w.Write(b)
}
type gzipResponseWriterCloseNotifier struct {
*gzipResponseWriter
}
func (rw *gzipResponseWriterCloseNotifier) CloseNotify() <-chan bool {
return rw.ResponseWriter.(http.CloseNotifier).CloseNotify()
}
func newGzipResponseWriter(rw negroni.ResponseWriter, w *gzip.Writer) negroni.ResponseWriter {
wr := &gzipResponseWriter{w: w, ResponseWriter: rw}
if _, ok := rw.(http.CloseNotifier); ok {
return &gzipResponseWriterCloseNotifier{gzipResponseWriter: wr}
}
return wr
}
// handler struct contains the ServeHTTP method
type handler struct {
pool sync.Pool
}
// Gzip returns a handler which will handle the Gzip compression in ServeHTTP.
// Valid values for level are identical to those in the compress/gzip package.
func Gzip(level int) *handler {
h := &handler{}
h.pool.New = func() interface{} {
gz, err := gzip.NewWriterLevel(ioutil.Discard, level)
if err != nil {
panic(err)
}
return gz
}
return h
}
// ServeHTTP wraps the http.ResponseWriter with a gzip.Writer.
func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
// Skip compression if the client doesn't accept gzip encoding.
if !strings.Contains(r.Header.Get(headerAcceptEncoding), encodingGzip) {
next(w, r)
return
}
// Skip compression if client attempt WebSocket connection
if len(r.Header.Get(headerSecWebSocketKey)) > 0 {
next(w, r)
return
}
// Retrieve gzip writer from the pool. Reset it to use the ResponseWriter.
// This allows us to re-use an already allocated buffer rather than
// allocating a new buffer for every request.
// We defer g.pool.Put here so that the gz writer is returned to the
// pool if any thing after here fails for some reason (functions in
// next could potentially panic, etc)
gz := h.pool.Get().(*gzip.Writer)
defer h.pool.Put(gz)
gz.Reset(w)
// Wrap the original http.ResponseWriter with negroni.ResponseWriter
// and create the gzipResponseWriter.
nrw := negroni.NewResponseWriter(w)
grw := newGzipResponseWriter(nrw, gz)
// Call the next handler supplying the gzipResponseWriter instead of
// the original.
next(grw, r)
gz.Close()
}