Advertise Flusher and Hijacker only when the underlying writer has them

README.md

@@ -216,7 +216,13 @@ the type sniffed from the first chunk of the body. By default the common textual
 for responses the handler already encoded (`Content-Encoding` set), and for partial responses (206 or a
 `Content-Range`), whose offsets describe the uncompressed representation. `Content-Length` is dropped when the
 body is compressed, interim 1xx responses pass through without becoming the final status, and `Flush` and
-`Hijack` pass through so streaming responses and protocol upgrades keep working.
+`Hijack` pass through so streaming responses and protocol upgrades keep working. The wrapper offers those two
+only when the writer beneath it does, so composing with `Timeout`, which offers neither, does not leave a
+handler with a `Flush` that silently does nothing.
+
+One deviation is worth knowing about: sniffing means the status cannot be sent until the body arrives, so when
+a handler calls `WriteHeader` without setting `Content-Type`, headers it changes before the first `Write` still
+reach the client, where `net/http` would have ignored them.
 
 ### RealIP middleware
 

gzip.go

@@ -29,6 +29,12 @@ var gzPool = sync.Pool{
 
 // gzipResponseWriter defers the compression decision until the response content type is known,
 // either from the header the handler set or sniffed from the first chunk of the body.
+//
+// One consequence is worth knowing about: when the handler calls WriteHeader without a Content-Type,
+// the status cannot be sent yet, because the body has to be sniffed first. Headers changed between
+// that call and the first Write therefore still reach the client, where net/http would have ignored
+// them. Handlers that mutate headers after WriteHeader are relying on a no-op, so this is more
+// permissive rather than wrong, but it is a difference from the bare ResponseWriter.
 type gzipResponseWriter struct {
 	http.ResponseWriter
 
@@ -55,8 +61,9 @@ func (w *gzipResponseWriter) WriteHeader(status int) {
 	w.status = status
 
 	// with a content type in hand the decision can be made right away, otherwise it waits for the
-	// first Write so the body can be sniffed
-	if ctype := w.Header().Get("Content-Type"); ctype != "" {
+	// first Write so the body can be sniffed. 101 cannot wait: the upgrade sequence carries no body
+	// and usually no content type, and the Hijack that follows would leave the status unsent
+	if ctype := w.Header().Get("Content-Type"); ctype != "" || status == http.StatusSwitchingProtocols {
 		w.decide(ctype)
 		w.commit()
 	}
@@ -141,8 +148,8 @@ func (w *gzipResponseWriter) close(finished bool) {
 	w.gz = nil
 }
 
-// Flush pushes buffered data out, keeping streaming responses working through the compressor
-func (w *gzipResponseWriter) Flush() {
+// flush pushes buffered data out, keeping streaming responses working through the compressor
+func (w *gzipResponseWriter) flush() {
 	if w.hijacked {
 		return
 	}
@@ -162,19 +169,59 @@ func (w *gzipResponseWriter) Flush() {
 	}
 }
 
-// Hijack passes through to the underlying writer for protocol upgrades
-func (w *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+// hijack passes through to the underlying writer for protocol upgrades
+func (w *gzipResponseWriter) hijack() (net.Conn, *bufio.ReadWriter, error) {
 	h, ok := w.ResponseWriter.(http.Hijacker)
 	if !ok {
 		return nil, nil, fmt.Errorf("http.Hijacker not supported")
 	}
+	// finish the stream first, whatever the handler already wrote has to reach the wire before
+	// the connection changes hands
+	if w.gz != nil {
+		_ = w.gz.Close()
+		gzPool.Put(w.gz)
+		w.gz = nil
+	}
 	conn, rw, err := h.Hijack()
 	if err == nil {
 		w.hijacked = true
 	}
 	return conn, rw, err
 }
 
+// the wrapper must offer exactly the optional interfaces the underlying writer has, otherwise a
+// handler's type assertion succeeds and the call then does nothing, which is how http.TimeoutHandler
+// (offering neither) would silently lose a Flush
+type gzipFlusher struct{ *gzipResponseWriter }
+
+func (w gzipFlusher) Flush() { w.flush() }
+
+type gzipHijacker struct{ *gzipResponseWriter }
+
+func (w gzipHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) { return w.hijack() }
+
+type gzipFlushHijacker struct{ *gzipResponseWriter }
+
+func (w gzipFlushHijacker) Flush() { w.flush() }
+
+func (w gzipFlushHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) { return w.hijack() }
+
+// wrapGzipWriter picks the variant matching the capabilities of the writer underneath
+func wrapGzipWriter(gw *gzipResponseWriter) http.ResponseWriter {
+	_, isFlusher := gw.ResponseWriter.(http.Flusher)
+	_, isHijacker := gw.ResponseWriter.(http.Hijacker)
+
+	switch {
+	case isFlusher && isHijacker:
+		return gzipFlushHijacker{gw}
+	case isFlusher:
+		return gzipFlusher{gw}
+	case isHijacker:
+		return gzipHijacker{gw}
+	}
+	return gw
+}
+
 // Unwrap exposes the underlying writer to http.ResponseController
 func (w *gzipResponseWriter) Unwrap() http.ResponseWriter {
 	return w.ResponseWriter
@@ -204,7 +251,7 @@ func Gzip(contentTypes ...string) func(http.Handler) http.Handler {
 			finished := false
 			defer func() { gw.close(finished) }()
 
-			next.ServeHTTP(gw, r)
+			next.ServeHTTP(wrapGzipWriter(gw), r)
 			finished = true
 		})
 	}

gzip_test.go

@@ -6,11 +6,13 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"strconv"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -449,7 +451,7 @@ func TestGzipWriterInterfaces(t *testing.T) {
 
 	t.Run("hijack unsupported by the underlying writer", func(t *testing.T) {
 		gw := &gzipResponseWriter{ResponseWriter: httptest.NewRecorder(), gzCts: gzDefaultContentTypes}
-		_, _, err := gw.Hijack()
+		_, _, err := gw.hijack()
 		assert.Error(t, err)
 	})
 
@@ -472,7 +474,7 @@ func TestGzipWriterInterfaces(t *testing.T) {
 	t.Run("flush without a body commits the status", func(t *testing.T) {
 		rec := httptest.NewRecorder()
 		gw := &gzipResponseWriter{ResponseWriter: rec, gzCts: gzDefaultContentTypes}
-		gw.Flush()
+		gw.flush()
 		gw.close(true)
 		assert.Equal(t, http.StatusOK, rec.Code)
 	})
@@ -736,3 +738,96 @@ func TestGzipNoSniffWhenAlreadyEncoded(t *testing.T) {
 	assert.Empty(t, rec.Header().Get("Content-Type"), "no content type should be guessed from encoded bytes")
 	assert.Equal(t, "br", rec.Header().Get("Content-Encoding"))
 }
+
+func TestGzipWriterCapabilitiesMatchUnderlying(t *testing.T) {
+	body := strings.Repeat("stream me. ", 40)
+
+	tbl := []struct {
+		name         string
+		wrap         func(http.Handler) http.Handler
+		wantFlusher  bool
+		wantHijacker bool
+	}{
+		{
+			name:         "plain server writer",
+			wrap:         func(h http.Handler) http.Handler { return h },
+			wantFlusher:  true,
+			wantHijacker: true,
+		},
+		{
+			// http.TimeoutHandler offers neither, so the wrapper must not claim them either
+			name:         "behind Timeout",
+			wrap:         Timeout(time.Minute),
+			wantFlusher:  false,
+			wantHijacker: false,
+		},
+	}
+
+	for _, tt := range tbl {
+		t.Run(tt.name, func(t *testing.T) {
+			var gotFlusher, gotHijacker bool
+			handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+				_, gotFlusher = w.(http.Flusher)
+				_, gotHijacker = w.(http.Hijacker)
+				w.Header().Set("Content-Type", "text/plain")
+				_, err := w.Write([]byte(body))
+				require.NoError(t, err)
+			})
+
+			ts := httptest.NewServer(tt.wrap(Gzip()(handler)))
+			defer ts.Close()
+
+			req, err := http.NewRequest("GET", ts.URL+"/x", http.NoBody)
+			require.NoError(t, err)
+			req.Header.Set("Accept-Encoding", "gzip")
+
+			resp, err := http.DefaultTransport.RoundTrip(req)
+			require.NoError(t, err)
+			defer resp.Body.Close()
+			_, err = io.ReadAll(resp.Body)
+			require.NoError(t, err)
+
+			assert.Equal(t, tt.wantFlusher, gotFlusher, "Flusher must be advertised only when it works")
+			assert.Equal(t, tt.wantHijacker, gotHijacker, "Hijacker must be advertised only when it works")
+		})
+	}
+}
+
+func TestGzipSwitchingProtocolsWithoutContentType(t *testing.T) {
+	// the plain upgrade sequence: 101 with no content type, then the handler takes the connection.
+	// the status has to be on the wire before Hijack, or it is never sent at all
+	handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Upgrade", "foo")
+		w.Header().Set("Connection", "Upgrade")
+		w.WriteHeader(http.StatusSwitchingProtocols)
+
+		hj, ok := w.(http.Hijacker)
+		require.True(t, ok, "upgrade needs a Hijacker")
+		conn, buf, err := hj.Hijack()
+		require.NoError(t, err)
+		defer conn.Close()
+
+		_, err = buf.WriteString("raw-protocol-bytes")
+		require.NoError(t, err)
+		require.NoError(t, buf.Flush())
+	})
+
+	ts := httptest.NewServer(Gzip()(handler))
+	defer ts.Close()
+
+	conn, err := net.Dial("tcp", strings.TrimPrefix(ts.URL, "http://"))
+	require.NoError(t, err)
+	defer conn.Close()
+
+	_, err = fmt.Fprint(conn, "GET /upgrade HTTP/1.1\r\nHost: example.com\r\n"+
+		"Accept-Encoding: gzip\r\nUpgrade: foo\r\nConnection: Upgrade\r\n\r\n")
+	require.NoError(t, err)
+	require.NoError(t, conn.SetReadDeadline(time.Now().Add(5*time.Second)))
+
+	got, err := io.ReadAll(conn)
+	require.NoError(t, err)
+
+	assert.Contains(t, string(got), "101 Switching Protocols", "the 101 status line must reach the client")
+	assert.Contains(t, string(got), "raw-protocol-bytes")
+	assert.NotContains(t, string(got), "Content-Encoding: gzip", "an upgraded connection must not be gzipped")
+}