fix address extraction from anonymous set (#23)

* fix address extraction from anonymous set

* add tests for `mineAddress` functionality

---------

Co-authored-by: Ivan Andreyko <[email protected]>

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

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
 			}

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))
+		}
+	}
+}