cmd/natc: move address storage behind an interface

Adds IPPool and moves all IP address management concerns behind that.

Updates #14667

Signed-off-by: Fran Bull <fran@tailscale.com>

cmd/natc/ippool/ippool.go

@@ -0,0 +1,127 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+// ippool implements IP address storage, creation, and retrieval for cmd/natc
+package ippool
+
+import (
+	"errors"
+	"log"
+	"math/big"
+	"net/netip"
+	"sync"
+
+	"github.com/gaissmai/bart"
+	"go4.org/netipx"
+	"tailscale.com/syncs"
+	"tailscale.com/tailcfg"
+	"tailscale.com/util/dnsname"
+	"tailscale.com/util/mak"
+)
+
+var ErrNoIPsAvailable = errors.New("no IPs available")
+
+type IPPool struct {
+	perPeerMap syncs.Map[tailcfg.NodeID, *perPeerState]
+	IPSet      *netipx.IPSet
+	V6ULA      netip.Prefix
+}
+
+func (ipp *IPPool) DomainForIP(from tailcfg.NodeID, addr netip.Addr) (string, bool) {
+	ps, ok := ipp.perPeerMap.Load(from)
+	if !ok {
+		log.Printf("handleTCPFlow: no perPeerState for %v", from)
+		return "", false
+	}
+	domain, ok := ps.domainForIP(addr)
+	if !ok {
+		log.Printf("handleTCPFlow: no domain for IP %v\n", addr)
+		return "", false
+	}
+	return domain, ok
+}
+
+func (ipp *IPPool) IPForDomain(from tailcfg.NodeID, domain string) ([]netip.Addr, error) {
+	npps := &perPeerState{
+		ipset: ipp.IPSet,
+		v6ULA: ipp.V6ULA,
+	}
+	ps, _ := ipp.perPeerMap.LoadOrStore(from, npps)
+	return ps.ipForDomain(domain)
+}
+
+// perPeerState holds the state for a single peer.
+type perPeerState struct {
+	v6ULA netip.Prefix
+	ipset *netipx.IPSet
+
+	mu           sync.Mutex
+	addrInUse    *big.Int
+	domainToAddr map[string][]netip.Addr
+	addrToDomain *bart.Table[string]
+}
+
+// domainForIP returns the domain name assigned to the given IP address and
+// whether it was found.
+func (ps *perPeerState) domainForIP(ip netip.Addr) (_ string, ok bool) {
+	ps.mu.Lock()
+	defer ps.mu.Unlock()
+	if ps.addrToDomain == nil {
+		return "", false
+	}
+	return ps.addrToDomain.Lookup(ip)
+}
+
+// ipForDomain assigns a pair of unique IP addresses for the given domain and
+// returns them. The first address is an IPv4 address and the second is an IPv6
+// address. If the domain already has assigned addresses, it returns them.
+func (ps *perPeerState) ipForDomain(domain string) ([]netip.Addr, error) {
+	fqdn, err := dnsname.ToFQDN(domain)
+	if err != nil {
+		return nil, err
+	}
+	domain = fqdn.WithoutTrailingDot()
+
+	ps.mu.Lock()
+	defer ps.mu.Unlock()
+	if addrs, ok := ps.domainToAddr[domain]; ok {
+		return addrs, nil
+	}
+	addrs := ps.assignAddrsLocked(domain)
+	if addrs == nil {
+		return nil, ErrNoIPsAvailable
+	}
+	return addrs, nil
+}
+
+// unusedIPv4Locked returns an unused IPv4 address from the available ranges.
+func (ps *perPeerState) unusedIPv4Locked() netip.Addr {
+	if ps.addrInUse == nil {
+		ps.addrInUse = big.NewInt(0)
+	}
+	return allocAddr(ps.ipset, ps.addrInUse)
+}
+
+// assignAddrsLocked assigns a pair of unique IP addresses for the given domain
+// and returns them. The first address is an IPv4 address and the second is an
+// IPv6 address. It does not check if the domain already has assigned addresses.
+// ps.mu must be held.
+func (ps *perPeerState) assignAddrsLocked(domain string) []netip.Addr {
+	if ps.addrToDomain == nil {
+		ps.addrToDomain = &bart.Table[string]{}
+	}
+	v4 := ps.unusedIPv4Locked()
+	if !v4.IsValid() {
+		return nil
+	}
+	as16 := ps.v6ULA.Addr().As16()
+	as4 := v4.As4()
+	copy(as16[12:], as4[:])
+	v6 := netip.AddrFrom16(as16)
+	addrs := []netip.Addr{v4, v6}
+	mak.Set(&ps.domainToAddr, domain, addrs)
+	for _, a := range addrs {
+		ps.addrToDomain.Insert(netip.PrefixFrom(a, a.BitLen()), domain)
+	}
+	return addrs
+}

