db: advance allocator cursor under lock during IP backfill

Reading prev4/prev6 unlocked raced Next; not advancing made backfill O(n^2).

hscontrol/db/ip.go

@@ -170,11 +170,36 @@ func (i *IPAllocator) Next() (*netip.Addr, *netip.Addr, error) {
 
 var ErrCouldNotAllocateIP = errors.New("failed to allocate IP")
 
-func (i *IPAllocator) nextLocked(prev netip.Addr, prefix *netip.Prefix) (*netip.Addr, error) {
+// allocateNext4 allocates the next IPv4 under i.mu, advancing prev4 so a run of
+// allocations (e.g. BackfillNodeIPs) does not rescan already-issued addresses,
+// and so prev4 is read under the lock rather than in the caller's frame.
+func (i *IPAllocator) allocateNext4() (*netip.Addr, error) {
 	i.mu.Lock()
 	defer i.mu.Unlock()
 
-	return i.next(prev, prefix)
+	ret, err := i.next(i.prev4, i.prefix4)
+	if err != nil {
+		return nil, err
+	}
+
+	i.prev4 = *ret
+
+	return ret, nil
+}
+
+// allocateNext6 mirrors allocateNext4 for the IPv6 prefix.
+func (i *IPAllocator) allocateNext6() (*netip.Addr, error) {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	ret, err := i.next(i.prev6, i.prefix6)
+	if err != nil {
+		return nil, err
+	}
+
+	i.prev6 = *ret
+
+	return ret, nil
 }
 
 func (i *IPAllocator) next(prev netip.Addr, prefix *netip.Prefix) (*netip.Addr, error) {
@@ -323,7 +348,7 @@ func (db *HSDatabase) BackfillNodeIPs(i *IPAllocator) ([]string, error) {
 			changed := false
 			// IPv4 prefix is set, but node ip is missing, alloc
 			if i.prefix4 != nil && node.IPv4 == nil {
-				ret4, err := i.nextLocked(i.prev4, i.prefix4)
+				ret4, err := i.allocateNext4()
 				if err != nil {
 					return fmt.Errorf("allocating IPv4 for node(%d): %w", node.ID, err)
 				}
@@ -336,7 +361,7 @@ func (db *HSDatabase) BackfillNodeIPs(i *IPAllocator) ([]string, error) {
 
 			// IPv6 prefix is set, but node ip is missing, alloc
 			if i.prefix6 != nil && node.IPv6 == nil {
-				ret6, err := i.nextLocked(i.prev6, i.prefix6)
+				ret6, err := i.allocateNext6()
 				if err != nil {
 					return fmt.Errorf("allocating IPv6 for node(%d): %w", node.ID, err)
 				}

hscontrol/db/ip_backfill_race_test.go

@@ -0,0 +1,52 @@
+package db
+
+import (
+	"net/netip"
+	"sync"
+	"testing"
+
+	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/stretchr/testify/require"
+)
+
+// TestAllocatorConcurrentNextAndBackfillNoRace exercises the registration path
+// (Next) concurrently with the backfill allocation path (allocateNext4/6) on
+// the same allocator. Backfill used to read prev4/prev6 in the caller's frame
+// without the lock, racing Next's writes; both must now take i.mu. Run with
+// -race.
+func TestAllocatorConcurrentNextAndBackfillNoRace(t *testing.T) {
+	p4 := netip.MustParsePrefix("100.64.0.0/10")
+	p6 := netip.MustParsePrefix("fd7a:115c:a1e0::/48")
+
+	alloc, err := NewIPAllocator(nil, &p4, &p6, types.IPAllocationStrategySequential)
+	require.NoError(t, err)
+
+	const iterations = 2000
+
+	var wg sync.WaitGroup
+
+	wg.Go(func() {
+		for range iterations {
+			_, _, err := alloc.Next()
+			if err != nil {
+				return
+			}
+		}
+	})
+
+	wg.Go(func() {
+		for range iterations {
+			_, err := alloc.allocateNext4()
+			if err != nil {
+				return
+			}
+
+			_, err = alloc.allocateNext6()
+			if err != nil {
+				return
+			}
+		}
+	})
+
+	wg.Wait()
+}