allocators: Create a bitmap allocator for IPv4

Signed-off-by: Anatole Denis <anatole@unverle.fr>

plugins/allocators/bitmap/bitmap_ipv4.go

@@ -0,0 +1,125 @@
+// Copyright 2018-present the CoreDHCP Authors. All rights reserved
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+package bitmap
+
+// This allocator handles IPv4 assignments with a similar logic to the base bitmap, but a simpler
+// implementation due to the ability to just use uint32 for IPv4 addresses
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"net"
+	"sync"
+
+	"github.com/coredhcp/coredhcp/plugins/allocators"
+	"github.com/willf/bitset"
+)
+
+var (
+	// ErrNoAddrAvail is returned when we can't allocate an IP
+	// because there's no unallocated space left
+	ErrNoAddrAvail = errors.New("No address available to allocate")
+
+	errNotInRange = errors.New("IP outside of allowed range")
+
+	errInvalidIP = errors.New("Invalid IP passed as input")
+)
+
+// IPv4Allocator allocates IPv4 addresses, tracking utilization with a bitmap
+type IPv4Allocator struct {
+	start uint32
+	end   uint32
+
+	bitmap *bitset.BitSet
+	l      sync.Mutex
+}
+
+func (a *IPv4Allocator) toIP(offset uint32) net.IP {
+	if offset > a.end-a.start {
+		return nil
+	}
+
+	r := make(net.IP, 4)
+	binary.BigEndian.PutUint32(r, a.start+offset)
+	return r
+}
+
+func (a *IPv4Allocator) toOffset(ip net.IP) (uint, error) {
+	if ip.To4() == nil {
+		return 0, errInvalidIP
+	}
+
+	intIP := binary.BigEndian.Uint32(ip.To4())
+	if intIP < a.start || intIP > a.end {
+		return 0, errNotInRange
+	}
+
+	return uint(intIP - a.start), nil
+}
+
+// Allocate reserves an IP for a client
+func (a *IPv4Allocator) Allocate(hint net.IPNet) (n net.IPNet, err error) {
+	n.Mask = net.CIDRMask(32, 32)
+
+	// This is just a hint, ignore any error with it
+	hintOffset, _ := a.toOffset(hint.IP)
+
+	a.l.Lock()
+	defer a.l.Unlock()
+
+	var next uint
+	// First try the exact match
+	if a.bitmap.Test(hintOffset) {
+		next = hintOffset
+	} else {
+		// Then any available address
+		avail, ok := a.bitmap.NextClear(0)
+		if !ok {
+			return n, ErrNoAddrAvail
+		}
+		next = avail
+	}
+
+	a.bitmap.Set(next)
+	n.IP = a.toIP(uint32(next))
+	return
+}
+
+// Free releases the given IP
+func (a *IPv4Allocator) Free(n net.IPNet) error {
+	offset, err := a.toOffset(n.IP)
+	if err != nil {
+		return errNotInRange
+	}
+
+	a.l.Lock()
+	defer a.l.Unlock()
+
+	if !a.bitmap.Test(uint(offset)) {
+		return &allocators.ErrDoubleFree{Loc: n}
+	}
+	a.bitmap.Clear(offset)
+	return nil
+}
+
+// NewIPv4Allocator creates a new allocator suitable for giving out IPv4 addresses
+func NewIPv4Allocator(start, end net.IP) (*IPv4Allocator, error) {
+	if start.To4() == nil || end.To4() == nil {
+		return nil, fmt.Errorf("Invalid IPv4s given to create the allocator: [%s,%s]", start, end)
+	}
+
+	alloc := IPv4Allocator{
+		start: binary.BigEndian.Uint32(start.To4()),
+		end:   binary.BigEndian.Uint32(end.To4()),
+	}
+
+	if alloc.start > alloc.end {
+		return nil, errors.New("No IPs in the given range to allocate")
+	}
+	alloc.bitmap = bitset.New(uint(alloc.end - alloc.start + 1))
+
+	return &alloc, nil
+}

plugins/allocators/bitmap/bitmap_ipv4_test.go

@@ -0,0 +1,54 @@
+// Copyright 2018-present the CoreDHCP Authors. All rights reserved
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+package bitmap
+
+import (
+	"net"
+	"testing"
+)
+
+func getv4Allocator() *IPv4Allocator {
+	alloc, err := NewIPv4Allocator(net.IPv4(192, 0, 2, 0), net.IPv4(192, 0, 2, 255))
+	if err != nil {
+		panic(err)
+	}
+
+	return alloc
+}
+func Test4Alloc(t *testing.T) {
+	alloc := getv4Allocator()
+
+	net, err := alloc.Allocate(net.IPNet{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = alloc.Free(net)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = alloc.Free(net)
+	if err == nil {
+		t.Fatal("Expected DoubleFree error")
+	}
+}
+
+func Test4OutOfPool(t *testing.T) {
+	alloc := getv4Allocator()
+
+	hint := net.IPv4(198, 51, 100, 5)
+	res, err := alloc.Allocate(net.IPNet{IP: hint, Mask: net.CIDRMask(32, 32)})
+	if err != nil {
+		t.Fatalf("Failed to allocate with invalid hint: %v", err)
+	}
+	_, prefix, _ := net.ParseCIDR("192.0.2.0/24")
+	if !prefix.Contains(res.IP) {
+		t.Fatal("Obtained prefix outside of range: ", res)
+	}
+	if prefLen, totalLen := res.Mask.Size(); prefLen != 32 || totalLen != 32 {
+		t.Fatalf("Prefixes have wrong size %d/%d", prefLen, totalLen)
+	}
+}