Mark netstack as save and use it only in tests

- Adds a new flag which will enable netstack s/r. When the flag is not enabled,
there is no change in the existing behavior. The flag will be enabled only in
tests to verify the s/r functionality of netstack.
- Some additional fields in netstack were causing panic when netstack is
save/restored. Such fields are marked as 'save'/'nosave' accordingly to resolve
the panic.

PiperOrigin-RevId: 668566657

pkg/sentry/inet/inet.go

@@ -120,6 +120,9 @@ type Stack interface {
 	// SetPortRange sets the UDP and TCP IPv4 and IPv6 ephemeral port range
 	// (inclusive).
 	SetPortRange(start uint16, end uint16) error
+
+	// EnableSaveRestore enables netstack s/r.
+	EnableSaveRestore() error
 }
 
 // Interface contains information about a network interface.

pkg/sentry/inet/namespace.go

@@ -29,7 +29,7 @@ type Namespace struct {
 	inode *nsfs.Inode
 
 	// stack is the network stack implementation of this network namespace.
-	stack Stack `state:"nosave"`
+	stack Stack
 
 	// creator allows kernel to create new network stack for network namespaces.
 	// If nil, no networking will function if network is namespaced.

pkg/sentry/inet/test_stack.go

@@ -215,3 +215,9 @@ func (*TestStack) SetGROTimeout(NICID int32, timeout time.Duration) error {
 	// No-op.
 	return nil
 }
+
+// EnableSaveRestore implements Stack.
+func (*TestStack) EnableSaveRestore() error {
+	// No-op.
+	return nil
+}

pkg/sentry/kernel/kernel.go

@@ -767,7 +767,7 @@ func (k *Kernel) invalidateUnsavableMappings(ctx context.Context) error {
 }
 
 // LoadFrom returns a new Kernel loaded from args.
-func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pagesFile *fd.FD, timeReady chan struct{}, net inet.Stack, clocks sentrytime.Clocks, vfsOpts *vfs.CompleteRestoreOptions) error {
+func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pagesFile *fd.FD, timeReady chan struct{}, net inet.Stack, clocks sentrytime.Clocks, vfsOpts *vfs.CompleteRestoreOptions, saveRestoreNet bool) error {
 	loadStart := time.Now()
 
 	var (
@@ -830,9 +830,13 @@ func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pages
 		return mfLoadErr
 	}
 
-	// rootNetworkNamespace should be populated after loading the state file.
-	// Restore the root network stack.
-	k.rootNetworkNamespace.RestoreRootStack(net)
+	if !saveRestoreNet {
+		// rootNetworkNamespace and stack should be populated after
+		// loading the state file. Reset the stack before restoring the
+		// root network stack.
+		k.rootNetworkNamespace.ResetStack()
+		k.rootNetworkNamespace.RestoreRootStack(net)
+	}
 
 	k.Timekeeper().SetClocks(clocks, k.vdsoParams)
 

pkg/sentry/socket/hostinet/stack.go

@@ -421,3 +421,8 @@ func (*Stack) PortRange() (uint16, uint16) {
 func (*Stack) SetPortRange(uint16, uint16) error {
 	return linuxerr.EACCES
 }
+
+// EnableSaveRestore implements inet.Stack.EnableSaveRestore.
+func (*Stack) EnableSaveRestore() error {
+	return fmt.Errorf("s/r is not supported for hostinet")
+}

pkg/sentry/socket/netstack/BUILD

@@ -15,6 +15,9 @@ go_library(
         "stack.go",
         "tun.go",
     ],
+    imports = [
+        "gvisor.dev/gvisor/pkg/tcpip/stack",
+    ],
     visibility = [
         "//pkg/sentry:internal",
     ],

pkg/sentry/socket/netstack/netstack_state.go

@@ -17,6 +17,8 @@ package netstack
 import (
 	"context"
 	"time"
+
+	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
 func (s *sock) saveTimestamp() int64 {
@@ -30,3 +32,17 @@ func (s *sock) loadTimestamp(_ context.Context, nsec int64) {
 	defer s.readMu.Unlock()
 	s.timestamp = time.Unix(0, nsec)
 }
+
+func (s *Stack) saveStack() *stack.Stack {
+	if s.shouldSaveRestoreStack {
+		return s.Stack
+	}
+
+	// Netstack s/r is not enabled. Do not save netstack, during
+	// restore a new stack will be configured.
+	return nil
+}
+
+func (s *Stack) loadStack(_ context.Context, st *stack.Stack) {
+	s.Stack = st
+}

pkg/sentry/socket/netstack/save_restore.go

@@ -22,6 +22,11 @@ import (
 
 // afterLoad is invoked by stateify.
 func (s *Stack) afterLoad(ctx context.Context) {
+	if s.shouldSaveRestoreStack {
+		// This indicates that netstack s/r is enabled and the stack
+		// should not be replaced with the new stack from context.
+		return
+	}
 	s.Stack = stack.RestoreStackFromContext(ctx)
 	if s.Stack == nil {
 		panic("can't restore without netstack/tcpip/stack.Stack")

pkg/sentry/socket/netstack/stack.go

@@ -40,7 +40,14 @@ import (
 //
 // +stateify savable
 type Stack struct {
-	Stack *stack.Stack `state:"manual"`
+	Stack                  *stack.Stack `state:".(*stack.Stack)"`
+	shouldSaveRestoreStack bool
+}
+
+// EnableSaveRestore enables netstack s/r.
+func (s *Stack) EnableSaveRestore() error {
+	s.shouldSaveRestoreStack = true
+	return nil
 }
 
 // Destroy implements inet.Stack.Destroy.

pkg/sentry/state/state.go

@@ -149,7 +149,7 @@ type LoadOpts struct {
 }
 
 // Load loads the given kernel, setting the provided platform and stack.
-func (opts LoadOpts) Load(ctx context.Context, k *kernel.Kernel, timeReady chan struct{}, n inet.Stack, clocks time.Clocks, vfsOpts *vfs.CompleteRestoreOptions) error {
+func (opts LoadOpts) Load(ctx context.Context, k *kernel.Kernel, timeReady chan struct{}, n inet.Stack, clocks time.Clocks, vfsOpts *vfs.CompleteRestoreOptions, saveRestoreNet bool) error {
 	// Open the file.
 	r, m, err := statefile.NewReader(opts.Source, opts.Key)
 	if err != nil {
@@ -167,5 +167,5 @@ func (opts LoadOpts) Load(ctx context.Context, k *kernel.Kernel, timeReady chan
 	previousMetadata = m
 
 	// Restore the Kernel object graph.
-	return k.LoadFrom(ctx, r, pagesMetadata, opts.PagesFile, timeReady, n, clocks, vfsOpts)
+	return k.LoadFrom(ctx, r, pagesMetadata, opts.PagesFile, timeReady, n, clocks, vfsOpts, saveRestoreNet)
 }

pkg/sleep/sleep_unsafe.go

@@ -394,6 +394,7 @@ type Waker struct {
 	allWakersNext *Waker
 }
 
+// +stateify savable
 type wakerState struct {
 	asserted bool
 	other    *Sleeper

pkg/tcpip/link/fdbased/packet_dispatchers.go

@@ -43,7 +43,7 @@ type iovecBuffer struct {
 	// (skipsVnetHdr) then the first iovec points to a buffer for the vnet header
 	// which is stripped before the views are passed up the stack for further
 	// processing.
-	iovecs []unix.Iovec
+	iovecs []unix.Iovec `state:"nosave"`
 
 	// sizes is an array of buffer sizes for the underlying views. sizes is
 	// immutable.

pkg/tcpip/stack/addressable_endpoint_state.go

@@ -41,10 +41,12 @@ type AddressableEndpointState struct {
 	// AddressableEndpointState.mu
 	//   addressState.mu
 	mu addressableEndpointStateRWMutex `state:"nosave"`
+	// TODO(b/361075310): Enable s/r for the below fields.
+	//
 	// +checklocks:mu
-	endpoints map[tcpip.Address]*addressState
+	endpoints map[tcpip.Address]*addressState `state:"nosave"`
 	// +checklocks:mu
-	primary []*addressState
+	primary []*addressState `state:"nosave"`
 }
 
 // AddressableEndpointStateOptions contains options used to configure an

runsc/boot/loader.go

@@ -247,6 +247,10 @@ type Loader struct {
 
 	// +checklocks:mu
 	saveFDs []*fd.FD
+
+	// saveRestoreNet indicates if the saved network stack should be used
+	// during restore.
+	saveRestoreNet bool
 }
 
 // execID uniquely identifies a sentry process that is executed in a container.
@@ -526,6 +530,14 @@ func New(args Args) (*Loader, error) {
 		return nil, fmt.Errorf("creating network: %w", err)
 	}
 
+	// S/R is not supported for hostinet.
+	if l.root.conf.Network != config.NetworkHost && args.Conf.TestOnlySaveRestoreNetstack {
+		l.saveRestoreNet = true
+		if err := netns.Stack().EnableSaveRestore(); err != nil {
+			return nil, fmt.Errorf("enable s/r: %w", err)
+		}
+	}
+
 	if args.NumCPU == 0 {
 		args.NumCPU = runtime.NumCPU()
 	}

runsc/boot/restore.go

@@ -222,7 +222,7 @@ func (r *restorer) restore(l *Loader) error {
 
 	// Load the state.
 	loadOpts := state.LoadOpts{Source: r.stateFile, PagesMetadata: r.pagesMetadata, PagesFile: r.pagesFile}
-	if err := loadOpts.Load(ctx, l.k, nil, oldInetStack, time.NewCalibratedClocks(), &vfs.CompleteRestoreOptions{}); err != nil {
+	if err := loadOpts.Load(ctx, l.k, nil, oldInetStack, time.NewCalibratedClocks(), &vfs.CompleteRestoreOptions{}, l.saveRestoreNet); err != nil {
 		return err
 	}
 

runsc/config/config.go

@@ -370,6 +370,9 @@ type Config struct {
 
 	// TestOnlyAutosaveResume indicates save resume for syscall tests.
 	TestOnlyAutosaveResume bool `flag:"TESTONLY-autosave-resume"`
+
+	// TestOnlySaveRestoreNetstack indicates netstack should be saved and restored.
+	TestOnlySaveRestoreNetstack bool `flag:"TESTONLY-save-restore-netstack"`
 }
 
 func (c *Config) validate() error {

runsc/config/flags.go

@@ -143,6 +143,7 @@ func RegisterFlags(flagSet *flag.FlagSet) {
 	flagSet.Bool("TESTONLY-afs-syscall-panic", false, "TEST ONLY; do not ever use! Used for tests exercising gVisor panic reporting.")
 	flagSet.String("TESTONLY-autosave-image-path", "", "TEST ONLY; enable auto save for syscall tests and set path for state file.")
 	flagSet.Bool("TESTONLY-autosave-resume", false, "TEST ONLY; enable auto save and resume for syscall tests and set path for state file.")
+	flagSet.Bool("TESTONLY-save-restore-netstack", false, "TEST ONLY; enable save/restore for netstack.")
 }
 
 // overrideAllowlist lists all flags that can be changed using OCI