[fix] httpbp middleware doesn't flush chunked responses

httpbp/middlewares.go

@@ -394,9 +394,10 @@ func recordStatusCode(counters counterGenerator) Middleware {
 	return func(name string, next HandlerFunc) HandlerFunc {
 		counter := counters.Counter("baseplate.http." + name + ".response")
 		return func(ctx context.Context, w http.ResponseWriter, r *http.Request) (err error) {
-			wrapped := &statusCodeRecorder{ResponseWriter: w}
+			rec := &statusCodeRecorder{ResponseWriter: w}
+			wrapped := allowFlushHijack(w, rec)
 			defer func() {
-				code := wrapped.getCode(err)
+				code := rec.getCode(err)
 				counter.With("status", statusCodeFamily(code)).Add(1)
 			}()
 
@@ -453,9 +454,11 @@ func PrometheusServerMetrics(_ string) Middleware {
 			}
 			serverActiveRequests.With(activeRequestLabels).Inc()
 
-			wrapped := &responseRecorder{ResponseWriter: w}
+			rec := &responseRecorder{ResponseWriter: w}
+			wrapped := allowFlushHijack(w, rec)
+
 			defer func() {
-				code := errorCodeForMetrics(wrapped.responseCode, err)
+				code := errorCodeForMetrics(rec.responseCode, err)
 				success := isRequestSuccessful(code, err)
 
 				labels := prometheus.Labels{
@@ -465,7 +468,7 @@ func PrometheusServerMetrics(_ string) Middleware {
 				}
 				serverLatency.With(labels).Observe(time.Since(start).Seconds())
 				serverRequestSize.With(labels).Observe(float64(r.ContentLength))
-				serverResponseSize.With(labels).Observe(float64(wrapped.bytesWritten))
+				serverResponseSize.With(labels).Observe(float64(rec.bytesWritten))
 
 				totalRequestLabels := prometheus.Labels{
 					methodLabel:   method,
@@ -510,3 +513,49 @@ func (rr *responseRecorder) WriteHeader(code int) {
 	rr.ResponseWriter.WriteHeader(code)
 	rr.responseCode = code
 }
+
+type wrappedHijacker struct {
+	http.ResponseWriter
+	http.Hijacker
+}
+
+type wrappedFlusher struct {
+	http.ResponseWriter
+	http.Flusher
+}
+
+type wrappedFlushHijacker struct {
+	http.ResponseWriter
+	http.Flusher
+	http.Hijacker
+}
+
+func allowFlushHijack(original, rw http.ResponseWriter) http.ResponseWriter {
+	flusher, isFlusher := original.(http.Flusher)
+	hijacker, isHijacker := original.(http.Hijacker)
+	switch {
+	case isFlusher && isHijacker:
+		return &wrappedFlushHijacker{rw, flusher, hijacker}
+	case isFlusher:
+		return allowFlush(original, rw)
+	case isHijacker:
+		return allowHijack(original, rw)
+	default:
+		return rw
+	}
+
+}
+
+func allowFlush(original, rw http.ResponseWriter) http.ResponseWriter {
+	if flusher, isFlusher := original.(http.Flusher); isFlusher {
+		return &wrappedFlusher{rw, flusher}
+	}
+	return rw
+}
+
+func allowHijack(original, rw http.ResponseWriter) http.ResponseWriter {
+	if hijacker, isHijacker := original.(http.Hijacker); isHijacker {
+		return &wrappedHijacker{rw, hijacker}
+	}
+	return rw
+}

httpbp/middlewares_test.go

@@ -1,15 +1,18 @@
 package httpbp_test
 
 import (
+	"bufio"
 	"context"
 	"encoding/json"
 	"errors"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"sort"
 	"strings"
 	"testing"
 
+	"github.com/reddit/baseplate.go"
 	"github.com/reddit/baseplate.go/ecinterface"
 	"github.com/reddit/baseplate.go/httpbp"
 	"github.com/reddit/baseplate.go/log"
@@ -323,3 +326,88 @@ func TestSupportedMethods(t *testing.T) {
 		)
 	}
 }
+
+func TestMiddlewareResponseWrapping(t *testing.T) {
+	store := newSecretsStore(t)
+	defer store.Close()
+
+	bp := baseplate.NewTestBaseplate(baseplate.NewTestBaseplateArgs{
+		Config:          baseplate.Config{Addr: ":8080"},
+		Store:           store,
+		EdgeContextImpl: ecinterface.Mock(),
+	})
+
+	args := httpbp.ServerArgs{
+		Baseplate: bp,
+		Middlewares: []httpbp.Middleware{
+			func(name string, next httpbp.HandlerFunc) httpbp.HandlerFunc {
+				return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
+					if flusher, isFlusher := w.(http.Flusher); isFlusher {
+						flusher.Flush()
+					}
+
+					next(ctx, w, r)
+					return nil
+				}
+			},
+			func(name string, next httpbp.HandlerFunc) httpbp.HandlerFunc {
+				return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
+					if hijacker, isHijacker := w.(http.Hijacker); isHijacker {
+						hijacker.Hijack()
+					}
+
+					next(ctx, w, r)
+					return nil
+				}
+			},
+		},
+		Endpoints: map[httpbp.Pattern]httpbp.Endpoint{
+			"/test": {
+				Name:    "test",
+				Methods: []string{http.MethodGet},
+				Handle: func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
+					w.Write([]byte("endpoint"))
+					return nil
+				},
+			},
+		},
+	}
+
+	// register our middleware to the EndpointRegistry
+	args, err := args.SetupEndpoints()
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Test the a flushable response
+	t.Run("flushable-response", func(t *testing.T) {
+		r := httptest.NewRequest(http.MethodGet, "/test", nil)
+		w := httptest.NewRecorder()
+		args.EndpointRegistry.ServeHTTP(w, r)
+
+		if !w.Flushed {
+			t.Error("expected http response to be flushed")
+		}
+	})
+
+	t.Run("hijackable-response", func(tt *testing.T) {
+		r := httptest.NewRequest(http.MethodGet, "/test", nil)
+		w := &hijackableResponseRecorder{httptest.NewRecorder(), false}
+		args.EndpointRegistry.ServeHTTP(w, r)
+
+		if !w.Hijacked {
+			t.Error("expected http response to be hijacked")
+		}
+	})
+}
+
+type hijackableResponseRecorder struct {
+	*httptest.ResponseRecorder
+	Hijacked bool
+}
+
+func (h *hijackableResponseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	h.Hijacked = true
+	return nil, nil, nil
+}