feature/conn25: move addrAssignments to their own file

Updates tailscale/corp#39975

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

feature/conn25/addrAssignments.go

@@ -0,0 +1,92 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+package conn25
+
+import (
+	"errors"
+	"net/netip"
+	"time"
+
+	"tailscale.com/tstime"
+	"tailscale.com/util/dnsname"
+	"tailscale.com/util/mak"
+)
+
+// domainDst is a key for looking up an existing address assignment by the
+// DNS response domain and destination IP pair.
+type domainDst struct {
+	domain dnsname.FQDN
+	dst    netip.Addr
+}
+
+// addrAssignments is the collection of addrs assigned by this client
+// supporting lookup by magic IP, transit IP or domain+dst, or to lookup all
+// transit IPs associated with a given connector (identified by its node key).
+// byConnKey stores netip.Prefix versions of the transit IPs for use in the
+// WireGuard hooks.
+type addrAssignments struct {
+	byMagicIP   map[netip.Addr]addrs
+	byTransitIP map[netip.Addr]addrs
+	byDomainDst map[domainDst]addrs
+	clock       tstime.Clock
+}
+
+const defaultExpiry = 48 * time.Hour
+
+func (a *addrAssignments) insert(as addrs) error {
+	return a.insertWithExpiry(as, defaultExpiry)
+}
+
+func (a *addrAssignments) insertWithExpiry(as addrs, d time.Duration) error {
+	if !as.expiresAt.IsZero() {
+		return errors.New("expiresAt already set")
+	}
+	now := a.clock.Now()
+	as.expiresAt = now.Add(d)
+	// we don't expect for addresses to be reused before expiry
+	if existing, ok := a.byMagicIP[as.magic]; ok {
+		if !existing.expiresAt.Before(now) {
+			return errors.New("byMagicIP key exists")
+		}
+	}
+	ddst := domainDst{domain: as.domain, dst: as.dst}
+	if existing, ok := a.byDomainDst[ddst]; ok {
+		if !existing.expiresAt.Before(now) {
+			return errors.New("byDomainDst key exists")
+		}
+	}
+	if existing, ok := a.byTransitIP[as.transit]; ok {
+		if !existing.expiresAt.Before(now) {
+			return errors.New("byTransitIP key exists")
+		}
+	}
+	mak.Set(&a.byMagicIP, as.magic, as)
+	mak.Set(&a.byTransitIP, as.transit, as)
+	mak.Set(&a.byDomainDst, ddst, as)
+	return nil
+}
+
+func (a *addrAssignments) lookupByDomainDst(domain dnsname.FQDN, dst netip.Addr) (addrs, bool) {
+	v, ok := a.byDomainDst[domainDst{domain: domain, dst: dst}]
+	if !ok || v.expiresAt.Before(a.clock.Now()) {
+		return addrs{}, false
+	}
+	return v, true
+}
+
+func (a *addrAssignments) lookupByMagicIP(mip netip.Addr) (addrs, bool) {
+	v, ok := a.byMagicIP[mip]
+	if !ok || v.expiresAt.Before(a.clock.Now()) {
+		return addrs{}, false
+	}
+	return v, true
+}
+
+func (a *addrAssignments) lookupByTransitIP(tip netip.Addr) (addrs, bool) {
+	v, ok := a.byTransitIP[tip]
+	if !ok || v.expiresAt.Before(a.clock.Now()) {
+		return addrs{}, false
+	}
+	return v, true
+}

feature/conn25/addrAssignments_test.go

@@ -0,0 +1,66 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+package conn25
+
+import (
+	"net/netip"
+	"testing"
+	"time"
+
+	"tailscale.com/tstest"
+)
+
+func TestAssignmentsExpire(t *testing.T) {
+	clock := tstest.NewClock(tstest.ClockOpts{Start: time.Now()})
+	assignments := addrAssignments{clock: clock}
+	as := addrs{
+		dst:     netip.MustParseAddr("0.0.0.1"),
+		magic:   netip.MustParseAddr("0.0.0.2"),
+		transit: netip.MustParseAddr("0.0.0.3"),
+		app:     "a",
+		domain:  "example.com.",
+	}
+	err := assignments.insert(as)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Time has not passed since the insert, the assignment should be returned.
+	foundAs, ok := assignments.lookupByMagicIP(as.magic)
+	if !ok {
+		t.Fatal("expected to find")
+	}
+	if foundAs.dst != as.dst {
+		t.Fatalf("want %v; got %v", as.dst, foundAs.dst)
+	}
+	// and we cannot insert over the addresses
+	err = assignments.insert(as)
+	if err == nil {
+		t.Fatal("expected an error but got nil")
+	}
+	// After a time greater than the default expiry passes, the assignment should
+	// not be returned.
+	clock.Advance(defaultExpiry * 2)
+	foundAsAfter, okAfter := assignments.lookupByMagicIP(as.magic)
+	if okAfter {
+		t.Fatal("expected not to find (expired)")
+	}
+	if foundAsAfter.isValid() {
+		t.Fatal("expected zero val")
+	}
+	// Now we can reuse the addresses
+	err = assignments.insert(as)
+	if err != nil {
+		t.Fatal(err)
+	}
+	foundAs, ok = assignments.lookupByMagicIP(as.magic)
+	if !ok {
+		t.Fatal("expected to find")
+	}
+	if foundAs.dst != as.dst {
+		t.Fatalf("want %v; got %v", as.dst, foundAs.dst)
+	}
+	if !foundAs.expiresAt.After(clock.Now()) {
+		t.Fatalf("expected foundAs to expire after now")
+	}
+}

feature/conn25/conn25.go

@@ -1247,84 +1247,6 @@ func (c addrs) isValid() bool {
 	return c.dst.IsValid()
 }
 
