checklocks: fix TCP atomic state access

pkg/tcpip/transport/tcp/endpoint.go

@@ -289,14 +289,16 @@ type sndQueueInfo struct {
 	TCPSndBufState
 }
 
-// CloneState clones sq into other. It is not thread safe
+// CloneState clones sq into other, which must be exclusively owned by the caller.
+//
+// +checklocks:sq.sndQueueMu
 func (sq *sndQueueInfo) CloneState(other *TCPSndBufState) {
 	other.SndBufSize = sq.SndBufSize
 	other.SndBufUsed = sq.SndBufUsed
 	other.SndClosed = sq.SndClosed
 	other.PacketTooBigCount = sq.PacketTooBigCount
 	other.SndMTU = sq.SndMTU
-	other.AutoTuneSndBufDisabled = atomicbitops.FromUint32(sq.AutoTuneSndBufDisabled.RacyLoad())
+	other.AutoTuneSndBufDisabled.Store(sq.AutoTuneSndBufDisabled.Load())
 }
 
 // Endpoint represents a TCP endpoint. This struct serves as the interface
@@ -409,8 +411,9 @@ type Endpoint struct {
 	// methods.
 	state atomicbitops.Uint32 `state:".(EndpointState)"`
 
-	// connectionDirectionState holds current state of send and receive,
-	// accessed atomically
+	// connectionDirectionState records whether sending and receiving are closed.
+	//
+	// +checkatomic
 	connectionDirectionState atomicbitops.Uint32
 
 	// origEndpointState is only used during a restore phase to save the
@@ -3095,14 +3098,15 @@ func (e *Endpoint) maxReceiveBufferSize() int {
 	return rs.Max
 }
 
-// directionState returns the close state of send and receive part of the endpoint
+// connDirectionState returns the send and receive close state of the endpoint.
 func (e *Endpoint) connDirectionState() connDirectionState {
 	return connDirectionState(e.connectionDirectionState.Load())
 }
 
-// updateDirectionState updates the close state of send and receive part of the endpoint
-func (e *Endpoint) updateConnDirectionState(state connDirectionState) connDirectionState {
-	return connDirectionState(e.connectionDirectionState.Swap(uint32(e.connDirectionState() | state)))
+// updateConnDirectionState adds closed directions to the endpoint's state.
+// Passing connDirectionStateOpen leaves the state unchanged.
+func (e *Endpoint) updateConnDirectionState(state connDirectionState) {
+	atomicbitops.OrUint32(&e.connectionDirectionState, uint32(state))
 }
 
 // rcvWndScaleForHandshake computes the receive window scale to offer to the

pkg/tcpip/transport/tcp/state.go

@@ -417,6 +417,8 @@ type TCPSndBufState struct {
 
 	// AutoTuneSndBufDisabled indicates that the auto tuning of send buffer
 	// is disabled.
+	//
+	// +checkatomic
 	AutoTuneSndBufDisabled atomicbitops.Uint32
 }
 

pkg/tcpip/transport/tcp/test/e2e/tcp_test.go

@@ -5798,6 +5798,14 @@ func TestTCPEndpointProbe(t *testing.T) {
 	c.CreateConnected(context.TestInitialSequenceNumber, 30000, -1 /* epRcvBuf */)
 	port = c.Port // c.Port is set during CreateConnected.
 
+	// The socket option callback can run concurrently with the probe without
+	// holding the endpoint or send queue mutex.
+	var wg sync.WaitGroup
+	defer wg.Wait()
+	wg.Go(func() {
+		_ = c.EP.(*tcp.Endpoint).OnSetSendBufferSize(4096)
+	})
+
 	data := []byte{1, 2, 3}
 	iss := seqnum.Value(context.TestInitialSequenceNumber).Add(1)
 	c.SendPacket(data, &context.Headers{