policy/v2: evaluate policy tests from compiled snapshots

Parse filters once and index SSH rules by source. Evaluate tests against
user and node snapshots before taking the write lock.

hscontrol/policy/v2/policy.go

@@ -673,10 +673,11 @@ func (pm *PolicyManager) SetPolicy(polB []byte) (bool, error) {
 		return false, fmt.Errorf("parsing policy: %w", err)
 	}
 
-	pm.mu.Lock()
-	defer pm.mu.Unlock()
+	pm.mu.RLock()
+	users, nodes := pm.users, pm.nodes
+	pm.mu.RUnlock()
 
-	err = validateUserReferences(pol, pm.users)
+	err = validateUserReferences(pol, users)
 	if err != nil {
 		return false, fmt.Errorf("validating policy user references: %w", err)
 	}
@@ -686,17 +687,25 @@ func (pm *PolicyManager) SetPolicy(polB []byte) (bool, error) {
 	// they fail, return without mutating the live PolicyManager so the
 	// failed write does not knock the running config offline.
 	//
+	// They run before the write lock is taken: how long they take is up
+	// to the policy author, and every map response and node write needs
+	// the lock. A user or node that changes meanwhile counts as changed
+	// right after the write.
+	//
 	// Aggregate ACL and SSH test failures via multierr so operators
 	// see both classes in a single response instead of having to
 	// fix-and-retry to discover the second one.
 	testErr := multierr.New(
-		evaluateTests(pol, pm.users, pm.nodes),
-		evaluateSSHTests(pol, pm.users, pm.nodes),
+		evaluateTests(pol, users, nodes),
+		evaluateSSHTests(pol, users, nodes),
 	)
 	if testErr != nil {
 		return false, testErr
 	}
 
+	pm.mu.Lock()
+	defer pm.mu.Unlock()
+
 	// Log policy metadata for debugging
 	log.Debug().
 		Int("policy.bytes", len(polB)).

hscontrol/policy/v2/policytester_compat_test.go

@@ -32,7 +32,6 @@ import (
 	"github.com/juanfont/headscale/hscontrol/types/testcapture"
 	"github.com/stretchr/testify/require"
 	tsapi "tailscale.com/client/tailscale/v2"
-	"tailscale.com/tailcfg"
 )
 
 // knownPolicyTesterDivergences lists checks where headscale's evaluator
@@ -284,9 +283,7 @@ func testerBreakdown(pm *PolicyManager) []string {
 		}
 	}
 
