diff --git a/hscontrol/db/ip.go b/hscontrol/db/ip.go index bdfce709..014172ed 100644 --- a/hscontrol/db/ip.go +++ b/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) } diff --git a/hscontrol/db/ip_backfill_race_test.go b/hscontrol/db/ip_backfill_race_test.go new file mode 100644 index 00000000..2c3e4b99 --- /dev/null +++ b/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() +}