Docker-DCO-1.1-Signed-off-by: Michael Crosby <michael@crosbymichael.com> (github: crosbymichael)

2026-10-04 07:48:44 +00:00 · 2014-01-23 00:28:35 -08:00
parent 79bb8212e0
commit 648357ffdf
3 changed files with 0 additions and 0 deletions
@@ -0,0 +1,229 @@
+package ipallocator
+
+import (
+	"encoding/binary"
+	"errors"
+	"github.com/dotcloud/docker/pkg/netlink"
+	"net"
+	"sync"
+)
+
+type networkSet map[iPNet]*iPSet
+
+type iPNet struct {
+	IP   string
+	Mask string
+}
+
+var (
+	ErrNetworkAlreadyAllocated = errors.New("requested network overlaps with existing network")
+	ErrNetworkAlreadyRegisterd = errors.New("requested network is already registered")
+	ErrNoAvailableIps          = errors.New("no available ips on network")
+	ErrIPAlreadyAllocated      = errors.New("ip already allocated")
+
+	lock         = sync.Mutex{}
+	allocatedIPs = networkSet{}
+	availableIPS = networkSet{}
+)
+
+func RegisterNetwork(network *net.IPNet) error {
+	lock.Lock()
+	defer lock.Unlock()
+
+	routes, err := netlink.NetworkGetRoutes()
+	if err != nil {
+		return err
+	}
+
+	if err := checkRouteOverlaps(routes, network); err != nil {
+		return err
+	}
+
+	if err := checkExistingNetworkOverlaps(network); err != nil {
+		return err
+	}
+	n := newIPNet(network)
+
+	allocatedIPs[n] = &iPSet{}
+	availableIPS[n] = &iPSet{}
+
+	return nil
+}
+
+func RequestIP(network *net.IPNet, ip *net.IP) (*net.IP, error) {
+	lock.Lock()
+	defer lock.Unlock()
+
+	if ip == nil {
+		next, err := getNextIp(network)
+		if err != nil {
+			return nil, err
+		}
+		return next, nil
+	}
+
+	if err := registerIP(network, ip); err != nil {
+		return nil, err
+	}
+	return ip, nil
+}
+
+func ReleaseIP(network *net.IPNet, ip *net.IP) error {
+	lock.Lock()
+	defer lock.Unlock()
+
+	var (
+		first, _  = networkRange(network)
+		base      = ipToInt(&first)
+		n         = newIPNet(network)
+		existing  = allocatedIPs[n]
+		available = availableIPS[n]
+		i         = ipToInt(ip)
+		pos       = i - base
+	)
+
+	existing.Remove(int(pos))
+	available.Push(int(pos))
+
+	return nil
+}
+
+func getNextIp(network *net.IPNet) (*net.IP, error) {
+	var (
+		n         = newIPNet(network)
+		ownIP     = ipToInt(&network.IP)
+		available = availableIPS[n]
+		allocated = allocatedIPs[n]
+
+		first, _ = networkRange(network)
+		base     = ipToInt(&first)
+
+		pos = int32(available.Pop())
+	)
+
+	// We pop and push the position not the ip
+	if pos != 0 {
+		ip := intToIP(int32(base + pos))
+		allocated.Push(int(pos))
+
+		return ip, nil
+	}
+
+	var (
+		size = int(networkSize(network.Mask))
+		max  = int32(size - 2) // size -1 for the broadcast address, -1 for the gateway address
+	)
+
+	if pos = int32(allocated.PullBack()); pos == 0 {
+		pos = 1
+	}
+
+	for i := int32(0); i < max; i++ {
+		next := int32(base + pos)
+		pos = pos%max + 1
+
+		if next == ownIP {
+			continue
+		}
+
+		ip := intToIP(next)
+		allocated.Push(int(pos))
+
+		return ip, nil
+	}
+	return nil, ErrNoAvailableIps
+}
+
+func registerIP(network *net.IPNet, ip *net.IP) error {
+	existing := allocatedIPs[newIPNet(network)]
+	// checking position not ip
+	if existing.Exists(int(ipToInt(ip))) {
+		return ErrIPAlreadyAllocated
+	}
+	return nil
+}
+
+func checkRouteOverlaps(networks []netlink.Route, toCheck *net.IPNet) error {
+	for _, network := range networks {
+		if network.IPNet != nil && networkOverlaps(toCheck, network.IPNet) {
+			return ErrNetworkAlreadyAllocated
+		}
+	}
+	return nil
+}
+
+// Detects overlap between one IPNet and another
+func networkOverlaps(netX *net.IPNet, netY *net.IPNet) bool {
+	if firstIP, _ := networkRange(netX); netY.Contains(firstIP) {
+		return true
+	}
+	if firstIP, _ := networkRange(netY); netX.Contains(firstIP) {
+		return true
+	}
+	return false
+}
+
+func checkExistingNetworkOverlaps(network *net.IPNet) error {
+	for existing := range allocatedIPs {
+		if newIPNet(network) == existing {
+			return ErrNetworkAlreadyRegisterd
+		}
+
+		ex := newNetIPNet(existing)
+		if networkOverlaps(network, ex) {
+			return ErrNetworkAlreadyAllocated
+		}
+	}
+	return nil
+}
+
+// Calculates the first and last IP addresses in an IPNet
+func networkRange(network *net.IPNet) (net.IP, net.IP) {
+	var (
+		netIP   = network.IP.To4()
+		firstIP = netIP.Mask(network.Mask)
+		lastIP  = net.IPv4(0, 0, 0, 0).To4()
+	)
+
+	for i := 0; i < len(lastIP); i++ {
+		lastIP[i] = netIP[i] | ^network.Mask[i]
+	}
+	return firstIP, lastIP
+}
+
+func newIPNet(network *net.IPNet) iPNet {
+	return iPNet{
+		IP:   string(network.IP),
+		Mask: string(network.Mask),
+	}
+}
+
+func newNetIPNet(network iPNet) *net.IPNet {
+	return &net.IPNet{
+		IP:   []byte(network.IP),
+		Mask: []byte(network.Mask),
+	}
+}
+
+// Converts a 4 bytes IP into a 32 bit integer
+func ipToInt(ip *net.IP) int32 {
+	return int32(binary.BigEndian.Uint32(ip.To4()))
+}
+
+// Converts 32 bit integer into a 4 bytes IP address
+func intToIP(n int32) *net.IP {
+	b := make([]byte, 4)
+	binary.BigEndian.PutUint32(b, uint32(n))
+	ip := net.IP(b)
+	return &ip
+}
+
+// Given a netmask, calculates the number of available hosts
+func networkSize(mask net.IPMask) int32 {
+	m := net.IPv4Mask(0, 0, 0, 0)
+	for i := 0; i < net.IPv4len; i++ {
+		m[i] = ^mask[i]
+	}
+
+	return int32(binary.BigEndian.Uint32(m)) + 1
+}
@@ -0,0 +1,163 @@
+package ipallocator
+
+import (
+	"fmt"
+	"net"
+	"testing"
+)
+
+func reset() {
+	allocatedIPs = networkSet{}
+	availableIPS = networkSet{}
+}
+
+func TestRegisterNetwork(t *testing.T) {
+	defer reset()
+	network := &net.IPNet{
+		IP:   []byte{192, 168, 0, 1},
+		Mask: []byte{255, 255, 255, 0},
+	}
+
+	if err := RegisterNetwork(network); err != nil {
+		t.Fatal(err)
+	}
+
+	n := newIPNet(network)
+	if _, exists := allocatedIPs[n]; !exists {
+		t.Fatal("IPNet should exist in allocated IPs")
+	}
+
+	if _, exists := availableIPS[n]; !exists {
+		t.Fatal("IPNet should exist in available IPs")
+	}
+}
+
+func TestRegisterTwoNetworks(t *testing.T) {
+	defer reset()
+	network := &net.IPNet{
+		IP:   []byte{192, 168, 0, 1},
+		Mask: []byte{255, 255, 255, 0},
+	}
+
+	if err := RegisterNetwork(network); err != nil {
+		t.Fatal(err)
+	}
+
+	network2 := &net.IPNet{
+		IP:   []byte{10, 1, 42, 1},
+		Mask: []byte{255, 255, 255, 0},
+	}
+
+	if err := RegisterNetwork(network2); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestRegisterNetworkThatExists(t *testing.T) {
+	defer reset()
+	network := &net.IPNet{
+		IP:   []byte{192, 168, 0, 1},
+		Mask: []byte{255, 255, 255, 0},
+	}
+
+	if err := RegisterNetwork(network); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := RegisterNetwork(network); err != ErrNetworkAlreadyRegisterd {
+		t.Fatalf("Expected error of %s got %s", ErrNetworkAlreadyRegisterd, err)
+	}
+}
+
+func TestRequestNewIps(t *testing.T) {
+	defer reset()
+	network := &net.IPNet{
+		IP:   []byte{192, 168, 0, 1},
+		Mask: []byte{255, 255, 255, 0},
+	}
+
+	if err := RegisterNetwork(network); err != nil {
+		t.Fatal(err)
+	}
+
+	for i := 2; i < 10; i++ {
+		ip, err := RequestIP(network, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if expected := fmt.Sprintf("192.168.0.%d", i); ip.String() != expected {
+			t.Fatalf("Expected ip %s got %s", expected, ip.String())
+		}
+	}
+}
+
+func TestReleaseIp(t *testing.T) {
+	defer reset()
+	network := &net.IPNet{
+		IP:   []byte{192, 168, 0, 1},
+		Mask: []byte{255, 255, 255, 0},
+	}
+
+	if err := RegisterNetwork(network); err != nil {
+		t.Fatal(err)
+	}
+
+	ip, err := RequestIP(network, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := ReleaseIP(network, ip); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestGetReleasedIp(t *testing.T) {
+	defer reset()
+	network := &net.IPNet{
+		IP:   []byte{192, 168, 0, 1},
+		Mask: []byte{255, 255, 255, 0},
+	}
+
+	if err := RegisterNetwork(network); err != nil {
+		t.Fatal(err)
+	}
+
+	ip, err := RequestIP(network, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	value := ip.String()
+	if err := ReleaseIP(network, ip); err != nil {
+		t.Fatal(err)
+	}
+
+	ip, err = RequestIP(network, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if ip.String() != value {
+		t.Fatalf("Expected to receive same ip %s got %s", value, ip.String())
+	}
+}
+
+func TestRequesetSpecificIp(t *testing.T) {
+	defer reset()
+	network := &net.IPNet{
+		IP:   []byte{192, 168, 0, 1},
+		Mask: []byte{255, 255, 255, 0},
+	}
+
+	if err := RegisterNetwork(network); err != nil {
+		t.Fatal(err)
+	}
+
+	ip := net.ParseIP("192.168.1.5")
+
+	if _, err := RequestIP(network, &ip); err != nil {
+		t.Fatal(err)
+	}
+}
@@ -0,0 +1,84 @@
+package ipallocator
+
+import (
+	"sort"
+	"sync"
+)
+
+// iPSet is a thread-safe sorted set and a stack.
+type iPSet struct {
+	sync.RWMutex
+	set []int
+}
+
+// Push takes a string and adds it to the set. If the elem aready exists, it has no effect.
+func (s *iPSet) Push(elem int) {
+	s.RLock()
+	for _, e := range s.set {
+		if e == elem {
+			s.RUnlock()
+			return
+		}
+	}
+	s.RUnlock()
+
+	s.Lock()
+	s.set = append(s.set, elem)
+	// Make sure the list is always sorted
+	sort.Ints(s.set)
+	s.Unlock()
+}
+
+// Pop is an alias to PopFront()
+func (s *iPSet) Pop() int {
+	return s.PopFront()
+}
+
+// Pop returns the first elemen from the list and removes it.
+// If the list is empty, it returns 0
+func (s *iPSet) PopFront() int {
+	s.RLock()
+
+	for i, e := range s.set {
+		ret := e
+		s.RUnlock()
+		s.Lock()
+		s.set = append(s.set[:i], s.set[i+1:]...)
+		s.Unlock()
+		return ret
+	}
+	s.RUnlock()
+
+	return 0
+}
+
+// PullBack retrieve the last element of the list.
+// The element is not removed.
+// If the list is empty, an empty element is returned.
+func (s *iPSet) PullBack() int {
+	if len(s.set) == 0 {
+		return 0
+	}
+	return s.set[len(s.set)-1]
+}
+
+// Exists checks if the given element present in the list.
+func (s *iPSet) Exists(elem int) bool {
+	for _, e := range s.set {
+		if e == elem {
+			return true
+		}
+	}
+	return false
+}
+
+// Remove removes an element from the list.
+// If the element is not found, it has no effect.
+func (s *iPSet) Remove(elem int) {
+	for i, e := range s.set {
+		if e == elem {
+			s.set = append(s.set[:i], s.set[i+1:]...)
+			return
+		}
+	}
+}