Merge pull request #66 from wathiede/flusher

[feature] http.Flusher/http.CloseNotifier support for CompressHandler.

Author: elithrar

compress.go

@@ -16,6 +16,8 @@ type compressResponseWriter struct {
 	io.Writer
 	http.ResponseWriter
 	http.Hijacker
+	http.Flusher
+	http.CloseNotifier
 }
 
 func (w *compressResponseWriter) WriteHeader(c int) {
@@ -37,6 +39,21 @@ func (w *compressResponseWriter) Write(b []byte) (int, error) {
 	return w.Writer.Write(b)
 }
 
+type flusher interface {
+	Flush() error
+}
+
+func (w *compressResponseWriter) Flush() {
+	// Flush compressed data if compressor supports it.
+	if f, ok := w.Writer.(flusher); ok {
+		f.Flush()
+	}
+	// Flush HTTP response.
+	if w.Flusher != nil {
+		w.Flusher.Flush()
+	}
+}
+
 // CompressHandler gzip compresses HTTP responses for clients that support it
 // via the 'Accept-Encoding' header.
 func CompressHandler(h http.Handler) http.Handler {
@@ -70,10 +87,22 @@ func CompressHandlerLevel(h http.Handler, level int) http.Handler {
 					h = nil
 				}
 
+				f, fok := w.(http.Flusher)
+				if !fok {
+					f = nil
+				}
+
+				cn, cnok := w.(http.CloseNotifier)
+				if !cnok {
+					cn = nil
+				}
+
 				w = &compressResponseWriter{
 					Writer:         gw,
 					ResponseWriter: w,
 					Hijacker:       h,
+					Flusher:        f,
+					CloseNotifier:  cn,
 				}
 
 				break L
@@ -89,10 +118,22 @@ func CompressHandlerLevel(h http.Handler, level int) http.Handler {
 					h = nil
 				}
 
+				f, fok := w.(http.Flusher)
+				if !fok {
+					f = nil
+				}
+
+				cn, cnok := w.(http.CloseNotifier)
+				if !cnok {
+					cn = nil
+				}
+
 				w = &compressResponseWriter{
 					Writer:         fw,
 					ResponseWriter: w,
 					Hijacker:       h,
+					Flusher:        f,
+					CloseNotifier:  cn,
 				}
 
 				break L

compress_test.go

@@ -5,7 +5,9 @@
 package handlers
 
 import (
+	"bufio"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"strconv"
@@ -88,3 +90,65 @@ func TestCompressHandlerGzipDeflate(t *testing.T) {
 		t.Fatalf("wrong content type, got %s want %s", w.HeaderMap.Get("Content-Type"), "text/plain; charset=utf-8")
 	}
 }
+
+type fullyFeaturedResponseWriter struct{}
+
+// Header/Write/WriteHeader implement the http.ResponseWriter interface.
+func (fullyFeaturedResponseWriter) Header() http.Header {
+	return http.Header{}
+}
+func (fullyFeaturedResponseWriter) Write([]byte) (int, error) {
+	return 0, nil
+}
+func (fullyFeaturedResponseWriter) WriteHeader(int) {}
+
+// Flush implements the http.Flusher interface.
+func (fullyFeaturedResponseWriter) Flush() {}
+
+// Hijack implements the http.Hijacker interface.
+func (fullyFeaturedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return nil, nil, nil
+}
+
+// CloseNotify implements the http.CloseNotifier interface.
+func (fullyFeaturedResponseWriter) CloseNotify() <-chan bool {
+	return nil
+}
+
+func TestCompressHandlerPreserveInterfaces(t *testing.T) {
+	// Compile time validation fullyFeaturedResponseWriter implements all the
+	// interfaces we're asserting in the test case below.
+	var (
+		_ http.Flusher       = fullyFeaturedResponseWriter{}
+		_ http.CloseNotifier = fullyFeaturedResponseWriter{}
+		_ http.Hijacker      = fullyFeaturedResponseWriter{}
+	)
+	var h http.Handler = http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		comp := r.Header.Get("Accept-Encoding")
+		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)
+		}
+		if _, ok := rw.(http.CloseNotifier); !ok {
+			t.Errorf("ResponseWriter lost http.CloseNotifier interface for %q", comp)
+		}
+		if _, ok := rw.(http.Hijacker); !ok {
+			t.Errorf("ResponseWriter lost http.Hijacker interface for %q", comp)
+		}
+	})
+	h = CompressHandler(h)
+	var (
+		rw fullyFeaturedResponseWriter
+	)
+	r, err := http.NewRequest("GET", "/", nil)
+	if err != nil {
+		t.Fatalf("Failed to create test request: %v", err)
+	}
+	r.Header.Set("Accept-Encoding", "gzip")
+	h.ServeHTTP(rw, r)
+
+	r.Header.Set("Accept-Encoding", "deflate")
+	h.ServeHTTP(rw, r)
+}