From f61753e73768ccc1d39c7dbba0003540566b67bf Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Fri, 5 Jun 2026 13:39:08 +0000
Subject: [PATCH] db: bound IP allocation scan so exhausted prefixes error out

Random strategy previously spun forever under the allocator mutex.
---
 hscontrol/db/ip.go                        | 46 ++++++++++++--------
 hscontrol/db/ip_random_exhaustion_test.go | 52 +++++++++++++++++++++++
 2 files changed, 80 insertions(+), 18 deletions(-)
 create mode 100644 hscontrol/db/ip_random_exhaustion_test.go

diff --git a/hscontrol/db/ip.go b/hscontrol/db/ip.go
index b369df65..b2033595 100644
--- a/hscontrol/db/ip.go
+++ b/hscontrol/db/ip.go
@@ -200,30 +200,40 @@ func (i *IPAllocator) next(prev netip.Addr, prefix *netip.Prefix) (*netip.Addr,
 		return nil, err
 	}
 
+	// Walk forward from the starting address until a free, non-reserved
+	// address inside the prefix is found. The random strategy only picks the
+	// starting point at random and then scans deterministically: this keeps
+	// the loop finite, so an exhausted prefix returns ErrCouldNotAllocateIP
+	// instead of re-drawing in-prefix addresses forever under i.mu.
+	start := ip
 	for {
-		if !prefix.Contains(ip) {
-			return nil, ErrCouldNotAllocateIP
+		if prefix.Contains(ip) && !set.Contains(ip) && !isTailscaleReservedIP(ip) {
+			i.usedIPs.Add(ip)
+
+			return &ip, nil
 		}
 
-		// Check if the IP has already been allocated
-		// or if it is a IP reserved by Tailscale.
-		if set.Contains(ip) || isTailscaleReservedIP(ip) {
-			switch i.strategy {
-			case types.IPAllocationStrategySequential:
-				ip = ip.Next()
-			case types.IPAllocationStrategyRandom:
-				ip, err = randomNext(*prefix)
-				if err != nil {
-					return nil, fmt.Errorf("getting random IP: %w", err)
-				}
+		ip = ip.Next()
+
+		switch i.strategy {
+		case types.IPAllocationStrategySequential:
+			// Sequential allocation never wraps: walking past the end of
+			// the prefix means the pool is exhausted.
+			if !prefix.Contains(ip) {
+				return nil, ErrCouldNotAllocateIP
+			}
+		case types.IPAllocationStrategyRandom:
+			// Random allocation wraps within the prefix so every address is
+			// examined exactly once; returning to the start means the prefix
+			// is exhausted.
+			if !prefix.Contains(ip) {
+				ip = prefix.Masked().Addr()
 			}
 
-			continue
+			if ip == start {
+				return nil, ErrCouldNotAllocateIP
+			}
 		}
-
-		i.usedIPs.Add(ip)
-
-		return &ip, nil
 	}
 }
 
diff --git a/hscontrol/db/ip_random_exhaustion_test.go b/hscontrol/db/ip_random_exhaustion_test.go
new file mode 100644
index 00000000..350ad9ba
--- /dev/null
+++ b/hscontrol/db/ip_random_exhaustion_test.go
@@ -0,0 +1,52 @@
+package db
+
+import (
+	"errors"
+	"net/netip"
+	"testing"
+	"time"
+
+	"github.com/juanfont/headscale/hscontrol/types"
+)
+
+// TestIPAllocatorRandomExhaustionReturnsError ensures the random allocation
+// strategy terminates when its prefix is exhausted. randomNext only ever
+// returns in-prefix addresses, so the allocation loop's exhaustion exit must
+// not depend on producing an out-of-prefix candidate; otherwise Next() spins
+// forever under the allocator mutex, wedging all registration and IP release.
+//
+// 100.64.0.0/30 has four addresses: .0 (network) and .3 (broadcast) are
+// reserved, leaving .1 and .2. After two allocations the pool is exhausted and
+// the third Next() must return ErrCouldNotAllocateIP promptly.
+func TestIPAllocatorRandomExhaustionReturnsError(t *testing.T) {
+	prefix4 := netip.MustParsePrefix("100.64.0.0/30")
+
+	alloc, err := NewIPAllocator(nil, &prefix4, nil, types.IPAllocationStrategyRandom)
+	if err != nil {
+		t.Fatalf("NewIPAllocator: %v", err)
+	}
+
+	for i := range 2 {
+		_, _, err := alloc.Next()
+		if err != nil {
+			t.Fatalf("Next() #%d unexpectedly failed: %v", i+1, err)
+		}
+	}
+
+	done := make(chan error, 1)
+
+	go func() {
+		_, _, err := alloc.Next()
+		done <- err
+	}()
+
+	select {
+	case err := <-done:
+		if !errors.Is(err, ErrCouldNotAllocateIP) {
+			t.Fatalf("expected ErrCouldNotAllocateIP on exhausted prefix, got: %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Next() on an exhausted random-strategy prefix did not return " +
+			"within 2s; it is spinning forever under the allocator mutex")
+	}
+}