cmd/natc/ippool/ippool_test.go

@@ -0,0 +1,129 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package ippool
+
+import (
+	"errors"
+	"fmt"
+	"net/netip"
+	"slices"
+	"testing"
+
+	"go4.org/netipx"
+	"tailscale.com/tailcfg"
+	"tailscale.com/util/must"
+)
+
+func TestIPPoolExhaustion(t *testing.T) {
+	smallPrefix := netip.MustParsePrefix("100.64.1.0/30") // Only 4 IPs: .0, .1, .2, .3
+	var ipsb netipx.IPSetBuilder
+	ipsb.AddPrefix(smallPrefix)
+	addrPool := must.Get(ipsb.IPSet())
+	v6ULA := netip.MustParsePrefix("fd7a:115c:a1e0:a99c:0001::/80")
+	pool := IPPool{V6ULA: v6ULA, IPSet: addrPool}
+
+	assignedIPs := make(map[netip.Addr]string)
+
+	domains := []string{"a.example.com", "b.example.com", "c.example.com", "d.example.com", "e.example.com"}
+
+	var errs []error
+
+	from := tailcfg.NodeID(12345)
+
+	for i := 0; i < 5; i++ {
+		for _, domain := range domains {
+			addrs, err := pool.IPForDomain(from, domain)
+			if err != nil {
+				errs = append(errs, fmt.Errorf("failed to get IP for domain %q: %w", domain, err))
+				continue
+			}
+
+			for _, addr := range addrs {
+				if d, ok := assignedIPs[addr]; ok {
+					if d != domain {
+						t.Errorf("IP %s reused for domain %q, previously assigned to %q", addr, domain, d)
+					}
+				} else {
+					assignedIPs[addr] = domain
+				}
+			}
+		}
+	}
+
+	for addr, domain := range assignedIPs {
+		if addr.Is4() && !smallPrefix.Contains(addr) {
+			t.Errorf("IP %s for domain %q not in expected range %s", addr, domain, smallPrefix)
+		}
+		if addr.Is6() && !v6ULA.Contains(addr) {
+			t.Errorf("IP %s for domain %q not in expected range %s", addr, domain, v6ULA)
+		}
+	}
+
+	// expect one error for each iteration with the 5th domain
+	if len(errs) != 5 {
+		t.Errorf("Expected 5 errors, got %d: %v", len(errs), errs)
+	}
+	for _, err := range errs {
+		if !errors.Is(err, ErrNoIPsAvailable) {
+			t.Errorf("generateDNSResponse() error = %v, want ErrNoIPsAvailable", err)
+		}
+	}
+}
+
+func TestIPPool(t *testing.T) {
+	var ipsb netipx.IPSetBuilder
+	ipsb.AddPrefix(netip.MustParsePrefix("100.64.1.0/24"))
+	addrPool := must.Get(ipsb.IPSet())
+	pool := IPPool{
+		V6ULA: netip.MustParsePrefix("fd7a:115c:a1e0:a99c:0001::/80"),
+		IPSet: addrPool,
+	}
+	from := tailcfg.NodeID(12345)
+	addrs, err := pool.IPForDomain(from, "example.com")
+	if err != nil {
+		t.Fatalf("ipForDomain() error = %v", err)
+	}
+
+	if len(addrs) != 2 {
+		t.Fatalf("ipForDomain() returned %d addresses, want 2", len(addrs))
+	}
+
+	v4 := addrs[0]
+	v6 := addrs[1]
+
+	if !v4.Is4() {
+		t.Errorf("First address is not IPv4: %s", v4)
+	}
+
+	if !v6.Is6() {
+		t.Errorf("Second address is not IPv6: %s", v6)
+	}
+
+	if !addrPool.Contains(v4) {
+		t.Errorf("IPv4 address %s not in range %s", v4, addrPool)
+	}
+
+	domain, ok := pool.DomainForIP(from, v4)
+	if !ok {
+		t.Errorf("domainForIP(%s) not found", v4)
+	} else if domain != "example.com" {
+		t.Errorf("domainForIP(%s) = %s, want %s", v4, domain, "example.com")
+	}
+
+	domain, ok = pool.DomainForIP(from, v6)
+	if !ok {
+		t.Errorf("domainForIP(%s) not found", v6)
+	} else if domain != "example.com" {
+		t.Errorf("domainForIP(%s) = %s, want %s", v6, domain, "example.com")
+	}
+
+	addrs2, err := pool.IPForDomain(from, "example.com")
+	if err != nil {
+		t.Fatalf("ipForDomain() second call error = %v", err)
+	}
+
+	if !slices.Equal(addrs, addrs2) {
+		t.Errorf("ipForDomain() second call = %v, want %v", addrs2, addrs)
+	}
+}

cmd/natc/ippool/ipx.go

