Changes to support netstack save restore.

- Added a new Stats() method in inet.Stack to get the saved stats
during restore.
- Mark stack.nic, tcpip.Route and stack.addressState structs as "nosave".
These fields should not be saved because the IP addresses and routes can
change during restore and new configuration of routes and IP addresses will be
extracted from the restore spec and initialized in the saved stack.
- Changes in Restore() method in icmp, udp, tcp, packet and raw endpoint files
to support save restore of these endpoints. These changes are flag guarded by
the TESTONLY-save-restore-netstack flag.

PiperOrigin-RevId: 707639274

pkg/sentry/inet/inet.go

@@ -137,6 +137,9 @@ type Stack interface {
 
 	// IsSaveRestoreEnabled returns true when netstack s/r is enabled.
 	IsSaveRestoreEnabled() bool
+
+	// Stats returns the network stats.
+	Stats() tcpip.Stats
 }
 
 // Interface contains information about a network interface.

pkg/sentry/inet/test_stack.go

@@ -235,3 +235,9 @@ func (*TestStack) IsSaveRestoreEnabled() bool {
 	// No-op.
 	return false
 }
+
+// Stats implements Stack.
+func (*TestStack) Stats() tcpip.Stats {
+	// No-op.
+	return tcpip.Stats{}
+}

pkg/sentry/kernel/timekeeper_state.go

@@ -26,6 +26,11 @@ func (t *Timekeeper) beforeSave() {
 		panic("pauseUpdates must be called before Save")
 	}
 
