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 golang#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>

Author: neild

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.

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
 }

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