@@ -0,0 +1,130 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package ippool
+
+import (
+	"math/big"
+	"math/bits"
+	"math/rand/v2"
+	"net/netip"
+
+	"go4.org/netipx"
+)
+
+func addrLessOrEqual(a, b netip.Addr) bool {
+	if a.Less(b) {
+		return true
+	}
+	if a == b {
+		return true
+	}
+	return false
+}
+
+// indexOfAddr returns the index of addr in ipset, or -1 if not found.
+func indexOfAddr(addr netip.Addr, ipset *netipx.IPSet) int {
+	var base int // offset of the current range
+	for _, r := range ipset.Ranges() {
+		if addr.Less(r.From()) {
+			return -1
+		}
+		numFrom := v4ToNum(r.From())
+		if addrLessOrEqual(addr, r.To()) {
+			numInRange := int(v4ToNum(addr) - numFrom)
+			return base + numInRange
+		}
+		numTo := v4ToNum(r.To())
+		base += int(numTo-numFrom) + 1
+	}
+	return -1
+}
+
+// addrAtIndex returns the address at the given index in ipset, or an empty
+// address if index is out of range.
+func addrAtIndex(index int, ipset *netipx.IPSet) netip.Addr {
+	if index < 0 {
+		return netip.Addr{}
+	}
+	var base int // offset of the current range
+	for _, r := range ipset.Ranges() {
+		numFrom := v4ToNum(r.From())
+		numTo := v4ToNum(r.To())
+		if index <= base+int(numTo-numFrom) {
+			return numToV4(uint32(int(numFrom) + index - base))
+		}
+		base += int(numTo-numFrom) + 1
+	}
+	return netip.Addr{}
+}
+
+// TODO(golang/go#9455): once we have uint128 we can easily implement for all addrs.
+
+// v4ToNum returns a uint32 representation of the IPv4 address. If addr is not
+// an IPv4 address, this function will panic.
+func v4ToNum(addr netip.Addr) uint32 {
+	addr = addr.Unmap()
+	if !addr.Is4() {
+		panic("only IPv4 addresses are supported by v4ToNum")
+	}
+	b := addr.As4()
+	var o uint32
+	o = o<<8 | uint32(b[0])
+	o = o<<8 | uint32(b[1])
+	o = o<<8 | uint32(b[2])
+	o = o<<8 | uint32(b[3])
+	return o
+}
+
+func numToV4(i uint32) netip.Addr {
+	var addr [4]byte
+	addr[0] = byte((i >> 24) & 0xff)
+	addr[1] = byte((i >> 16) & 0xff)
+	addr[2] = byte((i >> 8) & 0xff)
+	addr[3] = byte(i & 0xff)
+	return netip.AddrFrom4(addr)
+}
+
+// allocAddr returns an address in ipset that is not already marked allocated in allocated.
+func allocAddr(ipset *netipx.IPSet, allocated *big.Int) netip.Addr {
+	// first try to allocate a random IP from each range, if we land on one.
+	var base uint32 // index offset of the current range
+	for _, r := range ipset.Ranges() {
+		numFrom := v4ToNum(r.From())
+		numTo := v4ToNum(r.To())
+		randInRange := rand.N(numTo - numFrom)
+		randIndex := base + randInRange
+		if allocated.Bit(int(randIndex)) == 0 {
+			allocated.SetBit(allocated, int(randIndex), 1)
+			return numToV4(numFrom + randInRange)
+		}
+		base += numTo - numFrom + 1
+	}
+
+	// fall back to seeking a free bit in the allocated set
+	index := -1
+	for i, word := range allocated.Bits() {
+		zbi := leastZeroBit(uint(word))
+		if zbi == -1 {
+			continue
+		}
+		index = i*bits.UintSize + zbi
+		allocated.SetBit(allocated, index, 1)
+		break
+	}
+	if index == -1 {
+		return netip.Addr{}
+	}
+	return addrAtIndex(index, ipset)
+}
+
+// leastZeroBit returns the index of the least significant zero bit in the given uint, or -1
+// if all bits are set.
+func leastZeroBit(n uint) int {
+	notN := ^n
+	rightmostBit := notN & -notN
+	if rightmostBit == 0 {
+		return -1
+	}
+	return bits.TrailingZeros(rightmostBit)
+}

cmd/natc/ippool/ipx_test.go

