policy/v2: resolve each group once per policy compile

Group.resolve walked users and nodes for every member at every mention.
Resolve all groups once on a copy of the policy and compile from it.

hscontrol/policy/v2/policy.go

@@ -40,6 +40,12 @@ type PolicyManager struct {
 	// nodesByID indexes nodes; see [PolicyManager.cacheableLocked].
 	nodesByID map[types.NodeID]types.NodeView
 
+	// resolvedPol is pol with its tags indexed and its groups resolved
+	// against users and nodes, see [Policy.withNodesByTag] and
+	// [Policy.withResolvedGroups]. Readers that resolve aliases on demand
+	// go through it; pol itself is what is hashed.
+	resolvedPol *Policy
+
 	filterHash deephash.Sum
 	filter     []tailcfg.FilterRule
 	matchers   []matcher.Match
@@ -255,8 +261,10 @@ func NewPolicyManager(b []byte, users []types.User, nodes views.Slice[types.Node
 // pm is written, so a caller that restores its own input on error leaves pm
 // as it was, and a retry with the same input recompiles.
 func (pm *PolicyManager) updateLocked() (bool, error) {
-	// One tag index serves every resolution in this compile.
-	pol := pm.pol.withNodesByTag(pm.nodes)
+	// One tag index and one resolution of every group serve this whole
+	// compile; each step below would otherwise resolve them again for
+	// every rule naming them.
+	pol := pm.pol.withNodesByTag(pm.nodes).withResolvedGroups(pm.users, pm.nodes)
 
 	// Compile all grants once. Both global and per-node filter
 	// rules are derived from these compiled grants.
@@ -297,6 +305,7 @@ func (pm *PolicyManager) updateLocked() (bool, error) {
 		}
 	}
 
+	pm.resolvedPol = pol
 	pm.compiledGrants = grants
 	pm.userNodeIdx = buildUserNodeIndex(pm.nodes)
 	pm.needsPerNodeFilter = hasPerNodeGrants(pm.compiledGrants)
@@ -501,7 +510,7 @@ func (pm *PolicyManager) sshRulesLocked(baseURL string) ([]compiledSSHRule, erro
 	defer pm.sshRulesMu.Unlock()
 
 	if pm.sshRules == nil {
-		rules, err := pm.pol.compileSSHRules(baseURL, pm.users, pm.nodes)
+		rules, err := pm.resolvedPol.compileSSHRules(baseURL, pm.users, pm.nodes)
 		if err != nil {
 			return nil, err
 		}
@@ -561,7 +570,7 @@ func (pm *PolicyManager) SSHCheckParams(
 		}
 
 		// Resolve sources and check if src node matches.
-		srcIPs, err := rule.Sources.Resolve(pm.pol, pm.users, pm.nodes)
+		srcIPs, err := rule.Sources.Resolve(pm.resolvedPol, pm.users, pm.nodes)
 		if err != nil || srcIPs == nil {
 			continue
 		}
@@ -588,7 +597,7 @@ func (pm *PolicyManager) SSHCheckParams(
 
 			hasOtherDests = true
 
-			dstIPs, err := dst.Resolve(pm.pol, pm.users, pm.nodes)
+			dstIPs, err := dst.Resolve(pm.resolvedPol, pm.users, pm.nodes)
 			if err != nil || dstIPs == nil {
 				continue
 			}
@@ -1511,7 +1520,7 @@ func (pm *PolicyManager) ViaRoutesForPeer(viewer, peer types.NodeView) types.Via
 		resolved[i] = true
 
 		for _, src := range grants[i].Sources {
-			ips, err := src.Resolve(pm.pol, pm.users, pm.nodes)
+			ips, err := src.Resolve(pm.resolvedPol, pm.users, pm.nodes)
 			if err != nil {
 				continue
 			}
@@ -1525,7 +1534,7 @@ func (pm *PolicyManager) ViaRoutesForPeer(viewer, peer types.NodeView) types.Via
 
 		if viewerMatchesGrant[i] {
 			resolvedDstPrefixes[i], grantHasAutoGroupInternet[i] = resolveViaDestinations(
-				pm.pol, pm.users, pm.nodes, grants[i].Destinations,
+				pm.resolvedPol, pm.users, pm.nodes, grants[i].Destinations,
 			)
 		}
 

hscontrol/policy/v2/policy_test.go

@@ -4422,3 +4422,155 @@ func BenchmarkNodeCapMapsConstruction(b *testing.B) {
 		}
 	}
 }
+
+// TestWithResolvedGroups checks that the copy answers for a group exactly
+// as resolving it member by member does, without looking at users and
+// nodes again, and that the policy it was made from is left untouched.
+func TestWithResolvedGroups(t *testing.T) {
+	users := types.Users{
+		{ID: 1, Name: "alice"},
+		{ID: 2, Name: "bob"},
+		{ID: 3, Name: "twin"},
+		{ID: 4, Name: "twin"},
+	}
+	nodes := types.Nodes{
+		node("alice1", "100.64.0.1", "fd7a:115c:a1e0::1", users[0]),
+		node("bob1", "100.64.0.2", "fd7a:115c:a1e0::2", users[1]),
+	}.ViewSlice()
+
+	pol := &Policy{
+		Groups: Groups{
+			"group:dev":       Usernames{"alice@", "bob@", "missing@"},
+			"group:ambiguous": Usernames{"alice@", "twin@"},
+			"group:empty":     Usernames{},
+		},
+	}
+
+	resolved := pol.withResolvedGroups(users, nodes)
+	require.Nil(t, pol.resolvedGroups,
+		"the parsed policy is hashed and must stay as is")
+	require.Len(t, resolved.resolvedGroups, len(pol.Groups))
+
+	for g := range pol.Groups {
+		want, wantErr := g.resolve(pol, users, nodes)
+		got, gotErr := g.resolve(resolved, nil, views.Slice[types.NodeView]{})
+
+		require.Equal(t, want.Prefixes(), got.Prefixes(), g)
+
+		if wantErr != nil {
+			require.EqualError(t, gotErr, wantErr.Error(), g)
+		} else {
+			require.NoError(t, gotErr, g)
+		}
+	}
+
+	require.ErrorIs(t,
+		resolved.resolvedGroups["group:ambiguous"].err, ErrMultipleUsersFound)
+
+	unknown := Group("group:unknown")
+	got, err := unknown.resolve(resolved, users, nodes)
+	require.NoError(t, err)
+	require.Empty(t, got.Prefixes())
+
+	require.Nil(t, (*Policy)(nil).withResolvedGroups(users, nodes))
+}
+
+// TestResolvedGroupsFollowUsersAndNodes checks that readers resolving
+// aliases on demand through the groups resolved by updateLocked see the
+// users and nodes of the latest SetUsers, SetNodes and SetPolicy.
+func TestResolvedGroupsFollowUsersAndNodes(t *testing.T) {
+	alice := types.User{ID: 1, Name: "alice"}
+	bob := types.User{ID: 2, Name: "bob"}
+
+	srv := node("srv", "100.64.0.10", "fd7a:115c:a1e0::10", alice)
+	srv.ID = 10
+	srv.Tags = []string{"tag:srv"}
+	alice1 := node("alice1", "100.64.0.1", "fd7a:115c:a1e0::1", alice)
+	alice1.ID = 1
+	bob1 := node("bob1", "100.64.0.3", "fd7a:115c:a1e0::3", bob)
+	bob1.ID = 2
+
+	pol := func(members string) []byte {
+		return []byte(`{
+			"groups": {"group:dev": [` + members + `]},
+			"tagOwners": {"tag:srv": ["group:dev"]},
+			"acls": [{"action": "accept", "src": ["group:dev"], "dst": ["tag:srv:22"]}],
+			"ssh": [{"action": "accept", "src": ["group:dev"], "dst": ["tag:srv"], "users": ["root"]}]
+		}`)
+	}
+
+	pm, err := NewPolicyManager(
+		pol(`"alice@", "bob@"`),
+		types.Users{alice},
+		types.Nodes{srv, alice1}.ViewSlice(),
+	)
+	require.NoError(t, err)
+
+	sshSources := func() []string {
+		t.Helper()
+
+		sshPol, err := pm.SSHPolicy("", srv.View())
+		require.NoError(t, err)
+
+		var ips []string
+
+		for _, rule := range sshPol.Rules {
+			for _, principal := range rule.Principals {
+				ips = append(ips, principal.NodeIP)
+			}
+		}
+
+		slices.Sort(ips)
+
+		return ips
+	}
+
+	filterSources := func() []string {
+		t.Helper()
+
+		filter, _ := pm.Filter()
+		require.Len(t, filter, 1)
+
+		ips := slices.Clone(filter[0].SrcIPs)
+		slices.Sort(ips)
+
+		return ips
+	}
+
+	aliceIPs := []string{"100.64.0.1", "fd7a:115c:a1e0::1"}
+	bothIPs := []string{
+		"100.64.0.1", "100.64.0.3", "fd7a:115c:a1e0::1", "fd7a:115c:a1e0::3",
+	}
+
+	require.Equal(t, aliceIPs, filterSources())
+	require.Equal(t, aliceIPs, sshSources())
+
+	_, _, err = pm.SetUsers(types.Users{alice, bob})
+	require.NoError(t, err)
+
+	_, err = pm.SetNodes(types.Nodes{srv, alice1, bob1}.ViewSlice())
+	require.NoError(t, err)
+	require.Equal(t, bothIPs, filterSources())
+	require.Equal(t, bothIPs, sshSources())
+
+	_, err = pm.SetPolicy(pol(`"alice@"`))
+	require.NoError(t, err)
+	require.Equal(t, aliceIPs, filterSources())
+	require.Equal(t, aliceIPs, sshSources())
+}
+
+// TestResolvedPolicyKeepsTagIndex checks that the copy readers resolve
+// through carries both the tag index and the resolved groups.
+func TestResolvedPolicyKeepsTagIndex(t *testing.T) {
+	users := types.Users{{ID: 1, Name: "alice"}}
+	nodes := types.Nodes{node("a", "100.64.0.1", "fd7a:115c:a1e0::1", users[0])}
+
+	pm, err := NewPolicyManager([]byte(`{
+		"groups": {"group:a": ["alice@"]},
+		"tagOwners": {"tag:x": ["alice@"]},
+		"acls": [{"action": "accept", "src": ["group:a"], "dst": ["tag:x:*"]}]
+	}`), users, nodes.ViewSlice())
+	require.NoError(t, err)
+	require.NotNil(t, pm.resolvedPol.nodesByTag)
+	require.NotNil(t, pm.resolvedPol.resolvedGroups)
+}

hscontrol/policy/v2/sshtest.go

@@ -126,7 +126,7 @@ func (pm *PolicyManager) RunSSHTests() error {
 	defer pm.mu.RUnlock()
 
 	cache := make(map[types.NodeID]*tailcfg.SSHPolicy)
-	results := runSSHPolicyTests(pm.pol, pm.users, pm.nodes, cache)
+	results := runSSHPolicyTests(pm.resolvedPol, pm.users, pm.nodes, cache)
 
 	return wrapTestResult(errSSHPolicyTestsFailed, results.AllPassed, results.Errors)
 }
@@ -141,6 +141,7 @@ func evaluateSSHTests(
 		return nil
 	}
 
+	pol = pol.withResolvedGroups(users, nodes)
 	cache := make(map[types.NodeID]*tailcfg.SSHPolicy)
 	results := runSSHPolicyTests(pol, users, nodes, cache)
 

hscontrol/policy/v2/test.go

@@ -218,7 +218,7 @@ func (pm *PolicyManager) RunTests() error {
 	pm.mu.RLock()
 	defer pm.mu.RUnlock()
 
-	results := runPolicyTests(pm.pol, pm.filter, pm.users, pm.nodes)
+	results := runPolicyTests(pm.resolvedPol, pm.filter, pm.users, pm.nodes)
 
 	return wrapTestResult(errPolicyTestsFailed, results.AllPassed, results.Errors)
 }
@@ -231,6 +231,7 @@ func evaluateTests(pol *Policy, users []types.User, nodes views.Slice[types.Node
 		return nil
 	}
 
+	pol = pol.withResolvedGroups(users, nodes)
 	grants := pol.compileGrants(users, nodes)
 
 	var filter []tailcfg.FilterRule

hscontrol/policy/v2/types.go

@@ -519,6 +519,13 @@ func (g *Group) resolve(p *Policy, users types.Users, nodes views.Slice[types.No
 		errs []error
 	)
 
+	// A policy compiled through [Policy.withResolvedGroups] has already
+	// resolved every group once; answer from that instead of walking
+	// users and nodes again for each rule that names the group.
+	if r, ok := p.resolvedGroups[*g]; ok {
+		return r.ips, r.err
+	}
+
 	for _, user := range p.Groups[*g] {
 		uips, err := user.resolve(nil, users, nodes)
 		if err != nil {
@@ -1958,6 +1965,11 @@ type Policy struct {
 	// callers using it should panic if not
 	validated bool `json:"-"`
 
+	// resolvedGroups is only set on the copy returned by
+	// [Policy.withResolvedGroups]; a parsed policy never carries it,
+	// as the policy is hashed for change detection.
+	resolvedGroups map[Group]resolvedGroup `json:"-"`
+
 	Groups              Groups             `json:"groups,omitempty"`
 	Hosts               Hosts              `json:"hosts,omitempty"`
 	TagOwners           TagOwners          `json:"tagOwners,omitempty"`
@@ -2001,6 +2013,41 @@ func (pol *Policy) withNodesByTag(nodes views.Slice[types.NodeView]) *Policy {
 // MarshalJSON is deliberately not implemented for [Policy].
 // We use the default JSON marshalling behavior provided by the Go runtime.
 
+// resolvedGroup is what [Group.resolve] returns for one group.
+type resolvedGroup struct {
+	ips *netipx.IPSet
+	err error
+}
+
+// withResolvedGroups returns a shallow copy of pol on which every group
+// is already resolved against users and nodes, so that [Group.resolve]
+// is a map lookup. Resolving a group walks users and nodes once per
+// member, and a group is typically named by many rules; resolving it per
+// mention makes compilation quadratic in the policy size. The copy is
+// only valid for the users and nodes it was built from, and must not be
+// hashed in place of pol.
+func (pol *Policy) withResolvedGroups(
+	users types.Users,
+	nodes views.Slice[types.NodeView],
+) *Policy {
+	if pol == nil || len(pol.Groups) == 0 {
+		return pol
+	}
+
+	cp := *pol
+	cp.resolvedGroups = nil
+
+	resolved := make(map[Group]resolvedGroup, len(pol.Groups))
+	for g := range pol.Groups {
+		ips, err := g.resolve(&cp, users, nodes)
+		resolved[g] = resolvedGroup{ips: ips, err: err}
+	}
+
+	cp.resolvedGroups = resolved
+
+	return &cp
+}
+
 var (
 	// TODO(kradalby): Add these checks for tagOwners and autoApprovers.
 	autogroupForSrc       = []AutoGroup{AutoGroupMember, AutoGroupTagged, AutoGroupDangerAll}