Refactor WriterProxy out of web/middleware

It turns out WriterProxy is pretty generally useful, especially when
defining custom http loggers. Expose it in a util package so that other
packages can use it.

web/middleware/logger.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/zenazn/goji/web"
+	"github.com/zenazn/goji/web/util"
 )
 
 // Logger is a middleware that logs the start and end of each request, along
@@ -21,11 +22,14 @@ func Logger(c *web.C, h http.Handler) http.Handler {
 
 		printStart(reqID, r)
 
-		lw := wrapWriter(w)
+		lw := util.WrapWriter(w)
 
 		t1 := time.Now()
 		h.ServeHTTP(lw, r)
-		lw.maybeWriteHeader()
+
+		if lw.Status() == 0 {
+			lw.WriteHeader(http.StatusOK)
+		}
 		t2 := time.Now()
 
 		printEnd(reqID, lw, t2.Sub(t1))
@@ -49,23 +53,24 @@ func printStart(reqID string, r *http.Request) {
 	log.Print(buf.String())
 }
 
-func printEnd(reqID string, w writerProxy, dt time.Duration) {
+func printEnd(reqID string, w util.WriterProxy, dt time.Duration) {
 	var buf bytes.Buffer
 
 	if reqID != "" {
 		cW(&buf, bBlack, "[%s] ", reqID)
 	}
 	buf.WriteString("Returning ")
-	if w.status() < 200 {
-		cW(&buf, bBlue, "%03d", w.status())
-	} else if w.status() < 300 {
-		cW(&buf, bGreen, "%03d", w.status())
-	} else if w.status() < 400 {
-		cW(&buf, bCyan, "%03d", w.status())
-	} else if w.status() < 500 {
-		cW(&buf, bYellow, "%03d", w.status())
+	status := w.Status()
+	if status < 200 {
+		cW(&buf, bBlue, "%03d", status)
+	} else if status < 300 {
+		cW(&buf, bGreen, "%03d", status)
+	} else if status < 400 {
+		cW(&buf, bCyan, "%03d", status)
+	} else if status < 500 {
+		cW(&buf, bYellow, "%03d", status)
 	} else {
-		cW(&buf, bRed, "%03d", w.status())
+		cW(&buf, bRed, "%03d", status)
 	}
 	buf.WriteString(" in ")
 	if dt < 500*time.Millisecond {

web/middleware/writer_proxy.go

@@ -1,83 +0,0 @@
-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}
-	}
-	return &bw
-}
-
-type writerProxy interface {
-	http.ResponseWriter
-	maybeWriteHeader()
-	status() int
-}
-
-type basicWriter struct {
-	http.ResponseWriter
-	wroteHeader bool
-	code        int
-}
-
-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()
-	return b.ResponseWriter.Write(buf)
-}
-func (b *basicWriter) maybeWriteHeader() {
-	if !b.wroteHeader {
-		b.WriteHeader(http.StatusOK)
-	}
-}
-func (b *basicWriter) status() int {
-	return b.code
-}
-func (b *basicWriter) Unwrap() http.ResponseWriter {
-	return b.ResponseWriter
-}
-
-type fancyWriter struct {
-	basicWriter
-}
-
-func (f *fancyWriter) CloseNotify() <-chan bool {
-	cn := f.basicWriter.ResponseWriter.(http.CloseNotifier)
-	return cn.CloseNotify()
-}
-func (f *fancyWriter) Flush() {
-	fl := f.basicWriter.ResponseWriter.(http.Flusher)
-	fl.Flush()
-}
-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{}

web/util/writer_proxy.go

@@ -0,0 +1,115 @@
+package util
+
+import (
+	"bufio"
+	"io"
+	"net"
+	"net/http"
+)
+
+// WriterProxy is a proxy around an http.ResponseWriter that allows you to hook
+// into various parts of the response process.
+type WriterProxy interface {
+	http.ResponseWriter
+	// Status returns the HTTP status of the request, or 0 if one has not
+	// yet been sent.
+	Status() 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
+}
+
+// WrapWriter wraps an http.ResponseWriter into a proxy that allows you to hook
+// into various parts of the response process.
+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}
+	}
+	return &bw
+}
+
+// basicWriter wraps a http.ResponseWriter that implements the minimal
+// http.ResponseWriter interface.
+type basicWriter struct {
+	http.ResponseWriter
+	wroteHeader bool
+	code        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
+		}
+	}
+	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) Tee(w io.Writer) {
+	b.tee = w
+}
+func (b *basicWriter) Unwrap() http.ResponseWriter {
+	return b.ResponseWriter
+}
+
+// fancyWriter is a 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 fancyWriter struct {
+	basicWriter
+}
+
+func (f *fancyWriter) CloseNotify() <-chan bool {
+	cn := f.basicWriter.ResponseWriter.(http.CloseNotifier)
+	return cn.CloseNotify()
+}
+func (f *fancyWriter) Flush() {
+	fl := f.basicWriter.ResponseWriter.(http.Flusher)
+	fl.Flush()
+}
+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{}