Support ResponseController read and write deadlines on HTTP/2 streams

h2_server_frames_test.go

@@ -5,7 +5,10 @@ package http_test
 import (
 	"bytes"
 	"crypto/tls"
+	"errors"
+	"fmt"
 	"io"
+	"os"
 	"testing"
 	"time"
 
@@ -111,3 +114,80 @@ func TestH2ServerCancelsHandlerOnStreamError(t *testing.T) {
 		})
 	}
 }
+
+// A write deadline must release a handler whose write is blocked on flow
+// control, and cancel its request context.
+func TestH2ServerWriteDeadlineUnblocksFlowControlledWrite(t *testing.T) {
+	for _, srv := range h2Servers {
+		t.Run(srv.name, func(t *testing.T) {
+			errc, done := make(chan error, 1), make(chan struct{})
+			// The client grants no flow-control window, so the handler's write blocks.
+			c := newH2Client(t, srv.configure, func(w ResponseWriter, r *Request) {
+				defer close(done)
+				if err := NewResponseController(w).SetWriteDeadline(time.Now().Add(50 * time.Millisecond)); err != nil {
+					errc <- err
+					return
+				}
+				_, err := w.Write(make([]byte, 32<<10))
+				errc <- err
+				<-r.Context().Done()
+			}, http2.Setting{ID: http2.SettingInitialWindowSize, Val: 0})
+			c.writeRequest(1, true)
+			select {
+			case err := <-errc:
+				if err == nil || errors.Is(err, ErrNotSupported) {
+					t.Fatalf("blocked Write returned %v after the write deadline, want a write error", err)
+				}
+			case <-time.After(5 * time.Second):
+				t.Fatal("blocked Write did not return after the write deadline")
+			}
+			waitFor(t, done, "handler context was not canceled after the write deadline")
+		})
+	}
+}
+
+// A read deadline, in the future or already past, must release a handler
+// waiting for a request body that never comes.
+func TestH2ServerReadDeadline(t *testing.T) {
+	for _, srv := range h2Servers {
+		for _, d := range []time.Duration{50 * time.Millisecond, -time.Second} {
+			t.Run(fmt.Sprintf("%s/%v", srv.name, d), func(t *testing.T) {
+				errc := make(chan error, 1)
+				c := newH2Client(t, srv.configure, func(w ResponseWriter, r *Request) {
+					if err := NewResponseController(w).SetReadDeadline(time.Now().Add(d)); err != nil {
+						errc <- err
+						return
+					}
+					_, err := io.ReadAll(r.Body)
+					errc <- err
+				})
+				c.writeRequest(1, false) // the body never comes
+				select {
+				case err := <-errc:
+					if !errors.Is(err, os.ErrDeadlineExceeded) {
+						t.Fatalf("reading the body returned %v, want os.ErrDeadlineExceeded", err)
+					}
+				case <-time.After(5 * time.Second):
+					t.Fatal("reading the body did not return after the read deadline")
+				}
+			})
+		}
+	}
+}
+
+// A read deadline on a request without a body must not panic when it fires
+// (golang.org/issue/58237).
+func TestH2ServerReadDeadlineWithoutBody(t *testing.T) {
+	for _, srv := range h2Servers {
+		t.Run(srv.name, func(t *testing.T) {
+			done := make(chan struct{})
+			c := newH2Client(t, srv.configure, func(w ResponseWriter, r *Request) {
+				defer close(done)
+				NewResponseController(w).SetReadDeadline(time.Now().Add(time.Millisecond))
+				time.Sleep(10 * time.Millisecond)
+			})
+			c.writeRequest(1, true)
+			waitFor(t, done, "handler did not finish")
+		})
+	}
+}

http2/server.go

