initial commit

alloc.go

@@ -0,0 +1,95 @@
+package netalloc
+
+import (
+	"errors"
+	"net"
+
+	"github.com/mikioh/ipaddr"
+)
+
+var (
+	ErrPrefixFull = errors.New("no more ip addresses available in prefix")
+)
+
+type Allocator interface {
+	Alloc() (*net.IPAddr, error)
+	AllocCIDR() (cidr string, err error)
+	Free(*net.IPAddr) error
+	FreeCIDR(cidr string) error
+	SetStore(store Store)
+}
+
+type GenericAlloc struct {
+	store  Store
+	prefix *net.IPNet
+}
+
+func NewGenericAlloc(store Store, prefix string) (Allocator, error) {
+	_, pfx, err := net.ParseCIDR(prefix)
+	if err != nil {
+		return nil, err
+	}
+	return &GenericAlloc{
+		store:  store,
+		prefix: pfx,
+	}, nil
+}
+
+func NewMemAlloc(prefix string) (Allocator, error) {
+	return NewGenericAlloc(NewMemStore(), prefix)
+}
+
+func (a *GenericAlloc) Alloc() (ip *net.IPAddr, err error) {
+	a.store.Lock()
+	defer a.store.Unlock()
+	c := ipaddr.NewCursor([]ipaddr.Prefix{*ipaddr.NewPrefix(a.prefix)})
+	pos := c.Next()
+	last := c.Last()
+	for ; pos != nil && pos != last; pos = c.Next() {
+		net := &net.IPAddr{IP: pos.IP}
+		cidr := CIDR(a.prefix, net)
+		ok, e := a.store.Exists(cidr)
+		if e != nil {
+			return nil, e
+		}
+		if !ok {
+			a.store.Add(cidr)
+			ip = net
+			return
+		}
+	}
+	err = ErrPrefixFull
+	return
+}
+
+func (a *GenericAlloc) AllocCIDR() (cidr string, err error) {
+	ip, err := a.Alloc()
+	return CIDR(a.prefix, ip), err
+}
+
+func (a *GenericAlloc) Free(ip *net.IPAddr) (err error) {
+	a.store.Lock()
+	defer a.store.Unlock()
+	cidr := CIDR(a.prefix, ip)
+	a.store.Delete(cidr)
+	return
+}
+
+func (a *GenericAlloc) FreeCIDR(cidr string) (err error) {
+	ip, _, err := net.ParseCIDR(cidr)
+	if err != nil {
+		return
+	}
+	return a.Free(&net.IPAddr{IP: ip})
+}
+
+func (a *GenericAlloc) SetStore(store Store) {
+	a.store = store
+}
+
+func CIDR(prefix *net.IPNet, ip *net.IPAddr) string {
+	n := new(net.IPNet)
+	*n = *prefix
+	n.IP = ip.IP
+	return n.String()
+}