-
Notifications
You must be signed in to change notification settings - Fork 1
/
iplookuptree.go
160 lines (131 loc) · 3.92 KB
/
iplookuptree.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
package iplookuptree
import (
"encoding/binary"
"math"
"net"
)
// bitslen contstant must be a power of 2. It indicates how much space
// will be taken by a tree and maximum number of hops (treenode accesses). When bitslen is 4,
// the maximum number of hops will be 32 / bitslen and one node takes
// 1<< bitslen * (sizeof SID and *treenode). So current constant (4) will make maximum 8 hops and
// every node consumes 256 bytes.
const bitslen = 4
// SID type contains a list of corresponding service indexes.
// Every nth bit indicates nth service. So 0x1 stores (0),
// 0x2 stores (1), 0x3 stores (0, 1) services and etc.
// So you can handle up to 64 services.
// 0 indicates that there are no service.
// Example: service has index 6, then its SID representation will be 1<<6
type SID uint64
type Tree struct {
root *treenode
}
type treenode struct {
srvs [1<<bitslen]SID
ptrs [1<<bitslen]*treenode
}
// isEmpty checks if node is empty. Used to determine whether to delete node
// Reminder: don't accidentally remove root node
func (node *treenode) isEmpty() bool {
for i := 0; i < 1<<bitslen; i++ {
if node.srvs[i] != 0 || node.ptrs[i] != nil {
return false
}
}
return true
}
// New creates new IP subnet tree. It only works with IP v4
func New() *Tree {
tree := &Tree{&treenode{}}
return tree
}
// Add adds ip subnet to the tree.
// service should contain only one true bit (anyway it works well with multiple bits).
// ipnet.IP must be of the length 4 (IP v4).
// It is up to you to handle this.
// This method does not compress subnets - if you put 1.1.1.1/24 and 1.1.1.1/23
// of the same service, it will hold both subnets.
func (tree *Tree) Add(service SID, ipnet net.IPNet) {
node := tree.root
prefixLen, _ := ipnet.Mask.Size()
curLen := bitslen
for i := 0; i < 32 / bitslen; i++ {
if curLen >= prefixLen {
start := getSubstring(ipnet.IP, uint8(i))
end := start + (1<<uint(curLen - prefixLen)) - 1
for j := start; j <= end; j++ {
node.srvs[j] = node.srvs[j] | service
}
break
}
ind := getSubstring(ipnet.IP, uint8(i))
if node.ptrs[ind] != nil {
node = node.ptrs[ind]
} else {
node.ptrs[ind] = &treenode{}
node = node.ptrs[ind]
}
curLen += bitslen
}
}
// Get returns SID which corresponds to this ip v4
// ipv4 must be of length 4 and it is up to you to
// handle this.
func (tree *Tree) Get(ipv4 []byte) SID {
var ans SID
cur := tree.root
for i := 0; i < 32 / bitslen; i++ {
ind := getSubstring(ipv4, uint8(i))
ans = ans | cur.srvs[ind]
if cur = cur.ptrs[ind]; cur == nil {
break
}
}
return ans
}
// getSubstring is helper function that returns substring of
// bits placed in range [index * bitslen, index * bitslen + bitslen)
func getSubstring(ipv4 []byte, index uint8) uint32 {
var ans = binary.BigEndian.Uint32(ipv4)
ans = ans <<(bitslen * index)
ans = ans >>(32 - bitslen)
return ans
}
// Remove removes subnet from the tree. It works properly if you add
// and remove the same subnet. However, it will lead to undefined behaviour, if
// you remove subnet which was not added before.
func (tree *Tree) Remove(service SID, ipnet net.IPNet) {
reversedService := math.MaxUint64 ^ service
node := tree.root
path := make([]*treenode, 32 / bitslen, 32 / bitslen)
indpath := make([]uint32, 32 / bitslen, 32 / bitslen)
pathlen := 0
prefixLen, _ := ipnet.Mask.Size()
curLen := bitslen
for i := 0; i < 32 / bitslen; i++ {
path[pathlen] = node
pathlen++
if curLen >= prefixLen {
start := getSubstring(ipnet.IP, uint8(i))
end := start + (1<<uint(curLen - prefixLen)) - 1
for j := start; j <= end; j++ {
node.srvs[j] = node.srvs[j] & reversedService
}
break
}
ind := getSubstring(ipnet.IP, uint8(i))
indpath[pathlen] = ind
if node.ptrs[ind] != nil {
node = node.ptrs[ind]
}
curLen += bitslen
}
for i := pathlen - 1; i > 0; i-- {
node = path[i]
parent := path[i - 1]
ind := indpath[i]
if node.isEmpty() {
parent.ptrs[ind] = nil
}
}
}