personal_interest_mirrors/nebula
1
0
Fork 0
mirror of https://github.com/slackhq/nebula.git synced 2025-12-30 18:48:28 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

Use generics for CIDRTrees to avoid casting issues (#1004)

Browse Source
This commit is contained in:
Nate Brown
2023-11-02 17:05:08 -05:00
committed by GitHub
parent a44e1b8b05
commit 5181cb0474
21 changed files with 264 additions and 247 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
cidr/tree4.go
View File

@@ -6,35 +6,36 @@ import (
"github.com/slackhq/nebula/iputil"
)
type Node struct {
left *Node
right *Node
parent *Node
value interface{}
type Node[T any] struct {
left *Node[T]
right *Node[T]
parent *Node[T]
hasValue bool
value T
}
type entry struct {
type entry[T any] struct {
CIDR *net.IPNet
Value *interface{}
Value T
}
type Tree4 struct {
root *Node
list []entry
type Tree4[T any] struct {
root *Node[T]
list []entry[T]
}
const (
startbit = iputil.VpnIp(0x80000000)
)
func NewTree4() *Tree4 {
tree := new(Tree4)
tree.root = &Node{}
tree.list = []entry{}
func NewTree4[T any]() *Tree4[T] {
tree := new(Tree4[T])
tree.root = &Node[T]{}
tree.list = []entry[T]{}
return tree
}
func (tree *Tree4) AddCIDR(cidr *net.IPNet, val interface{}) {
func (tree *Tree4[T]) AddCIDR(cidr *net.IPNet, val T) {
bit := startbit
node := tree.root
next := tree.root
@@ -68,14 +69,15 @@ func (tree *Tree4) AddCIDR(cidr *net.IPNet, val interface{}) {
}
}
tree.list = append(tree.list, entry{CIDR: cidr, Value: &val})
tree.list = append(tree.list, entry[T]{CIDR: cidr, Value: val})
node.value = val
node.hasValue = true
return
}
// Build up the rest of the tree we don't already have
for bit&mask != 0 {
next = &Node{}
next = &Node[T]{}
next.parent = node
if ip&bit != 0 {
@@ -90,17 +92,18 @@ func (tree *Tree4) AddCIDR(cidr *net.IPNet, val interface{}) {
// Final node marks our cidr, set the value
node.value = val
tree.list = append(tree.list, entry{CIDR: cidr, Value: &val})
node.hasValue = true
tree.list = append(tree.list, entry[T]{CIDR: cidr, Value: val})
}
// Contains finds the first match, which may be the least specific
func (tree *Tree4) Contains(ip iputil.VpnIp) (value interface{}) {
func (tree *Tree4[T]) Contains(ip iputil.VpnIp) (ok bool, value T) {
bit := startbit
node := tree.root
for node != nil {
if node.value != nil {
return node.value
if node.hasValue {
return true, node.value
}
if ip&bit != 0 {
@@ -113,17 +116,18 @@ func (tree *Tree4) Contains(ip iputil.VpnIp) (value interface{}) {
}
return value
return false, value
}
// MostSpecificContains finds the most specific match
func (tree *Tree4) MostSpecificContains(ip iputil.VpnIp) (value interface{}) {
func (tree *Tree4[T]) MostSpecificContains(ip iputil.VpnIp) (ok bool, value T) {
bit := startbit
node := tree.root
for node != nil {
if node.value != nil {
if node.hasValue {
value = node.value
ok = true
}
if ip&bit != 0 {
@@ -135,11 +139,12 @@ func (tree *Tree4) MostSpecificContains(ip iputil.VpnIp) (value interface{}) {
bit >>= 1
}
return value
return ok, value
}
// Match finds the most specific match
func (tree *Tree4) Match(ip iputil.VpnIp) (value interface{}) {
// TODO this is exact match
func (tree *Tree4[T]) Match(ip iputil.VpnIp) (ok bool, value T) {
bit := startbit
node := tree.root
lastNode := node
@@ -157,11 +162,12 @@ func (tree *Tree4) Match(ip iputil.VpnIp) (value interface{}) {
if bit == 0 && lastNode != nil {
value = lastNode.value
ok = true
}
return value
return ok, value
}
// List will return all CIDRs and their current values. Do not modify the contents!
func (tree *Tree4) List() []entry {
func (tree *Tree4[T]) List() []entry[T] {
return tree.list
}

118
cidr/tree4_test.go
View File

@@ -9,7 +9,7 @@ import (
)
func TestCIDRTree_List(t *testing.T) {
tree := NewTree4()
tree := NewTree4[string]()
tree.AddCIDR(Parse("1.0.0.0/16"), "1")
tree.AddCIDR(Parse("1.0.0.0/8"), "2")
tree.AddCIDR(Parse("1.0.0.0/16"), "3")
@@ -17,13 +17,13 @@ func TestCIDRTree_List(t *testing.T) {
list := tree.List()
assert.Len(t, list, 2)
assert.Equal(t, "1.0.0.0/8", list[0].CIDR.String())
assert.Equal(t, "2", *list[0].Value)
assert.Equal(t, "2", list[0].Value)
assert.Equal(t, "1.0.0.0/16", list[1].CIDR.String())
assert.Equal(t, "4", *list[1].Value)
assert.Equal(t, "4", list[1].Value)
}
func TestCIDRTree_Contains(t *testing.T) {
tree := NewTree4()
tree := NewTree4[string]()
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")
@@ -33,35 +33,43 @@ func TestCIDRTree_Contains(t *testing.T) {
tree.AddCIDR(Parse("254.0.0.0/4"), "5")
tests := []struct {
Found bool
Result interface{}
IP string
}{
{"1", "1.0.0.0"},
{"1", "1.255.255.255"},
{"2", "2.1.0.0"},
{"2", "2.1.255.255"},
{"3", "3.1.1.0"},
{"3", "3.1.1.255"},
{"4a", "4.1.1.255"},
{"4a", "4.1.1.1"},
{"5", "240.0.0.0"},
{"5", "255.255.255.255"},
{nil, "239.0.0.0"},
{nil, "4.1.2.2"},
{true, "1", "1.0.0.0"},
{true, "1", "1.255.255.255"},
{true, "2", "2.1.0.0"},
{true, "2", "2.1.255.255"},
{true, "3", "3.1.1.0"},
{true, "3", "3.1.1.255"},
{true, "4a", "4.1.1.255"},
{true, "4a", "4.1.1.1"},
{true, "5", "240.0.0.0"},
{true, "5", "255.255.255.255"},
{false, "", "239.0.0.0"},
{false, "", "4.1.2.2"},
}
for _, tt := range tests {
assert.Equal(t, tt.Result, tree.Contains(iputil.Ip2VpnIp(net.ParseIP(tt.IP))))
ok, r := tree.Contains(iputil.Ip2VpnIp(net.ParseIP(tt.IP)))
assert.Equal(t, tt.Found, ok)
assert.Equal(t, tt.Result, r)
}
tree = NewTree4()
tree = NewTree4[string]()
tree.AddCIDR(Parse("1.1.1.1/0"), "cool")
assert.Equal(t, "cool", tree.Contains(iputil.Ip2VpnIp(net.ParseIP("0.0.0.0"))))
assert.Equal(t, "cool", tree.Contains(iputil.Ip2VpnIp(net.ParseIP("255.255.255.255"))))
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 TestCIDRTree_MostSpecificContains(t *testing.T) {
tree := NewTree4()
tree := NewTree4[string]()
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")
@@ -71,59 +79,75 @@ func TestCIDRTree_MostSpecificContains(t *testing.T) {
tree.AddCIDR(Parse("254.0.0.0/4"), "5")
tests := []struct {
Found bool
Result interface{}
IP string
}{
{"1", "1.0.0.0"},
{"1", "1.255.255.255"},
{"2", "2.1.0.0"},
{"2", "2.1.255.255"},
{"3", "3.1.1.0"},
{"3", "3.1.1.255"},
{"4a", "4.1.1.255"},
{"4b", "4.1.1.2"},
{"4c", "4.1.1.1"},
{"5", "240.0.0.0"},
{"5", "255.255.255.255"},
{nil, "239.0.0.0"},
{nil, "4.1.2.2"},
{true, "1", "1.0.0.0"},
{true, "1", "1.255.255.255"},
{true, "2", "2.1.0.0"},
{true, "2", "2.1.255.255"},
{true, "3", "3.1.1.0"},
{true, "3", "3.1.1.255"},
{true, "4a", "4.1.1.255"},
{true, "4b", "4.1.1.2"},
{true, "4c", "4.1.1.1"},
{true, "5", "240.0.0.0"},
{true, "5", "255.255.255.255"},
{false, "", "239.0.0.0"},
{false, "", "4.1.2.2"},
}
for _, tt := range tests {
assert.Equal(t, tt.Result, tree.MostSpecificContains(iputil.Ip2VpnIp(net.ParseIP(tt.IP))))
ok, r := tree.MostSpecificContains(iputil.Ip2VpnIp(net.ParseIP(tt.IP)))
assert.Equal(t, tt.Found, ok)
assert.Equal(t, tt.Result, r)
}
tree = NewTree4()
tree = NewTree4[string]()
tree.AddCIDR(Parse("1.1.1.1/0"), "cool")
assert.Equal(t, "cool", tree.MostSpecificContains(iputil.Ip2VpnIp(net.ParseIP("0.0.0.0"))))
assert.Equal(t, "cool", tree.MostSpecificContains(iputil.Ip2VpnIp(net.ParseIP("255.255.255.255"))))
ok, r := tree.MostSpecificContains(iputil.Ip2VpnIp(net.ParseIP("0.0.0.0")))
assert.True(t, ok)
assert.Equal(t, "cool", r)
ok, r = tree.MostSpecificContains(iputil.Ip2VpnIp(net.ParseIP("255.255.255.255")))
assert.True(t, ok)
assert.Equal(t, "cool", r)
}
func TestCIDRTree_Match(t *testing.T) {
tree := NewTree4()
tree := NewTree4[string]()
tree.AddCIDR(Parse("4.1.1.0/32"), "1a")
tree.AddCIDR(Parse("4.1.1.1/32"), "1b")
tests := []struct {
Found bool
Result interface{}
IP string
}{
{"1a", "4.1.1.0"},
{"1b", "4.1.1.1"},
{true, "1a", "4.1.1.0"},
{true, "1b", "4.1.1.1"},
}
for _, tt := range tests {
assert.Equal(t, tt.Result, tree.Match(iputil.Ip2VpnIp(net.ParseIP(tt.IP))))
ok, r := tree.Match(iputil.Ip2VpnIp(net.ParseIP(tt.IP)))
assert.Equal(t, tt.Found, ok)
assert.Equal(t, tt.Result, r)
}
tree = NewTree4()
tree = NewTree4[string]()
tree.AddCIDR(Parse("1.1.1.1/0"), "cool")
assert.Equal(t, "cool", tree.Contains(iputil.Ip2VpnIp(net.ParseIP("0.0.0.0"))))
assert.Equal(t, "cool", tree.Contains(iputil.Ip2VpnIp(net.ParseIP("255.255.255.255"))))
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) {
tree := NewTree4()
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")
@@ -145,7 +169,7 @@ func BenchmarkCIDRTree_Contains(b *testing.B) {
}
func BenchmarkCIDRTree_Match(b *testing.B) {
tree := NewTree4()
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")

44
cidr/tree6.go
View File

@@ -8,20 +8,20 @@ import (
const startbit6 = uint64(1 << 63)
type Tree6 struct {
root4 *Node
root6 *Node
type Tree6[T any] struct {
root4 *Node[T]
root6 *Node[T]
}
func NewTree6() *Tree6 {
tree := new(Tree6)
tree.root4 = &Node{}
tree.root6 = &Node{}
func NewTree6[T any]() *Tree6[T] {
tree := new(Tree6[T])
tree.root4 = &Node[T]{}
tree.root6 = &Node[T]{}
return tree
}
func (tree *Tree6) AddCIDR(cidr *net.IPNet, val interface{}) {
var node, next *Node
func (tree *Tree6[T]) AddCIDR(cidr *net.IPNet, val T) {
var node, next *Node[T]
cidrIP, ipv4 := isIPV4(cidr.IP)
if ipv4 {
@@ -56,7 +56,7 @@ func (tree *Tree6) AddCIDR(cidr *net.IPNet, val interface{}) {
// Build up the rest of the tree we don't already have
for bit&mask != 0 {
next = &Node{}
next = &Node[T]{}
next.parent = node
if ip&bit != 0 {
@@ -72,11 +72,12 @@ func (tree *Tree6) AddCIDR(cidr *net.IPNet, val interface{}) {
// Final node marks our cidr, set the value
node.value = val
node.hasValue = true
}
// Finds the most specific match
func (tree *Tree6) MostSpecificContains(ip net.IP) (value interface{}) {
var node *Node
func (tree *Tree6[T]) MostSpecificContains(ip net.IP) (ok bool, value T) {
var node *Node[T]
wholeIP, ipv4 := isIPV4(ip)
if ipv4 {
@@ -90,8 +91,9 @@ func (tree *Tree6) MostSpecificContains(ip net.IP) (value interface{}) {
bit := startbit
for node != nil {
if node.value != nil {
if node.hasValue {
value = node.value
ok = true
}
if bit == 0 {
@@ -108,16 +110,17 @@ func (tree *Tree6) MostSpecificContains(ip net.IP) (value interface{}) {
}
}
return value
return ok, value
}
func (tree *Tree6) MostSpecificContainsIpV4(ip iputil.VpnIp) (value interface{}) {
func (tree *Tree6[T]) MostSpecificContainsIpV4(ip iputil.VpnIp) (ok bool, value T) {
bit := startbit
node := tree.root4
for node != nil {
if node.value != nil {
if node.hasValue {
value = node.value
ok = true
}
if ip&bit != 0 {
@@ -129,10 +132,10 @@ func (tree *Tree6) MostSpecificContainsIpV4(ip iputil.VpnIp) (value interface{})
bit >>= 1
}
return value
return ok, value
}
func (tree *Tree6) MostSpecificContainsIpV6(hi, lo uint64) (value interface{}) {
func (tree *Tree6[T]) MostSpecificContainsIpV6(hi, lo uint64) (ok bool, value T) {
ip := hi
node := tree.root6
@@ -140,8 +143,9 @@ func (tree *Tree6) MostSpecificContainsIpV6(hi, lo uint64) (value interface{}) {
bit := startbit6
for node != nil {
if node.value != nil {
if node.hasValue {
value = node.value
ok = true
}
if bit == 0 {
@@ -160,7 +164,7 @@ func (tree *Tree6) MostSpecificContainsIpV6(hi, lo uint64) (value interface{}) {
ip = lo
}
return value
return ok, value
}
func isIPV4(ip net.IP) (net.IP, bool) {

73
cidr/tree6_test.go
View File

@@ -9,7 +9,7 @@ import (
)
func TestCIDR6Tree_MostSpecificContains(t *testing.T) {
tree := NewTree6()
tree := NewTree6[string]()
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")
@@ -22,53 +22,68 @@ func TestCIDR6Tree_MostSpecificContains(t *testing.T) {
tree.AddCIDR(Parse("1:2:0:4:5:0:0:0/96"), "6c")
tests := []struct {
Found bool
Result interface{}
IP string
}{
{"1", "1.0.0.0"},
{"1", "1.255.255.255"},
{"2", "2.1.0.0"},
{"2", "2.1.255.255"},
{"3", "3.1.1.0"},
{"3", "3.1.1.255"},
{"4a", "4.1.1.255"},
{"4b", "4.1.1.2"},
{"4c", "4.1.1.1"},
{"5", "240.0.0.0"},
{"5", "255.255.255.255"},
{"6a", "1:2:0:4:1:1:1:1"},
{"6b", "1:2:0:4:5:1:1:1"},
{"6c", "1:2:0:4:5:0:0:0"},
{nil, "239.0.0.0"},
{nil, "4.1.2.2"},
{true, "1", "1.0.0.0"},
{true, "1", "1.255.255.255"},
{true, "2", "2.1.0.0"},
{true, "2", "2.1.255.255"},
{true, "3", "3.1.1.0"},
{true, "3", "3.1.1.255"},
{true, "4a", "4.1.1.255"},
{true, "4b", "4.1.1.2"},
{true, "4c", "4.1.1.1"},
{true, "5", "240.0.0.0"},
{true, "5", "255.255.255.255"},
{true, "6a", "1:2:0:4:1:1:1:1"},
{true, "6b", "1:2:0:4:5:1:1:1"},
{true, "6c", "1:2:0:4:5:0:0:0"},
{false, "", "239.0.0.0"},
{false, "", "4.1.2.2"},
}
for _, tt := range tests {
assert.Equal(t, tt.Result, tree.MostSpecificContains(net.ParseIP(tt.IP)))
ok, r := tree.MostSpecificContains(net.ParseIP(tt.IP))
assert.Equal(t, tt.Found, ok)
assert.Equal(t, tt.Result, r)
}
tree = NewTree6()
tree = NewTree6[string]()
tree.AddCIDR(Parse("1.1.1.1/0"), "cool")
tree.AddCIDR(Parse("::/0"), "cool6")
assert.Equal(t, "cool", tree.MostSpecificContains(net.ParseIP("0.0.0.0")))
assert.Equal(t, "cool", tree.MostSpecificContains(net.ParseIP("255.255.255.255")))
assert.Equal(t, "cool6", tree.MostSpecificContains(net.ParseIP("::")))
assert.Equal(t, "cool6", tree.MostSpecificContains(net.ParseIP("1:2:3:4:5:6:7:8")))
ok, r := tree.MostSpecificContains(net.ParseIP("0.0.0.0"))
assert.True(t, ok)
assert.Equal(t, "cool", r)
ok, r = tree.MostSpecificContains(net.ParseIP("255.255.255.255"))
assert.True(t, ok)
assert.Equal(t, "cool", r)
ok, r = tree.MostSpecificContains(net.ParseIP("::"))
assert.True(t, ok)
assert.Equal(t, "cool6", r)
ok, r = tree.MostSpecificContains(net.ParseIP("1:2:3:4:5:6:7:8"))
assert.True(t, ok)
assert.Equal(t, "cool6", r)
}
func TestCIDR6Tree_MostSpecificContainsIpV6(t *testing.T) {
tree := NewTree6()
tree := NewTree6[string]()
tree.AddCIDR(Parse("1:2:0:4:5:0:0:0/64"), "6a")
tree.AddCIDR(Parse("1:2:0:4:5:0:0:0/80"), "6b")
tree.AddCIDR(Parse("1:2:0:4:5:0:0:0/96"), "6c")
tests := []struct {
Found bool
Result interface{}
IP string
}{
{"6a", "1:2:0:4:1:1:1:1"},
{"6b", "1:2:0:4:5:1:1:1"},
{"6c", "1:2:0:4:5:0:0:0"},
{true, "6a", "1:2:0:4:1:1:1:1"},
{true, "6b", "1:2:0:4:5:1:1:1"},
{true, "6c", "1:2:0:4:5:0:0:0"},
}
for _, tt := range tests {
@@ -76,6 +91,8 @@ func TestCIDR6Tree_MostSpecificContainsIpV6(t *testing.T) {
hi := binary.BigEndian.Uint64(ip[:8])
lo := binary.BigEndian.Uint64(ip[8:])
assert.Equal(t, tt.Result, tree.MostSpecificContainsIpV6(hi, lo))
ok, r := tree.MostSpecificContainsIpV6(hi, lo)
assert.Equal(t, tt.Found, ok)
assert.Equal(t, tt.Result, r)
}
}