policy: look up approved routes in a set when auto-approving

Avoid scanning the growing approval list for every announcement.
Use the sorted current routes when identifying additions.

hscontrol/policy/policy.go

@@ -8,6 +8,7 @@ import (
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/rs/zerolog/log"
+	"tailscale.com/util/set"
 )
 
 // ReduceRoutes returns a reduced list of routes for a given node that it can access.
@@ -36,11 +37,12 @@ func ApproveRoutesWithPolicy(pm PolicyManager, nv types.NodeView, currentApprove
 
 	// Start with ALL currently approved routes - we never remove approved routes
 	newApproved := slices.Clone(currentApproved)
+	approvedSet := set.SetOf(currentApproved)
 
 	// Then, check for new routes that can be auto-approved
 	for _, route := range announcedRoutes {
 		// Skip if already approved
-		if slices.Contains(newApproved, route) {
+		if approvedSet.Contains(route) {
 			continue
 		}
 
@@ -48,6 +50,7 @@ func ApproveRoutesWithPolicy(pm PolicyManager, nv types.NodeView, currentApprove
 		canApprove := pm.NodeCanApproveRoute(nv, route)
 		if canApprove {
 			newApproved = append(newApproved, route)
+			approvedSet.Add(route)
 		}
 	}
 
@@ -68,7 +71,10 @@ func ApproveRoutesWithPolicy(pm PolicyManager, nv types.NodeView, currentApprove
 		var added, kept []netip.Prefix
 
 		for _, route := range newApproved {
-			if !slices.Contains(sortedCurrent, route) {
+			_, exists := slices.BinarySearchFunc(
+				sortedCurrent, route, netip.Prefix.Compare,
+			)
+			if !exists {
 				added = append(added, route)
 			} else {
 				kept = append(kept, route)

hscontrol/policy/policy_autoapprove_test.go

@@ -339,3 +339,90 @@ func TestApproveRoutesWithPolicy_NilAndEmptyCases(t *testing.T) {
 		}
 	}
 }
+
+func TestApproveRoutesWithPolicyLargeRouteSet(t *testing.T) {
+	routes := make([]netip.Prefix, 16*1024)
+
+	ip := netip.MustParseAddr("10.0.0.0")
+	for i := range routes {
+		routes[i] = netip.PrefixFrom(ip, 32)
+		ip = ip.Next()
+	}
+
+	unannounced := netip.MustParsePrefix("192.0.2.0/24")
+	refused := netip.MustParsePrefix("203.0.113.0/24")
+
+	current := append(slices.Clone(routes[:len(routes)/2]), unannounced)
+	announced := append(slices.Clone(routes), routes[len(routes)-1], refused)
+	currentBefore := slices.Clone(current)
+	announcedBefore := slices.Clone(announced)
+	want := append(slices.Clone(routes), unannounced)
+	user := types.User{ID: 1, Name: "router"}
+	n := &types.Node{
+		ID:             1,
+		IPv4:           new(netip.MustParseAddr("100.64.0.1")),
+		UserID:         new(user.ID),
+		User:           new(user),
+		ApprovedRoutes: current,
+	}
+	pm, err := policyv2.NewPolicyManager([]byte(`{
+		"acls": [],
+		"autoApprovers": {"routes": {"10.0.0.0/8": ["router@"]}}
+	}`), []types.User{user}, types.Nodes{n}.ViewSlice())
+	require.NoError(t, err)
+
+	approved, changed := ApproveRoutesWithPolicy(pm, n.View(), current, announced)
+	require.True(t, changed)
+	require.Equal(t, want, approved)
+	require.Equal(t, currentBefore, current, "existing approvals must remain immutable")
+	require.Equal(t, announcedBefore, announced, "announcements must remain immutable")
+
+	again, changed := ApproveRoutesWithPolicy(pm, n.View(), approved, announced)
+	require.False(t, changed)
+	require.Equal(t, want, again)
+}
+
+func BenchmarkApproveRoutesWithPolicyLargeRouteSet(b *testing.B) {
+	routes := make([]netip.Prefix, 16*1024)
+
+	ip := netip.MustParseAddr("10.0.0.0")
+	for i := range routes {
+		routes[i] = netip.PrefixFrom(ip, 32)
+		ip = ip.Next()
+	}
+
+	user := types.User{ID: 1, Name: "router"}
+	n := &types.Node{
+		ID:     1,
+		IPv4:   new(netip.MustParseAddr("100.64.0.1")),
+		UserID: new(user.ID),
+		User:   new(user),
+	}
+	pm, err := policyv2.NewPolicyManager([]byte(`{
+		"acls": [],
+		"autoApprovers": {"routes": {"10.0.0.0/8": ["router@"]}}
+	}`), []types.User{user}, types.Nodes{n}.ViewSlice())
+	require.NoError(b, err)
+
+	for _, retained := range []int{0, len(routes) / 2, len(routes)} {
+		b.Run(fmt.Sprintf("retained=%d", retained), func(b *testing.B) {
+			current := routes[:retained]
+
+			var (
+				approved []netip.Prefix
+				changed  bool
+			)
+
+			b.ReportAllocs()
+
+			for b.Loop() {
+				approved, changed = ApproveRoutesWithPolicy(
+					pm, n.View(), current, routes)
+			}
+
+			if len(approved) != len(routes) || changed != (retained < len(routes)) {
+				b.Fatal("unexpected approval result")
+			}
+		})
+	}
+}