diff --git a/internal/httputils/wrap_writer.go b/internal/httputils/wrap_writer.go index 51bf910ff..caa91a716 100644 --- a/internal/httputils/wrap_writer.go +++ b/internal/httputils/wrap_writer.go @@ -48,6 +48,12 @@ func NewWrapResponseWriter(w http.ResponseWriter, protoMajor int) WrapResponseWr if fl && hj && rf { return &httpFancyWriter{bw} } + if fl && hj { + return &flushHijackWriter{bw} + } + if hj { + return &hijackWriter{bw} + } } if fl { return &flushWriter{bw} @@ -143,6 +149,36 @@ func (f *flushWriter) Flush() { var _ http.Flusher = &flushWriter{} +// flushHijackWriter is a HTTP writer that additionally satisfies http.Flusher +// and http.Hijacker, for writers that do not implement io.ReaderFrom. +type flushHijackWriter struct { + basicWriter +} + +func (f *flushHijackWriter) Flush() { + f.wroteHeader = true + f.ResponseWriter.(http.Flusher).Flush() +} + +func (f *flushHijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return f.ResponseWriter.(http.Hijacker).Hijack() +} + +var _ http.Flusher = &flushHijackWriter{} +var _ http.Hijacker = &flushHijackWriter{} + +// hijackWriter is a HTTP writer that additionally satisfies http.Hijacker, for +// writers that implement neither http.Flusher nor io.ReaderFrom. +type hijackWriter struct { + basicWriter +} + +func (f *hijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return f.ResponseWriter.(http.Hijacker).Hijack() +} + +var _ http.Hijacker = &hijackWriter{} + // httpFancyWriter is a HTTP writer that additionally satisfies // http.Flusher, http.Hijacker, and io.ReaderFrom. It exists for the common case // of wrapping the http.ResponseWriter that package http gives you, in order to diff --git a/internal/httputils/wrap_writer_test.go b/internal/httputils/wrap_writer_test.go index 2487cd3a3..2c889231d 100644 --- a/internal/httputils/wrap_writer_test.go +++ b/internal/httputils/wrap_writer_test.go @@ -37,6 +37,30 @@ func (c *CustomResponseWriter) ReadFrom(r io.Reader) (n int64, err error) { return } +// FlusherHijacker for testing a writer that supports http.Flusher and +// http.Hijacker, but not io.ReaderFrom, like most compressing middlewares. +type FlusherHijacker struct { + *httptest.ResponseRecorder +} + +func (c *FlusherHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return nil, nil, errors.New("hijack not supported in tests") +} + +func (c *FlusherHijacker) Flush() { + c.ResponseRecorder.Flush() +} + +// HijackerOnly for testing a writer that supports http.Hijacker, but neither +// http.Flusher nor io.ReaderFrom. +type HijackerOnly struct { + http.ResponseWriter +} + +func (c *HijackerOnly) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return nil, nil, errors.New("hijack not supported in tests") +} + func TestHttpFancyWriterRemembersWroteHeaderWhenFlushed(t *testing.T) { f := &httpFancyWriter{basicWriter: basicWriter{ResponseWriter: httptest.NewRecorder()}} f.Flush() @@ -95,6 +119,62 @@ func TestNewWrapResponseWriter(t *testing.T) { } } +func TestNewWrapResponseWriterKeepsHijacker(t *testing.T) { + // Flusher, Hijacker and io.ReaderFrom + w1 := NewWrapResponseWriter(&CustomResponseWriter{httptest.NewRecorder()}, 1) + if _, ok := w1.(*httpFancyWriter); !ok { + t.Fatalf("expected httpFancyWriter, got %T", w1) + } + if _, ok := w1.(http.Hijacker); !ok { + t.Fatalf("%T does not implement http.Hijacker", w1) + } + + // Flusher and Hijacker + w2 := NewWrapResponseWriter(&FlusherHijacker{httptest.NewRecorder()}, 1) + if _, ok := w2.(*flushHijackWriter); !ok { + t.Fatalf("expected flushHijackWriter, got %T", w2) + } + if _, ok := w2.(http.Hijacker); !ok { + t.Fatalf("%T does not implement http.Hijacker", w2) + } + + // Hijacker only + w3 := NewWrapResponseWriter(&HijackerOnly{httptest.NewRecorder()}, 1) + if _, ok := w3.(*hijackWriter); !ok { + t.Fatalf("expected hijackWriter, got %T", w3) + } + if _, ok := w3.(http.Hijacker); !ok { + t.Fatalf("%T does not implement http.Hijacker", w3) + } +} + +func TestFlushHijackWriterRemembersWroteHeaderWhenFlushed(t *testing.T) { + f := &flushHijackWriter{basicWriter{ResponseWriter: &FlusherHijacker{httptest.NewRecorder()}}} + f.Flush() + + if !f.wroteHeader { + t.Fatal("want Flush to have set wroteHeader=true") + } +} + +func TestFlushHijackWriterHijack(t *testing.T) { + f := &flushHijackWriter{basicWriter{ResponseWriter: &FlusherHijacker{httptest.NewRecorder()}}} + + _, _, err := f.Hijack() + if err == nil { + t.Fatal("expected error, got nil") + } +} + +func TestHijackWriterHijack(t *testing.T) { + f := &hijackWriter{basicWriter{ResponseWriter: &HijackerOnly{httptest.NewRecorder()}}} + + _, _, err := f.Hijack() + if err == nil { + t.Fatal("expected error, got nil") + } +} + func TestBasicWriterWriteHeader(t *testing.T) { rec := httptest.NewRecorder() bw := &basicWriter{ResponseWriter: rec}