Add checklocks annotations to tcp functions that should be locked.

PiperOrigin-RevId: 437100028

pkg/tcpip/transport/tcp/accept.go

@@ -181,8 +181,10 @@ func (l *listenContext) isCookieValid(id stack.TransportEndpointID, cookie seqnu
 }
 
 // createConnectingEndpoint creates a new endpoint in a connecting state, with
-// the connection parameters given by the arguments.
-func (l *listenContext) createConnectingEndpoint(s *segment, rcvdSynOpts header.TCPSynOptions, queue *waiter.Queue) (*endpoint, tcpip.Error) {
+// the connection parameters given by the arguments. The newly created endpoint
+// will be locked.
+// +checklocksacquire:n.mu
+func (l *listenContext) createConnectingEndpoint(s *segment, rcvdSynOpts header.TCPSynOptions, queue *waiter.Queue) (n *endpoint, _ tcpip.Error) {
 	// Create a new endpoint.
 	netProto := l.netProto
 	if netProto == 0 {
@@ -191,10 +193,11 @@ func (l *listenContext) createConnectingEndpoint(s *segment, rcvdSynOpts header.
 
 	route, err := l.stack.FindRoute(s.nicID, s.dstAddr, s.srcAddr, s.netProto, false /* multicastLoop */)
 	if err != nil {
-		return nil, err
+		return nil, err // +checklocksignore
 	}
 
-	n := newEndpoint(l.stack, l.protocol, netProto, queue)
+	n = newEndpoint(l.stack, l.protocol, netProto, queue)
+	n.mu.Lock()
 	n.ops.SetV6Only(l.v6Only)
 	n.TransportEndpointInfo.ID = s.id
 	n.boundNICID = s.nicID
@@ -225,19 +228,16 @@ func (l *listenContext) createConnectingEndpoint(s *segment, rcvdSynOpts header.
 // On success, a handshake h is returned with h.ep.mu held.
 //
 // Precondition: if l.listenEP != nil, l.listenEP.mu must be locked.
-func (l *listenContext) startHandshake(s *segment, opts header.TCPSynOptions, queue *waiter.Queue, owner tcpip.PacketOwner) (*handshake, tcpip.Error) {
+// +checklocksacquire:h.ep.mu
+func (l *listenContext) startHandshake(s *segment, opts header.TCPSynOptions, queue *waiter.Queue, owner tcpip.PacketOwner) (h *handshake, _ tcpip.Error) {
 	// Create new endpoint.
 	irs := s.sequenceNumber
 	isn := generateSecureISN(s.id, l.stack.Clock(), l.protocol.seqnumSecret)
 	ep, err := l.createConnectingEndpoint(s, opts, queue)
 	if err != nil {
-		return nil, err
+		return nil, err // +checklocksignore
 	}
 
-	// Lock the endpoint before registering to ensure that no out of
-	// band changes are possible due to incoming packets etc till
-	// the endpoint is done initializing.
-	ep.mu.Lock()
 	ep.owner = owner
 
 	// listenEP is nil when listenContext is used by tcp.Forwarder.
@@ -250,7 +250,7 @@ func (l *listenContext) startHandshake(s *segment, opts header.TCPSynOptions, qu
 			ep.mu.Unlock()
 			ep.Close()
 
-			return nil, &tcpip.ErrConnectionAborted{}
+			return nil, &tcpip.ErrConnectionAborted{} // +checklocksignore
 		}
 
 		// Propagate any inheritable options from the listening endpoint
@@ -261,7 +261,7 @@ func (l *listenContext) startHandshake(s *segment, opts header.TCPSynOptions, qu
 			ep.mu.Unlock()
 			ep.Close()
 
-			return nil, &tcpip.ErrConnectionAborted{}
+			return nil, &tcpip.ErrConnectionAborted{} // +checklocksignore
 		}
 
 		deferAccept = l.listenEP.deferAccept
@@ -281,13 +281,13 @@ func (l *listenContext) startHandshake(s *segment, opts header.TCPSynOptions, qu
 
 		ep.drainClosingSegmentQueue()
 
-		return nil, err
+		return nil, err // +checklocksignore
 	}
 
 	ep.isRegistered = true
 
 	// Initialize and start the handshake.
-	h := ep.newPassiveHandshake(isn, irs, opts, deferAccept)
+	h = ep.newPassiveHandshake(isn, irs, opts, deferAccept)
 	h.listenEP = l.listenEP
 	h.start()
 	return h, nil
@@ -661,8 +661,6 @@ func (e *endpoint) handleListenSegment(ctx *listenContext, s *segment) tcpip.Err
 			return err
 		}
 
-		n.mu.Lock()
-
 		// Propagate any inheritable options from the listening endpoint
 		// to the newly created endpoint.
 		e.propagateInheritableOptionsLocked(n)
@@ -711,6 +709,7 @@ func (e *endpoint) handleListenSegment(ctx *listenContext, s *segment) tcpip.Err
 			mss:                 rcvdSynOptions.MSS,
 			sampleRTTWithTSOnly: true,
 		}
+		h.ep.AssertLockHeld(n)
 		h.transitionToStateEstablishedLocked(s)
 
 		// Requeue the segment if the ACK completing the handshake has more info

pkg/tcpip/transport/tcp/connect.go

@@ -109,13 +109,16 @@ type handshake struct {
 	sampleRTTWithTSOnly bool
 }
 
-func (e *endpoint) newHandshake() *handshake {
-	h := &handshake{
+// +checklocks:e.mu
+// +checklocksacquire:h.ep.mu
+func (e *endpoint) newHandshake() (h *handshake) {
+	h = &handshake{
 		ep:          e,
 		active:      true,
 		rcvWnd:      seqnum.Size(e.initialReceiveWindow()),
 		rcvWndScale: e.rcvWndScaleForHandshake(),
 	}
+	h.ep.AssertLockHeld(e)
 	h.resetState()
 	// Store reference to handshake state in endpoint.
 	e.h = h
@@ -124,8 +127,10 @@ func (e *endpoint) newHandshake() *handshake {
 	return h
 }
 
-func (e *endpoint) newPassiveHandshake(isn, irs seqnum.Value, opts header.TCPSynOptions, deferAccept time.Duration) *handshake {
-	h := e.newHandshake()
+// +checklocks:e.mu
+// +checklocksacquire:h.ep.mu
+func (e *endpoint) newPassiveHandshake(isn, irs seqnum.Value, opts header.TCPSynOptions, deferAccept time.Duration) (h *handshake) {
+	h = e.newHandshake()
 	h.resetToSynRcvd(isn, irs, opts, deferAccept)
 	return h
 }
@@ -197,6 +202,7 @@ func (h *handshake) effectiveRcvWndScale() uint8 {
 
 // resetToSynRcvd resets the state of the handshake object to the SYN-RCVD
 // state.
+// +checklocks:h.ep.mu
 func (h *handshake) resetToSynRcvd(iss seqnum.Value, irs seqnum.Value, opts header.TCPSynOptions, deferAccept time.Duration) {
 	h.active = false
 	h.state = handshakeSynRcvd
@@ -227,6 +233,7 @@ func (h *handshake) checkAck(s *segment) bool {
 
 // synSentState handles a segment received when the TCP 3-way handshake is in
 // the SYN-SENT state.
+// +checklocks:h.ep.mu
 func (h *handshake) synSentState(s *segment) tcpip.Error {
 	// RFC 793, page 37, states that in the SYN-SENT state, a reset is
 	// acceptable if the ack field acknowledges the SYN.
@@ -314,6 +321,7 @@ func (h *handshake) synSentState(s *segment) tcpip.Error {
 
 // synRcvdState handles a segment received when the TCP 3-way handshake is in
 // the SYN-RCVD state.
+// +checklocks:h.ep.mu
 func (h *handshake) synRcvdState(s *segment) tcpip.Error {
 	if s.flags.Contains(header.TCPFlagRst) {
 		// RFC 793, page 37, states that in the SYN-RCVD state, a reset
@@ -424,6 +432,7 @@ func (h *handshake) synRcvdState(s *segment) tcpip.Error {
 	return nil
 }
 
+// +checklocks:h.ep.mu
 func (h *handshake) handleSegment(s *segment) tcpip.Error {
 	h.sndWnd = s.window
 	if !s.flags.Contains(header.TCPFlagSyn) && h.sndWndScale > 0 {
@@ -441,6 +450,7 @@ func (h *handshake) handleSegment(s *segment) tcpip.Error {
 
 // processSegments goes through the segment queue and processes up to
 // maxSegmentsPerWake (if they're available).
+// +checklocks:h.ep.mu
 func (h *handshake) processSegments() tcpip.Error {
 	for i := 0; i < maxSegmentsPerWake; i++ {
 		s := h.ep.segmentQueue.dequeue()
@@ -615,6 +625,7 @@ func (h *handshake) complete() tcpip.Error {
 // transitionToStateEstablisedLocked transitions the endpoint of the handshake
 // to an established state given the last segment received from peer. It also
 // initializes sender/receiver.
+// +checklocks:h.ep.mu
 func (h *handshake) transitionToStateEstablishedLocked(s *segment) {
 	// Transfer handshake state to TCP connection. We disable
 	// receive window scaling if the peer doesn't support it
@@ -950,7 +961,8 @@ func (e *endpoint) sendRaw(data buffer.VectorisedView, flags header.TCPFlags, se
 	return err
 }
 
-// Precondition: e.mu must be locked.
+// +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *endpoint) sendData(next *segment) {
 	// Initialize the next segment to write if it's currently nil.
 	if e.snd.writeNext == nil {
@@ -968,6 +980,7 @@ func (e *endpoint) sendData(next *segment) {
 // error code and sends a RST if and only if the error is not ErrConnectionReset
 // indicating that the connection is being reset due to receiving a RST. This
 // method must only be called from the protocol goroutine.
+// +checklocks:e.mu
 func (e *endpoint) resetConnectionLocked(err tcpip.Error) {
 	// Only send a reset if the connection is being aborted for a reason
 	// other than receiving a reset.
@@ -994,6 +1007,7 @@ func (e *endpoint) resetConnectionLocked(err tcpip.Error) {
 
 // completeWorkerLocked is called by the worker goroutine when it's about to
 // exit.
+// +checklocks:e.mu
 func (e *endpoint) completeWorkerLocked() {
 	// Worker is terminating(either due to moving to
 	// CLOSED or ERROR state, ensure we release all
@@ -1010,6 +1024,7 @@ func (e *endpoint) completeWorkerLocked() {
 // StateClose. This will ensure that no packet will be
 // delivered to this endpoint from the demuxer when the endpoint
 // is transitioned to StateClose.
+// +checklocks:e.mu
 func (e *endpoint) transitionToStateCloseLocked() {
 	s := e.EndpointState()
 	if s == StateClose {
@@ -1070,6 +1085,7 @@ func (e *endpoint) drainClosingSegmentQueue() {
 	}
 }
 
+// +checklocks:e.mu
 func (e *endpoint) handleReset(s *segment) (ok bool, err tcpip.Error) {
 	if e.rcv.acceptable(s.sequenceNumber, 0) {
 		// RFC 793, page 37 states that "in all states
@@ -1119,7 +1135,8 @@ func (e *endpoint) handleReset(s *segment) (ok bool, err tcpip.Error) {
 
 // handleSegments processes all inbound segments.
 //
-// Precondition: e.mu must be held.
+// +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *endpoint) handleSegmentsLocked(fastPath bool) tcpip.Error {
 	checkRequeue := true
 	for i := 0; i < maxSegmentsPerWake; i++ {
@@ -1159,7 +1176,7 @@ func (e *endpoint) handleSegmentsLocked(fastPath bool) tcpip.Error {
 	return nil
 }
 
-// Precondition: e.mu must be held.
+// +checklocks:e.mu
 func (e *endpoint) probeSegmentLocked() {
 	if fn := e.probe; fn != nil {
 		fn(e.completeStateLocked())
@@ -1169,7 +1186,9 @@ func (e *endpoint) probeSegmentLocked() {
 // handleSegment handles a given segment and notifies the worker goroutine if
 // if the connection should be terminated.
 //
-// Precondition: e.mu must be held.
+// +checklocks:e.mu
+// +checklocksalias:e.rcv.ep.mu=e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *endpoint) handleSegmentLocked(s *segment) (cont bool, err tcpip.Error) {
 	// Invoke the tcp probe if installed. The tcp probe function will update
 	// the TCPEndpointState after the segment is processed.
@@ -1243,6 +1262,8 @@ func (e *endpoint) handleSegmentLocked(s *segment) (cont bool, err tcpip.Error)
 // keepaliveTimerExpired is called when the keepaliveTimer fires. We send TCP
 // keepalive packets periodically when the connection is idle. If we don't hear
 // from the other side after a number of tries, we terminate the connection.
+// +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *endpoint) keepaliveTimerExpired() tcpip.Error {
 	userTimeout := e.userTimeout
 
@@ -1334,6 +1355,9 @@ func (e *endpoint) protocolMainLoopDone(closeTimer tcpip.Timer) {
 // handleWakeup handles a wakeup event while connected.
 //
 // +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
+// +checklocksalias:e.rcv.ep.mu=e.mu
+// +checklocksalias:e.snd.rc.snd.ep.mu=e.mu
 func (e *endpoint) handleWakeup(w, closeWaker *sleep.Waker, closeTimer *tcpip.Timer) tcpip.Error {
 	switch w {
 	case &e.sndQueueInfo.sndWaker:
@@ -1603,6 +1627,8 @@ loop:
 
 // handleTimeWaitSegments processes segments received during TIME_WAIT
 // state.
+// +checklocks:e.mu
+// +checklocksalias:e.rcv.ep.mu=e.mu
 func (e *endpoint) handleTimeWaitSegments() (extendTimeWait bool, reuseTW func()) {
 	checkRequeue := true
 	for i := 0; i < maxSegmentsPerWake; i++ {

pkg/tcpip/transport/tcp/dispatcher.go

@@ -115,7 +115,7 @@ func (p *processor) start(wg *sync.WaitGroup) {
 			// this should be fine as all normal shutdown states are handled by
 			// handleSegments and if the endpoint moves to a CLOSED/ERROR state
 			// then handleSegments is a noop.
-			if ep.EndpointState() == StateEstablished && ep.mu.TryLock() {
+			if ep.EndpointState() == StateEstablished && ep.TryLock() {
 				// If the endpoint is in a connected state then we do direct delivery
 				// to ensure low latency and avoid scheduler interactions.
 				switch err := ep.handleSegmentsLocked(true /* fastPath */); {
@@ -128,7 +128,7 @@ func (p *processor) start(wg *sync.WaitGroup) {
 				case !ep.segmentQueue.empty():
 					p.epQ.enqueue(ep)
 				}
-				ep.mu.Unlock() // +checklocksforce
+				ep.mu.Unlock()
 			} else {
 				ep.newSegmentWaker.Assert()
 			}

pkg/tcpip/transport/tcp/endpoint.go

@@ -665,7 +665,7 @@ func (e *endpoint) LockUser() {
 		// Try first if the sock is locked then check if it's owned
 		// by another user goroutine if not then we spin, otherwise
 		// we just go to sleep on the Lock() and wait.
-		if !e.mu.TryLock() {
+		if !e.TryLock() {
 			// If socket is owned by the user then just go to sleep
 			// as the lock could be held for a reasonably long time.
 			if atomic.LoadUint32(&e.ownedByUser) == 1 {
@@ -679,7 +679,7 @@ func (e *endpoint) LockUser() {
 			continue
 		}
 		atomic.StoreUint32(&e.ownedByUser, 1)
-		return // +checklocksforce
+		return
 	}
 }
 
@@ -743,11 +743,37 @@ func (e *endpoint) ResumeWork() {
 	e.mu.Unlock()
 }
 
+// AssertLockHeld forces the checklocks analyzer to consider e.mu held. This is
+// used in places where we know that e.mu is held, but checklocks does not,
+// which can happen when creating new locked objects. You must pass the known
+// locked endpoint to this function and it must be the same as the caller
+// endpoint.
+// TODO(b/226403629): Remove this function once checklocks understands local
+// variable locks.
+// +checklocks:locked.mu
+// +checklocksacquire:e.mu
+func (e *endpoint) AssertLockHeld(locked *endpoint) {
+	if e != locked {
+		panic("AssertLockHeld failed: locked endpoint != asserting endpoint")
+	}
+}
+
+// TryLock is a helper that calls TryLock on the endpoint's mutex and
+// adds the necessary checklocks annotations.
+// TODO(b/226403629): Remove this once checklocks understands TryLock.
+// +checklocksacquire:e.mu
+func (e *endpoint) TryLock() bool {
+	if e.mu.TryLock() {
+		return true // +checklocksforce
+	}
+	return false // +checklocksignore
+}
+
 // setEndpointState updates the state of the endpoint to state atomically. This
 // method is unexported as the only place we should update the state is in this
 // package but we allow the state to be read freely without holding e.mu.
 //
-// Precondition: e.mu must be held to call this method.
+// +checklocks:e.mu
 func (e *endpoint) setEndpointState(state EndpointState) {
 	oldstate := EndpointState(atomic.SwapUint32(&e.state, uint32(state)))
 	switch state {
@@ -1043,6 +1069,7 @@ func (e *endpoint) Close() {
 }
 
 // closeNoShutdown closes the endpoint without doing a full shutdown.
+// +checklocks:e.mu
 func (e *endpoint) closeNoShutdownLocked() {
 	// For listening sockets, we always release ports inline so that they
 	// are immediately available for reuse after Close() is called. If also
@@ -1124,6 +1151,7 @@ func (e *endpoint) closePendingAcceptableConnectionsLocked() {
 // cleanupLocked frees all resources associated with the endpoint. It is called
 // after Close() is called and the worker goroutine (if any) is done with its
 // work.
+// +checklocks:e.mu
 func (e *endpoint) cleanupLocked() {
 	// Close all endpoints that might have been accepted by TCP but not by
 	// the client.
@@ -1274,14 +1302,14 @@ func (e *endpoint) SetOwner(owner tcpip.PacketOwner) {
 	e.owner = owner
 }
 
-// Preconditions: e.mu must be held to call this function.
+// +checklocks:e.mu
 func (e *endpoint) hardErrorLocked() tcpip.Error {
 	err := e.hardError
 	e.hardError = nil
 	return err
 }
 
-// Preconditions: e.mu must be held to call this function.
+// +checklocks:e.mu
 func (e *endpoint) lastErrorLocked() tcpip.Error {
 	e.lastErrorMu.Lock()
 	defer e.lastErrorMu.Unlock()
@@ -1300,8 +1328,9 @@ func (e *endpoint) LastError() tcpip.Error {
 	return e.lastErrorLocked()
 }
 
-// LastErrorLocked reads and clears lastError with e.mu held.
+// LastErrorLocked reads and clears lastError.
 // Only to be used in tests.
+// +checklocks:e.mu
 func (e *endpoint) LastErrorLocked() tcpip.Error {
 	return e.lastErrorLocked()
 }
@@ -1376,7 +1405,7 @@ func (e *endpoint) Read(dst io.Writer, opts tcpip.ReadOptions) (tcpip.ReadResult
 // startRead checks that endpoint is in a readable state, and return the
 // inclusive range of segments that can be read.
 //
-// Precondition: e.rcvReadMu must be held.
+// +checklocks:e.rcvReadMu
 func (e *endpoint) startRead() (first, last *segment, err tcpip.Error) {
 	e.LockUser()
 	defer e.UnlockUser()
@@ -1427,7 +1456,7 @@ func (e *endpoint) startRead() (first, last *segment, err tcpip.Error) {
 // do this per segment read, hence this method conveniently returns the next
 // segment to read while holding the lock.
 //
-// Precondition: e.rcvReadMu must be held.
+// +checklocks:e.rcvReadMu
 func (e *endpoint) commitRead(done int) *segment {
 	e.LockUser()
 	defer e.UnlockUser()
@@ -1463,7 +1492,8 @@ func (e *endpoint) commitRead(done int) *segment {
 // and also returns the number of bytes that can be written at this
 // moment. If the endpoint is not writable then it returns an error
 // indicating the reason why it's not writable.
-// Caller must hold e.mu and e.sndQueueMu
+// +checklocks:e.mu
+// +checklocks:e.sndQueueInfo.sndQueueMu
 func (e *endpoint) isEndpointWritableLocked() (int, tcpip.Error) {
 	// The endpoint cannot be written to if it's not connected.
 	switch s := e.EndpointState(); {
@@ -1596,7 +1626,8 @@ func (e *endpoint) Write(p tcpip.Payloader, opts tcpip.WriteOptions) (int64, tcp
 
 // selectWindowLocked returns the new window without checking for shrinking or scaling
 // applied.
-// Precondition: e.mu and e.rcvQueueMu must be held.
+// +checklocks:e.mu
+// +checklocks:e.rcvQueueInfo.rcvQueueMu
 func (e *endpoint) selectWindowLocked(rcvBufSize int) (wnd seqnum.Size) {
 	wndFromAvailable := wndFromSpace(e.receiveBufferAvailableLocked(rcvBufSize))
 	maxWindow := wndFromSpace(rcvBufSize)
@@ -1619,6 +1650,7 @@ func (e *endpoint) selectWindowLocked(rcvBufSize int) (wnd seqnum.Size) {
 }
 
 // selectWindow invokes selectWindowLocked after acquiring e.rcvQueueMu.
+// +checklocks:e.mu
 func (e *endpoint) selectWindow() (wnd seqnum.Size) {
 	e.rcvQueueInfo.rcvQueueMu.Lock()
 	wnd = e.selectWindowLocked(int(e.ops.GetReceiveBufferSize()))
@@ -1640,7 +1672,8 @@ func (e *endpoint) selectWindow() (wnd seqnum.Size) {
 // above will be true if the new window is >= ACK threshold and false
 // otherwise.
 //
-// Precondition: e.mu and e.rcvQueueMu must be held.
+// +checklocks:e.mu
+// +checklocks:e.rcvQueueInfo.rcvQueueMu
 func (e *endpoint) windowCrossedACKThresholdLocked(deltaBefore int, rcvBufSize int) (crossed bool, above bool) {
 	newAvail := int(e.selectWindowLocked(rcvBufSize))
 	oldAvail := newAvail - deltaBefore
@@ -2104,6 +2137,7 @@ func (e *endpoint) GetSockOpt(opt tcpip.GettableSocketOption) tcpip.Error {
 
 // checkV4MappedLocked determines the effective network protocol and converts
 // addr to its canonical form.
+// +checklocks:e.mu
 func (e *endpoint) checkV4MappedLocked(addr tcpip.FullAddress) (tcpip.FullAddress, tcpip.NetworkProtocolNumber, tcpip.Error) {
 	unwrapped, netProto, err := e.TransportEndpointInfo.AddrNetProtoLocked(addr, e.ops.GetV6Only())
 	if err != nil {
@@ -2375,6 +2409,7 @@ func (e *endpoint) connect(addr tcpip.FullAddress, handshake bool, run bool) tcp
 			}
 		}
 		e.segmentQueue.mu.Unlock()
+		e.snd.ep.AssertLockHeld(e)
 		e.snd.updateMaxPayloadSize(int(e.route.MTU()), 0)
 		e.setEndpointState(StateEstablished)
 		// Set the new auto tuned send buffer size after entering
@@ -2421,6 +2456,7 @@ func (e *endpoint) Shutdown(flags tcpip.ShutdownFlags) tcpip.Error {
 	return e.shutdownLocked(flags)
 }
 
+// +checklocks:e.mu
 func (e *endpoint) shutdownLocked(flags tcpip.ShutdownFlags) tcpip.Error {
 	e.shutdownFlags |= flags
 	switch {
@@ -2635,6 +2671,7 @@ func (e *endpoint) Bind(addr tcpip.FullAddress) (err tcpip.Error) {
 	return e.bindLocked(addr)
 }
 
+// +checklocks:e.mu
 func (e *endpoint) bindLocked(addr tcpip.FullAddress) (err tcpip.Error) {
 	// Don't allow binding once endpoint is not in the initial state
 	// anymore. This is because once the endpoint goes into a connected or
@@ -2866,7 +2903,7 @@ func (e *endpoint) readyToRead(s *segment) {
 
 // receiveBufferAvailableLocked calculates how many bytes are still available
 // in the receive buffer.
-// rcvQueueMu must be held when this function is called.
+// +checklocks:e.rcvQueueInfo.rcvQueueMu
 func (e *endpoint) receiveBufferAvailableLocked(rcvBufSize int) int {
 	// We may use more bytes than the buffer size when the receive buffer
 	// shrinks.
@@ -2994,7 +3031,7 @@ func (e *endpoint) maxOptionSize() (size int) {
 // completeStateLocked makes a full copy of the endpoint and returns it. This is
 // used before invoking the probe.
 //
-// Precondition: e.mu must be held.
+// +checklocks:e.mu
 func (e *endpoint) completeStateLocked() stack.TCPEndpointState {
 	s := stack.TCPEndpointState{
 		TCPEndpointStateInner: e.TCPEndpointStateInner,

pkg/tcpip/transport/tcp/endpoint_state.go

@@ -193,6 +193,8 @@ func (e *endpoint) Resume(s *stack.Stack) {
 	}
 
 	bind := func() {
+		e.mu.Lock()
+		defer e.mu.Unlock()
 		addr, _, err := e.checkV4MappedLocked(tcpip.FullAddress{Addr: e.BindAddr, Port: e.TransportEndpointInfo.ID.LocalPort})
 		if err != nil {
 			panic("unable to parse BindAddr: " + err.String())

pkg/tcpip/transport/tcp/rack.go

@@ -186,6 +186,7 @@ func (s *sender) schedulePTO() {
 
 // probeTimerExpired is the same as TLP_send_probe() as defined in
 // https://tools.ietf.org/html/draft-ietf-tcpm-rack-08#section-7.5.2.
+// +checklocks:s.ep.mu
 func (s *sender) probeTimerExpired() tcpip.Error {
 	if !s.probeTimer.checkExpiration() {
 		return nil
@@ -384,6 +385,7 @@ func (rc *rackControl) detectLoss(rcvTime tcpip.MonotonicTime) int {
 
 // reorderTimerExpired will retransmit the segments which have not been acked
 // before the reorder timer expired.
+// +checklocks:rc.snd.ep.mu
 func (rc *rackControl) reorderTimerExpired() tcpip.Error {
 	// Check if the timer actually expired or if it's a spurious wake due
 	// to a previously orphaned runtime timer.
@@ -408,6 +410,7 @@ func (rc *rackControl) reorderTimerExpired() tcpip.Error {
 }
 
 // DoRecovery implements lossRecovery.DoRecovery.
+// +checklocks:rc.snd.ep.mu
 func (rc *rackControl) DoRecovery(_ *segment, fastRetransmit bool) {
 	snd := rc.snd
 	if fastRetransmit {

pkg/tcpip/transport/tcp/rcv.go

@@ -96,6 +96,7 @@ func (r *receiver) currentWindow() (curWnd seqnum.Size) {
 
 // getSendParams returns the parameters needed by the sender when building
 // segments to send.
+// +checklocks:r.ep.mu
 func (r *receiver) getSendParams() (RcvNxt seqnum.Value, rcvWnd seqnum.Size) {
 	newWnd := r.ep.selectWindow()
 	curWnd := r.currentWindow()
@@ -185,6 +186,8 @@ func (r *receiver) getSendParams() (RcvNxt seqnum.Value, rcvWnd seqnum.Size) {
 // nonZeroWindow is called when the receive window grows from zero to nonzero;
 // in such cases we may need to send an ack to indicate to our peer that it can
 // resume sending data.
+// +checklocks:r.ep.mu
+// +checklocksalias:r.ep.snd.ep.mu=r.ep.mu
 func (r *receiver) nonZeroWindow() {
 	// Immediately send an ack.
 	r.ep.snd.sendAck()
@@ -196,6 +199,8 @@ func (r *receiver) nonZeroWindow() {
 //
 // Returns true if the segment was consumed, false if it cannot be consumed
 // yet because of a missing segment.
+// +checklocks:r.ep.mu
+// +checklocksalias:r.ep.snd.ep.mu=r.ep.mu
 func (r *receiver) consumeSegment(s *segment, segSeq seqnum.Value, segLen seqnum.Size) bool {
 	if segLen > 0 {
 		// If the segment doesn't include the seqnum we're expecting to
@@ -347,6 +352,8 @@ func (r *receiver) updateRTT() {
 	r.ep.rcvQueueInfo.rcvQueueMu.Unlock()
 }
 
+// +checklocks:r.ep.mu
+// +checklocksalias:r.ep.snd.ep.mu=r.ep.mu
 func (r *receiver) handleRcvdSegmentClosing(s *segment, state EndpointState, closed bool) (drop bool, err tcpip.Error) {
 	r.ep.rcvQueueInfo.rcvQueueMu.Lock()
 	rcvClosed := r.ep.rcvQueueInfo.RcvClosed || r.closed
@@ -443,6 +450,8 @@ func (r *receiver) handleRcvdSegmentClosing(s *segment, state EndpointState, clo
 
 // handleRcvdSegment handles TCP segments directed at the connection managed by
 // r as they arrive. It is called by the protocol main loop.
+// +checklocks:r.ep.mu
+// +checklocksalias:r.ep.snd.ep.mu=r.ep.mu
 func (r *receiver) handleRcvdSegment(s *segment) (drop bool, err tcpip.Error) {
 	state := r.ep.EndpointState()
 	closed := r.ep.closed
@@ -524,6 +533,8 @@ func (r *receiver) handleRcvdSegment(s *segment) (drop bool, err tcpip.Error) {
 
 // handleTimeWaitSegment handles inbound segments received when the endpoint
 // has entered the TIME_WAIT state.
+// +checklocks:r.ep.mu
+// +checklocksalias:r.ep.snd.ep.mu=r.ep.mu
 func (r *receiver) handleTimeWaitSegment(s *segment) (resetTimeWait bool, newSyn bool) {
 	segSeq := s.sequenceNumber
 	segLen := seqnum.Size(s.data.Size())

pkg/tcpip/transport/tcp/reno_recovery.go

@@ -26,6 +26,7 @@ func newRenoRecovery(s *sender) *renoRecovery {
 	return &renoRecovery{s: s}
 }
 
+// +checklocks:rr.s.ep.mu
 func (rr *renoRecovery) DoRecovery(rcvdSeg *segment, fastRetransmit bool) {
 	ack := rcvdSeg.ackNumber
 	snd := rr.s

pkg/tcpip/transport/tcp/sack_recovery.go

@@ -30,6 +30,7 @@ func newSACKRecovery(s *sender) *sackRecovery {
 
 // handleSACKRecovery implements the loss recovery phase as described in RFC6675
 // section 5, step C.
+// +checklocks:sr.s.ep.mu
 func (sr *sackRecovery) handleSACKRecovery(limit int, end seqnum.Value) (dataSent bool) {
 	snd := sr.s
 	snd.SetPipe()
@@ -102,6 +103,7 @@ func (sr *sackRecovery) handleSACKRecovery(limit int, end seqnum.Value) (dataSen
 	return dataSent
 }
 
+// +checklocks:sr.s.ep.mu
 func (sr *sackRecovery) DoRecovery(rcvdSeg *segment, fastRetransmit bool) {
 	snd := sr.s
 	if fastRetransmit {

pkg/tcpip/transport/tcp/snd.go

@@ -165,6 +165,7 @@ type rtt struct {
 	stack.TCPRTTState
 }
 
+// +checklocks:ep.mu
 func newSender(ep *endpoint, iss, irs seqnum.Value, sndWnd seqnum.Size, mss uint16, sndWndScale int) *sender {
 	// The sender MUST reduce the TCP data length to account for any IP or
 	// TCP options that it is including in the packets that it sends.
@@ -210,8 +211,8 @@ func newSender(ep *endpoint, iss, irs seqnum.Value, sndWnd seqnum.Size, mss uint
 	s.reorderTimer.init(s.ep.stack.Clock(), &s.reorderWaker)
 	s.probeTimer.init(s.ep.stack.Clock(), &s.probeWaker)
 
+	s.ep.AssertLockHeld(ep)
 	s.updateMaxPayloadSize(int(ep.route.MTU()), 0)
-
 	// Initialize SACK Scoreboard after updating max payload size as we use
 	// the maxPayloadSize as the smss when determining if a segment is lost
 	// etc.
@@ -270,6 +271,7 @@ func (s *sender) initLossRecovery() lossRecovery {
 // MTU. If this is in response to "packet too big" control packets (indicated
 // by the count argument), it also reduces the number of outstanding packets and
 // attempts to retransmit the first packet above the MTU size.
+// +checklocks:s.ep.mu
 func (s *sender) updateMaxPayloadSize(mtu, count int) {
 	m := mtu - header.TCPMinimumSize
 
@@ -336,6 +338,7 @@ func (s *sender) updateMaxPayloadSize(mtu, count int) {
 }
 
 // sendAck sends an ACK segment.
+// +checklocks:s.ep.mu
 func (s *sender) sendAck() {
 	s.sendSegmentFromView(buffer.VectorisedView{}, header.TCPFlagAck, s.SndNxt)
 }
@@ -397,6 +400,7 @@ func (s *sender) updateRTO(rtt time.Duration) {
 }
 
 // resendSegment resends the first unacknowledged segment.
+// +checklocks:s.ep.mu
 func (s *sender) resendSegment() {
 	// Don't use any segments we already sent to measure RTT as they may
 	// have been affected by packets being lost.
@@ -427,6 +431,7 @@ func (s *sender) resendSegment() {
 // unacknowledged segments are assumed lost, and thus need to be resent.
 // Returns true if the connection is still usable, or false if the connection
 // is deemed lost.
+// +checklocks:s.ep.mu
 func (s *sender) retransmitTimerExpired() bool {
 	// Check if the timer actually expired or if it's a spurious wake due
 	// to a previously orphaned runtime timer.
@@ -717,6 +722,7 @@ func (s *sender) NextSeg(nextSegHint *segment) (nextSeg, hint *segment, rescueRt
 // other segments into this one or splits the specified segment based on the
 // lower of the specified limit value or the receivers window size specified by
 // end.
+// +checklocks:s.ep.mu
 func (s *sender) maybeSendSegment(seg *segment, limit int, end seqnum.Value) (sent bool) {
 	// We abuse the flags field to determine if we have already
 	// assigned a sequence number to this segment.
@@ -874,6 +880,8 @@ func (s *sender) maybeSendSegment(seg *segment, limit int, end seqnum.Value) (se
 	return true
 }
 
+// +checklocks:s.ep.mu
+// +checklocksalias:s.ep.rcv.ep.mu=s.ep.mu
 func (s *sender) sendZeroWindowProbe() {
 	ack, win := s.ep.rcv.getSendParams()
 	s.unackZeroWindowProbes++
@@ -937,6 +945,7 @@ func (s *sender) postXmit(dataSent bool, shouldScheduleProbe bool) {
 
 // sendData sends new data segments. It is called when data becomes available or
 // when the send window opens up.
+// +checklocks:s.ep.mu
 func (s *sender) sendData() {
 	limit := s.MaxPayloadSize
 	if s.gso {
@@ -1363,6 +1372,8 @@ func (s *sender) inRecovery() bool {
 
 // handleRcvdSegment is called when a segment is received; it is responsible for
 // updating the send-related state.
+// +checklocks:s.ep.mu
+// +checklocksalias:s.rc.snd.ep.mu=s.ep.mu
 func (s *sender) handleRcvdSegment(rcvdSeg *segment) {
 	// Check if we can extract an RTT measurement from this ack.
 	if !rcvdSeg.parsedOptions.TS && s.RTTMeasureSeqNum.LessThan(rcvdSeg.ackNumber) {
@@ -1629,6 +1640,7 @@ func (s *sender) handleRcvdSegment(rcvdSeg *segment) {
 }
 
 // sendSegment sends the specified segment.
+// +checklocks:s.ep.mu
 func (s *sender) sendSegment(seg *segment) tcpip.Error {
 	if seg.xmitCount > 0 {
 		s.ep.stack.Stats().TCP.Retransmits.Increment()
@@ -1661,6 +1673,8 @@ func (s *sender) sendSegment(seg *segment) tcpip.Error {
 
 // sendSegmentFromView sends a new segment containing the given payload, flags
 // and sequence number.
+// +checklocks:s.ep.mu
+// +checklocksalias:s.ep.rcv.ep.mu=s.ep.mu
 func (s *sender) sendSegmentFromView(data buffer.VectorisedView, flags header.TCPFlags, seq seqnum.Value) tcpip.Error {
 	s.LastSendTime = s.ep.stack.Clock().NowMonotonic()
 	if seq == s.RTTMeasureSeqNum {
@@ -1677,6 +1691,7 @@ func (s *sender) sendSegmentFromView(data buffer.VectorisedView, flags header.TC
 
 // maybeSendOutOfWindowAck sends an ACK if we are not being rate limited
 // currently.
+// +checklocks:s.ep.mu
 func (s *sender) maybeSendOutOfWindowAck(seg *segment) {
 	// Data packets are unlikely to be part of an ACK loop. So always send
 	// an ACK for a packet w/ data.