nebula/cidr/tree6.go

190 lines
2.8 KiB
Go
Raw Normal View History

package cidr
2021-03-18 19:37:24 -06:00
import (
"net"
"github.com/slackhq/nebula/iputil"
2021-03-18 19:37:24 -06:00
)
2021-04-01 09:23:31 -06:00
const startbit6 = uint64(1 << 63)
type Tree6[T any] struct {
root4 *Node[T]
root6 *Node[T]
2021-03-18 19:37:24 -06:00
}
func NewTree6[T any]() *Tree6[T] {
tree := new(Tree6[T])
tree.root4 = &Node[T]{}
tree.root6 = &Node[T]{}
2021-03-18 19:37:24 -06:00
return tree
}
func (tree *Tree6[T]) AddCIDR(cidr *net.IPNet, val T) {
var node, next *Node[T]
2021-03-18 19:37:24 -06:00
cidrIP, ipv4 := isIPV4(cidr.IP)
if ipv4 {
node = tree.root4
next = tree.root4
} else {
node = tree.root6
next = tree.root6
}
for i := 0; i < len(cidrIP); i += 4 {
ip := iputil.Ip2VpnIp(cidrIP[i : i+4])
mask := iputil.Ip2VpnIp(cidr.Mask[i : i+4])
2021-03-18 19:37:24 -06:00
bit := startbit
// Find our last ancestor in the tree
for bit&mask != 0 {
if ip&bit != 0 {
next = node.right
} else {
next = node.left
}
if next == nil {
break
}
bit = bit >> 1
node = next
}
// Build up the rest of the tree we don't already have
for bit&mask != 0 {
next = &Node[T]{}
2021-03-18 19:37:24 -06:00
next.parent = node
if ip&bit != 0 {
node.right = next
} else {
node.left = next
}
bit >>= 1
node = next
}
}
// Final node marks our cidr, set the value
node.value = val
node.hasValue = true
2021-03-18 19:37:24 -06:00
}
2021-04-01 09:23:31 -06:00
// Finds the most specific match
func (tree *Tree6[T]) MostSpecificContains(ip net.IP) (ok bool, value T) {
var node *Node[T]
2021-03-18 19:37:24 -06:00
wholeIP, ipv4 := isIPV4(ip)
if ipv4 {
node = tree.root4
} else {
node = tree.root6
}
for i := 0; i < len(wholeIP); i += 4 {
ip := iputil.Ip2VpnIp(wholeIP[i : i+4])
2021-03-18 19:37:24 -06:00
bit := startbit
for node != nil {
if node.hasValue {
2021-04-01 09:23:31 -06:00
value = node.value
ok = true
2021-03-18 19:37:24 -06:00
}
if bit == 0 {
break
}
if ip&bit != 0 {
node = node.right
} else {
node = node.left
}
bit >>= 1
}
}
return ok, value
2021-03-18 19:37:24 -06:00
}
func (tree *Tree6[T]) MostSpecificContainsIpV4(ip iputil.VpnIp) (ok bool, value T) {
2021-04-01 09:23:31 -06:00
bit := startbit
node := tree.root4
2021-03-18 19:37:24 -06:00
2021-04-01 09:23:31 -06:00
for node != nil {
if node.hasValue {
2021-04-01 09:23:31 -06:00
value = node.value
ok = true
2021-04-01 09:23:31 -06:00
}
if ip&bit != 0 {
node = node.right
} else {
node = node.left
}
bit >>= 1
2021-03-18 19:37:24 -06:00
}
return ok, value
2021-04-01 09:23:31 -06:00
}
func (tree *Tree6[T]) MostSpecificContainsIpV6(hi, lo uint64) (ok bool, value T) {
2021-04-01 09:23:31 -06:00
ip := hi
node := tree.root6
for i := 0; i < 2; i++ {
bit := startbit6
2021-03-18 19:37:24 -06:00
for node != nil {
if node.hasValue {
2021-03-18 19:37:24 -06:00
value = node.value
ok = true
2021-03-18 19:37:24 -06:00
}
if bit == 0 {
break
}
if ip&bit != 0 {
node = node.right
} else {
node = node.left
}
bit >>= 1
}
2021-04-01 09:23:31 -06:00
ip = lo
2021-03-18 19:37:24 -06:00
}
return ok, value
2021-03-18 19:37:24 -06:00
}
func isIPV4(ip net.IP) (net.IP, bool) {
if len(ip) == net.IPv4len {
return ip, true
}
if len(ip) == net.IPv6len && isZeros(ip[0:10]) && ip[10] == 0xff && ip[11] == 0xff {
return ip[12:16], true
}
return ip, false
}
func isZeros(p net.IP) bool {
for i := 0; i < len(p); i++ {
if p[i] != 0 {
return false
}
}
return true
}