diff --git a/go.mod b/go.mod index 45e5a70..2229659 100644 --- a/go.mod +++ b/go.mod @@ -12,6 +12,7 @@ require ( require ( github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect + github.com/google/go-cmp v0.6.0 // indirect github.com/kr/text v0.2.0 // indirect github.com/prometheus/client_model v0.6.0 // indirect github.com/prometheus/common v0.51.0 // indirect diff --git a/nftables.go b/nftables.go index e20b940..2c5f444 100644 --- a/nftables.go +++ b/nftables.go @@ -171,9 +171,17 @@ func (nft nftables) mineAddress(right gjson.Result) []string { if set.Exists() { var addresses []string // fmt.Printf("[prefix] %s\n", set.Get("#.prefix")) - for _, prefix := range set.Get("#.prefix").Array() { - // fmt.Printf("[prefix] %s\n", prefix) - addresses = append(addresses, nft.subnetToString(prefix)) + for _, el := range set.Array() { + switch el.Type { + case gjson.String: + addresses = append(addresses, el.String()) + case gjson.JSON: + if el.Get("prefix").Exists() { + addresses = append(addresses, nft.subnetToString(el.Get("prefix"))) + } + case gjson.False, gjson.Null, gjson.True: + // noop + } } return addresses } diff --git a/nftables_test.go b/nftables_test.go new file mode 100644 index 0000000..0d6988c --- /dev/null +++ b/nftables_test.go @@ -0,0 +1,52 @@ +package main + +import ( + "sort" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/tidwall/gjson" +) + +func TestMineAddress(t *testing.T) { + nft := nftables{} + + cases := []struct { + json string + want []string + }{ + { // plain string + json: `"8.8.8.8"`, + want: []string{"8.8.8.8"}, + }, + + { // anonymous ip addr only set + json: `{"set": ["8.8.8.8","10.96.0.10","21.21.0.242"]}}`, + want: []string{"10.96.0.10", "21.21.0.242", "8.8.8.8"}, + }, + + { // anonymous set with subnets only + json: `{"set": [{"prefix": {"addr": "10.10.0.0","len": 16}}, {"prefix": {"addr": "10.20.0.0","len": 16}}]}`, + want: []string{"10.10.0.0/16", "10.20.0.0/16"}, + }, + + { // anonymous mixed (ip addr and subnets) set + json: `{"set": ["8.8.8.8","10.96.0.10","21.21.0.242",{"prefix": {"addr": "127.0.0.0","len": 8}}]}`, + want: []string{"10.96.0.10", "127.0.0.0/8", "21.21.0.242", "8.8.8.8"}, + }, + + { // single subnet + json: `{"prefix": {"addr": "127.0.0.0","len": 8}}`, + want: []string{"127.0.0.0/8"}, + }, + } + for i, c := range cases { + json := gjson.Parse(c.json) + got := nft.mineAddress(json) + + sort.Strings(got) + if !cmp.Equal(got, c.want) { + t.Errorf("mineAddress case#%d failed:\n%v\n", i, cmp.Diff(c.want, got)) + } + } +}