-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
186 lines (151 loc) · 4.62 KB
/
main.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
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
package main
import (
"encoding/binary"
"errors"
"net/netip"
"slices"
)
var ErrNoFreePool = errors.New("no free address pools")
type Allocator struct {
pools []Pool
allocated []netip.Prefix
}
type Pool struct {
Prefix netip.Prefix
Size int
}
func NewAllocator(pools []Pool) *Allocator {
for i, p := range pools {
pools[i].Prefix = p.Prefix.Masked()
}
slices.SortFunc(pools, func(a, b Pool) int {
addr := a.Prefix.Addr().Compare(b.Prefix.Addr())
if addr < 0 || addr > 0 {
return addr
}
if a.Prefix.Bits() < b.Prefix.Bits() {
return -1
} else if a.Prefix.Bits() > b.Prefix.Bits() {
return 1
}
return 0
})
return &Allocator{
pools: pools,
allocated: []netip.Prefix{},
}
}
func (a *Allocator) Allocate() (netip.Prefix, error) {
var i, poolID int
var partialOverlap bool
nextPrefixAfter := func(prev, p netip.Prefix, targetSz int) netip.Prefix {
// TOOD: is there a simpler way? This code looks silly.
next := nextPrefix(netip.PrefixFrom(prev.Addr(), targetSz).Masked())
if p.Overlaps(next) {
return next
}
return netip.Prefix{}
}
for i < len(a.allocated) {
allocated := a.allocated[i]
if poolID >= len(a.pools) {
return netip.Prefix{}, ErrNoFreePool
}
p := a.pools[poolID]
if allocated.Overlaps(p.Prefix) {
i++
if allocated.Bits() <= p.Prefix.Bits() {
// The current 'allocated' prefix is bigger than the pool, thus
// the pool is fully overlapped.
partialOverlap = false
poolID++
continue
}
if lastAddr(allocated) == lastAddr(p.Prefix) {
// The last address of the current 'allocated' prefix is the
// same as the last address of the pool, it's fully overlapped.
// We can go to the next one.
partialOverlap = false
poolID++
continue
}
// This pool is partially overlapped. If the next iteration yields
// an 'allocated' prefix that don't overlap with the current pool,
// then might have found the right spot.
partialOverlap = true
continue
}
// Okay, so previous 'allocated' overlapped and current doesn't. Now
// the question is: is there enough space left between previous
// 'allocated' and the end of p?
if partialOverlap {
partialOverlap = false
// No need to check if 'i > 0' -- the lowest 'i' where 'partialOverlap'
// could be set is 1.
prevAlloc := a.allocated[i-1]
if next := nextPrefixAfter(prevAlloc, p.Prefix, p.Size); next != (netip.Prefix{}) {
a.allocated = slices.Insert(a.allocated, i, next)
return next, nil
}
// No luck -- nextPrefixAfter yielded an invalid prefix. There's
// not enough space left to use this pool.
poolID++
// We don't increment 'i' here, because we need to re-test the
// current 'allocated' against the next pool available.
continue
}
// If the pool doesn't overlap and has a binary value lower than the
// current 'allocated', we found the right spot.
if p.Prefix.Addr().Less(allocated.Addr()) {
copy(a.allocated[i+1:], a.allocated[i:])
a.allocated[i] = netip.PrefixFrom(p.Prefix.Addr(), p.Size)
return a.allocated[i], nil
}
i++
}
if poolID >= len(a.pools) {
return netip.Prefix{}, ErrNoFreePool
}
// We reached the end of 'allocated', but not the end of pools. Let's
// try two more times (once on the current 'p', and once on the next pool
// if any).
if partialOverlap {
p := a.pools[poolID]
prevAlloc := a.allocated[i-1]
if next := nextPrefixAfter(prevAlloc, p.Prefix, p.Size); next != (netip.Prefix{}) {
a.allocated = slices.Insert(a.allocated, i, next)
return next, nil
}
// No luck -- next yielded an invalid prefix. There's not enough
// space left to use this pool.
poolID++
}
// One last chance. Here we don't increment poolID since the last iteration
// on 'a.allocated' found either:
//
// - A full overlap, and incremented 'poolID'.
// - A partial overlap, and the previous 'if' incremented 'poolID'.
// - The current 'poolID' comes after the last 'allocated'.
//
// Hence, we're sure 'poolID' has never been subnetted yet.
if poolID < len(a.pools) {
p := a.pools[poolID]
a.allocated = append(a.allocated, netip.PrefixFrom(p.Prefix.Addr(), p.Size))
return a.allocated[i], nil
}
return netip.Prefix{}, ErrNoFreePool
}
func lastAddr(p netip.Prefix) netip.Addr {
return Add(p.Addr(), 1, uint(32-p.Bits())).Prev()
}
func nextPrefix(p netip.Prefix) netip.Prefix {
return netip.PrefixFrom(lastAddr(p).Next(), p.Bits())
}
// Add returns ip + (x << shift).
func Add(ip netip.Addr, x uint64, shift uint) netip.Addr {
a := ip.As4()
addr := binary.BigEndian.Uint32(a[:])
addr += uint32(x) << shift
binary.BigEndian.PutUint32(a[:], addr)
return netip.AddrFrom4(a)
}