diff --git a/graceful/middleware.go b/graceful/middleware.go index fedcf37..7f4683f 100644 --- a/graceful/middleware.go +++ b/graceful/middleware.go @@ -2,6 +2,7 @@ package graceful import ( "bufio" + "io" "net" "net/http" ) @@ -36,10 +37,11 @@ func Middleware(h http.Handler) http.Handler { _, cn := w.(http.CloseNotifier) _, fl := w.(http.Flusher) _, hj := w.(http.Hijacker) + _, rf := w.(io.ReaderFrom) bw := basicWriter{ResponseWriter: w} - if cn && fl && hj { + if cn && fl && hj && rf { h.ServeHTTP(&fancyWriter{bw}, r) } else { h.ServeHTTP(&bw, r) @@ -104,3 +106,15 @@ func (f *fancyWriter) Hijack() (c net.Conn, b *bufio.ReadWriter, e error) { return } +func (f *fancyWriter) ReadFrom(r io.Reader) (int64, error) { + rf := f.basicWriter.ResponseWriter.(io.ReaderFrom) + if !f.basicWriter.headerWritten { + f.basicWriter.maybeClose() + } + return rf.ReadFrom(r) +} + +var _ http.CloseNotifier = &fancyWriter{} +var _ http.Flusher = &fancyWriter{} +var _ http.Hijacker = &fancyWriter{} +var _ io.ReaderFrom = &fancyWriter{} diff --git a/web/middleware/logger.go b/web/middleware/logger.go index b8d590a..4fbb0e9 100644 --- a/web/middleware/logger.go +++ b/web/middleware/logger.go @@ -78,16 +78,3 @@ func printEnd(reqId string, w writerProxy, dt time.Duration) { log.Print(buf.String()) } - -func wrapWriter(w http.ResponseWriter) writerProxy { - _, cn := w.(http.CloseNotifier) - _, fl := w.(http.Flusher) - _, hj := w.(http.Hijacker) - - bw := basicWriter{ResponseWriter: w} - if cn && fl && hj { - return &fancyWriter{bw} - } else { - return &bw - } -} diff --git a/web/middleware/writer_proxy.go b/web/middleware/writer_proxy.go index 634bc4e..f9cfe94 100644 --- a/web/middleware/writer_proxy.go +++ b/web/middleware/writer_proxy.go @@ -2,10 +2,25 @@ package middleware import ( "bufio" + "io" "net" "net/http" ) +func wrapWriter(w http.ResponseWriter) writerProxy { + _, cn := w.(http.CloseNotifier) + _, fl := w.(http.Flusher) + _, hj := w.(http.Hijacker) + _, rf := w.(io.ReaderFrom) + + bw := basicWriter{ResponseWriter: w} + if cn && fl && hj && rf { + return &fancyWriter{bw} + } else { + return &bw + } +} + type writerProxy interface { http.ResponseWriter maybeWriteHeader() @@ -55,7 +70,13 @@ func (f *fancyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { hj := f.basicWriter.ResponseWriter.(http.Hijacker) return hj.Hijack() } +func (f *fancyWriter) ReadFrom(r io.Reader) (int64, error) { + rf := f.basicWriter.ResponseWriter.(io.ReaderFrom) + f.basicWriter.maybeWriteHeader() + return rf.ReadFrom(r) +} var _ http.CloseNotifier = &fancyWriter{} var _ http.Flusher = &fancyWriter{} var _ http.Hijacker = &fancyWriter{} +var _ io.ReaderFrom = &fancyWriter{}