[release-branch.go1.26] crypto/tls: avoid data race when canceling a QUICConn's Context

Methods on QUICConn are synchronous:
The connection state is expected to change only in reaction
to a user calling a QUICConn method, and the state change
should finish completely before the method returns.

The connection context provided to QUICConn.Start violates
this model, because canceling the context causes an
asynchronous state change.

Prior to CL 719040, this caused no problems because canceling
the context did not cause any user-visible state changes.
In particular, canceling the context did not cause any new
events to be immediately returned by QUICConn.NextEvent.

CL 719040 introduced a new error event. Now, canceling a
QUICConn's context causes a new connection event to be
generated.

Receiving this event causes a data race visible to the
race detector, but the core problem is not the data race
itself: It's that an asynchronous event (canceling the
connection context) causes an change to the connection
events.

Fix this race by reworking the handling of QUICConn
context cancellation a bit. We no longer react to
cancellation while control of the connection lies
with the user. We only process cancellation as
part of a user call, such as QUICConn.Close
or QUICConn.HandleData.

Fixes #77274

Change-Id: If2e0f73618c4852114e0931b6bd0cb0b6a6a6964
Reviewed-on: https://go-review.googlesource.com/c/go/+/742561
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Auto-Submit: Damien Neil <dneil@google.com>
Reviewed-by: Roland Shoemaker <roland@golang.org>
(cherry picked from commit d4febb45179fa99ee1d5783bcb693ed7ba14115c)
Reviewed-on: https://go-review.googlesource.com/c/go/+/742761
TryBot-Bypass: Dmitri Shuralyov <dmitshur@golang.org>
Auto-Submit: Michael Pratt <mpratt@google.com>
Reviewed-by: Michael Pratt <mpratt@google.com>

3 files changed, 29 insertions, 13 deletions

diff --git a/src/crypto/tls/conn.go b/src/crypto/tls/conn.go
index a840125a45..9c662ef8f6 100644
--- a/src/crypto/tls/conn.go
+++ b/src/crypto/tls/conn.go
@@ -1531,7 +1531,7 @@ func (c *Conn) handshakeContext(ctx context.Context) (ret error) {
 	defer cancel()
 
 	if c.quic != nil {
-		c.quic.cancelc = handshakeCtx.Done()
+		c.quic.ctx = handshakeCtx
 		c.quic.cancel = cancel
 	} else if ctx.Done() != nil {
 		// Close the connection if ctx is canceled before the function returns.
diff --git a/src/crypto/tls/quic.go b/src/crypto/tls/quic.go
index 76b7eb2cbd..c7d8508acb 100644
--- a/src/crypto/tls/quic.go
+++ b/src/crypto/tls/quic.go
@@ -162,7 +162,7 @@ type quicState struct {
 	started  bool
 	signalc  chan struct{}   // handshake data is available to be read
 	blockedc chan struct{}   // handshake is waiting for data, closed when done
-	cancelc  <-chan struct{} // handshake has been canceled
+	ctx      context.Context // handshake context
 	cancel   context.CancelFunc
 
 	waitingForDrain bool
@@ -261,10 +261,11 @@ func (q *QUICConn) NextEvent() QUICEvent {
 
 // Close closes the connection and stops any in-progress handshake.
 func (q *QUICConn) Close() error {
-	if q.conn.quic.cancel == nil {
+	if q.conn.quic.ctx == nil {
 		return nil // never started
 	}
 	q.conn.quic.cancel()
+	<-q.conn.quic.signalc
 	for range q.conn.quic.blockedc {
 		// Wait for the handshake goroutine to return.
 	}
@@ -511,20 +512,16 @@ func (c *Conn) quicWaitForSignal() error {
 	// Send on blockedc to notify the QUICConn that the handshake is blocked.
 	// Exported methods of QUICConn wait for the handshake to become blocked
 	// before returning to the user.
-	select {
-	case c.quic.blockedc <- struct{}{}:
-	case <-c.quic.cancelc:
-		return c.sendAlertLocked(alertCloseNotify)
-	}
+	c.quic.blockedc <- struct{}{}
 	// The QUICConn reads from signalc to notify us that the handshake may
 	// be able to proceed. (The QUICConn reads, because we close signalc to
 	// indicate that the handshake has completed.)
-	select {
-	case c.quic.signalc <- struct{}{}:
-		c.hand.Write(c.quic.readbuf)
-		c.quic.readbuf = nil
-	case <-c.quic.cancelc:
+	c.quic.signalc <- struct{}{}
+	if c.quic.ctx.Err() != nil {
+		// The connection has been canceled.
 		return c.sendAlertLocked(alertCloseNotify)
 	}
+	c.hand.Write(c.quic.readbuf)
+	c.quic.readbuf = nil
 	return nil
 }
diff --git a/src/crypto/tls/quic_test.go b/src/crypto/tls/quic_test.go
index bd0eaa4d47..0cf8b337e1 100644
--- a/src/crypto/tls/quic_test.go
+++ b/src/crypto/tls/quic_test.go
@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"reflect"
 	"strings"
+	"sync"
 	"testing"
 )
 
@@ -480,6 +481,24 @@ func TestQUICStartContextPropagation(t *testing.T) {
 	}
 }
 
+func TestQUICContextCancelation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	config := &QUICConfig{TLSConfig: testConfig.Clone()}
+	config.TLSConfig.MinVersion = VersionTLS13
+	cli := newTestQUICClient(t, config)
+	cli.conn.SetTransportParameters(nil)
+	srv := newTestQUICServer(t, config)
+	srv.conn.SetTransportParameters(nil)
+	// Verify that canceling the connection context concurrently does not cause any races.
+	// See https://go.dev/issue/77274.
+	var wg sync.WaitGroup
+	wg.Go(func() {
+		_ = runTestQUICConnection(ctx, cli, srv, nil)
+	})
+	wg.Go(cancel)
+	wg.Wait()
+}
+
 func TestQUICDelayedTransportParameters(t *testing.T) {
 	clientConfig := &QUICConfig{TLSConfig: testConfig.Clone()}
 	clientConfig.TLSConfig.MinVersion = VersionTLS13