Merge pull request #57 from sstrk/master

[agent] Adjust policy handling for IPv6 addresses
cloudandheat · Aug 1, 2024 · 0795732 · 0795732
2 parents 5a28338 + 7a3a99e
commit 0795732
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 3 deletions.
diff --git a/internal/agent/nftables_generator.go b/internal/agent/nftables_generator.go
@@ -34,8 +34,23 @@ import (
 	"github.com/cloudandheat/ch-k8s-lbaas/internal/model"
 )
 
+var funcMap = template.FuncMap{
+	// A very simplistic check if a string looks like an IPv4 address
+	"isIPv4Address": func (ipString string) bool {
+		return strings.Count(ipString, ":") == 0 && strings.Count(ipString, ".") == 3
+	},
+	// A very simplistic check if a string looks like an IPv6 address
+	"isIPv6Address": func (ipString string) bool {
+		return strings.Count(ipString, ":") >= 2
+	},
+	// Replace colons with dash
+	"replaceColons": func (ipString string) string {
+		return strings.ReplaceAll(ipString, ":", "-")
+	},
+}
+
 var (
-	nftablesTemplate = template.Must(template.New("nftables.conf").Parse(`
+	nftablesTemplate = template.Must(template.New("nftables.conf").Funcs(funcMap).Parse(`
 {{ $cfg := . }}
 
 {{- if $cfg.PartialReload }}
@@ -57,14 +72,14 @@ delete chain {{ $cfg.FilterTableType }} {{ $cfg.FilterTableName }} {{ $chain }}
 table {{ .FilterTableType }} {{ .FilterTableName }} {
 	chain {{ .FilterForwardChainName }} {
 		{{- range $dest := $cfg.PolicyAssignments }}
-		ct mark {{ $cfg.FWMarkBits | printf "0x%x" }} and {{ $cfg.FWMarkMask | printf "0x%x" }} ip daddr {{ $dest.Address }} goto {{ $cfg.PolicyPrefix }}POD-{{ $dest.Address }};
+		ct mark {{ $cfg.FWMarkBits | printf "0x%x" }} and {{ $cfg.FWMarkMask | printf "0x%x" }} {{if isIPv4Address $dest.Address }}ip{{else if isIPv6Address $dest.Address}}ip6{{end}} daddr {{ $dest.Address }} goto {{ $cfg.PolicyPrefix }}POD-{{replaceColons $dest.Address}};
 		{{- end }}
 		ct mark {{ $cfg.FWMarkBits | printf "0x%x" }} and {{ $cfg.FWMarkMask | printf "0x%x" }} accept;
 	}
 
 	# Using uppercase POD to prevent collisions with policy names like 'pod-x.x.x.x'
 	{{- range $pod := $cfg.PolicyAssignments }}
-	chain {{ $cfg.PolicyPrefix }}POD-{{ $pod.Address }} {
+	chain {{ $cfg.PolicyPrefix }}POD-{{replaceColons $pod.Address}} {
 		{{- range $pol := $pod.NetworkPolicies }}
 		jump {{ $cfg.PolicyPrefix }}{{ $pol }};
 		{{- end }}

diff --git a/internal/agent/nftables_generator_test.go b/internal/agent/nftables_generator_test.go
@@ -158,6 +158,18 @@ func TestNftablesStructuredConfigFromNonEmptyLBModel(t *testing.T) {
 					"block-range",
 				},
 			},
+			{
+				Address: "ff00::1",
+				NetworkPolicies: []string{
+					"allow-http",
+				},
+			},
+			{
+				Address: "ff00::2",
+				NetworkPolicies: []string{
+					"block-range",
+				},
+			},
 		},
 	}