From bcf520e027eaca77663021f7ce7e21631891ab7d Mon Sep 17 00:00:00 2001
From: Nayana Bidari <nybidari@google.com>
Date: Thu, 16 Jan 2025 00:23:52 -0800
Subject: [PATCH] S/R iptables

PiperOrigin-RevId: 716114923
---
 pkg/tcpip/stack/conntrack.go      |  2 +-
 pkg/tcpip/stack/iptables.go       | 14 ++++++++------
 pkg/tcpip/stack/iptables_types.go |  2 +-
 pkg/tcpip/stack/save_restore.go   |  1 +
 pkg/tcpip/stack/stack.go          |  5 +++--
 5 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/pkg/tcpip/stack/conntrack.go b/pkg/tcpip/stack/conntrack.go
index ba11e38124..57ce22820b 100644
--- a/pkg/tcpip/stack/conntrack.go
+++ b/pkg/tcpip/stack/conntrack.go
@@ -229,7 +229,7 @@ type ConnTrack struct {
 
 	// clock provides timing used to determine conntrack reapings.
 	clock tcpip.Clock
-	// TODO(b/341946753): Restore when netstack is savable.
+
 	rand *rand.Rand `state:"nosave"`
 
 	mu connTrackRWMutex `state:"nosave"`
diff --git a/pkg/tcpip/stack/iptables.go b/pkg/tcpip/stack/iptables.go
index c80b299fb9..d759fa61bf 100644
--- a/pkg/tcpip/stack/iptables.go
+++ b/pkg/tcpip/stack/iptables.go
@@ -15,7 +15,6 @@
 package stack
 
 import (
-	"context"
 	"fmt"
 	"math/rand"
 	"reflect"
@@ -606,15 +605,18 @@ func (it *IPTables) check(table Table, hook Hook, pkt *PacketBuffer, r *Route, a
 
 // beforeSave is invoked by stateify.
 func (it *IPTables) beforeSave() {
-	// Ensure the reaper exits cleanly.
-	it.reaper.Stop()
+	if it.reaper != nil {
+		// Ensure the reaper exits cleanly.
+		it.reaper.Stop()
+	}
 	// Prevent others from modifying the connection table.
 	it.connections.mu.Lock()
 }
 
-// afterLoad is invoked by stateify.
-func (it *IPTables) afterLoad(context.Context) {
-	it.startReaper(reaperDelay)
+func (it *IPTables) resume() {
+	if it.reaper != nil {
+		it.startReaper(reaperDelay)
+	}
 }
 
 // startReaper periodically reaps timed out connections.
diff --git a/pkg/tcpip/stack/iptables_types.go b/pkg/tcpip/stack/iptables_types.go
index 0c7ce686e5..43a8fc0b7f 100644
--- a/pkg/tcpip/stack/iptables_types.go
+++ b/pkg/tcpip/stack/iptables_types.go
@@ -82,7 +82,7 @@ const (
 type IPTables struct {
 	connections ConnTrack
 
-	reaper tcpip.Timer
+	reaper tcpip.Timer `state:"nosave"`
 
 	mu ipTablesRWMutex `state:"nosave"`
 	// v4Tables and v6tables map tableIDs to tables. They hold builtin
diff --git a/pkg/tcpip/stack/save_restore.go b/pkg/tcpip/stack/save_restore.go
index 838cf5f4fd..2139875849 100644
--- a/pkg/tcpip/stack/save_restore.go
+++ b/pkg/tcpip/stack/save_restore.go
@@ -26,4 +26,5 @@ import (
 func (s *Stack) afterLoad(context.Context) {
 	s.insecureRNG = rand.New(rand.NewSource(time.Now().UnixNano()))
 	s.secureRNG = cryptorand.RNGFrom(cryptorand.Reader)
+	s.tables.connections.rand = s.insecureRNG
 }
diff --git a/pkg/tcpip/stack/stack.go b/pkg/tcpip/stack/stack.go
index 472b6d4b03..ac0e7168f0 100644
--- a/pkg/tcpip/stack/stack.go
+++ b/pkg/tcpip/stack/stack.go
@@ -115,8 +115,7 @@ type Stack struct {
 	handleLocal bool
 
 	// tables are the iptables packet filtering and manipulation rules.
-	// TODO(gvisor.dev/issue/4595): S/R this field.
-	tables *IPTables `state:"nosave"`
+	tables *IPTables
 
 	// restoredEndpoints is a list of endpoints that need to be restored if the
 	// stack is being restored.
@@ -2005,6 +2004,7 @@ func (s *Stack) Restore() {
 	eps := s.restoredEndpoints
 	s.restoredEndpoints = nil
 	saveRestoreEnabled := s.saveRestoreEnabled
+	s.tables.resume()
 	s.mu.Unlock()
 	for _, e := range eps {
 		e.Restore(s)
@@ -2024,6 +2024,7 @@ func (s *Stack) Resume() {
 	s.mu.Lock()
 	eps := s.resumableEndpoints
 	s.resumableEndpoints = nil
+	s.tables.resume()
 	s.mu.Unlock()
 	for _, e := range eps {
 		e.Resume()