package middleware // The original work was derived from Goji's middleware, source: // https://github.com/zenazn/goji/tree/master/web/middleware import ( "bufio" "io" "net" "net/http" ) // WrapResponseWriter is a proxy around an http.ResponseWriter that allows you to hook // into various parts of the response process. type WrapResponseWriter interface { http.ResponseWriter // Status returns the HTTP status of the request, or 0 if one has not // yet been sent. Status() int // BytesWritten returns the total number of bytes sent to the client. BytesWritten() int // Tee causes the response body to be written to the given io.Writer in // addition to proxying the writes through. Only one io.Writer can be // tee'd to at once: setting a second one will overwrite the first. // Writes will be sent to the proxy before being written to this // io.Writer. It is illegal for the tee'd writer to be modified // concurrently with writes. Tee(io.Writer) // Unwrap returns the original proxied target. Unwrap() http.ResponseWriter } // basicWriter wraps a http.ResponseWriter that implements the minimal // http.ResponseWriter interface. type basicWriter struct { http.ResponseWriter wroteHeader bool code int bytes int tee io.Writer } func (b *basicWriter) WriteHeader(code int) { if !b.wroteHeader { b.code = code b.wroteHeader = true b.ResponseWriter.WriteHeader(code) } } func (b *basicWriter) Write(buf []byte) (int, error) { b.WriteHeader(http.StatusOK) n, err := b.ResponseWriter.Write(buf) if b.tee != nil { _, err2 := b.tee.Write(buf[:n]) // Prefer errors generated by the proxied writer. if err == nil { err = err2 } } b.bytes += n return n, err } func (b *basicWriter) maybeWriteHeader() { if !b.wroteHeader { b.WriteHeader(http.StatusOK) } } func (b *basicWriter) Status() int { return b.code } func (b *basicWriter) BytesWritten() int { return b.bytes } func (b *basicWriter) Tee(w io.Writer) { b.tee = w } func (b *basicWriter) Unwrap() http.ResponseWriter { return b.ResponseWriter } type flushWriter struct { basicWriter } func (f *flushWriter) Flush() { fl := f.basicWriter.ResponseWriter.(http.Flusher) fl.Flush() } var _ http.Flusher = &flushWriter{} // httpFancyWriter is a HTTP writer that additionally satisfies http.CloseNotifier, // http.Flusher, http.Hijacker, and io.ReaderFrom. It exists for the common case // of wrapping the http.ResponseWriter that package http gives you, in order to // make the proxied object support the full method set of the proxied object. type httpFancyWriter struct { basicWriter } func (f *httpFancyWriter) CloseNotify() <-chan bool { cn := f.basicWriter.ResponseWriter.(http.CloseNotifier) return cn.CloseNotify() } func (f *httpFancyWriter) Flush() { fl := f.basicWriter.ResponseWriter.(http.Flusher) fl.Flush() } func (f *httpFancyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { hj := f.basicWriter.ResponseWriter.(http.Hijacker) return hj.Hijack() } func (f *httpFancyWriter) ReadFrom(r io.Reader) (int64, error) { if f.basicWriter.tee != nil { return io.Copy(&f.basicWriter, r) } rf := f.basicWriter.ResponseWriter.(io.ReaderFrom) f.basicWriter.maybeWriteHeader() return rf.ReadFrom(r) } var _ http.CloseNotifier = &httpFancyWriter{} var _ http.Flusher = &httpFancyWriter{} var _ http.Hijacker = &httpFancyWriter{} var _ io.ReaderFrom = &httpFancyWriter{} // http2FancyWriter is a HTTP2 writer that additionally satisfies http.CloseNotifier, // http.Flusher, and io.ReaderFrom. It exists for the common case // of wrapping the http.ResponseWriter that package http gives you, in order to // make the proxied object support the full method set of the proxied object. type http2FancyWriter struct { basicWriter } func (f *http2FancyWriter) CloseNotify() <-chan bool { cn := f.basicWriter.ResponseWriter.(http.CloseNotifier) return cn.CloseNotify() } func (f *http2FancyWriter) Flush() { fl := f.basicWriter.ResponseWriter.(http.Flusher) fl.Flush() } var _ http.CloseNotifier = &http2FancyWriter{} var _ http.Flusher = &http2FancyWriter{}