Netstack s/r: Configure the new network config in the loaded stack.

- Refactor the scrapping of network config in {sandbox/boot}/network.go to
make is easy to configure in the loaded stack during restore.
- Remove ReplaceConfig at the inet.Stack level.
- Restore some fields correctly in tcpip.Stack.
- Store the required network config in the loader and use it to set new
configuration in the loaded stack during restore.

This CL doesn't change any behavior when netstack s/r is disabled.
When netstack s/r is enabled:
- the network configuration will be applied to the loaded stack from the spec
instead of using the newly created stack.
- Previously the network config was applied to the new network stack and then
copied to the loaded stack in the ReplaceConfig method. Now, the network config
is directly applied to the loaded stack with the exception when XDP is enabled.

PiperOrigin-RevId: 719364530

pkg/sentry/inet/inet.go

@@ -100,14 +100,6 @@ type Stack interface {
 	// Restore restarts the network stack after restore.
 	Restore()
 
-	// ReplaceConfig replaces the new network stack configuration to the
-	// loaded or saved network stack after restore.
-	// TODO(b/379115439): This method is a workaround to update netstack config
-	// during restore. It should be removed after a new method is added to
-	// extract the complete config from the spec and update it in the loaded
-	// stack during restore.
-	ReplaceConfig(st Stack)
-
 	// Destroy the network stack.
 	Destroy()
 

pkg/sentry/inet/test_stack.go

@@ -175,9 +175,6 @@ func (s *TestStack) Pause() {}
 // Restore implements Stack.
 func (s *TestStack) Restore() {}
 
-// ReplaceConfig implements Stack.
-func (s *TestStack) ReplaceConfig(_ Stack) {}
-
 // Resume implements Stack.
 func (s *TestStack) Resume() {}
 

pkg/sentry/kernel/kernel.go

@@ -79,7 +79,6 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 	"gvisor.dev/gvisor/pkg/state"
 	"gvisor.dev/gvisor/pkg/sync"
-	"gvisor.dev/gvisor/pkg/tcpip"
 )
 
 // IOUringEnabled is set to true when IO_URING is enabled. Added as a global to
@@ -836,24 +835,12 @@ func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pages
 
 	if saveRestoreNet {
 		log.Infof("netstack save restore is enabled")
-		s := k.rootNetworkNamespace.Stack()
-		if s == nil {
-			panic("inet.Stack cannot be nil when netstack s/r is enabled")
-		}
-		if net != nil {
-			s.ReplaceConfig(net)
-		}
-		s.Restore()
-	} else if net != nil {
-		net.Restore()
 	}
 
 	if err := k.vfs.CompleteRestore(ctx, vfsOpts); err != nil {
 		return vfs.PrependErrMsg("vfs.CompleteRestore() failed", err)
 	}
 
-	tcpip.AsyncLoading.Wait()
-
 	log.Infof("Overall load took [%s] after async work", time.Since(loadStart))
 
 	// Applications may size per-cpu structures based on k.applicationCores, so

pkg/sentry/socket/hostinet/stack.go

@@ -398,9 +398,6 @@ func (*Stack) Pause() {}
 // Restore implements inet.Stack.Restore.
 func (*Stack) Restore() {}
 
-// ReplaceConfig implements inet.Stack.ReplaceConfig.
-func (s *Stack) ReplaceConfig(_ inet.Stack) {}
-
 // Resume implements inet.Stack.Resume.
 func (*Stack) Resume() {}
 

pkg/sentry/socket/netstack/stack.go

