Fix connect() not properly setting endpoint state.

If connect does not set the state to Connecting before taking over the binding
of another endpoint in TIME_WAIT, a race can happen where an segment is
processed without the endpoint having a proper state, causing a panic.

PiperOrigin-RevId: 453109764

pkg/tcpip/transport/tcp/endpoint.go

@@ -2162,73 +2162,11 @@ func (e *endpoint) Connect(addr tcpip.FullAddress) tcpip.Error {
 	return err
 }
 
-// connect connects the endpoint to its peer.
+// registerEndpoint registers the endpoint with the provided address.
+//
 // +checklocks:e.mu
-func (e *endpoint) connect(addr tcpip.FullAddress, handshake bool) tcpip.Error {
-	connectingAddr := addr.Addr
-
-	addr, netProto, err := e.checkV4MappedLocked(addr)
-	if err != nil {
-		return err
-	}
-
-	if e.EndpointState().connected() {
-		// The endpoint is already connected. If caller hasn't been
-		// notified yet, return success.
-		if !e.isConnectNotified {
-			e.isConnectNotified = true
-			return nil
-		}
-		// Otherwise return that it's already connected.
-		return &tcpip.ErrAlreadyConnected{}
-	}
-
-	nicID := addr.NIC
-	switch e.EndpointState() {
-	case StateBound:
-		// If we're already bound to a NIC but the caller is requesting
-		// that we use a different one now, we cannot proceed.
-		if e.boundNICID == 0 {
-			break
-		}
-
-		if nicID != 0 && nicID != e.boundNICID {
-			return &tcpip.ErrNoRoute{}
-		}
-
-		nicID = e.boundNICID
-
-	case StateInitial:
-		// Nothing to do. We'll eventually fill-in the gaps in the ID (if any)
-		// when we find a route.
-
-	case StateConnecting, StateSynSent, StateSynRecv:
-		// A connection request has already been issued but hasn't completed
-		// yet.
-		return &tcpip.ErrAlreadyConnecting{}
-
-	case StateError:
-		if err := e.hardErrorLocked(); err != nil {
-			return err
-		}
-		return &tcpip.ErrConnectionAborted{}
-
-	default:
-		return &tcpip.ErrInvalidEndpointState{}
-	}
-
-	// Find a route to the desired destination.
-	r, err := e.stack.FindRoute(nicID, e.TransportEndpointInfo.ID.LocalAddress, addr.Addr, netProto, false /* multicastLoop */)
-	if err != nil {
-		return err
-	}
-	defer r.Release()
-
+func (e *endpoint) registerEndpoint(addr tcpip.FullAddress, netProto tcpip.NetworkProtocolNumber, nicID tcpip.NICID) tcpip.Error {
 	netProtos := []tcpip.NetworkProtocolNumber{netProto}
-	e.TransportEndpointInfo.ID.LocalAddress = r.LocalAddress()
-	e.TransportEndpointInfo.ID.RemoteAddress = r.RemoteAddress()
-	e.TransportEndpointInfo.ID.RemotePort = addr.Port
-
 	if e.TransportEndpointInfo.ID.LocalPort != 0 {
 		// The endpoint is bound to a port, attempt to register it.
 		err := e.stack.RegisterTransportEndpoint(netProtos, ProtocolNumber, e.TransportEndpointInfo.ID, e, e.boundPortFlags, e.boundBindToDevice)
@@ -2306,7 +2244,7 @@ func (e *endpoint) connect(addr tcpip.FullAddress, handshake bool) tcpip.Error {
 				// done yet) or the reservation was freed between the check above and
 				// the FindTransportEndpoint below. But rather than retry the same port
 				// we just skip it and move on.
-				transEP := e.stack.FindTransportEndpoint(netProto, ProtocolNumber, transEPID, r.NICID())
+				transEP := e.stack.FindTransportEndpoint(netProto, ProtocolNumber, transEPID, nicID)
 				if transEP == nil {
 					// ReservePort failed but there is no registered endpoint with
 					// demuxer. Which indicates there is at least some endpoint that has
@@ -2377,13 +2315,87 @@ func (e *endpoint) connect(addr tcpip.FullAddress, handshake bool) tcpip.Error {
 			return err
 		}
 	}
+	return nil
+}
+
+// connect connects the endpoint to its peer.
+// +checklocks:e.mu
+func (e *endpoint) connect(addr tcpip.FullAddress, handshake bool) tcpip.Error {
+	connectingAddr := addr.Addr
+
+	addr, netProto, err := e.checkV4MappedLocked(addr)
+	if err != nil {
+		return err
+	}
+
+	if e.EndpointState().connected() {
+		// The endpoint is already connected. If caller hasn't been
+		// notified yet, return success.
+		if !e.isConnectNotified {
+			e.isConnectNotified = true
+			return nil
+		}
+		// Otherwise return that it's already connected.
+		return &tcpip.ErrAlreadyConnected{}
+	}
+
+	nicID := addr.NIC
+	switch e.EndpointState() {
+	case StateBound:
+		// If we're already bound to a NIC but the caller is requesting
+		// that we use a different one now, we cannot proceed.
+		if e.boundNICID == 0 {
+			break
+		}
+
+		if nicID != 0 && nicID != e.boundNICID {
+			return &tcpip.ErrNoRoute{}
+		}
+
+		nicID = e.boundNICID
+
+	case StateInitial:
+		// Nothing to do. We'll eventually fill-in the gaps in the ID (if any)
+		// when we find a route.
+
+	case StateConnecting, StateSynSent, StateSynRecv:
+		// A connection request has already been issued but hasn't completed
+		// yet.
+		return &tcpip.ErrAlreadyConnecting{}
+
+	case StateError:
+		if err := e.hardErrorLocked(); err != nil {
+			return err
+		}
+		return &tcpip.ErrConnectionAborted{}
+
+	default:
+		return &tcpip.ErrInvalidEndpointState{}
+	}
+
+	// Find a route to the desired destination.
+	r, err := e.stack.FindRoute(nicID, e.TransportEndpointInfo.ID.LocalAddress, addr.Addr, netProto, false /* multicastLoop */)
+	if err != nil {
+		return err
+	}
+	defer r.Release()
+
+	e.TransportEndpointInfo.ID.LocalAddress = r.LocalAddress()
+	e.TransportEndpointInfo.ID.RemoteAddress = r.RemoteAddress()
+	e.TransportEndpointInfo.ID.RemotePort = addr.Port
+
+	oldState := e.EndpointState()
+	e.setEndpointState(StateConnecting)
+	if err := e.registerEndpoint(addr, netProto, r.NICID()); err != nil {
+		e.setEndpointState(oldState)
+		return err
+	}
 
 	e.isRegistered = true
-	e.setEndpointState(StateConnecting)
 	r.Acquire()
 	e.route = r
 	e.boundNICID = nicID
-	e.effectiveNetProtos = netProtos
+	e.effectiveNetProtos = []tcpip.NetworkProtocolNumber{netProto}
 	e.connectingAddress = connectingAddr
 
 	e.initGSO()

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

@@ -4982,6 +4982,105 @@ func TestReusePort(t *testing.T) {
 	}
 }
 
+func TestTimeWaitAssassination(t *testing.T) {
+	var wg sync.WaitGroup
+	defer wg.Wait()
+	// We need to run this test lots of times because it triggers a very rare race
+	// condition in segment processing.
+	initalTestPort := 1024
+	testRuns := 25
+	for port := initalTestPort; port < initalTestPort+testRuns; port++ {
+		wg.Add(1)
+		go func(port uint16) {
+			defer wg.Done()
+			c := context.New(t, e2e.DefaultMTU)
+			defer c.Cleanup()
+
+			twReuse := tcpip.TCPTimeWaitReuseOption(tcpip.TCPTimeWaitReuseGlobal)
+			if err := c.Stack().SetTransportProtocolOption(tcp.ProtocolNumber, &twReuse); err != nil {
+				t.Errorf("s.TransportProtocolOption(%v, %v) = %v", tcp.ProtocolNumber, &twReuse, err)
+			}
+
+			if err := c.Stack().SetPortRange(port, port); err != nil {
+				t.Errorf("got s.SetPortRange(%d, %d) = %s, want = nil", port, port, err)
+			}
+
+			iss := seqnum.Value(context.TestInitialSequenceNumber)
+			c.CreateConnected(context.TestInitialSequenceNumber, 30000, -1)
+			c.EP.Close()
+
+			checker.IPv4(t, c.GetPacket(), checker.TCP(
+				checker.SrcPort(port),
+				checker.DstPort(context.TestPort),
+				checker.TCPSeqNum(uint32(c.IRS+1)),
+				checker.TCPAckNum(uint32(iss)+1),
+				checker.TCPFlags(header.TCPFlagFin|header.TCPFlagAck)))
+
+			finHeaders := &context.Headers{
+				SrcPort: context.TestPort,
+				DstPort: port,
+				Flags:   header.TCPFlagAck | header.TCPFlagFin,
+				SeqNum:  iss + 1,
+				AckNum:  c.IRS + 2,
+			}
+
+			c.SendPacket(nil, finHeaders)
+
+			// c.EP is in TIME_WAIT. We must allow for a second to pass before the
+			// new endpoint is allowed to take over the old endpoint's binding.
+			time.Sleep(time.Second)
+
+			seq := iss + 1
+			ack := c.IRS + 2
+
+			var wg sync.WaitGroup
+			defer wg.Wait()
+
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				// The new endpoint will take over the binding.
+				c.Create(-1)
+				timeout := time.After(5 * time.Second)
+			connect:
+				for {
+					select {
+					case <-timeout:
+						break connect
+					default:
+						err := c.EP.Connect(tcpip.FullAddress{Addr: context.TestAddr, Port: context.TestPort})
+						// It can take some extra time for the port to be available.
+						if _, ok := err.(*tcpip.ErrNoPortAvailable); ok {
+							continue connect
+						}
+						if _, ok := err.(*tcpip.ErrConnectStarted); !ok {
+							t.Errorf("Unexpected return value from Connect: %v", err)
+						}
+						break connect
+					}
+				}
+			}()
+
+			// If the new endpoint does not properly transition to connecting before
+			// taking over the port reservation, sending acks will cause the processor
+			// to panic 1-5% of the time.
+			for i := 0; i < 5; i++ {
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					c.SendPacket(nil, &context.Headers{
+						SrcPort: context.TestPort,
+						DstPort: port,
+						Flags:   header.TCPFlagAck,
+						SeqNum:  seq,
+						AckNum:  ack,
+					})
+				}()
+			}
+		}(uint16(port))
+	}
+}
+
 func checkRecvBufferSize(t *testing.T, ep tcpip.Endpoint, v int) {
 	t.Helper()