new IP allocator and add postgres to integration tests. (#1756)

hscontrol/db/ip.go

@@ -0,0 +1,148 @@
+package db
+
+import (
+	"errors"
+	"fmt"
+	"net/netip"
+	"sync"
+
+	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/juanfont/headscale/hscontrol/util"
+	"go4.org/netipx"
+	"gorm.io/gorm"
+)
+
+// IPAllocator is a singleton responsible for allocating
+// IP addresses for nodes and making sure the same
+// address is not handed out twice. There can only be one
+// and it needs to be created before any other database
+// writes occur.
+type IPAllocator struct {
+	mu sync.Mutex
+
+	prefix4 netip.Prefix
+	prefix6 netip.Prefix
+
+	// Previous IPs handed out
+	prev4 netip.Addr
+	prev6 netip.Addr
+
+	// Set of all IPs handed out.
+	// This might not be in sync with the database,
+	// but it is more conservative. If saves to the
+	// database fails, the IP will be allocated here
+	// until the next restart of Headscale.
+	usedIPs netipx.IPSetBuilder
+}
+
+// NewIPAllocator returns a new IPAllocator singleton which
+// can be used to hand out unique IP addresses within the
+// provided IPv4 and IPv6 prefix. It needs to be created
+// when headscale starts and needs to finish its read
+// transaction before any writes to the database occur.
+func NewIPAllocator(db *HSDatabase, prefix4, prefix6 netip.Prefix) (*IPAllocator, error) {
+	var addressesSlices []string
+
+	if db != nil {
+		db.Read(func(rx *gorm.DB) error {
+			return rx.Model(&types.Node{}).Pluck("ip_addresses", &addressesSlices).Error
+		})
+	}
+
+	var ips netipx.IPSetBuilder
+
+	// Add network and broadcast addrs to used pool so they
+	// are not handed out to nodes.
+	network4, broadcast4 := util.GetIPPrefixEndpoints(prefix4)
+	network6, broadcast6 := util.GetIPPrefixEndpoints(prefix6)
+	ips.Add(network4)
+	ips.Add(broadcast4)
+	ips.Add(network6)
+	ips.Add(broadcast6)
+
+	// Fetch all the IP Addresses currently handed out from the Database
+	// and add them to the used IP set.
+	for _, slice := range addressesSlices {
+		var machineAddresses types.NodeAddresses
+		err := machineAddresses.Scan(slice)
+		if err != nil {
+			return nil, fmt.Errorf(
+				"parsing IPs from database %v: %w", machineAddresses,
+				err,
+			)
+		}
+
+		for _, ip := range machineAddresses {
+			ips.Add(ip)
+		}
+	}
+
+	// Build the initial IPSet to validate that we can use it.
+	_, err := ips.IPSet()
+	if err != nil {
+		return nil, fmt.Errorf(
+			"building initial IP Set: %w",
+			err,
+		)
+	}
+
+	return &IPAllocator{
+		usedIPs: ips,
+
+		prefix4: prefix4,
+		prefix6: prefix6,
+
+		// Use network as starting point, it will be used to call .Next()
+		// TODO(kradalby): Could potentially take all the IPs loaded from
+		// the database into account to start at a more "educated" location.
+		prev4: network4,
+		prev6: network6,
+	}, nil
+}
+
+func (i *IPAllocator) Next() (types.NodeAddresses, error) {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	v4, err := i.next(i.prev4, i.prefix4)
+	if err != nil {
+		return nil, fmt.Errorf("allocating IPv4 address: %w", err)
+	}
+
+	v6, err := i.next(i.prev6, i.prefix6)
+	if err != nil {
+		return nil, fmt.Errorf("allocating IPv6 address: %w", err)
+	}
+
+	return types.NodeAddresses{*v4, *v6}, nil
+}
+
+var ErrCouldNotAllocateIP = errors.New("failed to allocate IP")
+
+func (i *IPAllocator) next(prev netip.Addr, prefix netip.Prefix) (*netip.Addr, error) {
+	// Get the first IP in our prefix
+	ip := prev.Next()
+
+	// TODO(kradalby): maybe this can be done less often.
+	set, err := i.usedIPs.IPSet()
+	if err != nil {
+		return nil, err
+	}
+
+	for {
+		if !prefix.Contains(ip) {
+			return nil, ErrCouldNotAllocateIP
+		}
+
+		// Check if the IP has already been allocated.
+		if set.Contains(ip) {
+			ip = ip.Next()
+
+			continue
+		}
+
+		i.usedIPs.Add(ip)
+
+		return &ip, nil
+	}
+}