diff --git a/README.md b/README.md index 54999a4..b139dd5 100644 --- a/README.md +++ b/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 diff --git a/gzip.go b/gzip.go index 40d2625..47a5016 100644 --- a/gzip.go +++ b/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,12 +169,19 @@ 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 @@ -175,6 +189,39 @@ func (w *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { 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 }) } diff --git a/gzip_test.go b/gzip_test.go index 3a9d738..fe0b5cf 100644 --- a/gzip_test.go +++ b/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") +}