chore(internal/iptables): refactor egress network policy rules

internal/iptables/iptables.go

@@ -26,7 +26,7 @@ func (it *Iptables) Sync(state agent.State) error {
 	dns := state.Server.Status.Dns
 	peers := state.Peers
 
-	cfg := GenerateIptableRulesFromPeers(wgHostName, dns, peers)
+	cfg := GenerateIPTableRulesFromPeers(wgHostName, dns, peers)
 
 	err := ApplyRules(cfg)
 
@@ -37,7 +37,7 @@ func (it *Iptables) Sync(state agent.State) error {
 	return nil
 }
 
-func GenerateIptableRulesFromNetworkPolicies(policies v1alpha1.EgressNetworkPolicies, peerIp string, kubeDnsIp string, wgServerIp string) string {
+func GenerateIPTableRulesFromNetworkPolicies(policies v1alpha1.EgressNetworkPolicies, peerIp string, kubeDnsIp string, wgServerIp string) string {
 	peerChain := strings.ReplaceAll(peerIp, ".", "-")
 
 	rules := []string{
@@ -61,7 +61,7 @@ func GenerateIptableRulesFromNetworkPolicies(policies v1alpha1.EgressNetworkPoli
 	}
 
 	for _, policy := range policies {
-		rules = append(rules, EgressNetworkPolicyToIpTableRules(policy, peerChain)...)
+		rules = append(rules, EgressNetworkPolicyToIPTableRules(policy, peerChain)...)
 	}
 
 	// if policies are defined impose an implicit deny all
@@ -75,80 +75,67 @@ func GenerateIptableRulesFromNetworkPolicies(policies v1alpha1.EgressNetworkPoli
 	return strings.Join(rules, "\n")
 }
 
-func GenerateIptableRulesFromPeers(wgHostName string, dns string, peers []v1alpha1.WireguardPeer) string {
-	var rules []string
-
-	var natTableRules = `
-*nat
+func GenerateIPTableRulesFromPeers(wgHostName string, dns string, peers []v1alpha1.WireguardPeer) string {
+	const natTableRules = `*nat
 :PREROUTING ACCEPT [0:0]
 :INPUT ACCEPT [0:0]
 :OUTPUT ACCEPT [0:0]
 :POSTROUTING ACCEPT [0:0]
 -A POSTROUTING -s 10.8.0.0/24 -o eth0 -j MASQUERADE
 COMMIT`
 
+	var rules []string
 	for _, peer := range peers {
-
-		//tc(peer.Spec.DownloadSpeed, peer.Spec.UploadSpeed)
-		rules = append(rules, GenerateIptableRulesFromNetworkPolicies(peer.Spec.EgressNetworkPolicies, peer.Spec.Address, dns, wgHostName))
+		rules = append(rules, GenerateIPTableRulesFromNetworkPolicies(peer.Spec.EgressNetworkPolicies, peer.Spec.Address, dns, wgHostName))
 	}
 
-	var filterTableRules = fmt.Sprintf(`
-*filter
+	filterTableRules := fmt.Sprintf(`*filter
 :INPUT ACCEPT [0:0]
 :FORWARD ACCEPT [0:0]
 :OUTPUT ACCEPT [0:0]
 %s
-COMMIT
-`, strings.Join(rules, "\n"))
+COMMIT`, strings.Join(rules, "\n"))
 
 	return fmt.Sprintf("%s\n%s", natTableRules, filterTableRules)
 }
 
-func EgressNetworkPolicyToIpTableRules(policy v1alpha1.EgressNetworkPolicy, peerChain string) []string {
+func EgressNetworkPolicyToIPTableRules(policy v1alpha1.EgressNetworkPolicy, peerChain string) []string {
+	if policy.Protocol == "" && policy.To.Port != 0 {
+		tcpPolicy := policy
+		tcpPolicy.Protocol = "TCP"
 
-	var rules []string
+		udpPolicy := policy
+		udpPolicy.Protocol = "UDP"
 
-	if policy.Protocol == "" && policy.To.Port != 0 {
-		policy.Protocol = "TCP"
-		rules = append(rules, EgressNetworkPolicyToIpTableRules(policy, peerChain)[0])
-		policy.Protocol = "UDP"
-		rules = append(rules, EgressNetworkPolicyToIpTableRules(policy, peerChain)[0])
-		return rules
+		return []string{
+			EgressNetworkPolicyToIPTableRule(tcpPolicy, peerChain),
+			EgressNetworkPolicyToIPTableRule(udpPolicy, peerChain),
+		}
 	}
+	return []string{EgressNetworkPolicyToIPTableRule(policy, peerChain)}
+}
 
-	// customer rules
-	var rulePeerChain = "-A " + peerChain
-	var ruleAction = string("-j " + v1alpha1.EgressNetworkPolicyActionDeny)
-	var ruleProtocol = ""
-	var ruleDestIp = ""
-	var ruleDestPort = ""
+func EgressNetworkPolicyToIPTableRule(policy v1alpha1.EgressNetworkPolicy, peerChain string) string {
+	opts := []string{fmt.Sprintf("-A %s", peerChain)}
 
 	if policy.To.Ip != "" {
-		ruleDestIp = "-d " + policy.To.Ip
+		opts = append(opts, fmt.Sprintf("-d %s", policy.To.Ip))
 	}
 
 	if policy.Protocol != "" {
-		ruleProtocol = "-p " + strings.ToUpper(string(policy.Protocol))
+		opts = append(opts, fmt.Sprintf("-p %s", strings.ToUpper(string(policy.Protocol))))
 	}
 
 	if policy.To.Port != 0 {
-		ruleDestPort = "--dport " + fmt.Sprint(policy.To.Port)
+		opts = append(opts, fmt.Sprintf("--dport %d", policy.To.Port))
 	}
 
+	action := v1alpha1.EgressNetworkPolicyActionDeny
 	if policy.Action != "" {
-		ruleAction = "-j " + strings.ToUpper(string(policy.Action))
-	}
-
-	var options = []string{rulePeerChain, ruleDestIp, ruleProtocol, ruleDestPort, ruleAction}
-	var filteredOptions []string
-	for _, option := range options {
-		if len(option) != 0 {
-			filteredOptions = append(filteredOptions, option)
-		}
+		action = policy.Action
 	}
-	rules = append(rules, strings.Join(filteredOptions, " "))
 
-	return rules
+	opts = append(opts, fmt.Sprintf("-j %s", strings.ToUpper(string(action))))
 
+	return strings.Join(opts, " ")
 }

internal/iptables/iptables_test.go

@@ -1,13 +1,12 @@
 package iptables
 
 import (
-	"github.com/jodevsa/wireguard-operator/pkg/api/v1alpha1"
 	"testing"
-)
 
-// test helpers
+	"github.com/jodevsa/wireguard-operator/pkg/api/v1alpha1"
+)
 
-func TestIptableRules(t *testing.T) {
+func TestIPTableRules(t *testing.T) {
 	tests := []struct {
 		name                 string
 		peerIp               string
@@ -83,7 +82,7 @@ func TestIptableRules(t *testing.T) {
 -A 10-8-0-11 -d 10.7.0.1 -p icmp -j ACCEPT
 -A 10-8-0-11 -d 10.8.0.11 -j ACCEPT
 -A 10-8-0-11 -d 100.64.0.21 -p UDP --dport 53 -j ACCEPT
--A 10-8-0-11 -j Reject
+-A 10-8-0-11 -j REJECT
 -A 10-8-0-11 -j REJECT --reject-with icmp-port-unreachable
 # end of rules for peer 10.8.0.11`,
 		},
@@ -113,7 +112,7 @@ func TestIptableRules(t *testing.T) {
 
 		t.Run(test.name, func(t *testing.T) {
 
-			rules := GenerateIptableRulesFromNetworkPolicies(test.networkPolicies, test.peerIp, test.kubeDnsIp, test.wgServerIp)
+			rules := GenerateIPTableRulesFromNetworkPolicies(test.networkPolicies, test.peerIp, test.kubeDnsIp, test.wgServerIp)
 			if rules != test.expectedIptableRules {
 				t.Errorf("got %s, want %s", rules, test.expectedIptableRules)
 			}