-	cache := make(map[types.NodeID]*tailcfg.SSHPolicy)
-
-	for _, res := range runSSHPolicyTests(pm.pol, pm.users, pm.nodes, cache).Results {
+	for _, res := range runSSHPolicyTests(pm.pol, pm.users, pm.nodes).Results {
 		for range res.Errors {
 			out = append(out, res.Src+": error")
 		}

hscontrol/policy/v2/sshtest.go

@@ -155,8 +155,7 @@ func (pm *PolicyManager) RunSSHTests() error {
 	pm.mu.RLock()
 	defer pm.mu.RUnlock()
 
-	cache := make(map[types.NodeID]*tailcfg.SSHPolicy)
-	results := runSSHPolicyTests(pm.resolvedPol, pm.users, pm.nodes, cache)
+	results := runSSHPolicyTests(pm.resolvedPol, pm.users, pm.nodes)
 
 	return wrapTestResult(errSSHPolicyTestsFailed, results.AllPassed, results)
 }
@@ -172,20 +171,27 @@ func evaluateSSHTests(
 	}
 
 	pol = pol.withResolvedGroups(users, nodes)
-	cache := make(map[types.NodeID]*tailcfg.SSHPolicy)
-	results := runSSHPolicyTests(pol, users, nodes, cache)
+	results := runSSHPolicyTests(pol, users, nodes)
 
 	return wrapTestResult(errSSHPolicyTestsFailed, results.AllPassed, results)
 }
 
-// runSSHPolicyTests evaluates every sshTests entry. The cache is keyed
-// by dst [types.NodeID] so repeat destinations only compile once per pass.
+// runSSHPolicyTests evaluates every sshTests entry. The SSH rules are
+// compiled once per pass and each dst's rules are selected once.
 func runSSHPolicyTests(
 	pol *Policy,
 	users []types.User,
 	nodes views.Slice[types.NodeView],
-	cache map[types.NodeID]*tailcfg.SSHPolicy,
 ) SSHPolicyTestResults {
+	// baseURL is empty because reachability only checks for the presence
+	// of [tailcfg.SSHAction.HoldAndDelegate], not its value.
+	rules, err := pol.compileSSHRules("", users, nodes)
+	cache := &sshTestCache{
+		rules:  rules,
+		err:    err,
+		byNode: make(map[types.NodeID]sshRulesBySrc),
+	}
+
 	results := SSHPolicyTestResults{
 		AllPassed: true,
 		Results:   make([]SSHPolicyTestResult, 0, len(pol.SSHTests)),
@@ -210,7 +216,7 @@ func runSSHPolicyTest(
 	pol *Policy,
 	users []types.User,
 	nodes views.Slice[types.NodeView],
-	cache map[types.NodeID]*tailcfg.SSHPolicy,
+	cache *sshTestCache,
 ) SSHPolicyTestResult {
 	srcLabel := ""
 	if test.Src != nil {
@@ -265,11 +271,14 @@ func runSSHPolicyTest(
 		{test.Deny, assertDeny},
 		{test.Check, assertCheck},
 	} {
+		// Each distinct name and alias is evaluated once and reported at
+		// every occurrence, as Tailscale SaaS repeats them.
+		seen := make(map[sshTestKey]sshTestOutcome)
+
 		for _, user := range g.users {
 			evaluateAssertion(
-				pol, users, nodes, cache,
-				srcAddrs, dsts, user.String(),
-				g.kind, &res,
+				cache, srcAddrs, dsts, user.String(),
+				g.kind, seen, &res,
 			)
 		}
 	}
@@ -289,28 +298,24 @@ const (
 // entry, under the alias as written: the report goes back to the policy
 // author, who may not be allowed to list nodes. An alias passes when every
 // node it selects does, for every src address; one selecting no node is not
-// reached. Each distinct alias is evaluated once and reported at every
-// occurrence, as Tailscale SaaS reports them.
+// reached. seen holds the outcomes of one assertion kind already
+// evaluated.
 func evaluateAssertion(
-	pol *Policy,
-	users []types.User,
-	nodes views.Slice[types.NodeView],
-	cache map[types.NodeID]*tailcfg.SSHPolicy,
+	cache *sshTestCache,
 	srcAddrs []netip.Addr,
 	dsts []sshTestDst,
 	user string,
 	kind sshAssertion,
+	seen map[sshTestKey]sshTestOutcome,
 	res *SSHPolicyTestResult,
 ) {
-	type outcome struct{ passed, notReached, failed bool }
-
-	seen := make(map[string]outcome, len(dsts))
-
 	for _, dst := range dsts {
-		o, ok := seen[dst.label]
+		key := sshTestKey{user: user, dst: dst.label}
+
+		o, ok := seen[key]
 		if !ok {
 			passed, notReached, err := sshTestAliasOutcome(
-				pol, users, nodes, cache, srcAddrs, dst.nodes, user, kind,
+				cache, srcAddrs, dst.nodes, user, kind,
 			)
 			if err != nil {
 				res.Passed = false
@@ -318,8 +323,8 @@ func evaluateAssertion(
 					fmt.Sprintf("compiling SSH policy for %s: %v", dst.label, err))
 			}
 
-			o = outcome{passed, notReached, err != nil}
-			seen[dst.label] = o
+			o = sshTestOutcome{passed, notReached, err != nil}
+			seen[key] = o
 		}
 
 		if o.failed {
@@ -351,15 +356,19 @@ func evaluateAssertion(
 	}
 }
 
+// sshTestKey is one login name checked against one dst alias.
+type sshTestKey struct{ user, dst string }
+
+// sshTestOutcome is what evaluating an [sshTestKey] found; failed means it
+// could not be evaluated.
+type sshTestOutcome struct{ passed, notReached, failed bool }
+
 // sshTestAliasOutcome reports whether user's assertion holds on every node
 // of one dst alias from every src address, and whether a failing node was
 // not reached at all. All src addresses must agree: the first
 // counter-example fails the node.
 func sshTestAliasOutcome(
-	pol *Policy,
-	users []types.User,
-	nodes views.Slice[types.NodeView],
-	cache map[types.NodeID]*tailcfg.SSHPolicy,
+	cache *sshTestCache,
 	srcAddrs []netip.Addr,
 	dstNodes []types.NodeView,
 	user string,
@@ -372,13 +381,13 @@ func sshTestAliasOutcome(
 	passed, notReached := true, false
 
 	for _, dst := range dstNodes {
-		dstPol, err := compiledSSHPolicy(pol, users, nodes, cache, dst)
+		dstRules, err := compiledSSHPolicy(cache, dst)
 		if err != nil {
 			return false, false, err
 		}
 
 		for _, srcAddr := range srcAddrs {
-			a, c := reachability(dstPol, srcAddr, user)
+			a, c := reachability(dstRules[srcAddr], user)
 
 			var holds bool
 
@@ -609,50 +618,64 @@ func prefixesToIPSet(prefixes []netip.Prefix) (*netipx.IPSet, error) {
 	return b.IPSet()
 }
 
-// compiledSSHPolicy returns the per-node compiled [tailcfg.SSHPolicy], caching
-// on miss. baseURL is empty because reachability only checks for the
-// presence of [tailcfg.SSHAction.HoldAndDelegate], not its value.
-func compiledSSHPolicy(
-	pol *Policy,
-	users []types.User,
-	nodes views.Slice[types.NodeView],
-	cache map[types.NodeID]*tailcfg.SSHPolicy,
-	node types.NodeView,
-) (*tailcfg.SSHPolicy, error) {
-	if sshPol, ok := cache[node.ID()]; ok {
-		return sshPol, nil
-	}
-
-	sshPol, err := pol.compileSSHPolicy("", users, node, nodes)
-	if err != nil {
-		return nil, err
-	}
-
-	cache[node.ID()] = sshPol
-
-	return sshPol, nil
+// sshTestCache holds one sshTests pass's compiled SSH rules and, per dst,
+// the rules selected from them.
+type sshTestCache struct {
+	rules  []compiledSSHRule
+	err    error
+	byNode map[types.NodeID]sshRulesBySrc
 }
 
-// reachability reports whether srcAddr can log in as user via:
-//
-//   - any matching rule (acceptHit, satisfies accept assertions)
-//   - a check-action rule (checkHit, satisfies check assertions)
-func reachability(
-	dstPolicy *tailcfg.SSHPolicy,
-	srcAddr netip.Addr,
-	user string,
-) (bool, bool) {
-	if dstPolicy == nil {
-		return false, false
+// sshRulesBySrc is one destination's SSH rules keyed by principal address,
+// so an assertion walks only the rules its source can match and every
+// principal is parsed once per destination.
+type sshRulesBySrc map[netip.Addr][]*tailcfg.SSHRule
+
+// compiledSSHPolicy returns node's SSH rules, selecting them from the
+// compiled rules on first use.
+func compiledSSHPolicy(
+	cache *sshTestCache,
+	node types.NodeView,
+) (sshRulesBySrc, error) {
+	if cache.err != nil {
+		return nil, cache.err
 	}
 
+	if rules, ok := cache.byNode[node.ID()]; ok {
+		return rules, nil
+	}
+
+	rules := make(sshRulesBySrc)
+
+	for _, rule := range sshPolicyForNode(cache.rules, node).Rules {
+		for _, p := range rule.Principals {
+			if p == nil {
+				continue
+			}
+
+			addr, err := netip.ParseAddr(p.NodeIP)
+			if err != nil {
+				continue
+			}
+
+			rules[addr] = append(rules[addr], rule)
+		}
+	}
+
+	cache.byNode[node.ID()] = rules
+
+	return rules, nil
+}
+
+// reachability reports whether a source with these rules can log in as
+// user via:
+//
+//   - any of its rules (acceptHit, satisfies accept assertions)
+//   - a check-action rule (checkHit, satisfies check assertions)
+func reachability(rules []*tailcfg.SSHRule, user string) (bool, bool) {
 	var acceptHit, checkHit bool
 
-	for _, rule := range dstPolicy.Rules {
-		if !principalContainsAddr(rule.Principals, srcAddr) {
-			continue
-		}
-
+	for _, rule := range rules {
 		if !sshUserMapAllows(rule.SSHUsers, user) {
 			continue
 		}
@@ -677,34 +700,6 @@ func reachability(
 	return acceptHit, checkHit
 }
 
-// principalContainsAddr reports whether any principal's [tailcfg.SSHPrincipal.NodeIP]
-// matches srcAddr exactly (the SSH compiler emits one principal per source IP).
-func principalContainsAddr(
-	principals []*tailcfg.SSHPrincipal,
-	srcAddr netip.Addr,
-) bool {
-	for _, p := range principals {
-		if p == nil {
-			continue
-		}
-
-		if p.NodeIP == "" {
-			continue
-		}
-
-		addr, err := netip.ParseAddr(p.NodeIP)
-		if err != nil {
-			continue
-		}
-
-		if addr == srcAddr {
-			return true
-		}
-	}
-
-	return false
-}
-
 // sshUserMapAllows reports whether [SSHUsers] permits user. The [SSHUsers]
 // wire shape (see filter.go compileSSHPolicy):
 //

hscontrol/policy/v2/sshtest_test.go

@@ -2,9 +2,11 @@ package v2
 
 import (
 	"encoding/json"
+	"fmt"
 	"net/netip"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/stretchr/testify/assert"
@@ -1149,3 +1151,55 @@ func TestSSHTestsReportNamesAliasNotNode(t *testing.T) {
 		}
 	}
 }
+
+// TestRunSSHTestsManyNamesAndPrincipals: login names, destination nodes
+// and rule principals are all the policy author's to multiply, so a check
+// must look up its source's rules, not walk every principal.
+func TestRunSSHTestsManyNamesAndPrincipals(t *testing.T) {
+	users := sshTestUsers()
+
+	nodes := make(types.Nodes, 0, 1000)
+
+	for i := range cap(nodes) {
+		n := &types.Node{
+			ID:       types.NodeID(i + 1), //nolint:gosec
+			Hostname: fmt.Sprintf("node-%d", i),
+			IPv4:     ap(fmt.Sprintf("100.64.%d.%d", i/250, i%250+1)),
+			User:     &users[0],
+			UserID:   &users[0].ID,
+		}
+		if i < 100 {
+			n.Tags = []string{"tag:server"}
+		}
+
+		nodes = append(nodes, n)
+	}
+
+	names := make([]string, 5000)
+	for i := range names {
+		names[i] = fmt.Sprintf(`"user%d"`, i)
+	}
+
+	pol, err := unmarshalPolicy(fmt.Appendf(nil, `{
+		"tagOwners": {"tag:server": ["alice@headscale.net"]},
+		"ssh": [{
+			"action": "accept",
+			"src":    ["autogroup:member"],
+			"dst":    ["tag:server"],
+			"users":  ["autogroup:nonroot"]
+		}],
+		"sshTests": [{"src": "100.64.3.250", "dst": ["tag:server"], "accept": [%s]}]
+	}`, strings.Join(names, ",")))
+	require.NoError(t, err)
+
+	done := make(chan error, 1)
+
+	go func() { done <- evaluateSSHTests(pol, users, nodes.ViewSlice()) }()
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(5 * time.Second):
+		t.Fatal("sshTests did not finish in time")
+	}
+}

hscontrol/policy/v2/test.go

@@ -359,8 +359,10 @@ func runPolicyTests(pol *Policy, filter []tailcfg.FilterRule, users []types.User
 		Results:   make([]PolicyTestResult, 0, len(pol.Tests)),
 	}
 
+	rules := parseTestRules(filter)
+
 	for _, test := range pol.Tests {
-		res := runPolicyTest(test, pol, filter, users, nodes)
+		res := runPolicyTest(test, pol, rules, users, nodes)
 		if !res.Passed {
 			results.AllPassed = false
 		}
@@ -372,7 +374,13 @@ func runPolicyTests(pol *Policy, filter []tailcfg.FilterRule, users []types.User
 }
 
 // runPolicyTest evaluates one [PolicyTest].
-func runPolicyTest(test PolicyTest, pol *Policy, filter []tailcfg.FilterRule, users []types.User, nodes views.Slice[types.NodeView]) PolicyTestResult {
+func runPolicyTest(
+	test PolicyTest,
+	pol *Policy,
+	rules []testRule,
+	users []types.User,
+	nodes views.Slice[types.NodeView],
+) PolicyTestResult {
 	res := PolicyTestResult{
 		Src:    test.Src,
 		Proto:  test.Proto,
@@ -396,7 +404,9 @@ func runPolicyTest(test PolicyTest, pol *Policy, filter []tailcfg.FilterRule, us
 
 	check := func(dsts []string, wantAllowed bool, ok, fail *[]string) {
 		for _, dst := range dsts {
-			allowed, err := evalReachability(srcPrefixes, dst, test.Proto, pol, filter, users, nodes, wantAllowed)
+			allowed, err := evalReachability(
+				srcPrefixes, dst, test.Proto, pol, rules, users, nodes, wantAllowed,
+			)
 			if err != nil {
 				res.Passed = false
 				res.Errors = append(res.Errors, fmt.Sprintf("error testing %q: %v", dst, err))
@@ -466,7 +476,7 @@ func resolveTestSource(src string, pol *Policy, users []types.User, nodes views.
 }
 
 // evalReachability reports whether traffic from every srcPrefix to dst (in
-// `host:port` form) is allowed by filter for the requested protocol. With
+// `host:port` form) is allowed by rules for the requested protocol. With
 // everyNode, as an accept asks, each node behind dst must be reached, as on
 // Tailscale SaaS; otherwise any address of dst will do, so a deny passes only
 // when none is reached.
@@ -479,7 +489,7 @@ func evalReachability(
 	dst string,
 	proto Protocol,
 	pol *Policy,
-	filter []tailcfg.FilterRule,
+	rules []testRule,
 	users []types.User,
 	nodes views.Slice[types.NodeView],
 	everyNode bool,
@@ -507,7 +517,7 @@ func evalReachability(
 	// the test to consider it allowed. A partial allow is a fail.
 	for _, src := range srcPrefixes {
 		for _, target := range targets {
-			if !srcReachesDst(src, target, awp.Ports, proto, filter) {
+			if !srcReachesDst(src, target, awp.Ports, proto, rules) {
 				return false, nil
 			}
 		}
@@ -567,7 +577,107 @@ func parseDestinationAlias(dst string) (*AliasWithPorts, error) {
 	return &awp, nil
 }
 
-// srcReachesDst walks the compiled filter rules and reports whether
+// testRule is a [tailcfg.FilterRule] with its addresses parsed. A run
+// compares every assertion with every rule, and the policy author chooses
+// how many there are of both, so the filter is parsed once per run rather
+// than once per comparison.
+type testRule struct {
+	srcs   testAddrs
+	protos []int
+	// dsts merges the rule's DstPorts entries by port range, so a
+	// destination is looked up once per range instead of once per entry.
+	dsts []testRuleDst
+}
+
+// testRuleDst holds every address a rule allows a port range on.
+type testRuleDst struct {
+	ips   testAddrs
+	ports tailcfg.PortRange
+}
+
+// testAddrs is the sorted, non-overlapping [netipx.IPSet.Ranges] of a set.
+type testAddrs []netipx.IPRange
+
+// overlapsPrefix is [netipx.IPSet.OverlapsPrefix] without walking the set.
+func (a testAddrs) overlapsPrefix(p netip.Prefix) bool {
+	r := netipx.RangeOfPrefix(p)
+
+	// The first range that does not end before r is the only candidate:
+	// the ones after it start later still.
+	i, _ := slices.BinarySearchFunc(
+		a, r.From(), func(x netipx.IPRange, from netip.Addr) int {
+			return x.To().Compare(from)
+		},
+	)
+
+	return i < len(a) && a[i].From().Compare(r.To()) <= 0
+}
+
+// parseTestRules parses filter for [srcReachesDst].
+// [tailcfg.FilterRule.SrcIPs] and [tailcfg.NetPortRange.IP] may be CIDR, single
+// addresses, IP ranges (`a-b`), or `*`; we use [util.ParseIPSet] to cover all
+// of those uniformly. Unparseable entries are skipped (the rule compiler emits
+// well-formed strings, so this is defence-in-depth, not error handling).
+func parseTestRules(filter []tailcfg.FilterRule) []testRule {
+	rules := make([]testRule, 0, len(filter))
+
+	for _, rule := range filter {
+		var srcs netipx.IPSetBuilder
+
+		for _, raw := range rule.SrcIPs {
+			set, err := util.ParseIPSet(raw, nil)
+			if err != nil {
+				continue
+			}
+
+			srcs.AddSet(set)
+		}
+
+		dsts := make(map[tailcfg.PortRange]*netipx.IPSetBuilder)
+
+		for _, dp := range rule.DstPorts {
+			set, err := util.ParseIPSet(dp.IP, nil)
+			if err != nil {
+				continue
+			}
+
+			if dsts[dp.Ports] == nil {
+				dsts[dp.Ports] = new(netipx.IPSetBuilder)
+			}
+
+			dsts[dp.Ports].AddSet(set)
+		}
+
+		parsed := testRule{
+			srcs:   buildTestAddrs(&srcs),
+			protos: rule.IPProto,
+			dsts:   make([]testRuleDst, 0, len(dsts)),
+		}
+
+		for ports, b := range dsts {
+			parsed.dsts = append(
+				parsed.dsts, testRuleDst{ips: buildTestAddrs(b), ports: ports},
+			)
+		}
+
+		rules = append(rules, parsed)
+	}
+
+	return rules
+}
+
+// buildTestAddrs returns the addresses added to b. The builder only fails on
+// invalid input, which [util.ParseIPSet] has already refused.
+func buildTestAddrs(b *netipx.IPSetBuilder) testAddrs {
+	set, err := b.IPSet()
+	if err != nil {
+		return nil
+	}
+
+	return set.Ranges()
+}
+
+// srcReachesDst walks the parsed filter rules and reports whether
 // traffic from src to any prefix in dstPrefixes on at least one of ports
 // (or any port when ports is empty) is allowed under proto.
 //
@@ -576,14 +686,20 @@ func parseDestinationAlias(dst string) (*AliasWithPorts, error) {
 // is omitted. The captured Tailscale matches show these four IANA
 // numbers explicitly when no proto is set, so a rule restricted to any
 // of them satisfies an empty-proto test.
-func srcReachesDst(src netip.Prefix, dstPrefixes []netip.Prefix, ports []tailcfg.PortRange, proto Protocol, filter []tailcfg.FilterRule) bool {
+func srcReachesDst(
+	src netip.Prefix,
+	dstPrefixes []netip.Prefix,
+	ports []tailcfg.PortRange,
+	proto Protocol,
+	rules []testRule,
+) bool {
 	requestedProtos := proto.toIANAProtocolNumbers()
 	if len(requestedProtos) == 0 {
 		requestedProtos = []int{ProtocolTCP, ProtocolUDP, ProtocolICMP, ProtocolIPv6ICMP}
 	}
 
-	for _, rule := range filter {
-		if !ruleMatchesSource(rule, src) {
+	for _, rule := range rules {
+		if !rule.srcs.overlapsPrefix(src) {
 			continue
 		}
 
@@ -599,37 +715,17 @@ func srcReachesDst(src netip.Prefix, dstPrefixes []netip.Prefix, ports []tailcfg
 	return false
 }
 
-// ruleMatchesSource reports whether the rule's source list contains src.
-// [tailcfg.FilterRule.SrcIPs] may be CIDR, single addresses, IP ranges (`a-b`),
-// or `*`; we use [util.ParseIPSet] to cover all of those uniformly. Unparseable
-// entries are skipped (the rule compiler emits well-formed strings, so this is
-// defence-in-depth, not error handling).
-func ruleMatchesSource(rule tailcfg.FilterRule, src netip.Prefix) bool {
-	for _, raw := range rule.SrcIPs {
-		set, err := util.ParseIPSet(raw, nil)
-		if err != nil {
-			continue
-		}
-
-		if set.OverlapsPrefix(src) {
-			return true
-		}
-	}
-
-	return false
-}
-
 // ruleMatchesProto reports whether the rule permits any of requestedProtos.
 // An unset [tailcfg.FilterRule.IPProto] means "any protocol" and matches
 // everything. requestedProtos is the per-test protocol set: a single proto
 // for an explicit [PolicyTest.Proto], or the default set when
 // [PolicyTest.Proto] is empty.
-func ruleMatchesProto(rule tailcfg.FilterRule, requestedProtos []int) bool {
-	if len(rule.IPProto) == 0 {
+func ruleMatchesProto(rule testRule, requestedProtos []int) bool {
+	if len(rule.protos) == 0 {
 		return true
 	}
 
-	for _, ruleProto := range rule.IPProto {
+	for _, ruleProto := range rule.protos {
 		if slices.Contains(requestedProtos, ruleProto) {
 			return true
 		}
@@ -639,15 +735,19 @@ func ruleMatchesProto(rule tailcfg.FilterRule, requestedProtos []int) bool {
 }
 
 // ruleAllowsAnyDest reports whether at least one destination prefix in
-// dstPrefixes is allowed by at least one of the rule's DstPorts entries
+// dstPrefixes is allowed by at least one of the rule's destinations
 // for at least one of ports (or any port when ports is empty).
-func ruleAllowsAnyDest(rule tailcfg.FilterRule, dstPrefixes []netip.Prefix, ports []tailcfg.PortRange) bool {
-	for _, dp := range rule.DstPorts {
-		if !destEntryMatchesPrefixes(dp, dstPrefixes) {
+func ruleAllowsAnyDest(
+	rule testRule,
+	dstPrefixes []netip.Prefix,
+	ports []tailcfg.PortRange,
+) bool {
+	for _, dst := range rule.dsts {
+		if !portsAllowed(ports, dst.ports) {
 			continue
 		}
 
-		if portsAllowed(ports, dp.Ports) {
+		if slices.ContainsFunc(dstPrefixes, dst.ips.overlapsPrefix) {
 			return true
 		}
 	}
@@ -655,17 +755,6 @@ func ruleAllowsAnyDest(rule tailcfg.FilterRule, dstPrefixes []netip.Prefix, port
 	return false
 }
 
-// destEntryMatchesPrefixes reports whether the rule's [tailcfg.NetPortRange.IP]
-// (CIDR, single IP, IP range, or "*") covers any prefix in dstPrefixes.
-func destEntryMatchesPrefixes(dp tailcfg.NetPortRange, dstPrefixes []netip.Prefix) bool {
-	set, err := util.ParseIPSet(dp.IP, nil)
-	if err != nil {
-		return false
-	}
-
-	return slices.ContainsFunc(dstPrefixes, set.OverlapsPrefix)
-}
-
 // portsAllowed reports whether at least one requested port is contained
 // in allowed. Empty requested means "any port".
 func portsAllowed(requested []tailcfg.PortRange, allowed tailcfg.PortRange) bool {

hscontrol/policy/v2/test_test.go

@@ -2,12 +2,16 @@ package v2
 
 import (
 	"fmt"
+	"net/netip"
+	"slices"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go4.org/netipx"
 )
 
 // policyTestUsers/policyTestNodes are reused across the test cases below to
@@ -288,6 +292,59 @@ func TestRunTests(t *testing.T) {
 			wantErrSub:  []string{"bob@headscale.net", "expected ALLOWED"},
 			wantNoErrIs: errPolicyTestsFailed,
 		},
+		{
+			// group:all resolves to one prefix per address; no single rule
+			// covers them all, but every one is covered by some rule.
+			name: "src-prefixes-allowed-by-different-rules",
+			policy: `{
+				"groups":    { "group:all": ["alice@headscale.net", "bob@headscale.net"] },
+				"tagOwners": { "tag:server": ["alice@headscale.net"] },
+				"acls": [
+					{ "action": "accept", "src": ["alice@headscale.net"], "dst": ["tag:server:22"] },
+					{ "action": "accept", "src": ["bob@headscale.net"],   "dst": ["tag:server:22"] }
+				],
+				"tests": [{
+					"src": "group:all",
+					"accept": ["tag:server:22"]
+				}]
+			}`,
+			wantPass: true,
+		},
+		{
+			name: "src-prefixes-allowed-on-different-ports",
+			policy: `{
+				"groups":    { "group:all": ["alice@headscale.net", "bob@headscale.net"] },
+				"tagOwners": { "tag:server": ["alice@headscale.net"] },
+				"acls": [
+					{ "action": "accept", "src": ["alice@headscale.net"], "dst": ["tag:server:22"] },
+					{ "action": "accept", "src": ["bob@headscale.net"],   "dst": ["tag:server:80"] }
+				],
+				"tests": [{
+					"src": "group:all",
+					"accept": ["tag:server:22"]
+				}]
+			}`,
+			wantErrSub:  []string{"group:all", "tag:server:22", "expected ALLOWED"},
+			wantNoErrIs: errPolicyTestsFailed,
+		},
+		{
+			// One rule, several addresses per port and several ports per
+			// address: an address is only allowed on its own ports.
+			name: "dst-address-and-port-stay-paired",
+			policy: `{
+				"acls": [{
+					"action": "accept",
+					"src": ["*"],
+					"dst": ["10.0.0.1:22", "10.0.0.3:22", "10.0.0.5:22", "10.0.0.2:80", "10.0.0.3:1000-2000"]
+				}],
+				"tests": [{
+					"src": "alice@headscale.net",
+					"accept": ["10.0.0.1:22", "10.0.0.3:22", "10.0.0.5:22", "10.0.0.2:80", "10.0.0.3:1500"],
+					"deny": ["10.0.0.2:22", "10.0.0.4:22", "10.0.0.6:22", "10.0.0.1:80", "10.0.0.1:1500", "10.0.0.3:80"]
+				}]
+			}`,
+			wantPass: true,
+		},
 		{
 			name: "no-tests-block-is-no-op",
 			policy: `{
@@ -323,6 +380,108 @@ func TestRunTests(t *testing.T) {
 	}
 }
 
+// TestTestAddrsOverlapsPrefix checks the lookup against the one it replaces.
+func TestTestAddrsOverlapsPrefix(t *testing.T) {
+	var b netipx.IPSetBuilder
+
+	for _, p := range []string{
+		"10.0.0.1/32", "10.0.0.3/32", "10.0.2.0/24", "100.64.0.0/10",
+		"fd7a:115c:a1e0::1/128", "fd7a:115c:a1e0:1::/64",
+	} {
+		b.AddPrefix(mp(p))
+	}
+
+	set, err := b.IPSet()
+	require.NoError(t, err)
+
+	addrs := testAddrs(set.Ranges())
+
+	for _, p := range []string{
+		"0.0.0.0/0", "::/0", "9.0.0.0/8", "10.0.0.0/32", "10.0.0.1/32", "10.0.0.2/32",
+		"10.0.0.2/31", "10.0.0.4/30", "10.0.0.0/24", "10.0.1.0/24", "10.0.2.7/32",
+		"10.0.0.0/8", "100.63.255.255/32", "100.100.100.100/32", "100.128.0.0/32",
+		"255.255.255.255/32", "::/128", "fd7a:115c:a1e0::/128", "fd7a:115c:a1e0::1/128",
+		"fd7a:115c:a1e0::/48", "fd7a:115c:a1e0:1::5/128", "fd7a:115c:a1e0:2::/64",
+	} {
+		assert.Equal(t, set.OverlapsPrefix(mp(p)), addrs.overlapsPrefix(mp(p)), p)
+	}
+
+	assert.False(t, addrs.overlapsPrefix(netip.Prefix{}))
+	assert.False(t, testAddrs(nil).overlapsPrefix(mp("0.0.0.0/0")))
+}
+
+// TestRunTestsManyAssertionsAndDestinations runs the shape that made a run
+// quadratic: every assertion matched only the last of the rule's
+// destinations, and each comparison parsed that destination again.
+func TestRunTestsManyAssertionsAndDestinations(t *testing.T) {
+	users := policyTestUsers()
+	nodes := policyTestNodes(users)
+
+	const count = 2600
+
+	dsts := make([]string, 0, count+1)
+	for i := range count {
+		dsts = append(dsts, fmt.Sprintf(`"10.%d.%d.1:1"`, i/250, i%250))
+	}
+
+	dsts = append(dsts, `"192.0.2.1:1"`)
+
+	accept := slices.Repeat([]string{`"192.0.2.1:1"`}, count)
+	deny := slices.Repeat([]string{`"192.0.2.1:2"`}, count)
+
+	policy := fmt.Sprintf(`{
+		"acls": [{"action": "accept", "src": ["*"], "dst": [%s]}],
+		"tests": [{"src": "alice@headscale.net", "accept": [%s], "deny": [%s]}]
+	}`, strings.Join(dsts, ","), strings.Join(accept, ","), strings.Join(deny, ","))
+
+	pol, err := unmarshalPolicy([]byte(policy))
+	require.NoError(t, err)
+
+	require.NoError(t, evaluateTests(pol, users, nodes.ViewSlice()))
+}
+
+// TestSetPolicyEvaluatesTestsUnlocked asserts that a policy's tests run
+// before SetPolicy takes the write lock, so map generation and node writes
+// do not wait for them: a reader holding the lock does not hold up the
+// evaluation, and a rejected policy never needs the write lock at all.
+func TestSetPolicyEvaluatesTestsUnlocked(t *testing.T) {
+	users := policyTestUsers()
+	nodes := policyTestNodes(users)
+
+	pm, err := NewPolicyManager([]byte(`{}`), users, nodes.ViewSlice())
+	require.NoError(t, err)
+
+	bad := `{
+		"tagOwners": { "tag:server": ["alice@headscale.net"] },
+		"acls": [{
+			"action": "accept",
+			"src": ["alice@headscale.net"],
+			"dst": ["tag:server:22"]
+		}],
+		"tests": [{
+			"src": "bob@headscale.net",
+			"accept": ["tag:server:22"]
+		}]
+	}`
+
+	pm.mu.RLock()
+	defer pm.mu.RUnlock()
+
+	done := make(chan error, 1)
+
+	go func() {
+		_, err := pm.SetPolicy([]byte(bad))
+		done <- err
+	}()
+
+	select {
+	case err := <-done:
+		require.ErrorIs(t, err, errPolicyTestsFailed)
+	case <-time.After(10 * time.Second):
+		t.Fatal("SetPolicy waited for the write lock before evaluating the tests")
+	}
+}
+
 // TestSetPolicyRejectsFailingTests asserts that SetPolicy is the user-write
 // boundary: a policy whose tests fail must be rejected without mutating the
 // live PolicyManager. NewPolicyManager (boot path) does not run tests.