Merge pull request #87 from basecamp/prevent-response-buffering

Ensure all response writers are flushable

Author: kevinmcconnell

internal/cacheable_response.go

@@ -75,6 +75,13 @@ func (c *CacheableResponse) WriteHeader(statusCode int) {
 	c.headersWritten = true
 }
 
+func (c *CacheableResponse) Flush() {
+	flusher, ok := c.responseWriter.(http.Flusher)
+	if ok {
+		flusher.Flush()
+	}
+}
+
 func (c *CacheableResponse) CacheStatus() (bool, time.Time) {
 	if c.stasher.Overflowed() {
 		return false, time.Time{}

internal/handler_test.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"bytes"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -280,6 +281,26 @@ func TestHandlerAddsXRequestStartHeader(t *testing.T) {
 	assert.Equal(t, http.StatusOK, w.Code)
 }
 
+func TestHandlerAllowsFlushingTheResponseBody(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		fmt.Fprintf(w, "data: event\n\n")
+		w.(http.Flusher).Flush()
+
+	}))
+	defer upstream.Close()
+
+	h := NewHandler(handlerOptions(upstream.URL))
+
+	w := httptest.NewRecorder()
+	r := httptest.NewRequest("GET", "/", nil)
+	h.ServeHTTP(w, r)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+	assert.Equal(t, "text/event-stream", w.Header().Get("Content-Type"))
+	assert.True(t, w.Flushed)
+}
+
 // Helpers
 
 func handlerOptions(targetUrl string) HandlerOptions {

internal/logging_middleware.go

@@ -87,3 +87,10 @@ func (r *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	}
 	return con, rw, err
 }
+
+func (r *responseWriter) Flush() {
+	flusher, ok := r.ResponseWriter.(http.Flusher)
+	if ok {
+		flusher.Flush()
+	}
+}

internal/sendfile_handler.go

@@ -79,6 +79,13 @@ func (w *sendfileWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	return hijacker.Hijack()
 }
 
+func (w *sendfileWriter) Flush() {
+	flusher, ok := w.w.(http.Flusher)
+	if ok {
+		flusher.Flush()
+	}
+}
+
 func (w *sendfileWriter) sendingFilename() string {
 	return w.w.Header().Get("X-Sendfile")
 }