Fix dial contexts cancelled while connections are still in use

common/dialer/tfo.go

@@ -20,6 +20,7 @@ import (
 type slowOpenConn struct {
 	dialer      *tfo.Dialer
 	ctx         context.Context
+	cancel      context.CancelFunc
 	network     string
 	destination M.Socksaddr
 	conn        atomic.Pointer[net.TCPConn]
@@ -39,9 +40,11 @@ func DialSlowContext(dialer *tfo.Dialer, ctx context.Context, network string, de
 			return dialer.Dialer.DialContext(ctx, network, destination.AddrString())
 		}
 	}
+	dialCtx, cancel := context.WithCancel(context.WithoutCancel(ctx))
 	return &slowOpenConn{
 		dialer:      dialer,
-		ctx:         ctx,
+		ctx:         dialCtx,
+		cancel:      cancel,
 		network:     network,
 		destination: destination,
 		create:      make(chan struct{}),
@@ -86,7 +89,14 @@ func (c *slowOpenConn) Write(b []byte) (n int, err error) {
 	if err != nil {
 		c.err = err
 	} else {
-		c.conn.Store(conn.(*net.TCPConn))
+		select {
+		case <-c.done:
+			conn.Close()
+			err = os.ErrClosed
+			c.err = err
+		default:
+			c.conn.Store(conn.(*net.TCPConn))
+		}
 	}
 	n = len(b)
 	close(c.create)
@@ -96,6 +106,7 @@ func (c *slowOpenConn) Write(b []byte) (n int, err error) {
 func (c *slowOpenConn) Close() error {
 	c.closeOnce.Do(func() {
 		close(c.done)
+		c.cancel()
 		conn := c.conn.Load()
 		if conn != nil {
 			conn.Close()

common/httpclient/http3_transport.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	stdTLS "crypto/tls"
 	"errors"
+	"io"
 	"net/http"
 	"sync"
 	"time"
@@ -172,41 +173,49 @@ func (t *http3FallbackTransport) roundTripHTTP3(request *http.Request) (*http.Re
 }
 
 func (t *http3FallbackTransport) roundTripHTTP3Race(request *http.Request, authority string) (*http.Response, error) {
-	ctx, cancel := context.WithCancel(request.Context())
-	defer cancel()
 	type result struct {
 		response *http.Response
 		err      error
 		h3       bool
 	}
-	results := make(chan result, 2)
-	startRoundTrip := func(request *http.Request, useH3 bool) {
-		request = request.WithContext(ctx)
-		var (
-			response *http.Response
-			err      error
-		)
-		if useH3 {
-			response, err = t.h3Transport.RoundTrip(request)
-		} else {
-			response, err = t.h2FallbackRoundTrip(request)
-		}
-		results <- result{response: response, err: err, h3: useH3}
+	var (
+		results  = make(chan result, 2)
+		cancels  []context.CancelFunc
+		received int
+	)
+	startRoundTrip := func(useH3 bool) {
+		ctx, cancel := context.WithCancel(request.Context())
+		cancels = append(cancels, cancel)
+		raceRequest := cloneRequestForRetry(request).WithContext(ctx)
+		go func() {
+			var (
+				response *http.Response
+				err      error
+			)
+			if useH3 {
+				response, err = t.h3Transport.RoundTrip(raceRequest)
+			} else {
+				response, err = t.h2FallbackRoundTrip(raceRequest)
+			}
+			results <- result{response: response, err: err, h3: useH3}
+		}()
 	}
-	goroutines := 1
-	received := 0
-	drainRemaining := func() {
-		cancel()
-		for range goroutines - received {
+	finish := func(winner int) {
+		for index, cancel := range cancels {
+			if index != winner {
+				cancel()
+			}
+		}
+		for range len(cancels) - received {
 			go func() {
 				loser := <-results
-				if loser.response != nil && loser.response.Body != nil {
+				if loser.response != nil {
 					loser.response.Body.Close()
 				}
 			}()
 		}
 	}
-	go startRoundTrip(cloneRequestForRetry(request), true)
+	startRoundTrip(true)
 	timer := time.NewTimer(t.fallbackDelay)
 	defer timer.Stop()
 	var (
@@ -216,39 +225,40 @@ func (t *http3FallbackTransport) roundTripHTTP3Race(request *http.Request, autho
 	for {
 		select {
 		case <-timer.C:
-			if goroutines == 1 {
-				goroutines++
-				go startRoundTrip(cloneRequestForRetry(request), false)
+			if len(cancels) == 1 {
+				startRoundTrip(false)
 			}
 		case raceResult := <-results:
 			received++
 			if raceResult.err == nil {
+				winner := 1
 				if raceResult.h3 {
+					winner = 0
 					t.clearH3Broken(authority)
 				}
-				drainRemaining()
+				finish(winner)
+				raceResult.response.Body = &cancelOnCloseBody{ReadCloser: raceResult.response.Body, cancel: cancels[winner]}
 				return raceResult.response, nil
 			}
 			if raceResult.h3 {
 				t.markH3Broken(authority)
 				h3Err = raceResult.err
-				if goroutines == 1 {
-					goroutines++
+				if len(cancels) == 1 {
 					if !timer.Stop() {
 						select {
 						case <-timer.C:
 						default:
 						}
 					}
-					go startRoundTrip(cloneRequestForRetry(request), false)
+					startRoundTrip(false)
 				}
 			} else {
 				fallbackErr = raceResult.err
 			}
-			if received < goroutines {
+			if received < len(cancels) {
 				continue
 			}
-			drainRemaining()
+			finish(-1)
 			switch {
 			case h3Err != nil && fallbackErr != nil:
 				return nil, E.Errors(h3Err, fallbackErr)
@@ -261,6 +271,17 @@ func (t *http3FallbackTransport) roundTripHTTP3Race(request *http.Request, autho
 	}
 }
 
+type cancelOnCloseBody struct {
+	io.ReadCloser
+	cancel context.CancelFunc
+}
+
+func (b *cancelOnCloseBody) Close() error {
+	err := b.ReadCloser.Close()
+	b.cancel()
+	return err
+}
+
 func (t *http3FallbackTransport) h2FallbackRoundTrip(request *http.Request) (*http.Response, error) {
 	if fallback, isFallback := t.h2Fallback.(*http2FallbackTransport); isFallback {
 		return fallback.roundTrip(request, true)