tree4.go 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203
  1. package cidr
  2. import (
  3. "net"
  4. "github.com/slackhq/nebula/iputil"
  5. )
  6. type Node[T any] struct {
  7. left *Node[T]
  8. right *Node[T]
  9. parent *Node[T]
  10. hasValue bool
  11. value T
  12. }
  13. type entry[T any] struct {
  14. CIDR *net.IPNet
  15. Value T
  16. }
  17. type Tree4[T any] struct {
  18. root *Node[T]
  19. list []entry[T]
  20. }
  21. const (
  22. startbit = iputil.VpnIp(0x80000000)
  23. )
  24. func NewTree4[T any]() *Tree4[T] {
  25. tree := new(Tree4[T])
  26. tree.root = &Node[T]{}
  27. tree.list = []entry[T]{}
  28. return tree
  29. }
  30. func (tree *Tree4[T]) AddCIDR(cidr *net.IPNet, val T) {
  31. bit := startbit
  32. node := tree.root
  33. next := tree.root
  34. ip := iputil.Ip2VpnIp(cidr.IP)
  35. mask := iputil.Ip2VpnIp(cidr.Mask)
  36. // Find our last ancestor in the tree
  37. for bit&mask != 0 {
  38. if ip&bit != 0 {
  39. next = node.right
  40. } else {
  41. next = node.left
  42. }
  43. if next == nil {
  44. break
  45. }
  46. bit = bit >> 1
  47. node = next
  48. }
  49. // We already have this range so update the value
  50. if next != nil {
  51. addCIDR := cidr.String()
  52. for i, v := range tree.list {
  53. if addCIDR == v.CIDR.String() {
  54. tree.list = append(tree.list[:i], tree.list[i+1:]...)
  55. break
  56. }
  57. }
  58. tree.list = append(tree.list, entry[T]{CIDR: cidr, Value: val})
  59. node.value = val
  60. node.hasValue = true
  61. return
  62. }
  63. // Build up the rest of the tree we don't already have
  64. for bit&mask != 0 {
  65. next = &Node[T]{}
  66. next.parent = node
  67. if ip&bit != 0 {
  68. node.right = next
  69. } else {
  70. node.left = next
  71. }
  72. bit >>= 1
  73. node = next
  74. }
  75. // Final node marks our cidr, set the value
  76. node.value = val
  77. node.hasValue = true
  78. tree.list = append(tree.list, entry[T]{CIDR: cidr, Value: val})
  79. }
  80. // Contains finds the first match, which may be the least specific
  81. func (tree *Tree4[T]) Contains(ip iputil.VpnIp) (ok bool, value T) {
  82. bit := startbit
  83. node := tree.root
  84. for node != nil {
  85. if node.hasValue {
  86. return true, node.value
  87. }
  88. if ip&bit != 0 {
  89. node = node.right
  90. } else {
  91. node = node.left
  92. }
  93. bit >>= 1
  94. }
  95. return false, value
  96. }
  97. // MostSpecificContains finds the most specific match
  98. func (tree *Tree4[T]) MostSpecificContains(ip iputil.VpnIp) (ok bool, value T) {
  99. bit := startbit
  100. node := tree.root
  101. for node != nil {
  102. if node.hasValue {
  103. value = node.value
  104. ok = true
  105. }
  106. if ip&bit != 0 {
  107. node = node.right
  108. } else {
  109. node = node.left
  110. }
  111. bit >>= 1
  112. }
  113. return ok, value
  114. }
  115. type eachFunc[T any] func(T) bool
  116. // EachContains will call a function, passing the value, for each entry until the function returns true or the search is complete
  117. // The final return value will be true if the provided function returned true
  118. func (tree *Tree4[T]) EachContains(ip iputil.VpnIp, each eachFunc[T]) bool {
  119. bit := startbit
  120. node := tree.root
  121. for node != nil {
  122. if node.hasValue {
  123. // If the each func returns true then we can exit the loop
  124. if each(node.value) {
  125. return true
  126. }
  127. }
  128. if ip&bit != 0 {
  129. node = node.right
  130. } else {
  131. node = node.left
  132. }
  133. bit >>= 1
  134. }
  135. return false
  136. }
  137. // GetCIDR returns the entry added by the most recent matching AddCIDR call
  138. func (tree *Tree4[T]) GetCIDR(cidr *net.IPNet) (ok bool, value T) {
  139. bit := startbit
  140. node := tree.root
  141. ip := iputil.Ip2VpnIp(cidr.IP)
  142. mask := iputil.Ip2VpnIp(cidr.Mask)
  143. // Find our last ancestor in the tree
  144. for node != nil && bit&mask != 0 {
  145. if ip&bit != 0 {
  146. node = node.right
  147. } else {
  148. node = node.left
  149. }
  150. bit = bit >> 1
  151. }
  152. if bit&mask == 0 && node != nil {
  153. value = node.value
  154. ok = node.hasValue
  155. }
  156. return ok, value
  157. }
  158. // List will return all CIDRs and their current values. Do not modify the contents!
  159. func (tree *Tree4[T]) List() []entry[T] {
  160. return tree.list
  161. }