--- /dev/null
+// Copyright 2022 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package http
+
+import (
+ "bufio"
+ "fmt"
+ "net"
+ "time"
+)
+
+// A ResponseController is used by an HTTP handler to control the response.
+//
+// A ResponseController may not be used after the Handler.ServeHTTP method has returned.
+type ResponseController struct {
+ rw ResponseWriter
+}
+
+// NewResponseController creates a ResponseController for a request.
+//
+// The ResponseWriter should be the original value passed to the Handler.ServeHTTP method,
+// or have an Unwrap method returning the original ResponseWriter.
+//
+// If the ResponseWriter implements any of the following methods, the ResponseController
+// will call them as appropriate:
+//
+// Flush()
+// FlushError() error // alternative Flush returning an error
+// Hijack() (net.Conn, *bufio.ReadWriter, error)
+// SetReadDeadline(deadline time.Time) error
+// SetWriteDeadline(deadline time.Time) error
+//
+// If the ResponseWriter does not support a method, ResponseController returns
+// an error matching ErrNotSupported.
+func NewResponseController(rw ResponseWriter) *ResponseController {
+ return &ResponseController{rw}
+}
+
+type rwUnwrapper interface {
+ Unwrap() ResponseWriter
+}
+
+// Flush flushes buffered data to the client.
+func (c *ResponseController) Flush() error {
+ rw := c.rw
+ for {
+ switch t := rw.(type) {
+ case interface{ FlushError() error }:
+ return t.FlushError()
+ case Flusher:
+ t.Flush()
+ return nil
+ case rwUnwrapper:
+ rw = t.Unwrap()
+ default:
+ return errNotSupported()
+ }
+ }
+}
+
+// Hijack lets the caller take over the connection.
+// See the Hijacker interface for details.
+func (c *ResponseController) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+ rw := c.rw
+ for {
+ switch t := rw.(type) {
+ case Hijacker:
+ return t.Hijack()
+ case rwUnwrapper:
+ rw = t.Unwrap()
+ default:
+ return nil, nil, errNotSupported()
+ }
+ }
+}
+
+// SetReadDeadline sets the deadline for reading the entire request, including the body.
+// Reads from the request body after the deadline has been exceeded will return an error.
+// A zero value means no deadline.
+//
+// Setting the read deadline after it has been exceeded will not extend it.
+func (c *ResponseController) SetReadDeadline(deadline time.Time) error {
+ rw := c.rw
+ for {
+ switch t := rw.(type) {
+ case interface{ SetReadDeadline(time.Time) error }:
+ return t.SetReadDeadline(deadline)
+ case rwUnwrapper:
+ rw = t.Unwrap()
+ default:
+ return errNotSupported()
+ }
+ }
+}
+
+// SetWriteDeadline sets the deadline for writing the response.
+// Writes to the response body after the deadline has been exceeded will not block,
+// but may succeed if the data has been buffered.
+// A zero value means no deadline.
+//
+// Setting the write deadline after it has been exceeded will not extend it.
+func (c *ResponseController) SetWriteDeadline(deadline time.Time) error {
+ rw := c.rw
+ for {
+ switch t := rw.(type) {
+ case interface{ SetWriteDeadline(time.Time) error }:
+ return t.SetWriteDeadline(deadline)
+ case rwUnwrapper:
+ rw = t.Unwrap()
+ default:
+ return errNotSupported()
+ }
+ }
+}
+
+// errNotSupported returns an error that Is ErrNotSupported,
+// but is not == to it.
+func errNotSupported() error {
+ return fmt.Errorf("%w", ErrNotSupported)
+}
--- /dev/null
+package http_test
+
+import (
+ "errors"
+ "fmt"
+ "io"
+ . "net/http"
+ "os"
+ "sync"
+ "testing"
+ "time"
+)
+
+func TestResponseControllerFlush(t *testing.T) { run(t, testResponseControllerFlush) }
+func testResponseControllerFlush(t *testing.T, mode testMode) {
+ if mode == http2Mode {
+ t.Skip("skip until h2_bundle.go is updated")
+ }
+ continuec := make(chan struct{})
+ cst := newClientServerTest(t, mode, HandlerFunc(func(w ResponseWriter, r *Request) {
+ ctl := NewResponseController(w)
+ w.Write([]byte("one"))
+ if err := ctl.Flush(); err != nil {
+ t.Errorf("ctl.Flush() = %v, want nil", err)
+ return
+ }
+ <-continuec
+ w.Write([]byte("two"))
+ }))
+
+ res, err := cst.c.Get(cst.ts.URL)
+ if err != nil {
+ t.Fatalf("unexpected connection error: %v", err)
+ }
+ defer res.Body.Close()
+
+ buf := make([]byte, 16)
+ n, err := res.Body.Read(buf)
+ close(continuec)
+ if err != nil || string(buf[:n]) != "one" {
+ t.Fatalf("Body.Read = %q, %v, want %q, nil", string(buf[:n]), err, "one")
+ }
+
+ got, err := io.ReadAll(res.Body)
+ if err != nil || string(got) != "two" {
+ t.Fatalf("Body.Read = %q, %v, want %q, nil", string(got), err, "two")
+ }
+}
+
+func TestResponseControllerHijack(t *testing.T) { run(t, testResponseControllerHijack) }
+func testResponseControllerHijack(t *testing.T, mode testMode) {
+ if mode == http2Mode {
+ t.Skip("skip until h2_bundle.go is updated")
+ }
+ const header = "X-Header"
+ const value = "set"
+ cst := newClientServerTest(t, mode, HandlerFunc(func(w ResponseWriter, r *Request) {
+ ctl := NewResponseController(w)
+ c, _, err := ctl.Hijack()
+ if mode == http2Mode {
+ if err == nil {
+ t.Errorf("ctl.Hijack = nil, want error")
+ }
+ w.Header().Set(header, value)
+ return
+ }
+ if err != nil {
+ t.Errorf("ctl.Hijack = _, _, %v, want _, _, nil", err)
+ return
+ }
+ fmt.Fprintf(c, "HTTP/1.0 200 OK\r\n%v: %v\r\nContent-Length: 0\r\n\r\n", header, value)
+ }))
+ res, err := cst.c.Get(cst.ts.URL)
+ if err != nil {
+ t.Fatal(err)
+ }
+ if got, want := res.Header.Get(header), value; got != want {
+ t.Errorf("response header %q = %q, want %q", header, got, want)
+ }
+}
+
+func TestResponseControllerSetPastWriteDeadline(t *testing.T) {
+ run(t, testResponseControllerSetPastWriteDeadline)
+}
+func testResponseControllerSetPastWriteDeadline(t *testing.T, mode testMode) {
+ if mode == http2Mode {
+ t.Skip("skip until h2_bundle.go is updated")
+ }
+ cst := newClientServerTest(t, mode, HandlerFunc(func(w ResponseWriter, r *Request) {
+ ctl := NewResponseController(w)
+ w.Write([]byte("one"))
+ if err := ctl.Flush(); err != nil {
+ t.Errorf("before setting deadline: ctl.Flush() = %v, want nil", err)
+ }
+ if err := ctl.SetWriteDeadline(time.Now()); err != nil {
+ t.Errorf("ctl.SetWriteDeadline() = %v, want nil", err)
+ }
+
+ w.Write([]byte("two"))
+ if err := ctl.Flush(); err == nil {
+ t.Errorf("after setting deadline: ctl.Flush() = nil, want non-nil")
+ }
+ // Connection errors are sticky, so resetting the deadline does not permit
+ // making more progress. We might want to change this in the future, but verify
+ // the current behavior for now. If we do change this, we'll want to make sure
+ // to do so only for writing the response body, not headers.
+ if err := ctl.SetWriteDeadline(time.Now().Add(1 * time.Hour)); err != nil {
+ t.Errorf("ctl.SetWriteDeadline() = %v, want nil", err)
+ }
+ w.Write([]byte("three"))
+ if err := ctl.Flush(); err == nil {
+ t.Errorf("after resetting deadline: ctl.Flush() = nil, want non-nil")
+ }
+ }))
+
+ res, err := cst.c.Get(cst.ts.URL)
+ if err != nil {
+ t.Fatalf("unexpected connection error: %v", err)
+ }
+ defer res.Body.Close()
+ b, _ := io.ReadAll(res.Body)
+ if string(b) != "one" {
+ t.Errorf("unexpected body: %q", string(b))
+ }
+}
+
+func TestResponseControllerSetFutureWriteDeadline(t *testing.T) {
+ run(t, testResponseControllerSetFutureWriteDeadline)
+}
+func testResponseControllerSetFutureWriteDeadline(t *testing.T, mode testMode) {
+ if mode == http2Mode {
+ t.Skip("skip until h2_bundle.go is updated")
+ }
+ errc := make(chan error, 1)
+ cst := newClientServerTest(t, mode, HandlerFunc(func(w ResponseWriter, r *Request) {
+ ctl := NewResponseController(w)
+ if err := ctl.SetWriteDeadline(time.Now().Add(1 * time.Millisecond)); err != nil {
+ t.Errorf("ctl.SetWriteDeadline() = %v, want nil", err)
+ }
+ _, err := io.Copy(w, neverEnding('a'))
+ errc <- err
+ }))
+
+ res, err := cst.c.Get(cst.ts.URL)
+ if err != nil {
+ t.Fatalf("unexpected connection error: %v", err)
+ }
+ defer res.Body.Close()
+ _, err = io.Copy(io.Discard, res.Body)
+ if err == nil {
+ t.Errorf("client reading from truncated request body: got nil error, want non-nil")
+ }
+ err = <-errc // io.Copy error
+ if !errors.Is(err, os.ErrDeadlineExceeded) {
+ t.Errorf("server timed out writing request body: got err %v; want os.ErrDeadlineExceeded", err)
+ }
+}
+
+func TestResponseControllerSetPastReadDeadline(t *testing.T) {
+ run(t, testResponseControllerSetPastReadDeadline)
+}
+func testResponseControllerSetPastReadDeadline(t *testing.T, mode testMode) {
+ if mode == http2Mode {
+ t.Skip("skip until h2_bundle.go is updated")
+ }
+ readc := make(chan struct{})
+ cst := newClientServerTest(t, mode, HandlerFunc(func(w ResponseWriter, r *Request) {
+ ctl := NewResponseController(w)
+ b := make([]byte, 3)
+ n, err := io.ReadFull(r.Body, b)
+ b = b[:n]
+ if err != nil || string(b) != "one" {
+ t.Errorf("before setting read deadline: Read = %v, %q, want nil, %q", err, string(b), "one")
+ return
+ }
+ if err := ctl.SetReadDeadline(time.Now()); err != nil {
+ t.Errorf("ctl.SetReadDeadline() = %v, want nil", err)
+ return
+ }
+ b, err = io.ReadAll(r.Body)
+ if err == nil || string(b) != "" {
+ t.Errorf("after setting read deadline: Read = %q, nil, want error", string(b))
+ }
+ close(readc)
+ // Connection errors are sticky, so resetting the deadline does not permit
+ // making more progress. We might want to change this in the future, but verify
+ // the current behavior for now.
+ if err := ctl.SetReadDeadline(time.Time{}); err != nil {
+ t.Errorf("ctl.SetReadDeadline() = %v, want nil", err)
+ return
+ }
+ b, err = io.ReadAll(r.Body)
+ if err == nil {
+ t.Errorf("after resetting read deadline: Read = %q, nil, want error", string(b))
+ }
+ }))
+
+ pr, pw := io.Pipe()
+ var wg sync.WaitGroup
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ pw.Write([]byte("one"))
+ <-readc
+ pw.Write([]byte("two"))
+ pw.Close()
+ }()
+ defer wg.Wait()
+ res, err := cst.c.Post(cst.ts.URL, "text/foo", pr)
+ if err == nil {
+ defer res.Body.Close()
+ }
+}
+
+func TestResponseControllerSetFutureReadDeadline(t *testing.T) {
+ run(t, testResponseControllerSetFutureReadDeadline)
+}
+func testResponseControllerSetFutureReadDeadline(t *testing.T, mode testMode) {
+ if mode == http2Mode {
+ t.Skip("skip until h2_bundle.go is updated")
+ }
+ respBody := "response body"
+ cst := newClientServerTest(t, mode, HandlerFunc(func(w ResponseWriter, req *Request) {
+ ctl := NewResponseController(w)
+ if err := ctl.SetReadDeadline(time.Now().Add(1 * time.Millisecond)); err != nil {
+ t.Errorf("ctl.SetReadDeadline() = %v, want nil", err)
+ }
+ _, err := io.Copy(io.Discard, req.Body)
+ if !errors.Is(err, os.ErrDeadlineExceeded) {
+ t.Errorf("server timed out reading request body: got err %v; want os.ErrDeadlineExceeded", err)
+ }
+ w.Write([]byte(respBody))
+ }))
+ pr, pw := io.Pipe()
+ res, err := cst.c.Post(cst.ts.URL, "text/apocryphal", pr)
+ if err != nil {
+ t.Fatal(err)
+ }
+ defer res.Body.Close()
+ got, err := io.ReadAll(res.Body)
+ if string(got) != respBody || err != nil {
+ t.Errorf("client read response body: %q, %v; want %q, nil", string(got), err, respBody)
+ }
+ pw.Close()
+}
+
+type wrapWriter struct {
+ ResponseWriter
+}
+
+func (w wrapWriter) Unwrap() ResponseWriter {
+ return w.ResponseWriter
+}
+
+func TestWrappedResponseController(t *testing.T) { run(t, testWrappedResponseController) }
+func testWrappedResponseController(t *testing.T, mode testMode) {
+ if mode == http2Mode {
+ t.Skip("skip until h2_bundle.go is updated")
+ }
+ cst := newClientServerTest(t, mode, HandlerFunc(func(w ResponseWriter, r *Request) {
+ ctl := NewResponseController(w)
+ if err := ctl.Flush(); err != nil {
+ t.Errorf("ctl.Flush() = %v, want nil", err)
+ }
+ if err := ctl.SetReadDeadline(time.Time{}); err != nil {
+ t.Errorf("ctl.SetReadDeadline() = %v, want nil", err)
+ }
+ if err := ctl.SetWriteDeadline(time.Time{}); err != nil {
+ t.Errorf("ctl.SetWriteDeadline() = %v, want nil", err)
+ }
+ }))
+ res, err := cst.c.Get(cst.ts.URL)
+ if err != nil {
+ t.Fatalf("unexpected connection error: %v", err)
+ }
+ defer res.Body.Close()
+}