From 14bbd3907d15f30dad6fb13022c7caf6f8aa0451 Mon Sep 17 00:00:00 2001 From: sagitchu <601096721@qq.com> Date: Sat, 6 Jun 2026 18:01:09 +0800 Subject: [PATCH] feat(nftables): render traffic counters --- .../internal/runtime/nftables/manager_test.go | 2 +- .../internal/runtime/nftables/renderer.go | 47 ++++++++-- .../runtime/nftables/renderer_test.go | 86 ++++++++++++++++++- go-backend/internal/runtime/nftables/types.go | 4 + 4 files changed, 128 insertions(+), 11 deletions(-) diff --git a/go-backend/internal/runtime/nftables/manager_test.go b/go-backend/internal/runtime/nftables/manager_test.go index e788687..7529530 100644 --- a/go-backend/internal/runtime/nftables/manager_test.go +++ b/go-backend/internal/runtime/nftables/manager_test.go @@ -37,7 +37,7 @@ func TestManagerReconcileAppliesRenderedScript(t *testing.T) { if len(runner.scripts) != 1 { t.Fatalf("expected 1 script, got %d", len(runner.scripts)) } - if !strings.Contains(runner.scripts[0], "flvx forward:42 tcp") { + if !strings.Contains(runner.scripts[0], `flvx forward:42 dnat tcp`) { t.Fatalf("script missing forward comment:\n%s", runner.scripts[0]) } if result.NodeID != 7 || result.Hashes[42] == "" { diff --git a/go-backend/internal/runtime/nftables/renderer.go b/go-backend/internal/runtime/nftables/renderer.go index 4f2a116..34f6998 100644 --- a/go-backend/internal/runtime/nftables/renderer.go +++ b/go-backend/internal/runtime/nftables/renderer.go @@ -15,14 +15,18 @@ func RenderTable(plan NodePlan) string { b.WriteString(" chain prerouting {\n") b.WriteString(" type nat hook prerouting priority dstnat; policy accept;\n") for _, rule := range sortedRules(plan.Rules) { + family := nftAddressFamily(rule.TargetHost) + dnatFamily := "" + if family != "" { + dnatFamily = family + " " + } for _, protocol := range normalizedProtocols(rule.Protocols) { - b.WriteString(fmt.Sprintf(" %s dport %d dnat %s to %s comment \"flvx forward:%d %s\"\n", + b.WriteString(fmt.Sprintf(" %s dport %d counter dnat %sto %s comment %q\n", protocol, rule.InPort, - dnatFamilyPrefix(rule.TargetHost), + dnatFamily, formatDNATTarget(rule.TargetHost, rule.TargetPort), - rule.ForwardID, - protocol, + counterComment(rule.ForwardID, CounterDirectionDNAT, protocol), )) } } @@ -35,11 +39,38 @@ func RenderTable(plan NodePlan) string { b.WriteString(" }\n\n") b.WriteString(" chain forward {\n") b.WriteString(" type filter hook forward priority filter; policy accept;\n") + for _, rule := range sortedRules(plan.Rules) { + family := nftAddressFamily(rule.TargetHost) + if family == "" { + continue + } + targetHost := strings.Trim(strings.TrimSpace(rule.TargetHost), "[]") + for _, protocol := range normalizedProtocols(rule.Protocols) { + b.WriteString(fmt.Sprintf(" %s daddr %s %s dport %d counter comment %q\n", + family, + targetHost, + protocol, + rule.TargetPort, + counterComment(rule.ForwardID, CounterDirectionToTarget, protocol), + )) + b.WriteString(fmt.Sprintf(" %s saddr %s %s sport %d counter comment %q\n", + family, + targetHost, + protocol, + rule.TargetPort, + counterComment(rule.ForwardID, CounterDirectionFromTarget, protocol), + )) + } + } b.WriteString(" }\n") b.WriteString("}\n") return b.String() } +func counterComment(forwardID int64, direction, protocol string) string { + return fmt.Sprintf("flvx forward:%d %s %s", forwardID, direction, protocol) +} + func RuleHash(rule Rule) string { protocols := normalizedProtocols(rule.Protocols) sum := sha256.Sum256([]byte(fmt.Sprintf("%d|%d|%s|%d|%s", @@ -100,14 +131,14 @@ func formatDNATTarget(host string, port int) string { return fmt.Sprintf("%s:%d", trimmed, port) } -func dnatFamilyPrefix(host string) string { +func nftAddressFamily(host string) string { trimmed := strings.Trim(strings.TrimSpace(host), "[]") ip := net.ParseIP(trimmed) if ip == nil { return "" } - if ip.To4() != nil { - return "ip" + if ip.To4() == nil { + return "ip6" } - return "ip6" + return "ip" } diff --git a/go-backend/internal/runtime/nftables/renderer_test.go b/go-backend/internal/runtime/nftables/renderer_test.go index 04b29ec..b566f80 100644 --- a/go-backend/internal/runtime/nftables/renderer_test.go +++ b/go-backend/internal/runtime/nftables/renderer_test.go @@ -23,8 +23,8 @@ func TestRenderTableIncludesDNATAndMasquerade(t *testing.T) { "table inet flvx", "type nat hook prerouting priority dstnat; policy accept;", "type nat hook postrouting priority srcnat; policy accept;", - "tcp dport 24000 dnat ip to 198.51.100.20:443 comment \"flvx forward:42 tcp\"", - "udp dport 24000 dnat ip to 198.51.100.20:443 comment \"flvx forward:42 udp\"", + "tcp dport 24000 counter dnat ip to 198.51.100.20:443 comment \"flvx forward:42 dnat tcp\"", + "udp dport 24000 counter dnat ip to 198.51.100.20:443 comment \"flvx forward:42 dnat udp\"", "masquerade comment \"flvx masquerade\"", } for _, part := range expectedParts { @@ -46,6 +46,88 @@ func TestRenderTableBracketsIPv6Target(t *testing.T) { } } +func TestRenderTableIncludesForwardAccountingCounters(t *testing.T) { + plan := NodePlan{ + NodeID: 7, + Rules: []Rule{{ + ForwardID: 42, + InPort: 12345, + TargetHost: "198.51.100.20", + TargetPort: 443, + Protocols: []string{"tcp", "udp"}, + }}, + } + + got := RenderTable(plan) + wantLines := []string{ + `tcp dport 12345 counter dnat ip to 198.51.100.20:443 comment "flvx forward:42 dnat tcp"`, + `udp dport 12345 counter dnat ip to 198.51.100.20:443 comment "flvx forward:42 dnat udp"`, + `ip daddr 198.51.100.20 tcp dport 443 counter comment "flvx forward:42 to-target tcp"`, + `ip saddr 198.51.100.20 tcp sport 443 counter comment "flvx forward:42 from-target tcp"`, + `ip daddr 198.51.100.20 udp dport 443 counter comment "flvx forward:42 to-target udp"`, + `ip saddr 198.51.100.20 udp sport 443 counter comment "flvx forward:42 from-target udp"`, + } + for _, want := range wantLines { + if !strings.Contains(got, want) { + t.Fatalf("RenderTable() missing %q\n%s", want, got) + } + } +} + +func TestRenderTableIncludesIPv6ForwardAccountingCounters(t *testing.T) { + plan := NodePlan{ + NodeID: 7, + Rules: []Rule{{ + ForwardID: 43, + InPort: 12346, + TargetHost: "2001:db8::20", + TargetPort: 8443, + Protocols: []string{"tcp"}, + }}, + } + + got := RenderTable(plan) + wantLines := []string{ + `tcp dport 12346 counter dnat ip6 to [2001:db8::20]:8443 comment "flvx forward:43 dnat tcp"`, + `ip6 daddr 2001:db8::20 tcp dport 8443 counter comment "flvx forward:43 to-target tcp"`, + `ip6 saddr 2001:db8::20 tcp sport 8443 counter comment "flvx forward:43 from-target tcp"`, + } + for _, want := range wantLines { + if !strings.Contains(got, want) { + t.Fatalf("RenderTable() missing %q\n%s", want, got) + } + } +} + +func TestRenderTablePreservesHostnameDNATAndSkipsAccountingCounters(t *testing.T) { + plan := NodePlan{ + NodeID: 7, + Rules: []Rule{{ + ForwardID: 44, + InPort: 12347, + TargetHost: "example.com", + TargetPort: 9443, + Protocols: []string{"tcp"}, + }}, + } + + got := RenderTable(plan) + want := `tcp dport 12347 counter dnat to example.com:9443 comment "flvx forward:44 dnat tcp"` + if !strings.Contains(got, want) { + t.Fatalf("RenderTable() missing %q\n%s", want, got) + } + unwantedLines := []string{ + `dnat ip to example.com`, + `ip daddr example.com`, + `ip saddr example.com`, + } + for _, unwanted := range unwantedLines { + if strings.Contains(got, unwanted) { + t.Fatalf("RenderTable() unexpectedly contains %q\n%s", unwanted, got) + } + } +} + func TestRuleHashIsStable(t *testing.T) { rule := Rule{ForwardID: 42, InPort: 24000, TargetHost: "198.51.100.20", TargetPort: 443, Protocols: []string{"tcp", "udp"}} if RuleHash(rule) != RuleHash(rule) { diff --git a/go-backend/internal/runtime/nftables/types.go b/go-backend/internal/runtime/nftables/types.go index 1df1c29..3985af3 100644 --- a/go-backend/internal/runtime/nftables/types.go +++ b/go-backend/internal/runtime/nftables/types.go @@ -7,6 +7,10 @@ const ( StatusPending = "pending" StatusApplied = "applied" StatusError = "error" + + CounterDirectionDNAT = "dnat" + CounterDirectionToTarget = "to-target" + CounterDirectionFromTarget = "from-target" ) type Target struct {