-// domainDst is a key for looking up an existing address assignment by the
-// DNS response domain and destination IP pair.
-type domainDst struct {
-	domain dnsname.FQDN
-	dst    netip.Addr
-}
-
-// addrAssignments is the collection of addrs assigned by this client
-// supporting lookup by magic IP, transit IP or domain+dst, or to lookup all
-// transit IPs associated with a given connector (identified by its node key).
-// byConnKey stores netip.Prefix versions of the transit IPs for use in the
-// WireGuard hooks.
-type addrAssignments struct {
-	byMagicIP   map[netip.Addr]addrs
-	byTransitIP map[netip.Addr]addrs
-	byDomainDst map[domainDst]addrs
-	clock       tstime.Clock
-}
-
-const defaultExpiry = 48 * time.Hour
-
-func (a *addrAssignments) insert(as addrs) error {
-	return a.insertWithExpiry(as, defaultExpiry)
-}
-
-func (a *addrAssignments) insertWithExpiry(as addrs, d time.Duration) error {
-	if !as.expiresAt.IsZero() {
-		return errors.New("expiresAt already set")
-	}
-	now := a.clock.Now()
-	as.expiresAt = now.Add(d)
-	// we don't expect for addresses to be reused before expiry
-	if existing, ok := a.byMagicIP[as.magic]; ok {
-		if !existing.expiresAt.Before(now) {
-			return errors.New("byMagicIP key exists")
-		}
-	}
-	ddst := domainDst{domain: as.domain, dst: as.dst}
-	if existing, ok := a.byDomainDst[ddst]; ok {
-		if !existing.expiresAt.Before(now) {
-			return errors.New("byDomainDst key exists")
-		}
-	}
-	if existing, ok := a.byTransitIP[as.transit]; ok {
-		if !existing.expiresAt.Before(now) {
-			return errors.New("byTransitIP key exists")
-		}
-	}
-	mak.Set(&a.byMagicIP, as.magic, as)
-	mak.Set(&a.byTransitIP, as.transit, as)
-	mak.Set(&a.byDomainDst, ddst, as)
-	return nil
-}
-
-func (a *addrAssignments) lookupByDomainDst(domain dnsname.FQDN, dst netip.Addr) (addrs, bool) {
-	v, ok := a.byDomainDst[domainDst{domain: domain, dst: dst}]
-	if !ok || v.expiresAt.Before(a.clock.Now()) {
-		return addrs{}, false
-	}
-	return v, true
-}
-
-func (a *addrAssignments) lookupByMagicIP(mip netip.Addr) (addrs, bool) {
-	v, ok := a.byMagicIP[mip]
-	if !ok || v.expiresAt.Before(a.clock.Now()) {
-		return addrs{}, false
-	}
-	return v, true
-}
-
-func (a *addrAssignments) lookupByTransitIP(tip netip.Addr) (addrs, bool) {
-	v, ok := a.byTransitIP[tip]
-	if !ok || v.expiresAt.Before(a.clock.Now()) {
-		return addrs{}, false
-	}
-	return v, true
-}
-
 // insertTransitConnMapping adds an entry to the byConnKey map
 // for the provided transitIP (as a prefix).
 // The provided transitIP must already be present in the byTransitIP map.

feature/conn25/conn25_test.go

@@ -25,7 +25,6 @@ import (
 	"tailscale.com/net/tstun"
 	"tailscale.com/tailcfg"
 	"tailscale.com/tsd"
-	"tailscale.com/tstest"
 	"tailscale.com/types/appctype"
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
@@ -2367,57 +2366,3 @@ func TestGetMagicRange(t *testing.T) {
 		}
 	}
 }
-
-func TestAssignmentsExpire(t *testing.T) {
-	clock := tstest.NewClock(tstest.ClockOpts{Start: time.Now()})
-	assignments := addrAssignments{clock: clock}
-	as := addrs{
-		dst:     netip.MustParseAddr("0.0.0.1"),
-		magic:   netip.MustParseAddr("0.0.0.2"),
-		transit: netip.MustParseAddr("0.0.0.3"),
-		app:     "a",
-		domain:  "example.com.",
-	}
-	err := assignments.insert(as)
-	if err != nil {
-		t.Fatal(err)
-	}
-	// Time has not passed since the insert, the assignment should be returned.
-	foundAs, ok := assignments.lookupByMagicIP(as.magic)
-	if !ok {
-		t.Fatal("expected to find")
-	}
-	if foundAs.dst != as.dst {
-		t.Fatalf("want %v; got %v", as.dst, foundAs.dst)
-	}
-	// and we cannot insert over the addresses
-	err = assignments.insert(as)
-	if err == nil {
-		t.Fatal("expected an error but got nil")
-	}
-	// After a time greater than the default expiry passes, the assignment should
-	// not be returned.
-	clock.Advance(defaultExpiry * 2)
-	foundAsAfter, okAfter := assignments.lookupByMagicIP(as.magic)
-	if okAfter {
-		t.Fatal("expected not to find (expired)")
-	}
-	if foundAsAfter.isValid() {
-		t.Fatal("expected zero val")
-	}
-	// Now we can reuse the addresses
-	err = assignments.insert(as)
-	if err != nil {
-		t.Fatal(err)
-	}
-	foundAs, ok = assignments.lookupByMagicIP(as.magic)
-	if !ok {
-		t.Fatal("expected to find")
-	}
-	if foundAs.dst != as.dst {
-		t.Fatalf("want %v; got %v", as.dst, foundAs.dst)
-	}
-	if !foundAs.expiresAt.After(clock.Now()) {
-		t.Fatalf("expected foundAs to expire after now")
-	}
-}