http3: ignore context after response when using DontCloseRequestStream (#3473)

http3/body.go

@@ -110,7 +110,9 @@ func (r *hijackableBody) requestDone() {
 	if r.reqDoneClosed || r.reqDone == nil {
 		return
 	}
-	close(r.reqDone)
+	if r.reqDone != nil {
+		close(r.reqDone)
+	}
 	r.reqDoneClosed = true
 }
 

http3/client.go

@@ -273,7 +273,9 @@ func (c *client) RoundTripOpt(req *http.Request, opt RoundTripOpt) (*http.Respon
 	// This go routine keeps running even after RoundTripOpt() returns.
 	// It is shut down when the application is done processing the body.
 	reqDone := make(chan struct{})
+	done := make(chan struct{})
 	go func() {
+		defer close(done)
 		select {
 		case <-req.Context().Done():
 			str.CancelWrite(quic.StreamErrorCode(errorRequestCanceled))
@@ -282,9 +284,14 @@ func (c *client) RoundTripOpt(req *http.Request, opt RoundTripOpt) (*http.Respon
 		}
 	}()
 
-	rsp, rerr := c.doRequest(req, str, opt, reqDone)
+	doneChan := reqDone
+	if opt.DontCloseRequestStream {
+		doneChan = nil
+	}
+	rsp, rerr := c.doRequest(req, str, opt, doneChan)
 	if rerr.err != nil { // if any error occurred
 		close(reqDone)
+		<-done
 		if rerr.streamErr != 0 { // if it was a stream error
 			str.CancelWrite(quic.StreamErrorCode(rerr.streamErr))
 		}
@@ -296,6 +303,10 @@ func (c *client) RoundTripOpt(req *http.Request, opt RoundTripOpt) (*http.Respon
 			c.conn.CloseWithError(quic.ApplicationErrorCode(rerr.connErr), reason)
 		}
 	}
+	if opt.DontCloseRequestStream {
+		close(reqDone)
+		<-done
+	}
 	return rsp, rerr.err
 }
 
@@ -326,7 +337,7 @@ func (c *client) sendRequestBody(str Stream, body io.ReadCloser) error {
 	return nil
 }
 
-func (c *client) doRequest(req *http.Request, str quic.Stream, opt RoundTripOpt, reqDone chan struct{}) (*http.Response, requestError) {
+func (c *client) doRequest(req *http.Request, str quic.Stream, opt RoundTripOpt, reqDone chan<- struct{}) (*http.Response, requestError) {
 	var requestGzip bool
 	if !c.opts.DisableCompression && req.Method != "HEAD" && req.Header.Get("Accept-Encoding") == "" && req.Header.Get("Range") == "" {
 		requestGzip = true

http3/client_test.go

@@ -913,6 +913,26 @@ var _ = Describe("Client", func() {
 				cancel()
 				Eventually(done).Should(BeClosed())
 			})
+
+			It("doesn't cancel a request if DontCloseRequestStream is set", func() {
+				rspBuf := bytes.NewBuffer(getResponse(404))
+
+				ctx, cancel := context.WithCancel(context.Background())
+				req := req.WithContext(ctx)
+				conn.EXPECT().HandshakeComplete().Return(handshakeCtx)
+				conn.EXPECT().OpenStreamSync(ctx).Return(str, nil)
+				conn.EXPECT().ConnectionState().Return(quic.ConnectionState{})
+				buf := &bytes.Buffer{}
+				str.EXPECT().Close().MaxTimes(1)
+
+				str.EXPECT().Write(gomock.Any()).DoAndReturn(buf.Write)
+				str.EXPECT().Read(gomock.Any()).DoAndReturn(rspBuf.Read).AnyTimes()
+				rsp, err := client.RoundTripOpt(req, RoundTripOpt{DontCloseRequestStream: true})
+				Expect(err).ToNot(HaveOccurred())
+				cancel()
+				_, err = io.ReadAll(rsp.Body)
+				Expect(err).ToNot(HaveOccurred())
+			})
 		})
 
 		Context("gzip compression", func() {

http3/roundtrip.go

@@ -84,6 +84,7 @@ type RoundTripOpt struct {
 	// If set true and no cached connection is available, RoundTripOpt will return ErrNoCachedConn.
 	OnlyCachedConn bool
 	// DontCloseRequestStream controls whether the request stream is closed after sending the request.
+	// If set, context cancellations have no effect after the response headers are received.
 	DontCloseRequestStream bool
 }