Use httpsnoop to wrap ResponseWriter. (#193)

Wrapping http.ResponseWriter is fraught with danger. Our compress
handler made sure to implement all the optional ResponseWriter
interfaces, but that made it implement them even if the underlying
writer did not. For example, if the underlying ResponseWriter was
_not_ an http.Hijacker, the compress writer nonetheless appeared to
implement http.Hijacker, but would panic if you called Hijack().

On the other hand, the logging handler checked for certain
combinations of optional interfaces and only implemented them as
appropriate. However, it didn't check for all optional interfaces or
all combinations, so most optional interfaces would still get lost.

Fix both problems by using httpsnoop to do the wrapping. It uses code
generation to ensure correctness, and it handles std lib changes like
the http.Pusher addition in Go 1.8.

Fixes #169.

Author: Muir Manders

compress.go

@@ -10,35 +10,30 @@ import (
 	"io"
 	"net/http"
 	"strings"
+
+	"github.com/felixge/httpsnoop"
 )
 
 const acceptEncoding string = "Accept-Encoding"
 
 type compressResponseWriter struct {
-	io.Writer
-	http.ResponseWriter
-	http.Hijacker
-	http.Flusher
-	http.CloseNotifier
-}
-
-func (w *compressResponseWriter) WriteHeader(c int) {
-	w.ResponseWriter.Header().Del("Content-Length")
-	w.ResponseWriter.WriteHeader(c)
+	compressor io.Writer
+	w          http.ResponseWriter
 }
 
-func (w *compressResponseWriter) Header() http.Header {
-	return w.ResponseWriter.Header()
+func (cw *compressResponseWriter) WriteHeader(c int) {
+	cw.w.Header().Del("Content-Length")
+	cw.w.WriteHeader(c)
 }
 
-func (w *compressResponseWriter) Write(b []byte) (int, error) {
-	h := w.ResponseWriter.Header()
+func (cw *compressResponseWriter) Write(b []byte) (int, error) {
+	h := cw.w.Header()
 	if h.Get("Content-Type") == "" {
 		h.Set("Content-Type", http.DetectContentType(b))
 	}
 	h.Del("Content-Length")
 
-	return w.Writer.Write(b)
+	return cw.compressor.Write(b)
 }
 
 type flusher interface {
@@ -47,12 +42,12 @@ type flusher interface {
 
 func (w *compressResponseWriter) Flush() {
 	// Flush compressed data if compressor supports it.
-	if f, ok := w.Writer.(flusher); ok {
+	if f, ok := w.compressor.(flusher); ok {
 		f.Flush()
 	}
 	// Flush HTTP response.
-	if w.Flusher != nil {
-		w.Flusher.Flush()
+	if f, ok := w.w.(http.Flusher); ok {
+		f.Flush()
 	}
 }
 
@@ -119,28 +114,22 @@ func CompressHandlerLevel(h http.Handler, level int) http.Handler {
 		w.Header().Set("Content-Encoding", encoding)
 		r.Header.Del(acceptEncoding)
 
-		hijacker, ok := w.(http.Hijacker)
-		if !ok { /* w is not Hijacker... oh well... */
-			hijacker = nil
+		cw := &compressResponseWriter{
+			w:          w,
+			compressor: encWriter,
 		}
 
-		flusher, ok := w.(http.Flusher)
-		if !ok {
-			flusher = nil
-		}
-
-		closeNotifier, ok := w.(http.CloseNotifier)
-		if !ok {
-			closeNotifier = nil
-		}
-
-		w = &compressResponseWriter{
-			Writer:         encWriter,
-			ResponseWriter: w,
-			Hijacker:       hijacker,
-			Flusher:        flusher,
-			CloseNotifier:  closeNotifier,
-		}
+		w = httpsnoop.Wrap(w, httpsnoop.Hooks{
+			Write: func(httpsnoop.WriteFunc) httpsnoop.WriteFunc {
+				return cw.Write
+			},
+			WriteHeader: func(httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
+				return cw.WriteHeader
+			},
+			Flush: func(httpsnoop.FlushFunc) httpsnoop.FlushFunc {
+				return cw.Flush
+			},
+		})
 
 		h.ServeHTTP(w, r)
 	})

compress_test.go

@@ -29,7 +29,6 @@ func compressedRequest(w *httptest.ResponseRecorder, compression string) {
 			acceptEncoding: []string{compression},
 		},
 	})
-
 }
 
 func TestCompressHandlerNoCompression(t *testing.T) {
@@ -165,6 +164,7 @@ type fullyFeaturedResponseWriter struct{}
 func (fullyFeaturedResponseWriter) Header() http.Header {
 	return http.Header{}
 }
+
 func (fullyFeaturedResponseWriter) Write([]byte) (int, error) {
 	return 0, nil
 }
@@ -193,9 +193,6 @@ func TestCompressHandlerPreserveInterfaces(t *testing.T) {
 	)
 	var h http.Handler = http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
 		comp := r.Header.Get(acceptEncoding)
-		if _, ok := rw.(*compressResponseWriter); !ok {
-			t.Fatalf("ResponseWriter wasn't wrapped by compressResponseWriter, got %T type", rw)
-		}
 		if _, ok := rw.(http.Flusher); !ok {
 			t.Errorf("ResponseWriter lost http.Flusher interface for %q", comp)
 		}
@@ -207,9 +204,7 @@ func TestCompressHandlerPreserveInterfaces(t *testing.T) {
 		}
 	})
 	h = CompressHandler(h)
-	var (
-		rw fullyFeaturedResponseWriter
-	)
+	var rw fullyFeaturedResponseWriter
 	r, err := http.NewRequest("GET", "/", nil)
 	if err != nil {
 		t.Fatalf("Failed to create test request: %v", err)
@@ -220,3 +215,32 @@ func TestCompressHandlerPreserveInterfaces(t *testing.T) {
 	r.Header.Set(acceptEncoding, "deflate")
 	h.ServeHTTP(rw, r)
 }
+
+type paltryResponseWriter struct{}
+
+func (paltryResponseWriter) Header() http.Header {
+	return http.Header{}
+}
+
+func (paltryResponseWriter) Write([]byte) (int, error) {
+	return 0, nil
+}
+func (paltryResponseWriter) WriteHeader(int) {}
+
+func TestCompressHandlerDoesntInventInterfaces(t *testing.T) {
+	var h http.Handler = http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		if _, ok := rw.(http.Hijacker); ok {
+			t.Error("ResponseWriter shouldn't implement http.Hijacker")
+		}
+	})
+
+	h = CompressHandler(h)
+
+	var rw paltryResponseWriter
+	r, err := http.NewRequest("GET", "/", nil)
+	if err != nil {
+		t.Fatalf("Failed to create test request: %v", err)
+	}
+	r.Header.Set(acceptEncoding, "gzip")
+	h.ServeHTTP(rw, r)
+}

go.mod

@@ -1,3 +1,5 @@
 module github.com/gorilla/handlers
 
 go 1.14
+
+require github.com/felixge/httpsnoop v1.0.1

go.sum

@@ -0,0 +1,2 @@
+github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
+github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=

handlers.go

@@ -51,10 +51,6 @@ type responseLogger struct {
 	size   int
 }
 
-func (l *responseLogger) Header() http.Header {
-	return l.w.Header()
-}
-
 func (l *responseLogger) Write(b []byte) (int, error) {
 	size, err := l.w.Write(b)
 	l.size += size
@@ -74,39 +70,16 @@ func (l *responseLogger) Size() int {
 	return l.size
 }
 
-func (l *responseLogger) Flush() {
-	f, ok := l.w.(http.Flusher)
-	if ok {
-		f.Flush()
-	}
-}
-
-type hijackLogger struct {
-	responseLogger
-}
-
-func (l *hijackLogger) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	h := l.responseLogger.w.(http.Hijacker)
-	conn, rw, err := h.Hijack()
-	if err == nil && l.responseLogger.status == 0 {
+func (l *responseLogger) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	conn, rw, err := l.w.(http.Hijacker).Hijack()
+	if err == nil && l.status == 0 {
 		// The status will be StatusSwitchingProtocols if there was no error and
 		// WriteHeader has not been called yet
-		l.responseLogger.status = http.StatusSwitchingProtocols
+		l.status = http.StatusSwitchingProtocols
 	}
 	return conn, rw, err
 }
 
-type closeNotifyWriter struct {
-	loggingResponseWriter
-	http.CloseNotifier
-}
-
-type hijackCloseNotifier struct {
-	loggingResponseWriter
-	http.Hijacker
-	http.CloseNotifier
-}
-
 // isContentType validates the Content-Type header matches the supplied
 // contentType. That is, its type and subtype match.
 func isContentType(h http.Header, contentType string) bool {

handlers_go18.go

@@ -9,6 +9,15 @@ import (
 	"testing"
 )
 
+// *httptest.ResponseRecorder doesn't implement Pusher, so wrap it.
+type pushRecorder struct {
+	*httptest.ResponseRecorder
+}
+
+func (pr pushRecorder) Push(target string, opts *http.PushOptions) error {
+	return nil
+}
+
 func TestLoggingHandlerWithPush(t *testing.T) {
 	handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		if _, ok := w.(http.Pusher); !ok {
@@ -18,7 +27,7 @@ func TestLoggingHandlerWithPush(t *testing.T) {
 	})
 
 	logger := LoggingHandler(ioutil.Discard, handler)
-	logger.ServeHTTP(httptest.NewRecorder(), newRequest("GET", "/"))
+	logger.ServeHTTP(pushRecorder{httptest.NewRecorder()}, newRequest("GET", "/"))
 }
 
 func TestCombinedLoggingHandlerWithPush(t *testing.T) {
@@ -30,5 +39,5 @@ func TestCombinedLoggingHandlerWithPush(t *testing.T) {
 	})
 
 	logger := CombinedLoggingHandler(ioutil.Discard, handler)
-	logger.ServeHTTP(httptest.NewRecorder(), newRequest("GET", "/"))
+	logger.ServeHTTP(pushRecorder{httptest.NewRecorder()}, newRequest("GET", "/"))
 }

handlers_go18_test.go

@@ -12,6 +12,8 @@ import (
 	"strconv"
 	"time"
 	"unicode/utf8"
+
+	"github.com/felixge/httpsnoop"
 )
 
 // Logging
@@ -39,10 +41,10 @@ type loggingHandler struct {
 
 func (h loggingHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	t := time.Now()
-	logger := makeLogger(w)
+	logger, w := makeLogger(w)
 	url := *req.URL
 
-	h.handler.ServeHTTP(logger, req)
+	h.handler.ServeHTTP(w, req)
 	if req.MultipartForm != nil {
 		req.MultipartForm.RemoveAll()
 	}
@@ -58,27 +60,16 @@ func (h loggingHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	h.formatter(h.writer, params)
 }
 
-func makeLogger(w http.ResponseWriter) loggingResponseWriter {
-	var logger loggingResponseWriter = &responseLogger{w: w, status: http.StatusOK}
-	if _, ok := w.(http.Hijacker); ok {
-		logger = &hijackLogger{responseLogger{w: w, status: http.StatusOK}}
-	}
-	h, ok1 := logger.(http.Hijacker)
-	c, ok2 := w.(http.CloseNotifier)
-	if ok1 && ok2 {
-		return hijackCloseNotifier{logger, h, c}
-	}
-	if ok2 {
-		return &closeNotifyWriter{logger, c}
-	}
-	return logger
-}
-
-type commonLoggingResponseWriter interface {
-	http.ResponseWriter
-	http.Flusher
-	Status() int
-	Size() int
+func makeLogger(w http.ResponseWriter) (*responseLogger, http.ResponseWriter) {
+	logger := &responseLogger{w: w, status: http.StatusOK}
+	return logger, httpsnoop.Wrap(w, httpsnoop.Hooks{
+		Write: func(httpsnoop.WriteFunc) httpsnoop.WriteFunc {
+			return logger.Write
+		},
+		WriteHeader: func(httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
+			return logger.WriteHeader
+		},
+	})
 }
 
 const lowerhex = "0123456789abcdef"
@@ -145,7 +136,6 @@ func appendQuoted(buf []byte, s string) []byte {
 		}
 	}
 	return buf
-
 }
 
 // buildCommonLogLine builds a log entry for req in Apache Common Log Format.
@@ -160,7 +150,6 @@ func buildCommonLogLine(req *http.Request, url url.URL, ts time.Time, status int
 	}
 
 	host, _, err := net.SplitHostPort(req.RemoteAddr)
-
 	if err != nil {
 		host = req.RemoteAddr
 	}