@@ -591,6 +591,7 @@ type stream struct {
 	resetQueued      bool        // RST_STREAM queued for write; set by sc.resetStream
 	gotTrailerHeader bool        // HEADER frame for trailers was seen
 	wroteHeaders     bool        // whether we wrote headers (not status 100)
+	readDeadline     *time.Timer // nil if unused
 	writeDeadline    *time.Timer // nil if unused
 
 	trailer    http.Header // accumulated trailers
@@ -905,6 +906,8 @@ func (sc *serverConn) serve() {
 				}
 			case *startPushRequest:
 				sc.startPush(v)
+			case func(*serverConn):
+				v(sc)
 			default:
 				panic(fmt.Sprintf("unexpected type %T", v))
 			}
@@ -1516,6 +1519,9 @@ func (sc *serverConn) closeStream(st *stream, err error) {
 		panic(fmt.Sprintf("invariant; can't close stream in state %v", st.state))
 	}
 	st.state = stateClosed
+	if st.readDeadline != nil {
+		st.readDeadline.Stop()
+	}
 	if st.writeDeadline != nil {
 		st.writeDeadline.Stop()
 	}
@@ -1773,6 +1779,16 @@ func (st *stream) copyTrailersToHandlerRequest() {
 	}
 }
 
+// onReadTimeout is run on its own goroutine (from time.AfterFunc)
+// when the stream's ReadTimeout has fired.
+func (st *stream) onReadTimeout() {
+	if st.body != nil {
+		// Wrap the ErrDeadlineExceeded to avoid callers depending on us
+		// returning the bare error.
+		st.body.CloseWithError(fmt.Errorf("%w", os.ErrDeadlineExceeded))
+	}
+}
+
 // onWriteTimeout is run on its own goroutine (from time.AfterFunc)
 // when the stream's WriteTimeout has fired.
 func (st *stream) onWriteTimeout() {
@@ -2578,6 +2594,58 @@ func (rws *responseWriterState) promoteUndeclaredTrailers() {
 	}
 }
 
+func (w *responseWriter) SetReadDeadline(deadline time.Time) error {
+	st := w.rws.stream
+	if !deadline.IsZero() && deadline.Before(time.Now()) {
+		// If we're setting a deadline in the past, reset the stream immediately
+		// so writes after SetWriteDeadline returns will fail.
+		st.onReadTimeout()
+		return nil
+	}
+	w.rws.conn.sendServeMsg(func(sc *serverConn) {
+		if st.readDeadline != nil {
+			if !st.readDeadline.Stop() {
+				// Deadline already exceeded, or stream has been closed.
+				return
+			}
+		}
+		if deadline.IsZero() {
+			st.readDeadline = nil
+		} else if st.readDeadline == nil {
+			st.readDeadline = time.AfterFunc(deadline.Sub(time.Now()), st.onReadTimeout)
+		} else {
+			st.readDeadline.Reset(deadline.Sub(time.Now()))
+		}
+	})
+	return nil
+}
+
+func (w *responseWriter) SetWriteDeadline(deadline time.Time) error {
+	st := w.rws.stream
+	if !deadline.IsZero() && deadline.Before(time.Now()) {
+		// If we're setting a deadline in the past, reset the stream immediately
+		// so writes after SetWriteDeadline returns will fail.
+		st.onWriteTimeout()
+		return nil
+	}
+	w.rws.conn.sendServeMsg(func(sc *serverConn) {
+		if st.writeDeadline != nil {
+			if !st.writeDeadline.Stop() {
+				// Deadline already exceeded, or stream has been closed.
+				return
+			}
+		}
+		if deadline.IsZero() {
+			st.writeDeadline = nil
+		} else if st.writeDeadline == nil {
+			st.writeDeadline = time.AfterFunc(deadline.Sub(time.Now()), st.onWriteTimeout)
+		} else {
+			st.writeDeadline.Reset(deadline.Sub(time.Now()))
+		}
+	})
+	return nil
+}
+
 func (w *responseWriter) Flush() {
 	rws := w.rws
 	if rws == nil {