67 lines
1.5 KiB
Go
67 lines
1.5 KiB
Go
package wg
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"errors"
|
|
"net/netip"
|
|
|
|
"github.com/hashicorp/go-set/v3"
|
|
)
|
|
|
|
type IPPool struct {
|
|
scope netip.Prefix
|
|
leasure *set.Set[netip.Addr]
|
|
cache netip.Addr
|
|
}
|
|
|
|
var ErrIPPoolExhausted = errors.New("IP Pool Exhausted")
|
|
|
|
// @param base_ip4: is used as the gateway address
|
|
func IPPool_Init(base_ip4 netip.Addr, cidr uint8) *IPPool {
|
|
scope := netip.PrefixFrom(base_ip4, int(cidr))
|
|
pool := &IPPool{leasure: set.New[netip.Addr](1 << (32 - cidr)), scope: scope}
|
|
pool.cache = pool.scope.Addr().Next()
|
|
|
|
//calculate broadcast address
|
|
var wildcard uint32 = (1 << (32 - cidr)) - 1
|
|
|
|
wildcard_mask := [4]byte{}
|
|
binary.BigEndian.PutUint32(wildcard_mask[:4], wildcard)
|
|
broadcast := base_ip4.As4()
|
|
broadcast[0] |= wildcard_mask[0]
|
|
broadcast[1] |= wildcard_mask[1]
|
|
broadcast[2] |= wildcard_mask[2]
|
|
broadcast[3] |= wildcard_mask[3]
|
|
|
|
broadcast_addr := netip.AddrFrom4(broadcast)
|
|
pool.Lease(broadcast_addr)
|
|
//gateway
|
|
pool.Lease(base_ip4)
|
|
|
|
return pool
|
|
}
|
|
func (p *IPPool) Next() (netip.Addr, error) {
|
|
cyclic := false
|
|
for ; ; p.cache = p.cache.Next() {
|
|
if p.leasure.Contains(p.cache) && p.cache.IsValid() {
|
|
continue
|
|
} else {
|
|
if p.cache.IsValid() && p.scope.Contains(p.cache) && !p.leasure.Contains(p.cache) {
|
|
return p.cache, nil
|
|
} else if !cyclic {
|
|
p.cache = p.scope.Masked().Addr()
|
|
cyclic = true
|
|
|
|
} else {
|
|
return p.cache, ErrIPPoolExhausted
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (p *IPPool) Lease(addr netip.Addr) {
|
|
p.leasure.Insert(addr)
|
|
}
|
|
func (p *IPPool) HasIP(addr netip.Addr) {
|
|
p.leasure.Contains(addr)
|
|
}
|