diff --git a/hscontrol/policy/policy.go b/hscontrol/policy/policy.go index 55730cd18..c3e6741ca 100644 --- a/hscontrol/policy/policy.go +++ b/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) diff --git a/hscontrol/policy/policy_autoapprove_test.go b/hscontrol/policy/policy_autoapprove_test.go index 3f229ccd4..9116d1a93 100644 --- a/hscontrol/policy/policy_autoapprove_test.go +++ b/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") + } + }) + } +}