@@ -23,6 +23,7 @@ import (
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/refs"
 	"gvisor.dev/gvisor/pkg/sentry/inet"
+	"gvisor.dev/gvisor/pkg/sentry/socket/netfilter"
 	"gvisor.dev/gvisor/pkg/sentry/socket/netlink/nlmsg"
 	"gvisor.dev/gvisor/pkg/syserr"
 	"gvisor.dev/gvisor/pkg/tcpip"
@@ -922,15 +923,8 @@ func (s *Stack) Pause() {
 
 // Restore implements inet.Stack.Restore.
 func (s *Stack) Restore() {
-	s.Stack.Restore()
-}
-
-// ReplaceConfig implements inet.Stack.ReplaceConfig.
-func (s *Stack) ReplaceConfig(st inet.Stack) {
-	if _, ok := st.(*Stack); !ok {
-		panic("netstack.Stack cannot be nil when netstack s/r is enabled")
-	}
-	s.Stack.ReplaceConfig(st.(*Stack).Stack)
+	tables := netfilter.DefaultLinuxTables
+	s.Stack.Restore(tables)
 }
 
 // Resume implements inet.Stack.Resume.

pkg/tcpip/stack/save_restore.go

@@ -20,10 +20,14 @@ import (
 	"time"
 
 	cryptorand "gvisor.dev/gvisor/pkg/rand"
+	"gvisor.dev/gvisor/pkg/tcpip"
 )
 
 // afterLoad is invoked by stateify.
 func (s *Stack) afterLoad(context.Context) {
 	s.insecureRNG = rand.New(rand.NewSource(time.Now().UnixNano()))
 	s.secureRNG = cryptorand.RNGFrom(cryptorand.Reader)
+	s.mu.Lock()
+	s.nics = make(map[tcpip.NICID]*nic)
+	s.mu.Unlock()
 }

pkg/tcpip/stack/stack.go

@@ -1998,13 +1998,15 @@ func (s *Stack) ReplaceConfig(st *Stack) {
 
 // Restore restarts the stack after a restore. This must be called after the
 // entire system has been restored.
-func (s *Stack) Restore() {
+func (s *Stack) Restore(fn func(clock tcpip.Clock, rand *rand.Rand) *IPTables) {
 	// RestoredEndpoint.Restore() may call other methods on s, so we can't hold
 	// s.mu while restoring the endpoints.
 	s.mu.Lock()
 	eps := s.restoredEndpoints
 	s.restoredEndpoints = nil
 	saveRestoreEnabled := s.saveRestoreEnabled
+	s.tables = fn(s.clock, s.insecureRNG)
+	s.icmpRateLimiter = NewICMPRateLimiter(s.clock)
 	s.mu.Unlock()
 	for _, e := range eps {
 		e.Restore(s)

runsc/boot/controller.go

@@ -120,6 +120,9 @@ const (
 
 	// ContMgrContainerRuntimeState returns the runtime state of a container.
 	ContMgrContainerRuntimeState = "containerManager.ContainerRuntimeState"
+
+	// ContMgrStoreNetworkConfig stores the network configuration in the loader.
+	ContMgrStoreNetworkConfig = "containerManager.StoreNetworkConfig"
 )
 
 const (
@@ -131,6 +134,9 @@ const (
 
 	// DebugStacks collects sandbox stacks for debugging.
 	DebugStacks = "debug.Stacks"
+
+	// NetworkSetupNetwork sets up network.
+	NetworkSetupNetwork = "Network.SetupNetwork"
 )
 
 // Profiling related commands (see pprof.go for more details).
@@ -943,3 +949,9 @@ func (cm *containerManager) ContainerRuntimeState(cid *string, state *ContainerR
 	*state = cm.l.containerRuntimeState(*cid)
 	return nil
 }
+
+// StoreNetworkConfig stores the network configuration.
+func (cm *containerManager) StoreNetworkConfig(netConf *NetworkConfig, _ *struct{}) error {
+	cm.l.netConf = netConf
+	return nil
+}

runsc/boot/loader.go

@@ -255,6 +255,9 @@ type Loader struct {
 	// saveRestoreNet indicates if the saved network stack should be used
 	// during restore.
 	saveRestoreNet bool
+
+	// netConf contains the network configuration required during restore.
+	netConf *NetworkConfig
 }
 
 // execID uniquely identifies a sentry process that is executed in a container.

runsc/boot/network.go

@@ -605,3 +605,27 @@ func ipMaskToAddressMask(ipMask net.IPMask) tcpip.AddressMask {
 	addr := ipToAddress(net.IP(ipMask))
 	return tcpip.MaskFromBytes(addr.AsSlice())
 }
+
+// NetworkConfig contains network configuration.
+type NetworkConfig struct {
+	Args     *CreateLinksAndRoutesArgs
+	InitArgs *InitPluginStackArgs
+	Network  config.NetworkType
+}
+
+// SetupNetwork sets up the network during start and restore.
+func (n *Network) SetupNetwork(netConf *NetworkConfig, _ *struct{}) error {
+	switch netConf.Network {
+	case config.NetworkNone, config.NetworkSandbox:
+		if err := n.CreateLinksAndRoutes(netConf.Args, nil); err != nil {
+			return err
+		}
+	case config.NetworkPlugin:
+		if err := n.InitPluginStack(netConf.InitArgs, nil); err != nil {
+			return err
+		}
+	default:
+		return fmt.Errorf("invalid network type: %v", netConf.Network)
+	}
+	return nil
+}

runsc/boot/restore.go

@@ -40,6 +40,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 	"gvisor.dev/gvisor/pkg/sentry/watchdog"
 	"gvisor.dev/gvisor/pkg/sync"
+	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"gvisor.dev/gvisor/runsc/boot/pprof"
 	"gvisor.dev/gvisor/runsc/config"
@@ -143,6 +144,22 @@ func createNetworkStackForRestore(l *Loader) (*stack.Stack, inet.Stack) {
 	return nil, hostinet.NewStack()
 }
 
+// IsXDP returns true if the XDP mode has to be enabled in network.
+func IsXDP(conf *config.Config) bool {
+	if conf.Network != config.NetworkSandbox {
+		return false
+	}
+	switch conf.XDP.Mode {
+	case config.XDPModeOff:
+	case config.XDPModeNS:
+	case config.XDPModeRedirect, config.XDPModeTunnel:
+		return true
+	default:
+		panic("invalid XDP mode configured")
+	}
+	return false
+}
+
 func (r *restorer) restore(l *Loader) error {
 	log.Infof("Starting to restore %d containers", len(r.containers))
 
@@ -311,6 +328,41 @@ func (r *restorer) restore(l *Loader) error {
 	// Release `l.mu` before calling into callbacks.
 	cu.Clean()
 
+	if l.saveRestoreNet {
+		// TODO(b/340617793): Delete when netstack s/r is enabled by default.
+		if oldStack == nil {
+			return fmt.Errorf("invalid network config")
+		}
+
+		curNetwork := l.k.RootNetworkNamespace().Stack()
+		if eps, ok := curNetwork.(*netstack.Stack); ok {
+			// TODO(b/340617793): Configure routes and devices in the loaded stack
+			// similar to non-XDP and remove ReplaceConfig.
+			if IsXDP(l.root.conf) {
+				eps.Stack.ReplaceConfig(oldStack)
+				l.k.RootNetworkNamespace().Stack().Restore()
+			}
+			oldStack.Destroy()
+
+			n := &Network{
+				Stack:  eps.Stack,
+				Kernel: l.k,
+			}
+			if err := n.SetupNetwork(l.netConf, nil); err != nil {
+				return fmt.Errorf("SetupNetwork failed with error: %v", err)
+			}
+			l.k.RootNetworkNamespace().Stack().Restore()
+		} else {
+			// Restore the network stack with a new hostinet stack.
+			// Save/Restore is not supported for hostinet.
+			l.k.RootNetworkNamespace().RestoreRootStack(hostinet.NewStack())
+		}
+	} else {
+		// TODO(b/340617793): Delete when netstack s/r is enabled by default.
+		l.k.RootNetworkNamespace().Stack().Restore()
+	}
+	tcpip.AsyncLoading.Wait()
+
 	// r.restoreDone() signals and waits for the sandbox to start.
 	if err := r.restoreDone(); err != nil {
 		return fmt.Errorf("restorer.restoreDone callback failed: %w", err)