policy: consolidate cache-invalidation and evaluation helpers

hscontrol/policy/v2/compiled.go

@@ -11,6 +11,7 @@ import (
 	"go4.org/netipx"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/views"
+	"tailscale.com/util/set"
 )
 
 // grantCategory classifies a grant by what per-node work it needs.
@@ -324,17 +325,15 @@ func (pol *Policy) compileOneGrant(
 		)
 	}
 
-	// Classify and store deferred self data.
-	switch {
-	case len(autogroupSelfDests) > 0:
+	// Classify and store deferred self data. The struct literal already
+	// initializes category to grantCategoryRegular (the zero value).
+	if len(autogroupSelfDests) > 0 {
 		cg.category = grantCategorySelf
 		cg.self = &selfGrantData{
 			resolvedSrcs:      resolvedSrcs,
 			internetProtocols: grant.InternetProtocols,
 			app:               grant.App,
 		}
-	default:
-		cg.category = grantCategoryRegular
 	}
 
 	return cg, nil
@@ -472,15 +471,12 @@ func buildSrcIPStrings(
 	// individual IPs from non-wildcard sources alongside the
 	// merged CGNAT ranges rather than absorbing them.
 	if hasWildcard && len(nonWildcardSrcs) > 0 {
-		seen := make(map[string]bool, len(srcIPStrs))
-		for _, s := range srcIPStrs {
-			seen[s] = true
-		}
+		seen := set.SetOf(srcIPStrs)
 
 		for _, ips := range nonWildcardSrcs {
 			for _, s := range ips.Strings() {
-				if !seen[s] {
-					seen[s] = true
+				if !seen.Contains(s) {
+					seen.Add(s)
 					srcIPStrs = append(srcIPStrs, s)
 				}
 			}
@@ -627,7 +623,7 @@ func collectRelayTargetIPs(grants []compiledGrant) (*netipx.IPSet, error) {
 // traffic through it must recompute when it goes offline. Returns nil when no
 // via grants exist.
 func collectViaTargetTags(grants []compiledGrant) map[Tag]struct{} {
-	var tags map[Tag]struct{}
+	tags := make(map[Tag]struct{})
 
 	for i := range grants {
 		if grants[i].via == nil {
@@ -635,14 +631,14 @@ func collectViaTargetTags(grants []compiledGrant) map[Tag]struct{} {
 		}
 
 		for _, t := range grants[i].via.viaTags {
-			if tags == nil {
-				tags = make(map[Tag]struct{})
-			}
-
 			tags[t] = struct{}{}
 		}
 	}
 
+	if len(tags) == 0 {
+		return nil
+	}
+
 	return tags
 }
 

hscontrol/policy/v2/filter.go

@@ -1,7 +1,6 @@
 package v2
 
 import (
-	"cmp"
 	"errors"
 	"fmt"
 	"net/netip"
@@ -23,13 +22,22 @@ var (
 	errSelfInSources = errors.New("autogroup:self cannot be used in sources")
 )
 
-// companionCaps maps certain well-known Tailscale capabilities to
+// companionCap pairs a well-known Tailscale capability with its
+// companion capability.
+type companionCap struct {
+	original  tailcfg.PeerCapability
+	companion tailcfg.PeerCapability
+}
+
+// companionCaps lists certain well-known Tailscale capabilities and
 // their companion capability. When a grant includes one of these
 // capabilities, Tailscale automatically generates an additional
 // [tailcfg.FilterRule] with the companion capability and a nil CapMap value.
-var companionCaps = map[tailcfg.PeerCapability]tailcfg.PeerCapability{
-	tailcfg.PeerCapabilityTaildrive: tailcfg.PeerCapabilityTaildriveSharer,
-	tailcfg.PeerCapabilityRelay:     tailcfg.PeerCapabilityRelayTarget,
+// The slice is ordered by the original capability name so that
+// generated companion rules are emitted deterministically.
+var companionCaps = []companionCap{
+	{tailcfg.PeerCapabilityTaildrive, tailcfg.PeerCapabilityTaildriveSharer},
+	{tailcfg.PeerCapabilityRelay, tailcfg.PeerCapabilityRelayTarget},
 }
 
 // companionCapGrantRules returns additional [tailcfg.FilterRule]s for any
@@ -48,38 +56,22 @@ func companionCapGrantRules(
 	srcPrefixes []netip.Prefix,
 	capMap tailcfg.PeerCapMap,
 ) []tailcfg.FilterRule {
-	// Process in deterministic order by original capability name.
-	type pair struct {
-		original  tailcfg.PeerCapability
-		companion tailcfg.PeerCapability
-	}
+	companions := make([]tailcfg.FilterRule, 0, len(companionCaps))
 
-	var pairs []pair
-
-	for cap, companion := range companionCaps {
-		if _, ok := capMap[cap]; ok {
-			pairs = append(pairs, pair{cap, companion})
-		}
-	}
-
-	slices.SortFunc(pairs, func(a, b pair) int {
-		return cmp.Compare(a.original, b.original)
-	})
-
-	companions := make([]tailcfg.FilterRule, 0, len(pairs))
-
-	for _, p := range pairs {
-		companions = append(companions, tailcfg.FilterRule{
-			SrcIPs: dstIPStrings,
-			CapGrant: []tailcfg.CapGrant{
-				{
-					Dsts: srcPrefixes,
-					CapMap: tailcfg.PeerCapMap{
-						p.companion: nil,
+	for _, c := range companionCaps {
+		if _, ok := capMap[c.original]; ok {
+			companions = append(companions, tailcfg.FilterRule{
+				SrcIPs: dstIPStrings,
+				CapGrant: []tailcfg.CapGrant{
+					{
+						Dsts: srcPrefixes,
+						CapMap: tailcfg.PeerCapMap{
+							c.companion: nil,
+						},
 					},
 				},
-			},
-		})
+			})
+		}
 	}
 
 	return companions
@@ -238,7 +230,7 @@ func checkPeriodFromRule(rule SSH) time.Duration {
 	}
 }
 
-func sshCheck(baseURL string, _ time.Duration) tailcfg.SSHAction {
+func sshCheck(baseURL string) tailcfg.SSHAction {
 	holdURL := baseURL + "/machine/ssh/action/$SRC_NODE_ID/to/$DST_NODE_ID?local_user=$LOCAL_USER"
 
 	return tailcfg.SSHAction{
@@ -302,7 +294,7 @@ func (pol *Policy) compileSSHPolicy(
 		case SSHActionAccept:
 			action = sshAccept
 		case SSHActionCheck:
-			action = sshCheck(baseURL, checkPeriodFromRule(rule))
+			action = sshCheck(baseURL)
 		default:
 			return nil, fmt.Errorf(
 				"parsing SSH policy, unknown action %q, index: %d: %w",
@@ -425,12 +417,7 @@ func (pol *Policy) compileSSHPolicy(
 					allPrincipals = append(allPrincipals, taggedPrincipals...)
 
 					if len(allPrincipals) > 0 {
-						rules = append(rules, &tailcfg.SSHRule{
-							Principals: allPrincipals,
-							SSHUsers:   baseUserMap,
-							Action:     &action,
-							AcceptEnv:  acceptEnv,
-						})
+						appendRules(allPrincipals, 0, false)
 					}
 				}
 			} else if hasLocalpart && slices.ContainsFunc(node.IPs(), srcIPs.Contains) {

hscontrol/policy/v2/policy.go

@@ -822,7 +822,7 @@ func (pm *PolicyManager) SetUsers(users []types.User) (bool, error) {
 
 	// If SSH policies exist, force a policy change when users are updated
 	// This ensures nodes get updated SSH policies even if other policy hashes didn't change
-	if pm.pol != nil && pm.pol.SSHs != nil && len(pm.pol.SSHs) > 0 {
+	if pm.pol != nil && len(pm.pol.SSHs) > 0 {
 		return true, nil
 	}
 
@@ -878,15 +878,23 @@ func (pm *PolicyManager) SetNodes(nodes views.Slice[types.NodeView]) (bool, erro
 	return false, nil
 }
 
+// nodeIDViewMap indexes a slice of node views by node ID. On duplicate IDs the
+// last view wins, matching the open-coded loops it replaces.
+func nodeIDViewMap(s views.Slice[types.NodeView]) map[types.NodeID]types.NodeView {
+	m := make(map[types.NodeID]types.NodeView, s.Len())
+	for _, n := range s.All() {
+		m[n.ID()] = n
+	}
+
+	return m
+}
+
 func (pm *PolicyManager) nodesHavePolicyAffectingChanges(newNodes views.Slice[types.NodeView]) bool {
 	if pm.nodes.Len() != newNodes.Len() {
 		return true
 	}
 
-	oldNodes := make(map[types.NodeID]types.NodeView, pm.nodes.Len())
-	for _, node := range pm.nodes.All() {
-		oldNodes[node.ID()] = node
-	}
+	oldNodes := nodeIDViewMap(pm.nodes)
 
 	for _, newNode := range newNodes.All() {
 		oldNode, exists := oldNodes[newNode.ID()]
@@ -1387,15 +1395,8 @@ func (pm *PolicyManager) DebugString() string {
 // the entire cache.
 func (pm *PolicyManager) invalidateAutogroupSelfCache(oldNodes, newNodes views.Slice[types.NodeView]) {
 	// Build maps for efficient lookup
-	oldNodeMap := make(map[types.NodeID]types.NodeView)
-	for _, node := range oldNodes.All() {
-		oldNodeMap[node.ID()] = node
-	}
-
-	newNodeMap := make(map[types.NodeID]types.NodeView)
-	for _, node := range newNodes.All() {
-		newNodeMap[node.ID()] = node
-	}
+	oldNodeMap := nodeIDViewMap(oldNodes)
+	newNodeMap := nodeIDViewMap(newNodes)
 
 	// Track which users are affected by changes.
 	// Tagged nodes don't participate in autogroup:self (identity is tag-based),
@@ -1476,53 +1477,29 @@ func (pm *PolicyManager) invalidateAutogroupSelfCache(oldNodes, newNodes views.S
 	// For autogroup:self, we need to clear all nodes belonging to affected users
 	// because autogroup:self rules depend on the entire user's device set.
 	for nodeID := range pm.filterRulesMap {
-		// Find the user for this cached node
+		// Find the user for this cached node using the already-built indexes.
+		node, ok := newNodeMap[nodeID]
+		if !ok {
+			node, ok = oldNodeMap[nodeID]
+		}
+
+		// Node not found in either old or new list, clear it.
+		if !ok {
+			delete(pm.filterRulesMap, nodeID)
+			delete(pm.matchersForNodeMap, nodeID)
+
+			continue
+		}
+
+		// Tagged nodes don't participate in autogroup:self, so their cache
+		// doesn't need user-based invalidation; leave nodeUserID at zero.
 		var nodeUserID types.UserID
-
-		found := false
-
-		// Check in new nodes first
-		for _, node := range newNodes.All() {
-			if node.ID() == nodeID {
-				// Tagged nodes don't participate in autogroup:self,
-				// so their cache doesn't need user-based invalidation.
-				if node.IsTagged() {
-					found = true
-					break
-				}
-
-				nodeUserID = node.TypedUserID()
-				found = true
-
-				break
-			}
+		if !node.IsTagged() {
+			nodeUserID = node.TypedUserID()
 		}
 
-		// If not found in new nodes, check old nodes
-		if !found {
-			for _, node := range oldNodes.All() {
-				if node.ID() == nodeID {
-					if node.IsTagged() {
-						found = true
-						break
-					}
-
-					nodeUserID = node.TypedUserID()
-					found = true
-
-					break
-				}
-			}
-		}
-
-		// If we found the user and they're affected, clear this cache entry
-		if found {
-			if _, affected := affectedUsers[nodeUserID]; affected {
-				delete(pm.filterRulesMap, nodeID)
-				delete(pm.matchersForNodeMap, nodeID)
-			}
-		} else {
-			// Node not found in either old or new list, clear it
+		// If the owning user is affected, clear this cache entry.
+		if _, affected := affectedUsers[nodeUserID]; affected {
 			delete(pm.filterRulesMap, nodeID)
 			delete(pm.matchersForNodeMap, nodeID)
 		}
@@ -1553,15 +1530,8 @@ func (pm *PolicyManager) invalidateNodeCache(newNodes views.Slice[types.NodeView
 // invalidateGlobalPolicyCache invalidates only nodes whose properties affecting
 // [policyutil.ReduceFilterRules] changed. For global policies, each node's filter is independent.
 func (pm *PolicyManager) invalidateGlobalPolicyCache(newNodes views.Slice[types.NodeView]) {
-	oldNodeMap := make(map[types.NodeID]types.NodeView)
-	for _, node := range pm.nodes.All() {
-		oldNodeMap[node.ID()] = node
-	}
-
-	newNodeMap := make(map[types.NodeID]types.NodeView)
-	for _, node := range newNodes.All() {
-		newNodeMap[node.ID()] = node
-	}
+	oldNodeMap := nodeIDViewMap(pm.nodes)
+	newNodeMap := nodeIDViewMap(newNodes)
 
 	// Invalidate nodes whose properties changed
 	for nodeID, newNode := range newNodeMap {

hscontrol/policy/v2/sshtest.go

@@ -134,11 +134,7 @@ func (pm *PolicyManager) RunSSHTests() error {
 	cache := make(map[types.NodeID]*tailcfg.SSHPolicy)
 	results := runSSHPolicyTests(pm.pol, pm.users, pm.nodes, cache)
 
-	if results.AllPassed {
-		return nil
-	}
-
-	return fmt.Errorf("%w:\n%s", errSSHPolicyTestsFailed, results.Errors())
+	return wrapTestResult(errSSHPolicyTestsFailed, results.AllPassed, results.Errors)
 }
 
 // evaluateSSHTests runs the block against pol without mutating live state.
@@ -154,11 +150,7 @@ func evaluateSSHTests(
 	cache := make(map[types.NodeID]*tailcfg.SSHPolicy)
 	results := runSSHPolicyTests(pol, users, nodes, cache)
 
-	if results.AllPassed {
-		return nil
-	}
-
-	return fmt.Errorf("%w:\n%s", errSSHPolicyTestsFailed, results.Errors())
+	return wrapTestResult(errSSHPolicyTestsFailed, results.AllPassed, results.Errors)
 }
 
 // runSSHPolicyTests evaluates every sshTests entry. The cache is keyed
@@ -251,28 +243,21 @@ func runSSHPolicyTest(
 		return res
 	}
 
-	for _, user := range test.Accept {
-		evaluateAssertion(
-			pol, users, nodes, cache,
-			srcAddrs, dstNodes, user.String(),
-			assertAccept, &res,
-		)
-	}
-
-	for _, user := range test.Deny {
-		evaluateAssertion(
-			pol, users, nodes, cache,
-			srcAddrs, dstNodes, user.String(),
-			assertDeny, &res,
-		)
-	}
-
-	for _, user := range test.Check {
-		evaluateAssertion(
-			pol, users, nodes, cache,
-			srcAddrs, dstNodes, user.String(),
-			assertCheck, &res,
-		)
+	for _, g := range []struct {
+		users []SSHUser
+		kind  sshAssertion
+	}{
+		{test.Accept, assertAccept},
+		{test.Deny, assertDeny},
+		{test.Check, assertCheck},
+	} {
+		for _, user := range g.users {
+			evaluateAssertion(
+				pol, users, nodes, cache,
+				srcAddrs, dstNodes, user.String(),
+				g.kind, &res,
+			)
+		}
 	}
 
 	return res

hscontrol/policy/v2/test.go

@@ -196,6 +196,17 @@ func (r PolicyTestResults) Errors() string {
 	return strings.Join(lines, "\n")
 }
 
+// wrapTestResult returns nil when every test passed, otherwise wraps the
+// rendered failure breakdown in sentinel. errs is passed uncalled so it is
+// only evaluated on the failure path.
+func wrapTestResult(sentinel error, allPassed bool, errs func() string) error {
+	if allPassed {
+		return nil
+	}
+
+	return fmt.Errorf("%w:\n%s", sentinel, errs())
+}
+
 // RunTests evaluates the policy's own `tests` block against the live compiled
 // filter and returns a wrapped error when any test fails. Callers that need
 // the per-test breakdown can call runPolicyTests directly.
@@ -208,11 +219,8 @@ func (pm *PolicyManager) RunTests() error {
 	defer pm.mu.Unlock()
 
 	results := runPolicyTests(pm.pol, pm.filter, pm.users, pm.nodes)
-	if results.AllPassed {
-		return nil
-	}
 
-	return fmt.Errorf("%w:\n%s", errPolicyTestsFailed, results.Errors())
+	return wrapTestResult(errPolicyTestsFailed, results.AllPassed, results.Errors)
 }
 
 // evaluateTests runs the `tests` block against a fresh compilation of pol.
@@ -233,11 +241,8 @@ func evaluateTests(pol *Policy, users []types.User, nodes views.Slice[types.Node
 	}
 
 	results := runPolicyTests(pol, filter, users, nodes)
-	if results.AllPassed {
-		return nil
-	}
 
-	return fmt.Errorf("%w:\n%s", errPolicyTestsFailed, results.Errors())
+	return wrapTestResult(errPolicyTestsFailed, results.AllPassed, results.Errors)
 }
 
 // runPolicyTests is the pure evaluation function: given a policy, the
@@ -285,39 +290,28 @@ func runPolicyTest(test PolicyTest, pol *Policy, filter []tailcfg.FilterRule, us
 		return res
 	}
 
-	for _, dst := range test.Accept {
-		allowed, err := evalReachability(srcPrefixes, dst, test.Proto, pol, filter, users, nodes)
-		if err != nil {
-			res.Passed = false
-			res.Errors = append(res.Errors, fmt.Sprintf("error testing %q: %v", dst, err))
+	check := func(dsts []string, wantAllowed bool, ok, fail *[]string) {
+		for _, dst := range dsts {
+			allowed, err := evalReachability(srcPrefixes, dst, test.Proto, pol, filter, users, nodes)
+			if err != nil {
+				res.Passed = false
+				res.Errors = append(res.Errors, fmt.Sprintf("error testing %q: %v", dst, err))
 
-			continue
-		}
+				continue
+			}
 
-		if allowed {
-			res.AcceptOK = append(res.AcceptOK, dst)
-		} else {
-			res.Passed = false
-			res.AcceptFail = append(res.AcceptFail, dst)
+			if allowed == wantAllowed {
+				*ok = append(*ok, dst)
+			} else {
+				res.Passed = false
+
+				*fail = append(*fail, dst)
+			}
 		}
 	}
 
-	for _, dst := range test.Deny {
-		allowed, err := evalReachability(srcPrefixes, dst, test.Proto, pol, filter, users, nodes)
-		if err != nil {
-			res.Passed = false
-			res.Errors = append(res.Errors, fmt.Sprintf("error testing %q: %v", dst, err))
-
-			continue
-		}
-
-		if !allowed {
-			res.DenyOK = append(res.DenyOK, dst)
-		} else {
-			res.Passed = false
-			res.DenyFail = append(res.DenyFail, dst)
-		}
-	}
+	check(test.Accept, true, &res.AcceptOK, &res.AcceptFail)
+	check(test.Deny, false, &res.DenyOK, &res.DenyFail)
 
 	return res
 }

hscontrol/policy/v2/types.go

@@ -20,6 +20,7 @@ import (
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/views"
 	"tailscale.com/util/multierr"
+	"tailscale.com/util/set"
 	"tailscale.com/util/slicesx"
 )
 
@@ -341,15 +342,15 @@ func (a Asterix) resolve(_ *Policy, _ types.Users, _ views.Slice[types.NodeView]
 // IPSet merges overlapping ranges (e.g. 10.0.0.0/8 absorbs
 // 10.33.0.0/16), but Tailscale preserves individual route entries.
 func approvedSubnetRoutes(nodes views.Slice[types.NodeView]) []string {
-	seen := make(map[string]bool)
+	seen := make(set.Set[string])
 
 	var routes []string
 
 	for _, node := range nodes.All() {
 		for _, route := range node.SubnetRoutes() {
 			s := route.String()
-			if !seen[s] {
-				seen[s] = true
+			if !seen.Contains(s) {
+				seen.Add(s)
 				routes = append(routes, s)
 			}
 		}