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/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.