From 4d9fdbbbcff1941c2981c4e750f036ac3e958016 Mon Sep 17 00:00:00 2001 From: Carl Jackson Date: Sat, 23 Aug 2014 17:04:35 -0700 Subject: [PATCH] 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 | 29 +++++---- web/middleware/writer_proxy.go | 83 ------------------------ web/util/writer_proxy.go | 115 +++++++++++++++++++++++++++++++++ 3 files changed, 132 insertions(+), 95 deletions(-) delete mode 100644 web/middleware/writer_proxy.go create mode 100644 web/util/writer_proxy.go diff --git a/web/middleware/logger.go b/web/middleware/logger.go index ca4773b..fce822e 100644 --- a/web/middleware/logger.go +++ b/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 { diff --git a/web/middleware/writer_proxy.go b/web/middleware/writer_proxy.go deleted file mode 100644 index 0142403..0000000 --- a/web/middleware/writer_proxy.go +++ /dev/null @@ -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{} diff --git a/web/util/writer_proxy.go b/web/util/writer_proxy.go new file mode 100644 index 0000000..99e7ce0 --- /dev/null +++ b/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{}