@@ -0,0 +1,150 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package ippool
+
+import (
+	"math"
+	"math/big"
+	"net/netip"
+	"testing"
+
+	"go4.org/netipx"
+	"tailscale.com/util/must"
+)
+
+func TestV4ToNum(t *testing.T) {
+	cases := []struct {
+		addr netip.Addr
+		num  uint32
+	}{
+		{netip.MustParseAddr("0.0.0.0"), 0},
+		{netip.MustParseAddr("255.255.255.255"), 0xffffffff},
+		{netip.MustParseAddr("8.8.8.8"), 0x08080808},
+		{netip.MustParseAddr("192.168.0.1"), 0xc0a80001},
+		{netip.MustParseAddr("10.0.0.1"), 0x0a000001},
+		{netip.MustParseAddr("172.16.0.1"), 0xac100001},
+		{netip.MustParseAddr("100.64.0.1"), 0x64400001},
+	}
+
+	for _, tc := range cases {
+		num := v4ToNum(tc.addr)
+		if num != tc.num {
+			t.Errorf("addrNum(%v) = %d, want %d", tc.addr, num, tc.num)
+		}
+		if numToV4(num) != tc.addr {
+			t.Errorf("numToV4(%d) = %v, want %v", num, numToV4(num), tc.addr)
+		}
+	}
+
+	func() {
+		defer func() {
+			if r := recover(); r == nil {
+				t.Fatal("expected panic")
+			}
+		}()
+
+		v4ToNum(netip.MustParseAddr("::1"))
+	}()
+}
+
+func TestAddrIndex(t *testing.T) {
+	builder := netipx.IPSetBuilder{}
+	builder.AddRange(netipx.MustParseIPRange("10.0.0.1-10.0.0.5"))
+	builder.AddRange(netipx.MustParseIPRange("192.168.0.1-192.168.0.10"))
+	ipset := must.Get(builder.IPSet())
+
+	indexCases := []struct {
+		addr  netip.Addr
+		index int
+	}{
+		{netip.MustParseAddr("10.0.0.1"), 0},
+		{netip.MustParseAddr("10.0.0.2"), 1},
+		{netip.MustParseAddr("10.0.0.3"), 2},
+		{netip.MustParseAddr("10.0.0.4"), 3},
+		{netip.MustParseAddr("10.0.0.5"), 4},
+		{netip.MustParseAddr("192.168.0.1"), 5},
+		{netip.MustParseAddr("192.168.0.5"), 9},
+		{netip.MustParseAddr("192.168.0.10"), 14},
+		{netip.MustParseAddr("172.16.0.1"), -1}, // Not in set
+	}
+
+	for _, tc := range indexCases {
+		index := indexOfAddr(tc.addr, ipset)
+		if index != tc.index {
+			t.Errorf("indexOfAddr(%v) = %d, want %d", tc.addr, index, tc.index)
+		}
+		if tc.index == -1 {
+			continue
+		}
+		addr := addrAtIndex(tc.index, ipset)
+		if addr != tc.addr {
+			t.Errorf("addrAtIndex(%d) = %v, want %v", tc.index, addr, tc.addr)
+		}
+	}
+}
+
+func TestAllocAddr(t *testing.T) {
+	builder := netipx.IPSetBuilder{}
+	builder.AddRange(netipx.MustParseIPRange("10.0.0.1-10.0.0.5"))
+	builder.AddRange(netipx.MustParseIPRange("192.168.0.1-192.168.0.10"))
+	ipset := must.Get(builder.IPSet())
+
+	allocated := new(big.Int)
+	for range 15 {
+		addr := allocAddr(ipset, allocated)
+		if !addr.IsValid() {
+			t.Errorf("allocAddr() = invalid, want valid")
+		}
+		if !ipset.Contains(addr) {
+			t.Errorf("allocAddr() = %v, not in set", addr)
+		}
+	}
+	addr := allocAddr(ipset, allocated)
+	if addr.IsValid() {
+		t.Errorf("allocAddr() = %v, want invalid", addr)
+	}
+	wantAddr := netip.MustParseAddr("10.0.0.2")
+	allocated.SetBit(allocated, indexOfAddr(wantAddr, ipset), 0)
+	addr = allocAddr(ipset, allocated)
+	if addr != wantAddr {
+		t.Errorf("allocAddr() = %v, want %v", addr, wantAddr)
+	}
+}
+
+func TestLeastZeroBit(t *testing.T) {
+	cases := []struct {
+		num  uint
+		want int
+	}{
+		{math.MaxUint, -1},
+		{0, 0},
+		{0b01, 1},
+		{0b11, 2},
+		{0b111, 3},
+		{math.MaxUint, -1},
+		{math.MaxUint - 1, 0},
+	}
+	if math.MaxUint == math.MaxUint64 {
+		cases = append(cases, []struct {
+			num  uint
+			want int
+		}{
+			{math.MaxUint >> 1, 63},
+		}...)
+	} else {
+		cases = append(cases, []struct {
+			num  uint
+			want int
+		}{
+			{math.MaxUint >> 1, 31},
+		}...)
+	}
+
+	for _, tc := range cases {
+		got := leastZeroBit(tc.num)
+		if got != tc.want {
+			t.Errorf("leastZeroBit(%b) = %d, want %d", tc.num, got, tc.want)
+		}
+	}
+}