+	if t.clocks == nil {
+		t.restored = nil
+		return
+	}
+
 	// N.B. we want the *offset* monotonic time.
 	var err error
 	if t.saveMonotonic, err = t.GetTime(time.Monotonic); err != nil {

pkg/sentry/socket/hostinet/stack.go

@@ -438,3 +438,8 @@ func (*Stack) EnableSaveRestore() error {
 func (s *Stack) IsSaveRestoreEnabled() bool {
 	return false
 }
+
+// Stats implements inet.Stack.Stats.
+func (s *Stack) Stats() tcpip.Stats {
+	return tcpip.Stats{}
+}

pkg/sentry/socket/netstack/stack.go

@@ -592,6 +592,7 @@ func (s *Stack) SetTCPRecovery(recovery inet.TCPLossRecovery) error {
 
 // Statistics implements inet.Stack.Statistics.
 func (s *Stack) Statistics(stat any, arg string) error {
+	netStats := s.Stats()
 	switch stats := stat.(type) {
 	case *inet.StatDev:
 		for _, ni := range s.Stack.NICInfo() {
@@ -622,7 +623,7 @@ func (s *Stack) Statistics(stat any, arg string) error {
 			break
 		}
 	case *inet.StatSNMPIP:
-		ip := Metrics.IP
+		ip := netStats.IP
 		// TODO(gvisor.dev/issue/969) Support stubbed stats.
 		*stats = inet.StatSNMPIP{
 			0,                          // Ip/Forwarding.
@@ -646,8 +647,8 @@ func (s *Stack) Statistics(stat any, arg string) error {
 			0,                               // Support Ip/FragCreates.
 		}
 	case *inet.StatSNMPICMP:
-		in := Metrics.ICMP.V4.PacketsReceived.ICMPv4PacketStats
-		out := Metrics.ICMP.V4.PacketsSent.ICMPv4PacketStats
+		in := netStats.ICMP.V4.PacketsReceived.ICMPv4PacketStats
+		out := netStats.ICMP.V4.PacketsSent.ICMPv4PacketStats
 		// TODO(gvisor.dev/issue/969) Support stubbed stats.
 		*stats = inet.StatSNMPICMP{
 			0, // Icmp/InMsgs.
@@ -679,7 +680,7 @@ func (s *Stack) Statistics(stat any, arg string) error {
 			out.InfoReply.Value(),                           // OutAddrMaskReps.
 		}
 	case *inet.StatSNMPTCP:
-		tcp := Metrics.TCP
+		tcp := netStats.TCP
 		// RFC 2012 (updates 1213):  SNMPv2-MIB-TCP.
 		*stats = inet.StatSNMPTCP{
 			1,                                     // RtoAlgorithm.
@@ -699,7 +700,7 @@ func (s *Stack) Statistics(stat any, arg string) error {
 			tcp.ChecksumErrors.Value(),            // InCsumErrors.
 		}
 	case *inet.StatSNMPUDP:
-		udp := Metrics.UDP
+		udp := netStats.UDP
 		// TODO(gvisor.dev/issue/969) Support stubbed stats.
 		*stats = inet.StatSNMPUDP{
 			udp.PacketsReceived.Value(),     // InDatagrams.
@@ -717,6 +718,11 @@ func (s *Stack) Statistics(stat any, arg string) error {
 	return nil
 }
 
+// Stats implements inet.Stack.Stats.
+func (s *Stack) Stats() tcpip.Stats {
+	return s.Stack.Stats()
+}
+
 // RouteTable implements inet.Stack.RouteTable.
 func (s *Stack) RouteTable() []inet.Route {
 	var routeTable []inet.Route

pkg/tcpip/stack/addressable_endpoint_state.go

@@ -738,8 +738,6 @@ func (a *AddressableEndpointState) Cleanup() {
 var _ AddressEndpoint = (*addressState)(nil)
 
 // addressState holds state for an address.
-//
-// +stateify savable
 type addressState struct {
 	addressableEndpointState *AddressableEndpointState
 	addr                     tcpip.AddressWithPrefix
@@ -750,7 +748,7 @@ type addressState struct {
 	//
 	// AddressableEndpointState.mu
 	//   addressState.mu
-	mu   addressStateRWMutex `state:"nosave"`
+	mu   addressStateRWMutex
 	refs addressStateRefs
 	// checklocks:mu
 	kind AddressKind

pkg/tcpip/stack/stack.go

@@ -1986,6 +1986,7 @@ func (s *Stack) ReplaceConfig(st *Stack) {
 		s.nics[id] = nic
 		_ = s.NextNICID()
 	}
+	s.tables = st.tables
 }
 
 // Restore restarts the stack after a restore. This must be called after the

pkg/tcpip/transport/icmp/endpoint.go

@@ -57,7 +57,7 @@ type endpoint struct {
 
 	// The following fields are initialized at creation time and are
 	// immutable.
-	stack       *stack.Stack `state:"manual"`
+	stack       *stack.Stack
 	transProto  tcpip.TransportProtocolNumber
 	waiterQueue *waiter.Queue
 	net         network.Endpoint

pkg/tcpip/transport/icmp/endpoint_state.go

@@ -36,7 +36,11 @@ func (p *icmpPacket) loadReceivedAt(_ context.Context, nsec int64) {
 
 // afterLoad is invoked by stateify.
 func (e *endpoint) afterLoad(ctx context.Context) {
-	stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e)
+	if e.stack.IsSaveRestoreEnabled() {
+		e.stack.RegisterRestoredEndpoint(e)
+	} else {
+		stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e)
+	}
 }
 
 // beforeSave is invoked by stateify.
@@ -50,6 +54,10 @@ func (e *endpoint) Restore(s *stack.Stack) {
 	e.thaw()
 
 	e.net.Resume(s)
+	if e.stack.IsSaveRestoreEnabled() {
+		e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits)
+		return
+	}
 
 	e.stack = s
 	e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits)

pkg/tcpip/transport/internal/network/endpoint.go

@@ -35,7 +35,7 @@ import (
 // +stateify savable
 type Endpoint struct {
 	// The following fields must only be set once then never changed.
-	stack       *stack.Stack `state:"manual"`
+	stack       *stack.Stack
 	ops         *tcpip.SocketOptions
 	netProto    tcpip.NetworkProtocolNumber
 	transProto  tcpip.TransportProtocolNumber
@@ -53,7 +53,7 @@ type Endpoint struct {
 	// +checklocks:mu
 	effectiveNetProto tcpip.NetworkProtocolNumber
 	// +checklocks:mu
-	connectedRoute *stack.Route `state:"manual"`
+	connectedRoute *stack.Route `state:"nosave"`
 	// +checklocks:mu
 	multicastMemberships map[multicastMembership]struct{}
 	// +checklocks:mu

pkg/tcpip/transport/packet/endpoint.go

@@ -63,7 +63,7 @@ type endpoint struct {
 
 	// The following fields are initialized at creation time and are
 	// immutable.
-	stack       *stack.Stack `state:"manual"`
+	stack       *stack.Stack
 	waiterQueue *waiter.Queue
 	cooked      bool
 	ops         tcpip.SocketOptions

pkg/tcpip/transport/packet/endpoint_state.go

@@ -43,12 +43,20 @@ func (ep *endpoint) beforeSave() {
 
 // afterLoad is invoked by stateify.
 func (ep *endpoint) afterLoad(ctx context.Context) {
+	if !ep.stack.IsSaveRestoreEnabled() {
+		ep.mu.Lock()
+		ep.stack = stack.RestoreStackFromContext(ctx)
+		ep.mu.Unlock()
+	}
+	ep.stack.RegisterRestoredEndpoint(ep)
+}
+
+// Restore implements tcpip.RestoredEndpoint.Restore.
+func (ep *endpoint) Restore(_ *stack.Stack) {
 	ep.mu.Lock()
 	defer ep.mu.Unlock()
 
-	ep.stack = stack.RestoreStackFromContext(ctx)
 	ep.ops.InitHandler(ep, ep.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits)
-
 	if err := ep.stack.RegisterPacketEndpoint(ep.boundNIC, ep.boundNetProto, ep); err != nil {
 		panic(fmt.Sprintf("RegisterPacketEndpoint(%d, %d, _): %s", ep.boundNIC, ep.boundNetProto, err))
 	}

pkg/tcpip/transport/raw/endpoint.go

@@ -73,7 +73,7 @@ type endpoint struct {
 
 	// The following fields are initialized at creation time and are
 	// immutable.
-	stack       *stack.Stack `state:"manual"`
+	stack       *stack.Stack
 	transProto  tcpip.TransportProtocolNumber
 	waiterQueue *waiter.Queue
 	associated  bool

pkg/tcpip/transport/raw/endpoint_state.go

@@ -16,7 +16,6 @@ package raw
 
 import (
 	"context"
-	"fmt"
 	"time"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
@@ -35,7 +34,11 @@ func (p *rawPacket) loadReceivedAt(_ context.Context, nsec int64) {
 
 // afterLoad is invoked by stateify.
 func (e *endpoint) afterLoad(ctx context.Context) {
-	stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e)
+	if e.stack.IsSaveRestoreEnabled() {
+		e.stack.RegisterRestoredEndpoint(e)
+	} else {
+		stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e)
+	}
 }
 
 // beforeSave is invoked by stateify.
@@ -46,16 +49,20 @@ func (e *endpoint) beforeSave() {
 
 // Restore implements tcpip.RestoredEndpoint.Restore.
 func (e *endpoint) Restore(s *stack.Stack) {
-	e.net.Resume(s)
-
 	e.setReceiveDisabled(false)
+	e.net.Resume(s)
+	if e.stack.IsSaveRestoreEnabled() {
+		e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits)
+		return
+	}
+
 	e.stack = s
 	e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits)
 
 	if e.associated {
 		netProto := e.net.NetProto()
 		if err := e.stack.RegisterRawTransportEndpoint(netProto, e.transProto, e); err != nil {
-			panic(fmt.Sprintf("e.stack.RegisterRawTransportEndpoint(%d, %d, _): %s", netProto, e.transProto, err))
+			panic("RegisterRawTransportEndpoint failed during restore")
 		}
 	}
 }

pkg/tcpip/transport/tcp/endpoint.go

@@ -418,7 +418,7 @@ type Endpoint struct {
 	isPortReserved    bool
 	isRegistered      bool
 	boundNICID        tcpip.NICID
-	route             *stack.Route `state:"manual"`
+	route             *stack.Route `state:"nosave"`
 	ipv4TTL           uint8
 	ipv6HopLimit      int16
 	isConnectNotified bool

pkg/tcpip/transport/tcp/endpoint_state.go

@@ -107,8 +107,9 @@ var connectingLoading sync.WaitGroup
 func (e *Endpoint) loadState(_ context.Context, epState EndpointState) {
 	// This is to ensure that the loading wait groups include all applicable
 	// endpoints before any asynchronous calls to the Wait() methods.
-	// For restore purposes we treat TimeWait like a connected endpoint.
-	if epState.connected() || epState == StateTimeWait {
+	// For restore purposes we treat all endpoints with state after
+	// StateEstablished and before StateClosed like connected endpoint.
+	if epState.connected() {
 		connectedLoading.Add(1)
 	}
 	switch {
@@ -159,21 +160,23 @@ func (e *Endpoint) Restore(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}, true /* bind */)
-		if err != nil {
-			panic("unable to parse BindAddr: " + err.String())
-		}
-		portRes := ports.Reservation{
-			Networks:     e.effectiveNetProtos,
-			Transport:    ProtocolNumber,
-			Addr:         addr.Addr,
-			Port:         addr.Port,
-			Flags:        e.boundPortFlags,
-			BindToDevice: e.boundBindToDevice,
-			Dest:         e.boundDest,
-		}
-		if ok := e.stack.ReserveTuple(portRes); !ok {
-			panic(fmt.Sprintf("unable to re-reserve tuple (%v, %q, %d, %+v, %d, %v)", e.effectiveNetProtos, addr.Addr, addr.Port, e.boundPortFlags, e.boundBindToDevice, e.boundDest))
+		if !saveRestoreEnabled {
+			addr, _, err := e.checkV4MappedLocked(tcpip.FullAddress{Addr: e.BindAddr, Port: e.TransportEndpointInfo.ID.LocalPort}, true /* bind */)
+			if err != nil {
+				panic("unable to parse BindAddr: " + err.String())
+			}
+			portRes := ports.Reservation{
+				Networks:     e.effectiveNetProtos,
+				Transport:    ProtocolNumber,
+				Addr:         addr.Addr,
+				Port:         addr.Port,
+				Flags:        e.boundPortFlags,
+				BindToDevice: e.boundBindToDevice,
+				Dest:         e.boundDest,
+			}
+			if ok := e.stack.ReserveTuple(portRes); !ok {
+				panic(fmt.Sprintf("unable to re-reserve tuple (%v, %q, %d, %+v, %d, %v)", e.effectiveNetProtos, addr.Addr, addr.Port, e.boundPortFlags, e.boundBindToDevice, e.boundDest))
+			}
 		}
 		e.isPortReserved = true
 
@@ -201,6 +204,10 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 		// Reset the scoreboard to reinitialize the sack information as
 		// we do not restore SACK information.
 		e.scoreboard.Reset()
+		if saveRestoreEnabled {
+			// Unregister the endpoint before registering again during Connect.
+			e.stack.UnregisterTransportEndpoint(e.effectiveNetProtos, header.TCPProtocolNumber, e.TransportEndpointInfo.ID, e, e.boundPortFlags, e.boundBindToDevice)
+		}
 		e.mu.Lock()
 		err := e.connect(tcpip.FullAddress{NIC: e.boundNICID, Addr: e.connectingAddress, Port: e.TransportEndpointInfo.ID.RemotePort}, false /* handshake */)
 		if _, ok := err.(*tcpip.ErrConnectStarted); !ok {
@@ -224,8 +231,8 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 		e.mu.Unlock()
 		connectedLoading.Done()
 	case epState == StateListen:
+		tcpip.AsyncLoading.Add(1)
 		if !saveRestoreEnabled {
-			tcpip.AsyncLoading.Add(1)
 			go func() {
 				connectedLoading.Wait()
 				bind()
@@ -244,14 +251,19 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 				tcpip.AsyncLoading.Done()
 			}()
 		} else {
-			e.LockUser()
-			// All endpoints will be moved to initial state after
-			// restore. Set endpoint to its originial listen state.
-			e.setEndpointState(StateListen)
-			// Initialize the listening context.
-			rcvWnd := seqnum.Size(e.receiveBufferAvailable())
-			e.listenCtx = newListenContext(e.stack, e.protocol, e, rcvWnd, e.ops.GetV6Only(), e.NetProto)
-			e.UnlockUser()
+			go func() {
+				connectedLoading.Wait()
+				e.LockUser()
+				// All endpoints will be moved to initial state after
+				// restore. Set endpoint to its originial listen state.
+				e.setEndpointState(StateListen)
+				// Initialize the listening context.
+				rcvWnd := seqnum.Size(e.receiveBufferAvailable())
+				e.listenCtx = newListenContext(e.stack, e.protocol, e, rcvWnd, e.ops.GetV6Only(), e.NetProto)
+				e.UnlockUser()
+				listenLoading.Done()
+				tcpip.AsyncLoading.Done()
+			}()
 		}
 	case epState == StateConnecting:
 		// Initial SYN hasn't been sent yet so initiate a connect.
@@ -268,26 +280,30 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 			tcpip.AsyncLoading.Done()
 		}()
 	case epState == StateSynSent || epState == StateSynRecv:
-		connectedLoading.Wait()
-		listenLoading.Wait()
-		// Initial SYN has been sent/received so we should bind the
-		// ports start the retransmit timer for the SYNs and let it
-		// naturally complete the connection.
-		bind()
-		e.mu.Lock()
-		defer e.mu.Unlock()
-		e.setEndpointState(epState)
-		r, err := e.stack.FindRoute(e.boundNICID, e.TransportEndpointInfo.ID.LocalAddress, e.TransportEndpointInfo.ID.RemoteAddress, e.effectiveNetProtos[0], false /* multicastLoop */)
-		if err != nil {
-			panic(fmt.Sprintf("FindRoute failed when restoring endpoint w/ ID: %+v", e.ID))
-		}
-		e.route = r
-		timer, err := newBackoffTimer(e.stack.Clock(), InitialRTO, MaxRTO, timerHandler(e, e.h.retransmitHandlerLocked))
-		if err != nil {
-			panic(fmt.Sprintf("newBackOffTimer(_, %s, %s, _) failed: %s", InitialRTO, MaxRTO, err))
-		}
-		e.h.retransmitTimer = timer
-		connectingLoading.Done()
+		tcpip.AsyncLoading.Add(1)
+		go func() {
+			connectedLoading.Wait()
+			listenLoading.Wait()
+			// Initial SYN has been sent/received so we should bind the
+			// ports start the retransmit timer for the SYNs and let it
+			// naturally complete the connection.
+			bind()
+			e.mu.Lock()
+			defer e.mu.Unlock()
+			e.setEndpointState(epState)
+			r, err := e.stack.FindRoute(e.boundNICID, e.TransportEndpointInfo.ID.LocalAddress, e.TransportEndpointInfo.ID.RemoteAddress, e.effectiveNetProtos[0], false /* multicastLoop */)
+			if err != nil {
+				panic(fmt.Sprintf("FindRoute failed when restoring endpoint w/ ID: %+v", e.ID))
+			}
+			e.route = r
+			timer, err := newBackoffTimer(e.stack.Clock(), InitialRTO, MaxRTO, timerHandler(e, e.h.retransmitHandlerLocked))
+			if err != nil {
+				panic(fmt.Sprintf("newBackOffTimer(_, %s, %s, _) failed: %s", InitialRTO, MaxRTO, err))
+			}
+			e.h.retransmitTimer = timer
+			connectingLoading.Done()
+			tcpip.AsyncLoading.Done()
+		}()
 	case epState == StateBound:
 		tcpip.AsyncLoading.Add(1)
 		go func() {

pkg/tcpip/transport/udp/endpoint.go

@@ -61,7 +61,7 @@ type endpoint struct {
 
 	// The following fields are initialized at creation time and do not
 	// change throughout the lifetime of the endpoint.
-	stack       *stack.Stack `state:"manual"`
+	stack       *stack.Stack
 	waiterQueue *waiter.Queue
 	net         network.Endpoint
 	stats       tcpip.TransportEndpointStats

pkg/tcpip/transport/udp/endpoint_state.go

@@ -16,7 +16,6 @@ package udp
 
 import (
 	"context"
-	"fmt"
 	"time"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
@@ -36,7 +35,11 @@ func (p *udpPacket) loadReceivedAt(_ context.Context, nsec int64) {
 
 // afterLoad is invoked by stateify.
 func (e *endpoint) afterLoad(ctx context.Context) {
-	stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e)
+	if e.stack.IsSaveRestoreEnabled() {
+		e.stack.RegisterRestoredEndpoint(e)
+	} else {
+		stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e)
+	}
 }
 
 // beforeSave is invoked by stateify.
@@ -53,7 +56,10 @@ func (e *endpoint) Restore(s *stack.Stack) {
 	defer e.mu.Unlock()
 
 	e.net.Resume(s)
-
+	if e.stack.IsSaveRestoreEnabled() {
+		e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits)
+		return
+	}
 	e.stack = s
 	e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits)
 
@@ -69,12 +75,12 @@ func (e *endpoint) Restore(s *stack.Stack) {
 		id.RemotePort = e.remotePort
 		id, e.boundBindToDevice, err = e.registerWithStack(e.effectiveNetProtos, id)
 		if err != nil {
-			panic(err)
+			panic("registering udp endpoint with the stack failed during restore")
 		}
 		e.localPort = id.LocalPort
 		e.remotePort = id.RemotePort
 	default:
-		panic(fmt.Sprintf("unhandled state = %s", state))
+		panic("unhandled state")
 	}
 }