Turn IPTablesSaveRestore into an interface

cloudnativelabs · Jan 23, 2023 · a7e5803 · a7e5803
1 parent 033444b
commit a7e5803
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 10 deletions.
diff --git a/pkg/controllers/netpol/network_policy_controller.go b/pkg/controllers/netpol/network_policy_controller.go
@@ -77,7 +77,7 @@ type NetworkPolicyController struct {
 	ipsetMutex                     *sync.Mutex
 
 	iptablesCmdHandlers map[v1core.IPFamily]utils.IPTablesHandler
-	iptablesSaveRestore map[v1core.IPFamily]*utils.IPTablesSaveRestore
+	iptablesSaveRestore map[v1core.IPFamily]utils.IPTablesSaveRestorer
 	filterTableRules    map[v1core.IPFamily]*bytes.Buffer
 	ipSetHandlers       map[v1core.IPFamily]utils.IPSetHandler
 	nodeIPs             map[v1core.IPFamily]net.IP
@@ -778,7 +778,7 @@ func NewNetworkPolicyController(clientset kubernetes.Interface,
 	}
 
 	npc.iptablesCmdHandlers = iptablesCmdHandlers
-	npc.iptablesSaveRestore = make(map[v1core.IPFamily]*utils.IPTablesSaveRestore, 2)
+	npc.iptablesSaveRestore = make(map[v1core.IPFamily]utils.IPTablesSaveRestorer, 2)
 	npc.filterTableRules = make(map[v1core.IPFamily]*bytes.Buffer, 2)
 	npc.ipSetHandlers = ipSetHandlers
 	npc.nodeIPs = make(map[v1core.IPFamily]net.IP, 2)

diff --git a/pkg/controllers/netpol/network_policy_controller_test.go b/pkg/controllers/netpol/network_policy_controller_test.go
@@ -198,7 +198,7 @@ func newUneventfulNetworkPolicyController(podInformer cache.SharedIndexInformer,
 	npc.syncPeriod = time.Hour
 
 	npc.iptablesCmdHandlers = make(map[v1.IPFamily]utils.IPTablesHandler)
-	npc.iptablesSaveRestore = make(map[v1.IPFamily]*utils.IPTablesSaveRestore)
+	npc.iptablesSaveRestore = make(map[v1.IPFamily]utils.IPTablesSaveRestorer)
 	npc.filterTableRules = make(map[v1.IPFamily]*bytes.Buffer)
 	npc.ipSetHandlers = make(map[v1.IPFamily]utils.IPSetHandler)
 	npc.nodeIPs = make(map[v1.IPFamily]net.IP)
@@ -891,13 +891,7 @@ func TestNetworkPolicyController(t *testing.T) {
 	for _, test := range testCases {
 		t.Run(test.name, func(t *testing.T) {
 			// TODO: Handle IPv6
-			iptablesHandlers := make(map[v1.IPFamily]utils.IPTablesHandler, 1)
-			iptablesHandlers[v1.IPv4Protocol] = newFakeIPTables(iptables.ProtocolIPv4)
-			ipSetHandlers := make(map[v1.IPFamily]utils.IPSetHandler, 1)
-			ipSetHandlers[v1.IPv4Protocol] = &fakeIPSet{}
-
-			_, err := NewNetworkPolicyController(client, test.config, podInformer, netpolInformer, nsInformer, &sync.Mutex{},
-				iptablesHandlers, ipSetHandlers)
+			_, err := NewNetworkPolicyController(client, test.config, podInformer, netpolInformer, nsInformer, &sync.Mutex{})
 			if err == nil && test.expectError {
 				t.Error("This config should have failed, but it was successful instead")
 			} else if err != nil {

diff --git a/pkg/utils/iptables.go b/pkg/utils/iptables.go
@@ -136,6 +136,11 @@ func Append(buffer *bytes.Buffer, chain string, rule []string) {
 	buffer.WriteString(ruleStr)
 }
 
+type IPTablesSaveRestorer interface {
+	SaveInto(table string, buffer *bytes.Buffer) error
+	Restore(table string, data []byte) error
+}
+
 type IPTablesSaveRestore struct {
 	saveCmd    string
 	restoreCmd string