You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
181 lines
4.5 KiB
Go
181 lines
4.5 KiB
Go
4 years ago
|
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"
|
||
|
)
|
||
|
|
||
|
// NewWrapResponseWriter wraps an http.ResponseWriter, returning a proxy that allows you to
|
||
|
// hook into various parts of the response process.
|
||
|
func NewWrapResponseWriter(w http.ResponseWriter, protoMajor int) WrapResponseWriter {
|
||
|
_, fl := w.(http.Flusher)
|
||
|
|
||
|
bw := basicWriter{ResponseWriter: w}
|
||
|
|
||
|
if protoMajor == 2 {
|
||
|
_, ps := w.(http.Pusher)
|
||
|
if fl && ps {
|
||
|
return &http2FancyWriter{bw}
|
||
|
}
|
||
|
} else {
|
||
|
_, hj := w.(http.Hijacker)
|
||
|
_, rf := w.(io.ReaderFrom)
|
||
|
if fl && hj && rf {
|
||
|
return &httpFancyWriter{bw}
|
||
|
}
|
||
|
}
|
||
|
if fl {
|
||
|
return &flushWriter{bw}
|
||
|
}
|
||
|
|
||
|
return &bw
|
||
|
}
|
||
|
|
||
|
// 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.maybeWriteHeader()
|
||
|
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() {
|
||
|
f.wroteHeader = true
|
||
|
fl := f.basicWriter.ResponseWriter.(http.Flusher)
|
||
|
fl.Flush()
|
||
|
}
|
||
|
|
||
|
var _ http.Flusher = &flushWriter{}
|
||
|
|
||
|
// httpFancyWriter is a HTTP writer that additionally satisfies
|
||
|
// 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) Flush() {
|
||
|
f.wroteHeader = true
|
||
|
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 *http2FancyWriter) Push(target string, opts *http.PushOptions) error {
|
||
|
return f.basicWriter.ResponseWriter.(http.Pusher).Push(target, opts)
|
||
|
}
|
||
|
|
||
|
func (f *httpFancyWriter) ReadFrom(r io.Reader) (int64, error) {
|
||
|
if f.basicWriter.tee != nil {
|
||
|
n, err := io.Copy(&f.basicWriter, r)
|
||
|
f.basicWriter.bytes += int(n)
|
||
|
return n, err
|
||
|
}
|
||
|
rf := f.basicWriter.ResponseWriter.(io.ReaderFrom)
|
||
|
f.basicWriter.maybeWriteHeader()
|
||
|
n, err := rf.ReadFrom(r)
|
||
|
f.basicWriter.bytes += int(n)
|
||
|
return n, err
|
||
|
}
|
||
|
|
||
|
var _ http.Flusher = &httpFancyWriter{}
|
||
|
var _ http.Hijacker = &httpFancyWriter{}
|
||
|
var _ http.Pusher = &http2FancyWriter{}
|
||
|
var _ io.ReaderFrom = &httpFancyWriter{}
|
||
|
|
||
|
// http2FancyWriter is a HTTP2 writer that additionally satisfies
|
||
|
// 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) Flush() {
|
||
|
f.wroteHeader = true
|
||
|
fl := f.basicWriter.ResponseWriter.(http.Flusher)
|
||
|
fl.Flush()
|
||
|
}
|
||
|
|
||
|
var _ http.Flusher = &http2FancyWriter{}
|