341 lines
9.4 KiB
Go
341 lines
9.4 KiB
Go
package nettrie
|
|
|
|
import (
|
|
"maps"
|
|
"net/netip"
|
|
"reflect"
|
|
"slices"
|
|
"testing"
|
|
)
|
|
|
|
func TestValueTrie_IPv4(t *testing.T) {
|
|
trie := NewValue[string]()
|
|
|
|
routes := map[string]string{
|
|
"0.0.0.0/0": "Default",
|
|
"10.0.0.0/8": "Private A",
|
|
"192.168.0.0/16": "Private B",
|
|
"192.168.1.0/24": "LAN",
|
|
"192.168.1.128/25": "Subnet",
|
|
"192.168.1.200/32": "Host",
|
|
}
|
|
|
|
for s, v := range routes {
|
|
trie.Insert(netip.MustParsePrefix(s), v)
|
|
}
|
|
|
|
testCases := []struct {
|
|
name string
|
|
lookupAddr string
|
|
expectedVal string
|
|
expectedOK bool
|
|
}{
|
|
{"Exact Host Match", "192.168.1.200", "Host", true},
|
|
{"Subnet Match", "192.168.1.201", "Subnet", true},
|
|
{"LAN Match", "192.168.1.50", "LAN", true},
|
|
{"Private B Match", "192.168.255.255", "Private B", true},
|
|
{"Private A Match", "10.255.255.255", "Private A", true},
|
|
{"Default Match", "8.8.8.8", "Default", true},
|
|
{"No Match", "203.0.113.1", "Default", true}, // Falls back to default
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
addr := netip.MustParseAddr(tc.lookupAddr)
|
|
val, ok := trie.Lookup(addr)
|
|
|
|
if ok != tc.expectedOK {
|
|
t.Errorf("expected ok=%v, got ok=%v", tc.expectedOK, ok)
|
|
}
|
|
if val != tc.expectedVal {
|
|
t.Errorf("expected value=%q, got %q", tc.expectedVal, val)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestValueTrie_IPv6(t *testing.T) {
|
|
trie := NewValue[string]()
|
|
|
|
routes := map[string]string{
|
|
"::/0": "Default",
|
|
"2001:db8::/32": "Global Unicast",
|
|
"2001:db8:acad::/48": "Academic",
|
|
"2001:db8:acad:1::/64": "CS Department",
|
|
"2001:db8:acad:1::1/128": "Host Route",
|
|
}
|
|
|
|
for s, v := range routes {
|
|
trie.Insert(netip.MustParsePrefix(s), v)
|
|
}
|
|
|
|
testCases := []struct {
|
|
name string
|
|
lookupAddr string
|
|
expectedVal string
|
|
expectedOK bool
|
|
}{
|
|
{"Exact Host Match", "2001:db8:acad:1::1", "Host Route", true},
|
|
{"CS Dept Match", "2001:db8:acad:1::2", "CS Department", true},
|
|
{"Academic Match", "2001:db8:acad:2::1", "Academic", true},
|
|
{"Global Unicast Match", "2001:db8:cafe::1", "Global Unicast", true},
|
|
{"Default Match", "2606:4700:4700::1111", "Default", true},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
addr := netip.MustParseAddr(tc.lookupAddr)
|
|
val, ok := trie.Lookup(addr)
|
|
|
|
if ok != tc.expectedOK {
|
|
t.Errorf("expected ok=%v, got ok=%v", tc.expectedOK, ok)
|
|
}
|
|
if val != tc.expectedVal {
|
|
t.Errorf("expected value=%q, got %q", tc.expectedVal, val)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestValueTrie_UpdateValue(t *testing.T) {
|
|
trie := NewValue[string]()
|
|
prefix := netip.MustParsePrefix("192.168.1.0/24")
|
|
addr := netip.MustParseAddr("192.168.1.1")
|
|
|
|
trie.Insert(prefix, "Initial Value")
|
|
val, _ := trie.Lookup(addr)
|
|
if val != "Initial Value" {
|
|
t.Fatalf("expected initial value, got %q", val)
|
|
}
|
|
|
|
trie.Insert(prefix, "Updated Value")
|
|
val, _ = trie.Lookup(addr)
|
|
if val != "Updated Value" {
|
|
t.Fatalf("expected updated value, got %q", val)
|
|
}
|
|
}
|
|
|
|
func TestValueTrie_Delete(t *testing.T) {
|
|
trie := NewValue[string]()
|
|
|
|
trie.Insert(netip.MustParsePrefix("10.0.0.0/8"), "A")
|
|
trie.Insert(netip.MustParsePrefix("10.0.0.0/24"), "B")
|
|
trie.Insert(netip.MustParsePrefix("10.0.1.0/24"), "C")
|
|
|
|
t.Run("Delete Non-Existent", func(t *testing.T) {
|
|
if trie.Delete(netip.MustParsePrefix("1.2.3.4/32")) {
|
|
t.Error("deleted a non-existent prefix")
|
|
}
|
|
})
|
|
|
|
t.Run("Delete Leaf Node", func(t *testing.T) {
|
|
// Delete 10.0.0.0/24, which is a leaf.
|
|
if !trie.Delete(netip.MustParsePrefix("10.0.0.0/24")) {
|
|
t.Fatal("failed to delete 10.0.0.0/24")
|
|
}
|
|
// Lookup should now resolve to the parent 10.0.0.0/8
|
|
val, ok := trie.Lookup(netip.MustParseAddr("10.0.0.1"))
|
|
if !ok || val != "A" {
|
|
t.Errorf("lookup failed after delete, got val=%q ok=%v, want val=\"A\" ok=true", val, ok)
|
|
}
|
|
})
|
|
|
|
t.Run("Delete and Merge", func(t *testing.T) {
|
|
// Insert a new prefix that causes a split
|
|
trie.Insert(netip.MustParsePrefix("10.0.0.0/8"), "A-updated")
|
|
|
|
// We now have 10.0.0.0/8 and 10.0.1.0/24. Deleting 10.0.1.0/24 should
|
|
// cause the split node to be removed and merged back into 10.0.0.0/8.
|
|
if !trie.Delete(netip.MustParsePrefix("10.0.1.0/24")) {
|
|
t.Fatal("failed to delete 10.0.1.0/24")
|
|
}
|
|
|
|
// A lookup for 10.0.1.1 should now match 10.0.0.0/8
|
|
val, ok := trie.Lookup(netip.MustParseAddr("10.0.1.1"))
|
|
if !ok || val != "A-updated" {
|
|
t.Errorf("lookup failed after merge, got val=%q ok=%v, want val=\"A-updated\" ok=true", val, ok)
|
|
}
|
|
})
|
|
|
|
t.Run("Delete Prefix Of Another", func(t *testing.T) {
|
|
// Delete 10.0.0.0/8
|
|
trie.Insert(netip.MustParsePrefix("10.0.0.0/8"), "A")
|
|
trie.Insert(netip.MustParsePrefix("10.1.0.0/16"), "D")
|
|
|
|
if !trie.Delete(netip.MustParsePrefix("10.0.0.0/8")) {
|
|
t.Fatal("failed to delete 10.0.0.0/8")
|
|
}
|
|
|
|
// Lookup for 10.0.0.1 should now fail (no default route)
|
|
_, ok := trie.Lookup(netip.MustParseAddr("10.0.0.1"))
|
|
if ok {
|
|
t.Error("lookup for 10.0.0.1 should have failed")
|
|
}
|
|
|
|
// Lookup for 10.1.0.1 should still succeed
|
|
val, ok := trie.Lookup(netip.MustParseAddr("10.1.0.1"))
|
|
if !ok || val != "D" {
|
|
t.Error("lookup for 10.1.0.1 failed unexpectedly")
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestValueTrie_Contains(t *testing.T) {
|
|
trie := NewValue[string]()
|
|
|
|
prefixes := []string{
|
|
"10.0.0.0/8",
|
|
"192.168.1.0/24",
|
|
"192.168.1.200/32",
|
|
"2001:db8::/32",
|
|
"2001:db8:acad:1::1/128",
|
|
}
|
|
|
|
for _, s := range prefixes {
|
|
trie.Insert(netip.MustParsePrefix(s), "present")
|
|
}
|
|
|
|
t.Run("ContainsPrefix", func(t *testing.T) {
|
|
testCases := []struct {
|
|
prefix string
|
|
want bool
|
|
}{
|
|
{"10.0.0.0/8", true},
|
|
{"192.168.1.200/32", true},
|
|
{"2001:db8::/32", true},
|
|
{"2001:db8:acad:1::1/128", true},
|
|
{"10.0.0.0/9", false}, // shorter parent, but not exact
|
|
{"192.168.1.0/25", false}, // non-existent child
|
|
{"172.16.0.0/12", false}, // completely different prefix
|
|
{"2001:db8::/48", false}, // non-existent child
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
p := netip.MustParsePrefix(tc.prefix)
|
|
if got := trie.ContainsPrefix(p); got != tc.want {
|
|
t.Errorf("ContainsPrefix(%q) = %v, want %v", tc.prefix, got, tc.want)
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("Contains", func(t *testing.T) {
|
|
testCases := []struct {
|
|
addr string
|
|
want bool
|
|
}{
|
|
{"192.168.1.200", true},
|
|
{"2001:db8:acad:1::1", true},
|
|
{"192.168.1.201", false}, // In /24 range, but not a /32 host route
|
|
{"10.0.0.1", false}, // In /8 range, but not a /32 host route
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
a := netip.MustParseAddr(tc.addr)
|
|
if got := trie.Contains(a); got != tc.want {
|
|
t.Errorf("Contains(%q) = %v, want %v", tc.addr, got, tc.want)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestValueTrie_Walk(t *testing.T) {
|
|
trie := NewValue[string]()
|
|
prefixes := map[string]string{
|
|
"10.0.0.0/8": "A",
|
|
"192.168.1.0/24": "B",
|
|
"2001:db8::/32": "C",
|
|
"172.16.0.0/12": "D",
|
|
"2001:db8:acad::/48": "E",
|
|
}
|
|
for s, v := range prefixes {
|
|
trie.Insert(netip.MustParsePrefix(s), v)
|
|
}
|
|
|
|
t.Run("Walk all prefixes", func(t *testing.T) {
|
|
walked := make(map[string]string)
|
|
trie.Walk(func(p netip.Prefix, v string) bool {
|
|
walked[p.String()] = v
|
|
return true
|
|
})
|
|
|
|
if !maps.Equal(walked, prefixes) {
|
|
t.Errorf("walked prefixes mismatch:\nexpected: %v\ngot: %v", prefixes, walked)
|
|
}
|
|
})
|
|
|
|
t.Run("Stop walk early", func(t *testing.T) {
|
|
count := 0
|
|
trie.Walk(func(p netip.Prefix, v string) bool {
|
|
count++
|
|
return count < 3 // Stop after visiting 3 prefixes
|
|
})
|
|
|
|
if count != 3 {
|
|
t.Errorf("expected walk to stop after 3 prefixes, but it visited %d", count)
|
|
}
|
|
})
|
|
|
|
t.Run("Stop walk between families", func(t *testing.T) {
|
|
stopTrie := NewValue[int]()
|
|
stopTrie.Insert(netip.MustParsePrefix("10.0.0.0/8"), 1)
|
|
stopTrie.Insert(netip.MustParsePrefix("2001:db8::/32"), 2)
|
|
|
|
count := 0
|
|
stopTrie.Walk(func(p netip.Prefix, v int) bool {
|
|
count++
|
|
return false
|
|
})
|
|
|
|
if count != 1 {
|
|
t.Errorf("expected walk to stop after 1 prefix, but it visited %d", count)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestValueTrie_Merge(t *testing.T) {
|
|
trieA := NewValue[string]()
|
|
trieA.Insert(netip.MustParsePrefix("10.0.0.0/8"), "net10")
|
|
trieA.Insert(netip.MustParsePrefix("192.168.1.0/24"), "lan_A")
|
|
trieA.Insert(netip.MustParsePrefix("2001:db8::/32"), "v6_A")
|
|
|
|
trieB := NewValue[string]()
|
|
trieB.Insert(netip.MustParsePrefix("10.1.0.0/16"), "net10_subnet")
|
|
trieB.Insert(netip.MustParsePrefix("192.168.1.0/24"), "lan_B_override") // Overlap
|
|
trieB.Insert(netip.MustParsePrefix("172.16.0.0/12"), "corp")
|
|
trieB.Insert(netip.MustParsePrefix("2001:db8:acad::/48"), "v6_B")
|
|
|
|
trieA.Merge(trieB)
|
|
|
|
expected := map[string]string{
|
|
"10.0.0.0/8": "net10",
|
|
"192.168.1.0/24": "lan_B_override", // Value should be from trieB
|
|
"2001:db8::/32": "v6_A",
|
|
"10.1.0.0/16": "net10_subnet",
|
|
"172.16.0.0/12": "corp",
|
|
"2001:db8:acad::/48": "v6_B",
|
|
}
|
|
|
|
actual := make(map[string]string)
|
|
trieA.Walk(func(p netip.Prefix, v string) bool {
|
|
actual[p.String()] = v
|
|
return true
|
|
})
|
|
|
|
if !reflect.DeepEqual(expected, actual) {
|
|
t.Errorf("Merge result incorrect.\nExpected: %v\nGot: %v", expected, actual)
|
|
}
|
|
|
|
// Verify trieB is unchanged by collecting its prefixes
|
|
bPrefixes := []string{}
|
|
trieB.Walk(func(p netip.Prefix, v string) bool {
|
|
bPrefixes = append(bPrefixes, p.String())
|
|
return true
|
|
})
|
|
expectedBPrefixes := []string{"10.1.0.0/16", "172.16.0.0/12", "192.168.1.0/24", "2001:db8:acad::/48"}
|
|
slices.Sort(bPrefixes)
|
|
slices.Sort(expectedBPrefixes)
|
|
if !reflect.DeepEqual(bPrefixes, expectedBPrefixes) {
|
|
t.Errorf("trieB was modified during merge.\nExpected: %v\nGot: %v", expectedBPrefixes, bPrefixes)
|
|
}
|
|
}
|