diff --git a/gzip.go b/gzip.go index 47a5016..ddf6e8e 100644 --- a/gzip.go +++ b/gzip.go @@ -175,18 +175,29 @@ func (w *gzipResponseWriter) hijack() (net.Conn, *bufio.ReadWriter, error) { 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 + // finish the stream first, whatever the handler already wrote has to reach the wire before the + // connection changes hands, and a failure there means truncated output rather than something to + // swallow behind a successful hijack if w.gz != nil { - _ = w.gz.Close() - gzPool.Put(w.gz) - w.gz = nil + if err := w.gz.Close(); err != nil { + return nil, nil, fmt.Errorf("finish gzip stream before hijack: %w", err) + } } + conn, rw, err := h.Hijack() - if err == nil { - w.hijacked = true + if err != nil { + // the connection was not taken over, so the writer stays attached and closed: a handler that + // carries on writing now gets an error instead of appending raw bytes to a body already + // advertised as gzip, and the deferred close still returns the writer to the pool + return nil, nil, err + } + + if w.gz != nil { + gzPool.Put(w.gz) + w.gz = nil } - return conn, rw, err + w.hijacked = true + return conn, rw, nil } // the wrapper must offer exactly the optional interfaces the underlying writer has, otherwise a diff --git a/gzip_test.go b/gzip_test.go index fe0b5cf..48c9088 100644 --- a/gzip_test.go +++ b/gzip_test.go @@ -1,6 +1,7 @@ package rest import ( + "bufio" "bytes" "compress/gzip" "errors" @@ -831,3 +832,85 @@ func TestGzipSwitchingProtocolsWithoutContentType(t *testing.T) { assert.Contains(t, string(got), "raw-protocol-bytes") assert.NotContains(t, string(got), "Content-Encoding: gzip", "an upgraded connection must not be gzipped") } + +// hijackStub is a ResponseWriter whose Hijack and Write can be made to fail on demand +type hijackStub struct { + http.ResponseWriter + hijackErr error + writeErr error + written bytes.Buffer +} + +func (h *hijackStub) Hijack() (net.Conn, *bufio.ReadWriter, error) { + if h.hijackErr != nil { + return nil, nil, h.hijackErr + } + return nil, nil, nil +} + +func (h *hijackStub) Write(b []byte) (int, error) { + if h.writeErr != nil { + return 0, h.writeErr + } + return h.written.Write(b) +} + +func TestGzipHijackErrorPath(t *testing.T) { + t.Run("failing stream close is reported, not swallowed", func(t *testing.T) { + stub := &hijackStub{ResponseWriter: httptest.NewRecorder()} + gw := &gzipResponseWriter{ResponseWriter: stub, gzCts: gzDefaultContentTypes} + gw.Header().Set("Content-Type", "text/plain") + + _, err := gw.Write([]byte(strings.Repeat("compress me. ", 40))) + require.NoError(t, err) + require.NotNil(t, gw.gz, "the stream has to be open for this case to mean anything") + + boom := errors.New("write failed") + stub.writeErr = boom // the flush inside gz.Close now fails + + _, _, err = gw.hijack() + require.Error(t, err, "a truncated stream must not be hidden behind a successful hijack") + assert.ErrorIs(t, err, boom) + assert.False(t, gw.hijacked) + }) + + t.Run("failed hijack leaves the stream attached so writes error instead of corrupting", func(t *testing.T) { + nope := errors.New("already hijacked") + stub := &hijackStub{ResponseWriter: httptest.NewRecorder(), hijackErr: nope} + gw := &gzipResponseWriter{ResponseWriter: stub, gzCts: gzDefaultContentTypes} + gw.Header().Set("Content-Type", "text/plain") + + _, err := gw.Write([]byte(strings.Repeat("compress me. ", 40))) + require.NoError(t, err) + + _, _, err = gw.hijack() + require.ErrorIs(t, err, nope) + + assert.False(t, gw.hijacked, "the connection was never taken over") + require.NotNil(t, gw.gz, "the writer stays attached so a further write cannot bypass it") + assert.Equal(t, "gzip", gw.Header().Get("Content-Encoding")) + + // the response already advertises gzip, so raw bytes must not be appended to it + _, err = gw.Write([]byte("RAW-AFTER-FAILED-HIJACK")) + require.Error(t, err, "writing after a failed hijack has to fail rather than corrupt the body") + assert.NotContains(t, stub.written.String(), "RAW-AFTER-FAILED-HIJACK") + + // the deferred close still returns the writer to the pool + gw.close(true) + assert.Nil(t, gw.gz) + }) + + t.Run("successful hijack releases the writer", func(t *testing.T) { + stub := &hijackStub{ResponseWriter: httptest.NewRecorder()} + gw := &gzipResponseWriter{ResponseWriter: stub, gzCts: gzDefaultContentTypes} + gw.Header().Set("Content-Type", "text/plain") + + _, err := gw.Write([]byte(strings.Repeat("compress me. ", 40))) + require.NoError(t, err) + + _, _, err = gw.hijack() + require.NoError(t, err) + assert.True(t, gw.hijacked) + assert.Nil(t, gw.gz, "the writer goes back to the pool once the connection is taken over") + }) +}