From 740dc367dbaec5661b0b8b5bd32f8fdcaea2864b Mon Sep 17 00:00:00 2001
From: Nayana Bidari <nybidari@google.com>
Date: Wed, 28 Aug 2024 12:45:30 -0700
Subject: [PATCH] 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                       |  3 +++
 pkg/sentry/inet/namespace.go                  |  2 +-
 pkg/sentry/inet/test_stack.go                 |  6 ++++++
 pkg/sentry/kernel/kernel.go                   | 12 ++++++++----
 pkg/sentry/socket/hostinet/stack.go           |  5 +++++
 pkg/sentry/socket/netstack/BUILD              |  3 +++
 pkg/sentry/socket/netstack/netstack_state.go  | 16 ++++++++++++++++
 pkg/sentry/socket/netstack/save_restore.go    |  5 +++++
 pkg/sentry/socket/netstack/stack.go           |  9 ++++++++-
 pkg/sentry/state/state.go                     |  4 ++--
 pkg/sleep/sleep_unsafe.go                     |  1 +
 pkg/tcpip/link/fdbased/packet_dispatchers.go  |  2 +-
 pkg/tcpip/stack/addressable_endpoint_state.go |  6 ++++--
 runsc/boot/loader.go                          | 12 ++++++++++++
 runsc/boot/restore.go                         |  2 +-
 runsc/config/config.go                        |  3 +++
 runsc/config/flags.go                         |  1 +
 17 files changed, 80 insertions(+), 12 deletions(-)

diff --git a/pkg/sentry/inet/inet.go b/pkg/sentry/inet/inet.go
index 350c8ce3b..ba7cd422a 100644
--- a/pkg/sentry/inet/inet.go
+++ b/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.
diff --git a/pkg/sentry/inet/namespace.go b/pkg/sentry/inet/namespace.go
index d83082366..0ffeee03d 100644
--- a/pkg/sentry/inet/namespace.go
+++ b/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.
diff --git a/pkg/sentry/inet/test_stack.go b/pkg/sentry/inet/test_stack.go
index d5d7031f5..4e3b4ddd1 100644
--- a/pkg/sentry/inet/test_stack.go
+++ b/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
+}
diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go
index e487d8749..05e71fb77 100644
--- a/pkg/sentry/kernel/kernel.go
+++ b/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)
 
diff --git a/pkg/sentry/socket/hostinet/stack.go b/pkg/sentry/socket/hostinet/stack.go
index 7b0fea80c..f91d4a013 100644
--- a/pkg/sentry/socket/hostinet/stack.go
+++ b/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")
+}
diff --git a/pkg/sentry/socket/netstack/BUILD b/pkg/sentry/socket/netstack/BUILD
index 61d2d08ca..02efa96e0 100644
--- a/pkg/sentry/socket/netstack/BUILD
+++ b/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",
     ],
diff --git a/pkg/sentry/socket/netstack/netstack_state.go b/pkg/sentry/socket/netstack/netstack_state.go
index 57d9dc310..ff35c6077 100644
--- a/pkg/sentry/socket/netstack/netstack_state.go
+++ b/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
+}
diff --git a/pkg/sentry/socket/netstack/save_restore.go b/pkg/sentry/socket/netstack/save_restore.go
index 9a054f81c..2c730ce89 100644
--- a/pkg/sentry/socket/netstack/save_restore.go
+++ b/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")
diff --git a/pkg/sentry/socket/netstack/stack.go b/pkg/sentry/socket/netstack/stack.go
index 0071ddd34..47277c63b 100644
--- a/pkg/sentry/socket/netstack/stack.go
+++ b/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.
diff --git a/pkg/sentry/state/state.go b/pkg/sentry/state/state.go
index 45ec19905..ffd465a32 100644
--- a/pkg/sentry/state/state.go
+++ b/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)
 }
diff --git a/pkg/sleep/sleep_unsafe.go b/pkg/sleep/sleep_unsafe.go
index eab682dc0..9dcd78c09 100644
--- a/pkg/sleep/sleep_unsafe.go
+++ b/pkg/sleep/sleep_unsafe.go
@@ -394,6 +394,7 @@ type Waker struct {
 	allWakersNext *Waker
 }
 
+// +stateify savable
 type wakerState struct {
 	asserted bool
 	other    *Sleeper
diff --git a/pkg/tcpip/link/fdbased/packet_dispatchers.go b/pkg/tcpip/link/fdbased/packet_dispatchers.go
index 6e091243d..e324090a7 100644
--- a/pkg/tcpip/link/fdbased/packet_dispatchers.go
+++ b/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.
diff --git a/pkg/tcpip/stack/addressable_endpoint_state.go b/pkg/tcpip/stack/addressable_endpoint_state.go
index c0290ca6a..bb2e0faf0 100644
--- a/pkg/tcpip/stack/addressable_endpoint_state.go
+++ b/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
diff --git a/runsc/boot/loader.go b/runsc/boot/loader.go
index 94f925b8a..acec716b4 100644
--- a/runsc/boot/loader.go
+++ b/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()
 	}
diff --git a/runsc/boot/restore.go b/runsc/boot/restore.go
index a1607c1a9..ccf311e8d 100644
--- a/runsc/boot/restore.go
+++ b/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
 	}
 
diff --git a/runsc/config/config.go b/runsc/config/config.go
index 0b5bbe5ad..6ec62982d 100644
--- a/runsc/config/config.go
+++ b/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 {
diff --git a/runsc/config/flags.go b/runsc/config/flags.go
index 2fd7a8fdc..0307de226 100644
--- a/runsc/config/flags.go
+++ b/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