policy/v2: compile SSH rules once and share them across nodes

SSHPolicy re-resolved every rule and built private principals per node.
Compile the rules once per policy, users and nodes; per node, only
select.

hscontrol/policy/v2/filter.go

@@ -254,29 +254,105 @@ func sshCheck(baseURL string) tailcfg.SSHAction {
 	}
 }
 
-//nolint:gocyclo // SSH compilation walks per-rule branches with intertwined autogroup:self handling
-func (pol *Policy) compileSSHPolicy(
-	baseURL string,
-	users types.Users,
-	node types.NodeView,
-	nodes views.Slice[types.NodeView],
-) (*tailcfg.SSHPolicy, error) {
-	// Never nil: clients read a nil SSHPolicy as "keep the previous
-	// rules", so revoked access would stay. SaaS sends "rules":[].
-	if pol == nil || len(pol.SSHs) == 0 {
-		return &tailcfg.SSHPolicy{Rules: []*tailcfg.SSHRule{}}, nil
+// compiledSSHRule is an [SSH] rule with everything that does not depend on
+// the node it is delivered to already resolved, as [compiledGrant] is for
+// grants. It is read-only once built: the SSH policies of all nodes point
+// into it, so what a rule expands to is kept once and not once per node.
+type compiledSSHRule struct {
+	srcs *sshSources
+
+	sshUsers        map[string]string
+	localpartByUser map[uint]string
+	action          *tailcfg.SSHAction
+	acceptEnv       []string
+
+	hasLocalpart bool
+	hasSelfDest  bool
+
+	// userRules holds the rules scoped to one source user's own nodes,
+	// for autogroup:self destinations and localpart self-access.
+	userRules map[uint][]*tailcfg.SSHRule
+
+	// dests holds the destinations other than autogroup:self, nil if
+	// there are none. Every node in it gets destRules.
+	dests     *netipx.IPSet
+	destRules []*tailcfg.SSHRule
+}
+
+// sshSources is what the sources of an [SSH] rule resolve to. It is
+// read-only once built: rules with the same source addresses share it.
+type sshSources struct {
+	ips              ResolvedAddresses
+	userIDs          []uint
+	principalsByUser map[uint][]*tailcfg.SSHPrincipal
+	taggedPrincipals []*tailcfg.SSHPrincipal
+
+	// allPrincipals merges user and tagged principals into a single
+	// list. Tagged principals preserve per-tag duplication (a node with
+	// N tags appears N times, matching SaaS behavior).
+	allPrincipals []*tailcfg.SSHPrincipal
+}
+
+// rules emits a common rule and, if the user has a localpart match, a
+// per-user localpart rule.
+func (r *compiledSSHRule) rules(
+	principals []*tailcfg.SSHPrincipal,
+	uid uint,
+	hasUID bool,
+) []*tailcfg.SSHRule {
+	rules := []*tailcfg.SSHRule{{
+		Principals: principals,
+		SSHUsers:   r.sshUsers,
+		Action:     r.action,
+		AcceptEnv:  r.acceptEnv,
+	}}
+
+	if hasUID {
+		if lp, ok := r.localpartByUser[uid]; ok {
+			rules = append(rules, &tailcfg.SSHRule{
+				Principals: principals,
+				SSHUsers:   map[string]string{lp: lp},
+				Action:     r.action,
+				AcceptEnv:  r.acceptEnv,
+			})
+		}
 	}
 
-	log.Trace().Caller().Msgf("compiling SSH policy for node %q", node.Hostname())
+	return rules
+}
 
-	rules := []*tailcfg.SSHRule{}
+// compileSSHRules resolves all SSH rules of the policy into
+// [compiledSSHRule] structs. Rules whose sources resolve to nothing are
+// left out.
+//
+//nolint:gocyclo // SSH compilation walks per-rule branches with intertwined autogroup:self handling
+func (pol *Policy) compileSSHRules(
+	baseURL string,
+	users types.Users,
+	nodes views.Slice[types.NodeView],
+) ([]compiledSSHRule, error) {
+	if pol == nil {
+		return nil, nil
+	}
+
+	compiled := make([]compiledSSHRule, 0, len(pol.SSHs))
+
+	// There is one principal per source address and the policy author
+	// decides how many rules there are, so rules whose sources resolve to
+	// the same addresses share their principals.
+	sources := make(map[string]*sshSources)
 
 	for index, rule := range pol.SSHs {
-		var autogroupSelfDests, otherDests []Alias
+		var otherDests []Alias
+
+		cr := compiledSSHRule{
+			hasLocalpart: rule.Users.ContainsLocalpart(),
+			acceptEnv:    rule.AcceptEnv,
+		}
 
 		for _, dst := range rule.Destinations {
 			if ag, ok := dst.(*AutoGroup); ok && ag.Is(AutoGroupSelf) {
-				autogroupSelfDests = append(autogroupSelfDests, dst)
+				cr.hasSelfDest = true
 			} else {
 				otherDests = append(otherDests, dst)
 			}
@@ -307,73 +383,64 @@ func (pol *Policy) compileSSHPolicy(
 			)
 		}
 
-		acceptEnv := rule.AcceptEnv
+		cr.action = &action
 
 		// Build the common userMap (always has at least a root entry).
 		const rootUser = "root"
 
-		baseUserMap := make(map[string]string, len(rule.Users))
+		cr.sshUsers = make(map[string]string, len(rule.Users))
 		if rule.Users.ContainsNonRoot() {
-			baseUserMap["*"] = "="
+			cr.sshUsers["*"] = "="
 		}
 
 		if rule.Users.ContainsRoot() {
-			baseUserMap[rootUser] = rootUser
+			cr.sshUsers[rootUser] = rootUser
 		} else {
-			baseUserMap[rootUser] = ""
+			cr.sshUsers[rootUser] = ""
 		}
 
 		for _, u := range rule.Users.NormalUsers() {
-			baseUserMap[u.String()] = u.String()
+			cr.sshUsers[u.String()] = u.String()
 		}
 
-		hasLocalpart := rule.Users.ContainsLocalpart()
-
-		var localpartByUser map[uint]string
-		if hasLocalpart {
-			localpartByUser = resolveLocalparts(
+		if cr.hasLocalpart {
+			cr.localpartByUser = resolveLocalparts(
 				rule.Users.LocalpartEntries(), users,
 			)
 		}
 
-		userIDs, principalsByUser, taggedPrincipals := groupSourcesByUser(
-			nodes, srcIPs,
-		)
+		srcKey := strings.Join(srcIPs.Strings(), ",")
 
-		// appendRules emits a common rule and, if the user has a
-		// localpart match, a per-user localpart rule.
-		appendRules := func(principals []*tailcfg.SSHPrincipal, uid uint, hasUID bool) {
-			rules = append(rules, &tailcfg.SSHRule{
-				Principals: principals,
-				SSHUsers:   baseUserMap,
-				Action:     &action,
-				AcceptEnv:  acceptEnv,
-			})
+		srcs, ok := sources[srcKey]
+		if !ok {
+			userIDs, byUser, tagged := groupSourcesByUser(nodes, srcIPs)
+			srcs = &sshSources{
+				ips:              srcIPs,
+				userIDs:          userIDs,
+				principalsByUser: byUser,
+				taggedPrincipals: tagged,
+			}
 
-			if hasUID {
-				if lp, ok := localpartByUser[uid]; ok {
-					rules = append(rules, &tailcfg.SSHRule{
-						Principals: principals,
-						SSHUsers:   map[string]string{lp: lp},
-						Action:     &action,
-						AcceptEnv:  acceptEnv,
-					})
-				}
+			for _, uid := range srcs.userIDs {
+				srcs.allPrincipals = append(
+					srcs.allPrincipals, srcs.principalsByUser[uid]...,
+				)
+			}
+
+			srcs.allPrincipals = append(srcs.allPrincipals, srcs.taggedPrincipals...)
+
+			sources[srcKey] = srcs
+		}
+
+		cr.srcs = srcs
+
+		if cr.hasSelfDest || cr.hasLocalpart {
+			cr.userRules = make(map[uint][]*tailcfg.SSHRule, len(srcs.userIDs))
+			for _, uid := range srcs.userIDs {
+				cr.userRules[uid] = cr.rules(srcs.principalsByUser[uid], uid, true)
 			}
 		}
 
-		// Handle autogroup:self destinations.
-		// Tagged nodes can't match autogroup:self.
-		if len(autogroupSelfDests) > 0 &&
-			!node.IsTagged() && node.User().Valid() {
-			uid := node.User().ID()
-
-			if principals := principalsByUser[uid]; len(principals) > 0 {
-				appendRules(principals, uid, true)
-			}
-		}
-
-		// Handle other destinations.
 		if len(otherDests) > 0 {
 			var dest netipx.IPSetBuilder
 
@@ -391,60 +458,78 @@ func (pol *Policy) compileSSHPolicy(
 				}
 			}
 
-			destSet, err := dest.IPSet()
+			cr.dests, err = dest.IPSet()
 			if err != nil {
 				return nil, err
 			}
 
-			if node.InIPSet(destSet) {
-				// Node is a destination — emit rules.
-				// When localpart entries exist, interleave common
-				// and localpart rules per source user to match
-				// Tailscale SaaS first-match-wins ordering.
-				if hasLocalpart {
-					for _, uid := range userIDs {
-						appendRules(principalsByUser[uid], uid, true)
-					}
+			// When localpart entries exist, interleave common
+			// and localpart rules per source user to match
+			// Tailscale SaaS first-match-wins ordering.
+			if cr.hasLocalpart {
+				for _, uid := range srcs.userIDs {
+					cr.destRules = append(cr.destRules, cr.userRules[uid]...)
+				}
 
-					if len(taggedPrincipals) > 0 {
-						appendRules(taggedPrincipals, 0, false)
-					}
-				} else {
-					// Merge user and tagged principals into a
-					// single list. Tagged principals preserve
-					// per-tag duplication (a node with N tags
-					// appears N times, matching SaaS behavior).
-					var allPrincipals []*tailcfg.SSHPrincipal
-					for _, uid := range userIDs {
-						allPrincipals = append(allPrincipals, principalsByUser[uid]...)
-					}
+				if len(srcs.taggedPrincipals) > 0 {
+					cr.destRules = append(
+						cr.destRules, cr.rules(srcs.taggedPrincipals, 0, false)...,
+					)
+				}
+			} else if len(srcs.allPrincipals) > 0 {
+				cr.destRules = cr.rules(srcs.allPrincipals, 0, false)
+			}
+		}
 
-					allPrincipals = append(allPrincipals, taggedPrincipals...)
+		compiled = append(compiled, cr)
+	}
 
-					if len(allPrincipals) > 0 {
-						appendRules(allPrincipals, 0, false)
-					}
-				}
-			} else if hasLocalpart && slices.ContainsFunc(node.IPs(), srcIPs.Contains) {
-				// Self-access: source node not in destination set
-				// receives rules scoped to its own user.
-				if node.IsTagged() {
-					var builder netipx.IPSetBuilder
-
-					node.AppendToIPSet(&builder)
-
-					ipSet, err := builder.IPSet()
-					if err == nil && ipSet != nil {
-						if principals := ipSetToPrincipals(ipSet); len(principals) > 0 {
-							appendRules(principals, 0, false)
-						}
-					}
-				} else if node.User().Valid() {
-					uid := node.User().ID()
-					if principals := principalsByUser[uid]; len(principals) > 0 {
-						appendRules(principals, uid, true)
+	return compiled, nil
+}
+
+// sshPolicyForNode picks the rules that apply to node out of compiled.
+func sshPolicyForNode(
+	compiled []compiledSSHRule,
+	node types.NodeView,
+) *tailcfg.SSHPolicy {
+	// Never nil: clients read a nil SSHPolicy as "keep the previous
+	// rules", so revoked access would stay. SaaS sends "rules":[].
+	rules := []*tailcfg.SSHRule{}
+
+	for i := range compiled {
+		cr := &compiled[i]
+
+		// Handle autogroup:self destinations.
+		// Tagged nodes can't match autogroup:self.
+		if cr.hasSelfDest && !node.IsTagged() && node.User().Valid() {
+			rules = append(rules, cr.userRules[node.User().ID()]...)
+		}
+
+		// Handle other destinations.
+		if cr.dests == nil {
+			continue
+		}
+
+		if node.InIPSet(cr.dests) {
+			// Node is a destination — emit rules.
+			rules = append(rules, cr.destRules...)
+		} else if cr.hasLocalpart &&
+			slices.ContainsFunc(node.IPs(), cr.srcs.ips.Contains) {
+			// Self-access: source node not in destination set
+			// receives rules scoped to its own user.
+			if node.IsTagged() {
+				var builder netipx.IPSetBuilder
+
+				node.AppendToIPSet(&builder)
+
+				ipSet, err := builder.IPSet()
+				if err == nil && ipSet != nil {
+					if principals := ipSetToPrincipals(ipSet); len(principals) > 0 {
+						rules = append(rules, cr.rules(principals, 0, false)...)
 					}
 				}
+			} else if node.User().Valid() {
+				rules = append(rules, cr.userRules[node.User().ID()]...)
 			}
 		}
 	}
@@ -468,7 +553,26 @@ func (pol *Policy) compileSSHPolicy(
 
 	return &tailcfg.SSHPolicy{
 		Rules: rules,
-	}, nil
+	}
+}
+
+// compileSSHPolicy compiles the SSH policy of a single node.
+// [PolicyManager.SSHPolicy], which is asked for every node, keeps the
+// result of [Policy.compileSSHRules] instead.
+func (pol *Policy) compileSSHPolicy(
+	baseURL string,
+	users types.Users,
+	node types.NodeView,
+	nodes views.Slice[types.NodeView],
+) (*tailcfg.SSHPolicy, error) {
+	log.Trace().Caller().Msgf("compiling SSH policy for node %q", node.Hostname())
+
+	compiled, err := pol.compileSSHRules(baseURL, users, nodes)
+	if err != nil {
+		return nil, err
+	}
+
+	return sshPolicyForNode(compiled, node), nil
 }
 
 // ipSetToPrincipals converts an [netipx.IPSet] into SSH principals, one per address.

hscontrol/policy/v2/policy.go

@@ -64,6 +64,15 @@ type PolicyManager struct {
 	// Lazy map of SSH policies
 	sshPolicyMap *xsync.Map[types.NodeID, *tailcfg.SSHPolicy]
 
+	// sshRules are the SSH rules with sources and destinations
+	// pre-resolved, shared by every node's policy in sshPolicyMap so
+	// the principal lists exist once rather than once per rule and
+	// node. They are reset to nil wherever sshPolicyMap is cleared
+	// and filled by the first SSHPolicy call after that; sshRulesMu
+	// serialises that fill between readers of mu.
+	sshRulesMu sync.Mutex
+	sshRules   []compiledSSHRule
+
 	// compiledGrants are the grants with sources pre-resolved.
 	// The single source of truth for filter compilation. Both
 	// global and per-node filter rules are derived from these.
@@ -394,6 +403,7 @@ func (pm *PolicyManager) updateLocked() (bool, error) {
 		// policies for nodes that have changed. Particularly if the only difference is
 		// that nodes has been added or removed.
 		pm.sshPolicyMap.Clear()
+		pm.sshRules = nil
 		pm.filterRulesMap.Clear()
 		pm.matchersForNodeMap.Clear()
 	}
@@ -468,11 +478,13 @@ func (pm *PolicyManager) SSHPolicy(baseURL string, node types.NodeView) (*tailcf
 		return sshPol, nil
 	}
 
-	sshPol, err := pm.pol.compileSSHPolicy(baseURL, pm.users, node, pm.nodes)
+	rules, err := pm.sshRulesLocked(baseURL)
 	if err != nil {
 		return nil, fmt.Errorf("compiling SSH policy: %w", err)
 	}
 
+	sshPol := sshPolicyForNode(rules, node)
+
 	if pm.cacheableLocked(node) {
 		pm.sshPolicyMap.Store(node.ID(), sshPol)
 	}
@@ -480,6 +492,26 @@ func (pm *PolicyManager) SSHPolicy(baseURL string, node types.NodeView) (*tailcf
 	return sshPol, nil
 }
 
+// sshRulesLocked returns the compiled SSH rules, compiling them on first
+// access. It must be called with pm.mu held, for reading or writing;
+// writers reset pm.sshRules directly. The rules keep the baseURL of the
+// call that compiled them; [state.State] always passes the server URL.
+func (pm *PolicyManager) sshRulesLocked(baseURL string) ([]compiledSSHRule, error) {
+	pm.sshRulesMu.Lock()
+	defer pm.sshRulesMu.Unlock()
+
+	if pm.sshRules == nil {
+		rules, err := pm.pol.compileSSHRules(baseURL, pm.users, pm.nodes)
+		if err != nil {
+			return nil, err
+		}
+
+		pm.sshRules = rules
+	}
+
+	return pm.sshRules, nil
+}
+
 // SSHCheckParams resolves the SSH check period for src logging in to dst
 // as localUser by looking up the current policy. This avoids trusting URL
 // parameters that a client could tamper with. First-match wins across
@@ -989,6 +1021,7 @@ func (pm *PolicyManager) SetUsers(users []types.User) (bool, bool, error) {
 	// the self sources are outside the filter hash, so updateLocked can
 	// report no change while per-node results moved.
 	pm.sshPolicyMap.Clear()
+	pm.sshRules = nil
 	pm.filterRulesMap.Clear()
 	pm.matchersForNodeMap.Clear()
 
@@ -1071,6 +1104,7 @@ func (pm *PolicyManager) SetNodes(nodes views.Slice[types.NodeView]) (bool, erro
 		if !needsUpdate {
 			// This ensures fresh filter rules are generated for all nodes
 			pm.sshPolicyMap.Clear()
+			pm.sshRules = nil
 			pm.filterRulesMap.Clear()
 			pm.matchersForNodeMap.Clear()
 		}

hscontrol/policy/v2/policy_test.go

@@ -1053,6 +1053,52 @@ func TestSSHPolicyRemovalClearsRules(t *testing.T) {
 	}
 }
 
+// TestSSHPolicySharedBetweenNodes pins that what an SSH rule expands to is
+// kept once, not once per rule and destination node. There is one principal
+// per source address, so separate copies grow with rules × nodes².
+func TestSSHPolicySharedBetweenNodes(t *testing.T) {
+	users := types.Users{{Name: "user1", Email: "user1@headscale.net"}}
+	users[0].ID = 1
+
+	nodes := types.Nodes{
+		node("node1", "100.64.0.1", "fd7a:115c:a1e0::1", users[0]),
+		node("node2", "100.64.0.2", "fd7a:115c:a1e0::2", users[0]),
+		node("node3", "100.64.0.3", "fd7a:115c:a1e0::3", users[0]),
+	}
+	for i, n := range nodes {
+		n.ID = types.NodeID(i + 1) //nolint:gosec // small test value
+	}
+
+	pm, err := NewPolicyManager([]byte(`{"ssh": [
+		{"action": "accept", "src": ["autogroup:member"], "dst": ["autogroup:member"], "users": ["a"]},
+		{"action": "accept", "src": ["user1@headscale.net"], "dst": ["autogroup:member"], "users": ["b"]}
+	]}`), users, nodes[:2].ViewSlice())
+	require.NoError(t, err)
+
+	first, err := pm.SSHPolicy("", nodes[0].View())
+	require.NoError(t, err)
+	second, err := pm.SSHPolicy("", nodes[1].View())
+	require.NoError(t, err)
+
+	require.Len(t, first.Rules, 2)
+	require.Len(t, first.Rules[0].Principals, 4)
+	require.Equal(t, first, second)
+
+	require.Same(t, first.Rules[0], second.Rules[0], "between nodes")
+	require.Same(t, first.Rules[0].Principals[0], first.Rules[1].Principals[0],
+		"between rules whose sources resolve to the same addresses")
+
+	// The shared rules are dropped together with the per-node policies.
+	changed, err := pm.SetNodes(nodes.ViewSlice())
+	require.NoError(t, err)
+	require.True(t, changed)
+
+	after, err := pm.SSHPolicy("", nodes[0].View())
+	require.NoError(t, err)
+	require.Len(t, after.Rules, 2)
+	require.Len(t, after.Rules[0].Principals, 6)
+}
+
 // TestTagPropagationToPeerMap tests that when a node's tags change,
 // the peer map is correctly updated. This is a regression test for
 // https://github.com/juanfont/headscale/issues/2389