mirror of
https://github.com/Sagit-chu/flvx.git
synced 2026-10-07 02:06:38 +08:00
feat(nftables): parse traffic counters
This commit is contained in:
@@ -0,0 +1,119 @@
|
|||||||
|
package nftables
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
type CounterSample struct {
|
||||||
|
ForwardID int64
|
||||||
|
Direction string
|
||||||
|
Protocol string
|
||||||
|
Bytes uint64
|
||||||
|
Packets uint64
|
||||||
|
}
|
||||||
|
|
||||||
|
func ParseCounterComment(comment string) (CounterSample, bool) {
|
||||||
|
parts := strings.Split(comment, " ")
|
||||||
|
if len(parts) != 4 || parts[0] != "flvx" {
|
||||||
|
return CounterSample{}, false
|
||||||
|
}
|
||||||
|
if !strings.HasPrefix(parts[1], "forward:") {
|
||||||
|
return CounterSample{}, false
|
||||||
|
}
|
||||||
|
forwardText := strings.TrimPrefix(parts[1], "forward:")
|
||||||
|
forwardID, err := strconv.ParseInt(forwardText, 10, 64)
|
||||||
|
if err != nil || forwardID <= 0 {
|
||||||
|
return CounterSample{}, false
|
||||||
|
}
|
||||||
|
direction := parts[2]
|
||||||
|
if direction != CounterDirectionToTarget && direction != CounterDirectionFromTarget {
|
||||||
|
return CounterSample{}, false
|
||||||
|
}
|
||||||
|
protocol := parts[3]
|
||||||
|
if protocol != "tcp" && protocol != "udp" {
|
||||||
|
return CounterSample{}, false
|
||||||
|
}
|
||||||
|
return CounterSample{
|
||||||
|
ForwardID: forwardID,
|
||||||
|
Direction: direction,
|
||||||
|
Protocol: protocol,
|
||||||
|
}, true
|
||||||
|
}
|
||||||
|
|
||||||
|
func ParseCounterSamples(raw []byte) ([]CounterSample, error) {
|
||||||
|
var doc nftListTable
|
||||||
|
if err := json.Unmarshal(raw, &doc); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
samples := make([]CounterSample, 0)
|
||||||
|
for _, item := range doc.Nftables {
|
||||||
|
ruleRaw, ok := item["rule"]
|
||||||
|
if !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
var rule nftCounterRule
|
||||||
|
if err := json.Unmarshal(ruleRaw, &rule); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if rule.Table != "flvx" || rule.Chain != "forward" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
sample, ok, err := parseCounterRule(rule)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
samples = append(samples, sample)
|
||||||
|
}
|
||||||
|
return samples, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type nftListTable struct {
|
||||||
|
Nftables []map[string]json.RawMessage `json:"nftables"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type nftCounterRule struct {
|
||||||
|
Table string `json:"table"`
|
||||||
|
Chain string `json:"chain"`
|
||||||
|
Comment string `json:"comment"`
|
||||||
|
Expr []map[string]json.RawMessage `json:"expr"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type nftCounter struct {
|
||||||
|
Bytes uint64 `json:"bytes"`
|
||||||
|
Packets uint64 `json:"packets"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseCounterRule(rule nftCounterRule) (CounterSample, bool, error) {
|
||||||
|
var (
|
||||||
|
counter nftCounter
|
||||||
|
hasCounter bool
|
||||||
|
)
|
||||||
|
|
||||||
|
for _, expr := range rule.Expr {
|
||||||
|
if rawCounter, ok := expr["counter"]; ok {
|
||||||
|
if err := json.Unmarshal(rawCounter, &counter); err != nil {
|
||||||
|
return CounterSample{}, false, err
|
||||||
|
}
|
||||||
|
hasCounter = true
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !hasCounter {
|
||||||
|
return CounterSample{}, false, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
sample, ok := ParseCounterComment(rule.Comment)
|
||||||
|
if !ok {
|
||||||
|
return CounterSample{}, false, nil
|
||||||
|
}
|
||||||
|
sample.Bytes = counter.Bytes
|
||||||
|
sample.Packets = counter.Packets
|
||||||
|
return sample, true, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,135 @@
|
|||||||
|
package nftables
|
||||||
|
|
||||||
|
import "testing"
|
||||||
|
|
||||||
|
func TestParseCounterCommentAcceptsValidToTargetTCP(t *testing.T) {
|
||||||
|
sample, ok := ParseCounterComment("flvx forward:42 to-target tcp")
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("expected comment to parse")
|
||||||
|
}
|
||||||
|
if sample.ForwardID != 42 ||
|
||||||
|
sample.Direction != CounterDirectionToTarget ||
|
||||||
|
sample.Protocol != "tcp" {
|
||||||
|
t.Fatalf("unexpected sample: %+v", sample)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseCounterCommentRejectsDNAT(t *testing.T) {
|
||||||
|
if sample, ok := ParseCounterComment("flvx forward:42 dnat tcp"); ok {
|
||||||
|
t.Fatalf("expected dnat comment to be rejected, got %+v", sample)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseCounterSamplesParsesForwardBillableCounters(t *testing.T) {
|
||||||
|
raw := []byte(`{
|
||||||
|
"nftables": [
|
||||||
|
{"metainfo": {"json_schema_version": 1}},
|
||||||
|
{"rule": {
|
||||||
|
"family": "inet",
|
||||||
|
"table": "flvx",
|
||||||
|
"chain": "forward",
|
||||||
|
"handle": 10,
|
||||||
|
"comment": "flvx forward:42 to-target tcp",
|
||||||
|
"expr": [
|
||||||
|
{"match": {"left": {"payload": {"protocol": "ip", "field": "daddr"}}, "op": "==", "right": "198.51.100.20"}},
|
||||||
|
{"counter": {"packets": 7, "bytes": 4096}}
|
||||||
|
]
|
||||||
|
}},
|
||||||
|
{"rule": {
|
||||||
|
"family": "inet",
|
||||||
|
"table": "flvx",
|
||||||
|
"chain": "forward",
|
||||||
|
"handle": 11,
|
||||||
|
"comment": "flvx forward:42 from-target udp",
|
||||||
|
"expr": [
|
||||||
|
{"counter": {"packets": 9, "bytes": 8192}}
|
||||||
|
]
|
||||||
|
}},
|
||||||
|
{"rule": {
|
||||||
|
"family": "inet",
|
||||||
|
"table": "flvx",
|
||||||
|
"chain": "prerouting",
|
||||||
|
"handle": 12,
|
||||||
|
"comment": "flvx forward:42 dnat tcp",
|
||||||
|
"expr": [
|
||||||
|
{"counter": {"packets": 100, "bytes": 65536}}
|
||||||
|
]
|
||||||
|
}}
|
||||||
|
]
|
||||||
|
}`)
|
||||||
|
|
||||||
|
samples, err := ParseCounterSamples(raw)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ParseCounterSamples: %v", err)
|
||||||
|
}
|
||||||
|
if len(samples) != 2 {
|
||||||
|
t.Fatalf("expected 2 samples, got %d: %+v", len(samples), samples)
|
||||||
|
}
|
||||||
|
|
||||||
|
want := []CounterSample{
|
||||||
|
{ForwardID: 42, Direction: CounterDirectionToTarget, Protocol: "tcp", Bytes: 4096, Packets: 7},
|
||||||
|
{ForwardID: 42, Direction: CounterDirectionFromTarget, Protocol: "udp", Bytes: 8192, Packets: 9},
|
||||||
|
}
|
||||||
|
for i := range want {
|
||||||
|
if samples[i] != want[i] {
|
||||||
|
t.Fatalf("sample %d: expected %+v, got %+v", i, want[i], samples[i])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseCounterSamplesUsesRuleLevelComment(t *testing.T) {
|
||||||
|
raw := []byte(`{
|
||||||
|
"nftables": [
|
||||||
|
{"rule": {
|
||||||
|
"table": "flvx",
|
||||||
|
"chain": "forward",
|
||||||
|
"comment": "flvx forward:77 to-target udp",
|
||||||
|
"expr": [
|
||||||
|
{"counter": {"packets": 3, "bytes": 2048}}
|
||||||
|
]
|
||||||
|
}}
|
||||||
|
]
|
||||||
|
}`)
|
||||||
|
|
||||||
|
samples, err := ParseCounterSamples(raw)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ParseCounterSamples: %v", err)
|
||||||
|
}
|
||||||
|
if len(samples) != 1 {
|
||||||
|
t.Fatalf("expected 1 sample, got %d: %+v", len(samples), samples)
|
||||||
|
}
|
||||||
|
want := CounterSample{
|
||||||
|
ForwardID: 77,
|
||||||
|
Direction: CounterDirectionToTarget,
|
||||||
|
Protocol: "udp",
|
||||||
|
Bytes: 2048,
|
||||||
|
Packets: 3,
|
||||||
|
}
|
||||||
|
if samples[0] != want {
|
||||||
|
t.Fatalf("expected %+v, got %+v", want, samples[0])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseCounterSamplesMalformedJSONReturnsError(t *testing.T) {
|
||||||
|
if _, err := ParseCounterSamples([]byte(`{"nftables": [`)); err == nil {
|
||||||
|
t.Fatal("expected malformed JSON error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseCounterSamplesMalformedRuleJSONReturnsError(t *testing.T) {
|
||||||
|
raw := []byte(`{
|
||||||
|
"nftables": [
|
||||||
|
{"rule": {
|
||||||
|
"table": "flvx",
|
||||||
|
"chain": "forward",
|
||||||
|
"comment": "flvx forward:42 to-target tcp",
|
||||||
|
"expr": [
|
||||||
|
{"counter": {"packets": "bad", "bytes": 4096}}
|
||||||
|
]
|
||||||
|
}}
|
||||||
|
]
|
||||||
|
}`)
|
||||||
|
if _, err := ParseCounterSamples(raw); err == nil {
|
||||||
|
t.Fatal("expected malformed rule JSON error")
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user