/
gzip.go
126 lines (109 loc) · 3.38 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
package rest
import (
"bufio"
"compress/gzip"
"net"
"net/http"
"strings"
)
// GzipMiddleware is responsible for compressing the payload with gzip and setting the proper
// headers when supported by the client. It must be wrapped by TimerMiddleware for the
// compression time to be captured. And It must be wrapped by RecorderMiddleware for the
// compressed BYTES_WRITTEN to be captured.
type GzipMiddleware struct{}
// MiddlewareFunc makes GzipMiddleware implement the Middleware interface.
func (mw *GzipMiddleware) MiddlewareFunc(h HandlerFunc) HandlerFunc {
return func(w ResponseWriter, r *Request) {
// gzip support enabled
canGzip := strings.Contains(r.Header.Get("Accept-Encoding"), "gzip")
// client accepts gzip ?
writer := &gzipResponseWriter{w, false, canGzip, nil}
// call the handler with the wrapped writer
h(writer, r)
}
}
// Private responseWriter intantiated by the gzip middleware.
// It encodes the payload with gzip and set the proper headers.
// It implements the following interfaces:
// ResponseWriter
// http.ResponseWriter
// http.Flusher
// http.CloseNotifier
// http.Hijacker
type gzipResponseWriter struct {
ResponseWriter
wroteHeader bool
canGzip bool
gzipWriter *gzip.Writer
}
// Set the right headers for gzip encoded responses.
func (w *gzipResponseWriter) WriteHeader(code int) {
// Always set the Vary header, even if this particular request
// is not gzipped.
w.Header().Add("Vary", "Accept-Encoding")
if w.canGzip {
w.Header().Set("Content-Encoding", "gzip")
}
w.ResponseWriter.WriteHeader(code)
w.wroteHeader = true
}
// Make sure the local Write is called.
func (w *gzipResponseWriter) WriteJson(v interface{}) error {
b, err := w.EncodeJson(v)
if err != nil {
return err
}
_, err = w.Write(b)
if err != nil {
return err
}
return nil
}
// Make sure the local WriteHeader is called, and call the parent Flush.
// Provided in order to implement the http.Flusher interface.
func (w *gzipResponseWriter) Flush() {
if !w.wroteHeader {
w.WriteHeader(http.StatusOK)
}
flusher := w.ResponseWriter.(http.Flusher)
flusher.Flush()
}
// Call the parent CloseNotify.
// Provided in order to implement the http.CloseNotifier interface.
func (w *gzipResponseWriter) CloseNotify() <-chan bool {
notifier := w.ResponseWriter.(http.CloseNotifier)
return notifier.CloseNotify()
}
// Provided in order to implement the http.Hijacker interface.
func (w *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hijacker := w.ResponseWriter.(http.Hijacker)
return hijacker.Hijack()
}
// Make sure the local WriteHeader is called, and encode the payload if necessary.
// Provided in order to implement the http.ResponseWriter interface.
func (w *gzipResponseWriter) Write(b []byte) (int, error) {
if !w.wroteHeader {
w.WriteHeader(http.StatusOK)
}
writer := w.ResponseWriter.(http.ResponseWriter)
if w.canGzip {
// Write can be called multiple times for a given response.
// (see the streaming example:
// https://github.com/ant0ine/go-json-rest-examples/tree/master/streaming)
// The gzipWriter is instantiated only once, and flushed after
// each write.
if w.gzipWriter == nil {
w.gzipWriter = gzip.NewWriter(writer)
}
count, errW := w.gzipWriter.Write(b)
errF := w.gzipWriter.Flush()
if errW != nil {
return count, errW
}
if errF != nil {
return count, errF
}
return count, nil
}
return writer.Write(b)
}