From 704f1cafd37db02cb0a70d10253497f963ddb682 Mon Sep 17 00:00:00 2001
From: Igor Serganov <iserganov@tunein.com>
Date: Fri, 25 Sep 2026 09:27:13 -0700
Subject: [PATCH] db: release the IPv4 address when IPv6 allocation fails in
 Next

IPAllocator.Next took the allocator mutex once per address family.
When the IPv6 allocation failed, the IPv4 address had already been
added to usedIPs and prev4 had moved past it, so the address stayed
reserved until headscale restarted even though no node held it.

Allocate both families under a single critical section, and on an
IPv6 failure remove the IPv4 address from usedIPs and restore prev4.
---
 hscontrol/db/ip.go      | 21 +++++++++++++--
 hscontrol/db/ip_test.go | 58 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/hscontrol/db/ip.go b/hscontrol/db/ip.go
index fede11b7f..fab5aa054 100644
--- a/hscontrol/db/ip.go
+++ b/hscontrol/db/ip.go
@@ -138,22 +138,34 @@ func NewIPAllocator(
 }
 
 func (i *IPAllocator) Next() (*netip.Addr, *netip.Addr, error) {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
 	var (
 		err  error
 		ret4 *netip.Addr
 		ret6 *netip.Addr
 	)
 
+	prev4 := i.prev4
+
 	if i.prefix4 != nil {
-		ret4, err = i.allocateNext(&i.prev4, i.prefix4)
+		ret4, err = i.allocateNextLocked(&i.prev4, i.prefix4)
 		if err != nil {
 			return nil, nil, fmt.Errorf("allocating IPv4 address: %w", err)
 		}
 	}
 
 	if i.prefix6 != nil {
-		ret6, err = i.allocateNext(&i.prev6, i.prefix6)
+		ret6, err = i.allocateNextLocked(&i.prev6, i.prefix6)
 		if err != nil {
+			// ret4 never reaches a caller, so nothing would ever free it, and
+			// sequential allocation does not wrap back to it on its own.
+			if ret4 != nil {
+				i.usedIPs.Remove(*ret4)
+				i.prev4 = prev4
+			}
+
 			return nil, nil, fmt.Errorf("allocating IPv6 address: %w", err)
 		}
 	}
@@ -171,6 +183,11 @@ func (i *IPAllocator) allocateNext(prev *netip.Addr, prefix *netip.Prefix) (*net
 	i.mu.Lock()
 	defer i.mu.Unlock()
 
+	return i.allocateNextLocked(prev, prefix)
+}
+
+// allocateNextLocked is [IPAllocator.allocateNext] for a caller holding i.mu.
+func (i *IPAllocator) allocateNextLocked(prev *netip.Addr, prefix *netip.Prefix) (*netip.Addr, error) {
 	ret, err := i.next(*prev, prefix)
 	if err != nil {
 		return nil, err
diff --git a/hscontrol/db/ip_test.go b/hscontrol/db/ip_test.go
index cf450e38d..849840707 100644
--- a/hscontrol/db/ip_test.go
+++ b/hscontrol/db/ip_test.go
@@ -514,3 +514,61 @@ func TestIPAllocatorNextNoReservedIPs(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, na("100.115.94.0"), *nextChrome)
 }
+
+// TestIPAllocatorNextIPv6OnlyExhausted pins that an IPv6 failure with no IPv4
+// prefix has no IPv4 address to release.
+func TestIPAllocatorNextIPv6OnlyExhausted(t *testing.T) {
+	alloc, err := NewIPAllocator(nil, nil, mpp("fd7a:115c:a1e0::/127"), types.IPAllocationStrategySequential)
+	require.NoError(t, err)
+
+	_, _, err = alloc.Next()
+	require.ErrorIs(t, err, ErrCouldNotAllocateIP)
+}
+
+// TestIPAllocatorNextReleasesIPv4WhenIPv6Fails uses a /127 because both of its
+// addresses are reserved as network and broadcast, so IPv6 always fails. Both
+// /30 addresses must stay free afterwards; the prev4 restore only shows under
+// sequential, as random ignores prev4.
+func TestIPAllocatorNextReleasesIPv4WhenIPv6Fails(t *testing.T) {
+	tests := []struct {
+		name     string
+		strategy types.IPAllocationStrategy
+	}{
+		{
+			name:     "sequential",
+			strategy: types.IPAllocationStrategySequential,
+		},
+		{
+			name:     "random",
+			strategy: types.IPAllocationStrategyRandom,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			alloc, err := NewIPAllocator(
+				nil,
+				mpp("100.64.0.0/30"),
+				mpp("fd7a:115c:a1e0::/127"),
+				tt.strategy,
+			)
+			require.NoError(t, err)
+
+			got4, got6, err := alloc.Next()
+			require.ErrorIs(t, err, ErrCouldNotAllocateIP)
+			assert.Nil(t, got4)
+			assert.Nil(t, got6)
+
+			got := make([]netip.Addr, 0, 2)
+
+			for range 2 {
+				ip, err := alloc.allocateNext(&alloc.prev4, alloc.prefix4)
+				require.NoError(t, err)
+
+				got = append(got, *ip)
+			}
+
+			assert.ElementsMatch(t, []netip.Addr{na("100.64.0.1"), na("100.64.0.2")}, got)
+		})
+	}
+}