netstack: move TCP state to tcp package and cleanup probe

TCP probe support is implemented, unnecessarily, across the stack and and tcp
packages. It can live entirely in tcp. Additionally, it is only ever set at
initialization time, so support for dynamically adding/removing the probe isn't
necesary.

The probe is getting in the way of adding debugging for b/339664055.

PiperOrigin-RevId: 695793246

pkg/tcpip/stack/BUILD

@@ -252,7 +252,6 @@ go_library(
         "stack_mutex.go",
         "stack_options.go",
         "state_conn_mutex.go",
-        "tcp.go",
         "transport_demuxer.go",
         "transport_endpoints_mutex.go",
         "tuple_list.go",

pkg/tcpip/stack/stack.go

@@ -20,11 +20,11 @@
 package stack
 
 import (
+	"context"
 	"encoding/binary"
 	"fmt"
 	"io"
 	"math/rand"
-	"sync/atomic"
 	"time"
 
 	"golang.org/x/time/rate"
@@ -108,11 +108,6 @@ type Stack struct {
 
 	*ports.PortManager
 
-	// If not nil, then any new endpoints will have this probe function
-	// invoked everytime they receive a TCP segment.
-	// TODO(b/341946753): Restore them when netstack is savable.
-	tcpProbeFunc atomic.Value `state:"nosave"` // TCPProbeFunc
-
 	// clock is used to generate user-visible times.
 	clock tcpip.Clock
 
@@ -2139,41 +2134,6 @@ func (s *Stack) TransportProtocolInstance(num tcpip.TransportProtocolNumber) Tra
 	return nil
 }
 
-// AddTCPProbe installs a probe function that will be invoked on every segment
-// received by a given TCP endpoint. The probe function is passed a copy of the
-// TCP endpoint state before and after processing of the segment.
-//
-// NOTE: TCPProbe is added only to endpoints created after this call. Endpoints
-// created prior to this call will not call the probe function.
-//
-// Further, installing two different probes back to back can result in some
-// endpoints calling the first one and some the second one. There is no
-// guarantee provided on which probe will be invoked. Ideally this should only
-// be called once per stack.
-func (s *Stack) AddTCPProbe(probe TCPProbeFunc) {
-	s.tcpProbeFunc.Store(probe)
-}
-
-// GetTCPProbe returns the TCPProbeFunc if installed with AddTCPProbe, nil
-// otherwise.
-func (s *Stack) GetTCPProbe() TCPProbeFunc {
-	p := s.tcpProbeFunc.Load()
-	if p == nil {
-		return nil
-	}
-	return p.(TCPProbeFunc)
-}
-
-// RemoveTCPProbe removes an installed TCP probe.
-//
-// NOTE: This only ensures that endpoints created after this call do not
-// have a probe attached. Endpoints already created will continue to invoke
-// TCP probe.
-func (s *Stack) RemoveTCPProbe() {
-	// This must be TCPProbeFunc(nil) because atomic.Value.Store(nil) panics.
-	s.tcpProbeFunc.Store(TCPProbeFunc(nil))
-}
-
 // JoinGroup joins the given multicast group on the given NIC.
 func (s *Stack) JoinGroup(protocol tcpip.NetworkProtocolNumber, nicID tcpip.NICID, multicastAddr tcpip.Address) tcpip.Error {
 	s.mu.RLock()
@@ -2452,3 +2412,16 @@ func (s *Stack) IsSaveRestoreEnabled() bool {
 
 	return s.saveRestoreEnabled
 }
+
+// contextID is this package's type for context.Context.Value keys.
+type contextID int
+
+const (
+	// CtxRestoreStack is a Context.Value key for the stack to be used in restore.
+	CtxRestoreStack contextID = iota
+)
+
+// RestoreStackFromContext returns the stack to be used during restore.
+func RestoreStackFromContext(ctx context.Context) *Stack {
+	return ctx.Value(CtxRestoreStack).(*Stack)
+}

pkg/tcpip/transport/tcp/BUILD

@@ -66,6 +66,7 @@ go_library(
         "segment_state.go",
         "segment_unsafe.go",
         "snd.go",
+        "state.go",
         "tcp_endpoint_list.go",
         "tcp_segment_list.go",
         "tcp_segment_refs.go",

pkg/tcpip/transport/tcp/connect.go

@@ -1220,7 +1220,7 @@ func (e *Endpoint) handleSegmentsLocked() tcpip.Error {
 // +checklocks:e.mu
 func (e *Endpoint) probeSegmentLocked() {
 	if fn := e.probe; fn != nil {
-		var state stack.TCPEndpointState
+		var state TCPEndpointState
 		e.completeStateLocked(&state)
 		fn(&state)
 	}

pkg/tcpip/transport/tcp/cubic.go

@@ -19,7 +19,6 @@ import (
 	"time"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
-	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
 // effectivelyInfinity is an initialization value used for round-trip times
@@ -58,7 +57,7 @@ const (
 // See: https://tools.ietf.org/html/rfc8312.
 // +stateify savable
 type cubicState struct {
-	stack.TCPCubicState
+	TCPCubicState
 
 	// numCongestionEvents tracks the number of congestion events since last
 	// RTO.
@@ -72,7 +71,7 @@ type cubicState struct {
 func newCubicCC(s *sender) *cubicState {
 	now := s.ep.stack.Clock().NowMonotonic()
 	return &cubicState{
-		TCPCubicState: stack.TCPCubicState{
+		TCPCubicState: TCPCubicState{
 			T:    now,
 			Beta: 0.7,
 			C:    0.4,

pkg/tcpip/transport/tcp/cubic_test.go

@@ -40,7 +40,7 @@ func TestHyStartAckTrainOK(t *testing.T) {
 	iss := seqnum.Value(0)
 	snd := &sender{
 		ep: ep,
-		TCPSenderState: stack.TCPSenderState{
+		TCPSenderState: TCPSenderState{
 			SndUna:   iss + 1,
 			SndNxt:   iss + 1,
 			Ssthresh: InitialSsthresh,
@@ -121,7 +121,7 @@ func TestHyStartAckTrainTooSpread(t *testing.T) {
 	iss := seqnum.Value(0)
 	snd := &sender{
 		ep: ep,
-		TCPSenderState: stack.TCPSenderState{
+		TCPSenderState: TCPSenderState{
 			SndUna:   iss + 1,
 			SndNxt:   iss + 1,
 			Ssthresh: InitialSsthresh,
@@ -190,7 +190,7 @@ func TestHyStartDelayOK(t *testing.T) {
 	iss := seqnum.Value(0)
 	snd := &sender{
 		ep: ep,
-		TCPSenderState: stack.TCPSenderState{
+		TCPSenderState: TCPSenderState{
 			SndUna:   iss + 1,
 			SndNxt:   iss + 1,
 			Ssthresh: InitialSsthresh,
@@ -241,7 +241,7 @@ func TestHyStartDelay_BelowThresh(t *testing.T) {
 	iss := seqnum.Value(0)
 	snd := &sender{
 		ep: ep,
-		TCPSenderState: stack.TCPSenderState{
+		TCPSenderState: TCPSenderState{
 			SndUna:   iss + 1,
 			SndNxt:   iss + 1,
 			Ssthresh: InitialSsthresh,

pkg/tcpip/transport/tcp/endpoint.go

@@ -286,11 +286,11 @@ func (*Stats) IsEndpointStats() {}
 // +stateify savable
 type sndQueueInfo struct {
 	sndQueueMu sync.Mutex `state:"nosave"`
-	stack.TCPSndBufState
+	TCPSndBufState
 }
 
 // CloneState clones sq into other. It is not thread safe
-func (sq *sndQueueInfo) CloneState(other *stack.TCPSndBufState) {
+func (sq *sndQueueInfo) CloneState(other *TCPSndBufState) {
 	other.SndBufSize = sq.SndBufSize
 	other.SndBufUsed = sq.SndBufUsed
 	other.SndClosed = sq.SndClosed
@@ -340,7 +340,7 @@ func (sq *sndQueueInfo) CloneState(other *stack.TCPSndBufState) {
 //
 // +stateify savable
 type Endpoint struct {
-	stack.TCPEndpointStateInner
+	TCPEndpointStateInner
 	stack.TransportEndpointInfo
 	tcpip.DefaultSocketOptionsHandler
 
@@ -377,7 +377,7 @@ type Endpoint struct {
 	rcvQueueMu sync.Mutex `state:"nosave"`
 
 	// +checklocks:rcvQueueMu
-	stack.TCPRcvBufState
+	TCPRcvBufState
 
 	// rcvMemUsed tracks the total amount of memory in use by received segments
 	// held in rcvQueue, pendingRcvdSegments and the segment queue. This is used to
@@ -535,7 +535,7 @@ type Endpoint struct {
 
 	// probe if not nil is invoked on every received segment. It is passed
 	// a copy of the current state of the endpoint.
-	probe stack.TCPProbeFunc `state:"nosave"`
+	probe TCPProbeFunc `state:"nosave"`
 
 	// The following are only used to assist the restore run to re-connect.
 	connectingAddress tcpip.Address
@@ -843,7 +843,7 @@ func newEndpoint(s *stack.Stack, protocol *protocol, netProto tcpip.NetworkProto
 			TransProto: header.TCPProtocolNumber,
 		},
 		sndQueueInfo: sndQueueInfo{
-			TCPSndBufState: stack.TCPSndBufState{
+			TCPSndBufState: TCPSndBufState{
 				SndMTU: math.MaxInt32,
 			},
 		},
@@ -904,10 +904,7 @@ func newEndpoint(s *stack.Stack, protocol *protocol, netProto tcpip.NetworkProto
 		e.maxSynRetries = uint8(synRetries)
 	}
 
-	if p := s.GetTCPProbe(); p != nil {
-		e.probe = p
-	}
-
+	e.probe = protocol.probe
 	e.segmentQueue.ep = e
 
 	// TODO(https://gvisor.dev/issues/7493): Defer creating the timer until TCP connection becomes
@@ -3137,9 +3134,9 @@ func (e *Endpoint) maxOptionSize() (size int) {
 // used before invoking the probe.
 //
 // +checklocks:e.mu
-func (e *Endpoint) completeStateLocked(s *stack.TCPEndpointState) {
+func (e *Endpoint) completeStateLocked(s *TCPEndpointState) {
 	s.TCPEndpointStateInner = e.TCPEndpointStateInner
-	s.ID = stack.TCPEndpointID(e.TransportEndpointInfo.ID)
+	s.ID = TCPEndpointID(e.TransportEndpointInfo.ID)
 	s.SegTime = e.stack.Clock().NowMonotonic()
 	s.Receiver = e.rcv.TCPReceiverState
 	s.Sender = e.snd.TCPSenderState

pkg/tcpip/transport/tcp/protocol.go

@@ -109,6 +109,12 @@ type protocol struct {
 	synRetries                 uint8
 	dispatcher                 dispatcher
 
+	// probe, if not nil, will be invoked any time an endpoint receives a
+	// TCP segment.
+	//
+	// This is immutable after creation.
+	probe TCPProbeFunc `state:"nosave"`
+
 	// The following secrets are initialized once and stay unchanged after.
 	seqnumSecret   [16]byte
 	tsOffsetSecret [16]byte
@@ -520,18 +526,30 @@ func (*protocol) Parse(pkt *stack.PacketBuffer) bool {
 
 // NewProtocol returns a TCP transport protocol with Reno congestion control.
 func NewProtocol(s *stack.Stack) stack.TransportProtocol {
-	return newProtocol(s, ccReno)
+	return newProtocol(s, ccReno, nil)
+}
+
+// NewProtocolProbe returns a TCP transport protocol with Reno congestion
+// control and the given probe.
+//
+// The probe will be invoked on every segment received by TCP endpoints. The
+// probe function is passed a copy of the TCP endpoint state before and after
+// processing of the segment.
+func NewProtocolProbe(probe TCPProbeFunc) func(*stack.Stack) stack.TransportProtocol {
+	return func(s *stack.Stack) stack.TransportProtocol {
+		return newProtocol(s, ccReno, probe)
+	}
 }
 
 // NewProtocolCUBIC returns a TCP transport protocol with CUBIC congestion
 // control.
 //
 // TODO(b/345835636): Remove this and make CUBIC the default across the board.
 func NewProtocolCUBIC(s *stack.Stack) stack.TransportProtocol {
-	return newProtocol(s, ccCubic)
+	return newProtocol(s, ccCubic, nil)
 }
 
-func newProtocol(s *stack.Stack, cc string) stack.TransportProtocol {
+func newProtocol(s *stack.Stack, cc string, probe TCPProbeFunc) stack.TransportProtocol {
 	rng := s.SecureRNG()
 	var seqnumSecret [16]byte
 	var tsOffsetSecret [16]byte
@@ -567,6 +585,7 @@ func newProtocol(s *stack.Stack, cc string) stack.TransportProtocol {
 		recovery:                   tcpip.TCPRACKLossDetection,
 		seqnumSecret:               seqnumSecret,
 		tsOffsetSecret:             tsOffsetSecret,
+		probe:                      probe,
 	}
 	p.dispatcher.init(s.InsecureRNG(), runtime.GOMAXPROCS(0))
 	return &p

pkg/tcpip/transport/tcp/rack.go

@@ -19,7 +19,6 @@ import (
 
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/seqnum"
-	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
 const (
@@ -47,7 +46,7 @@ const (
 //
 // +stateify savable
 type rackControl struct {
-	stack.TCPRACKState
+	TCPRACKState
 
 	// exitedRecovery indicates if the connection is exiting loss recovery.
 	// This flag is set if the sender is leaving the recovery after

pkg/tcpip/transport/tcp/rcv.go

@@ -21,15 +21,14 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/seqnum"
-	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
 // receiver holds the state necessary to receive TCP segments and turn them
 // into a stream of bytes.
 //
 // +stateify savable
 type receiver struct {
-	stack.TCPReceiverState
+	TCPReceiverState
 	ep *Endpoint
 
 	// rcvWnd is the non-scaled receive window last advertised to the peer.
@@ -55,7 +54,7 @@ type receiver struct {
 func newReceiver(ep *Endpoint, irs seqnum.Value, rcvWnd seqnum.Size, rcvWndScale uint8) *receiver {
 	return &receiver{
 		ep: ep,
-		TCPReceiverState: stack.TCPReceiverState{
+		TCPReceiverState: TCPReceiverState{
 			RcvNxt:      irs + 1,
 			RcvAcc:      irs.Add(rcvWnd + 1),
 			RcvWndScale: rcvWndScale,

pkg/tcpip/transport/tcp/snd.go

@@ -99,7 +99,7 @@ type lossRecovery interface {
 //
 // +stateify savable
 type sender struct {
-	stack.TCPSenderState
+	TCPSenderState
 	ep *Endpoint
 
 	// lr is the loss recovery algorithm used by the sender.
@@ -187,7 +187,7 @@ type sender struct {
 type rtt struct {
 	sync.Mutex `state:"nosave"`
 
-	stack.TCPRTTState
+	TCPRTTState
 }
 
 // +checklocks:ep.mu
@@ -199,15 +199,15 @@ func newSender(ep *Endpoint, iss, irs seqnum.Value, sndWnd seqnum.Size, mss uint
 
 	s := &sender{
 		ep: ep,
-		TCPSenderState: stack.TCPSenderState{
+		TCPSenderState: TCPSenderState{
 			SndWnd:           sndWnd,
 			SndUna:           iss + 1,
 			SndNxt:           iss + 1,
 			RTTMeasureSeqNum: iss + 1,
 			LastSendTime:     ep.stack.Clock().NowMonotonic(),
 			MaxPayloadSize:   maxPayloadSize,
 			MaxSentAck:       irs + 1,
-			FastRecovery: stack.TCPFastRecoveryState{
+			FastRecovery: TCPFastRecoveryState{
 				// See: https://tools.ietf.org/html/rfc6582#section-3.2 Step 1.
 				Last:      iss,
 				HighRxt:   iss,