In the middle

cidr/tree4.go

@@ -142,15 +142,22 @@ func (tree *Tree4[T]) MostSpecificContains(ip iputil.VpnIp) (ok bool, value T) {
 	return ok, value
 }
 
-// Match finds the most specific match
-// TODO this is exact match
-func (tree *Tree4[T]) Match(ip iputil.VpnIp) (ok bool, value T) {
+type eachFunc[T any] func(T) bool
+
+// EachContains will call a function, passing the value, for each entry until the function returns false or the search is complete
+// The final return value will be true if the provided function returned true
+func (tree *Tree4[T]) EachContains(ip iputil.VpnIp, each eachFunc[T]) bool {
 	bit := startbit
 	node := tree.root
-	lastNode := node
 
 	for node != nil {
-		lastNode = node
+		if node.hasValue {
+			// If the each func returns true then we can exit the loop
+			if each(node.value) {
+				return true
+			}
+		}
+
 		if ip&bit != 0 {
 			node = node.right
 		} else {
@@ -160,10 +167,33 @@ func (tree *Tree4[T]) Match(ip iputil.VpnIp) (ok bool, value T) {
 		bit >>= 1
 	}
 
-	if bit == 0 && lastNode != nil {
-		value = lastNode.value
-		ok = true
+	return false
+}
+
+// GetCIDR returns the entry added by the most recent matching AddCIDR call
+func (tree *Tree4[T]) GetCIDR(cidr *net.IPNet) (ok bool, value T) {
+	bit := startbit
+	node := tree.root
+
+	ip := iputil.Ip2VpnIp(cidr.IP)
+	mask := iputil.Ip2VpnIp(cidr.Mask)
+
+	// Find our last ancestor in the tree
+	for node != nil && bit&mask != 0 {
+		if ip&bit != 0 {
+			node = node.right
+		} else {
+			node = node.left
+		}
+
+		bit = bit >> 1
 	}
+
+	if bit&mask == 0 && node != nil {
+		value = node.value
+		ok = node.hasValue
+	}
+
 	return ok, value
 }
 

cidr/tree4_test.go

@@ -115,35 +115,36 @@ func TestCIDRTree_MostSpecificContains(t *testing.T) {
 	assert.Equal(t, "cool", r)
 }
 
-func TestCIDRTree_Match(t *testing.T) {
+func TestTree4_GetCIDR(t *testing.T) {
 	tree := NewTree4[string]()
-	tree.AddCIDR(Parse("4.1.1.0/32"), "1a")
-	tree.AddCIDR(Parse("4.1.1.1/32"), "1b")
+	tree.AddCIDR(Parse("1.0.0.0/8"), "1")
+	tree.AddCIDR(Parse("2.1.0.0/16"), "2")
+	tree.AddCIDR(Parse("3.1.1.0/24"), "3")
+	tree.AddCIDR(Parse("4.1.1.0/24"), "4a")
+	tree.AddCIDR(Parse("4.1.1.1/32"), "4b")
+	tree.AddCIDR(Parse("4.1.2.1/32"), "4c")
+	tree.AddCIDR(Parse("254.0.0.0/4"), "5")
 
 	tests := []struct {
 		Found  bool
 		Result interface{}
-		IP     string
+		IPNet  *net.IPNet
 	}{
-		{true, "1a", "4.1.1.0"},
-		{true, "1b", "4.1.1.1"},
+		{true, "1", Parse("1.0.0.0/8")},
+		{true, "2", Parse("2.1.0.0/16")},
+		{true, "3", Parse("3.1.1.0/24")},
+		{true, "4a", Parse("4.1.1.0/24")},
+		{true, "4b", Parse("4.1.1.1/32")},
+		{true, "4c", Parse("4.1.2.1/32")},
+		{true, "5", Parse("254.0.0.0/4")},
+		{false, "", Parse("2.0.0.0/8")},
 	}
 
 	for _, tt := range tests {
-		ok, r := tree.Match(iputil.Ip2VpnIp(net.ParseIP(tt.IP)))
+		ok, r := tree.GetCIDR(tt.IPNet)
 		assert.Equal(t, tt.Found, ok)
 		assert.Equal(t, tt.Result, r)
 	}
-
-	tree = NewTree4[string]()
-	tree.AddCIDR(Parse("1.1.1.1/0"), "cool")
-	ok, r := tree.Contains(iputil.Ip2VpnIp(net.ParseIP("0.0.0.0")))
-	assert.True(t, ok)
-	assert.Equal(t, "cool", r)
-
-	ok, r = tree.Contains(iputil.Ip2VpnIp(net.ParseIP("255.255.255.255")))
-	assert.True(t, ok)
-	assert.Equal(t, "cool", r)
 }
 
 func BenchmarkCIDRTree_Contains(b *testing.B) {
@@ -167,25 +168,3 @@ func BenchmarkCIDRTree_Contains(b *testing.B) {
 		}
 	})
 }
-
-func BenchmarkCIDRTree_Match(b *testing.B) {
-	tree := NewTree4[string]()
-	tree.AddCIDR(Parse("1.1.0.0/16"), "1")
-	tree.AddCIDR(Parse("1.2.1.1/32"), "1")
-	tree.AddCIDR(Parse("192.2.1.1/32"), "1")
-	tree.AddCIDR(Parse("172.2.1.1/32"), "1")
-
-	ip := iputil.Ip2VpnIp(net.ParseIP("1.2.1.1"))
-	b.Run("found", func(b *testing.B) {
-		for i := 0; i < b.N; i++ {
-			tree.Match(ip)
-		}
-	})
-
-	ip = iputil.Ip2VpnIp(net.ParseIP("1.2.1.255"))
-	b.Run("not found", func(b *testing.B) {
-		for i := 0; i < b.N; i++ {
-			tree.Match(ip)
-		}
-	})
-}