wrap_writer.go
1package middleware
2
3// The original work was derived from Goji's middleware, source:
4// https://github.com/zenazn/goji/tree/master/web/middleware
5
6import (
7 "bufio"
8 "io"
9 "net"
10 "net/http"
11)
12
13// NewWrapResponseWriter wraps an http.ResponseWriter, returning a proxy that allows you to
14// hook into various parts of the response process.
15func NewWrapResponseWriter(w http.ResponseWriter, protoMajor int) WrapResponseWriter {
16 _, fl := w.(http.Flusher)
17
18 bw := basicWriter{ResponseWriter: w}
19
20 if protoMajor == 2 {
21 _, ps := w.(http.Pusher)
22 if fl && ps {
23 return &http2FancyWriter{bw}
24 }
25 } else {
26 _, hj := w.(http.Hijacker)
27 _, rf := w.(io.ReaderFrom)
28 if fl && hj && rf {
29 return &httpFancyWriter{bw}
30 }
31 if fl && hj {
32 return &flushHijackWriter{bw}
33 }
34 if hj {
35 return &hijackWriter{bw}
36 }
37 }
38
39 if fl {
40 return &flushWriter{bw}
41 }
42
43 return &bw
44}
45
46// WrapResponseWriter is a proxy around an http.ResponseWriter that allows you to hook
47// into various parts of the response process.
48type WrapResponseWriter interface {
49 http.ResponseWriter
50 // Status returns the HTTP status of the request, or 0 if one has not
51 // yet been sent.
52 Status() int
53 // BytesWritten returns the total number of bytes sent to the client.
54 BytesWritten() int
55 // Tee causes the response body to be written to the given io.Writer in
56 // addition to proxying the writes through. Only one io.Writer can be
57 // tee'd to at once: setting a second one will overwrite the first.
58 // Writes will be sent to the proxy before being written to this
59 // io.Writer. It is illegal for the tee'd writer to be modified
60 // concurrently with writes.
61 Tee(io.Writer)
62 // Unwrap returns the original proxied target.
63 Unwrap() http.ResponseWriter
64 // Discard causes all writes to the original ResponseWriter be discarded,
65 // instead writing only to the tee'd writer if it's set.
66 // The caller is responsible for calling WriteHeader and Write on the
67 // original ResponseWriter once the processing is done.
68 Discard()
69}
70
71// basicWriter wraps a http.ResponseWriter that implements the minimal
72// http.ResponseWriter interface.
73type basicWriter struct {
74 http.ResponseWriter
75 tee io.Writer
76 code int
77 bytes int
78 wroteHeader bool
79 discard bool
80}
81
82func (b *basicWriter) WriteHeader(code int) {
83 if code >= 100 && code <= 199 && code != http.StatusSwitchingProtocols {
84 if !b.discard {
85 b.ResponseWriter.WriteHeader(code)
86 }
87 } else if !b.wroteHeader {
88 b.code = code
89 b.wroteHeader = true
90 if !b.discard {
91 b.ResponseWriter.WriteHeader(code)
92 }
93 }
94}
95
96func (b *basicWriter) Write(buf []byte) (n int, err error) {
97 b.maybeWriteHeader()
98 if !b.discard {
99 n, err = b.ResponseWriter.Write(buf)
100 if b.tee != nil {
101 _, err2 := b.tee.Write(buf[:n])
102 // Prefer errors generated by the proxied writer.
103 if err == nil {
104 err = err2
105 }
106 }
107 } else if b.tee != nil {
108 n, err = b.tee.Write(buf)
109 } else {
110 n, err = io.Discard.Write(buf)
111 }
112 b.bytes += n
113 return n, err
114}
115
116func (b *basicWriter) maybeWriteHeader() {
117 if !b.wroteHeader {
118 b.WriteHeader(http.StatusOK)
119 }
120}
121
122func (b *basicWriter) Status() int {
123 return b.code
124}
125
126func (b *basicWriter) BytesWritten() int {
127 return b.bytes
128}
129
130func (b *basicWriter) Tee(w io.Writer) {
131 b.tee = w
132}
133
134func (b *basicWriter) Unwrap() http.ResponseWriter {
135 return b.ResponseWriter
136}
137
138func (b *basicWriter) Discard() {
139 b.discard = true
140}
141
142// flushWriter ...
143type flushWriter struct {
144 basicWriter
145}
146
147func (f *flushWriter) Flush() {
148 f.wroteHeader = true
149 fl := f.basicWriter.ResponseWriter.(http.Flusher)
150 fl.Flush()
151}
152
153var _ http.Flusher = &flushWriter{}
154
155// hijackWriter ...
156type hijackWriter struct {
157 basicWriter
158}
159
160func (f *hijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
161 hj := f.basicWriter.ResponseWriter.(http.Hijacker)
162 return hj.Hijack()
163}
164
165var _ http.Hijacker = &hijackWriter{}
166
167// flushHijackWriter ...
168type flushHijackWriter struct {
169 basicWriter
170}
171
172func (f *flushHijackWriter) Flush() {
173 f.wroteHeader = true
174 fl := f.basicWriter.ResponseWriter.(http.Flusher)
175 fl.Flush()
176}
177
178func (f *flushHijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
179 hj := f.basicWriter.ResponseWriter.(http.Hijacker)
180 return hj.Hijack()
181}
182
183var _ http.Flusher = &flushHijackWriter{}
184var _ http.Hijacker = &flushHijackWriter{}
185
186// httpFancyWriter is a HTTP writer that additionally satisfies
187// http.Flusher, http.Hijacker, and io.ReaderFrom. It exists for the common case
188// of wrapping the http.ResponseWriter that package http gives you, in order to
189// make the proxied object support the full method set of the proxied object.
190type httpFancyWriter struct {
191 basicWriter
192}
193
194func (f *httpFancyWriter) Flush() {
195 f.wroteHeader = true
196 fl := f.basicWriter.ResponseWriter.(http.Flusher)
197 fl.Flush()
198}
199
200func (f *httpFancyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
201 hj := f.basicWriter.ResponseWriter.(http.Hijacker)
202 return hj.Hijack()
203}
204
205func (f *http2FancyWriter) Push(target string, opts *http.PushOptions) error {
206 return f.basicWriter.ResponseWriter.(http.Pusher).Push(target, opts)
207}
208
209func (f *httpFancyWriter) ReadFrom(r io.Reader) (int64, error) {
210 if f.basicWriter.tee != nil {
211 n, err := io.Copy(&f.basicWriter, r)
212 f.basicWriter.bytes += int(n)
213 return n, err
214 }
215 rf := f.basicWriter.ResponseWriter.(io.ReaderFrom)
216 f.basicWriter.maybeWriteHeader()
217 n, err := rf.ReadFrom(r)
218 f.basicWriter.bytes += int(n)
219 return n, err
220}
221
222var _ http.Flusher = &httpFancyWriter{}
223var _ http.Hijacker = &httpFancyWriter{}
224var _ http.Pusher = &http2FancyWriter{}
225var _ io.ReaderFrom = &httpFancyWriter{}
226
227// http2FancyWriter is a HTTP2 writer that additionally satisfies
228// http.Flusher, and io.ReaderFrom. It exists for the common case
229// of wrapping the http.ResponseWriter that package http gives you, in order to
230// make the proxied object support the full method set of the proxied object.
231type http2FancyWriter struct {
232 basicWriter
233}
234
235func (f *http2FancyWriter) Flush() {
236 f.wroteHeader = true
237 fl := f.basicWriter.ResponseWriter.(http.Flusher)
238 fl.Flush()
239}
240
241var _ http.Flusher = &http2FancyWriter{}
