policy/v2: parse the compiled filter once per tests run

The tests evaluator parsed every SrcIPs and DstPorts entry again for
each assertion, source prefix and rule. Parse the filter once per run.

hscontrol/policy/v2/test.go

@@ -317,8 +317,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
 		}
@@ -330,7 +332,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,
@@ -354,7 +362,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))
@@ -424,7 +434,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.
@@ -437,7 +447,7 @@ func evalReachability(
 	dst string,
 	proto Protocol,
 	pol *Policy,
-	filter []tailcfg.FilterRule,
+	rules []testRule,
 	users []types.User,
 	nodes views.Slice[types.NodeView],
 	everyNode bool,
@@ -465,7 +475,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
 			}
 		}
@@ -525,7 +535,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.
 //
@@ -534,14 +644,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 {
+	for _, rule := range rules {
-		if !ruleMatchesSource(rule, src) {
+		if !rule.srcs.overlapsPrefix(src) {
 			continue
 		}
 
@@ -557,37 +673,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 {
+func ruleMatchesProto(rule testRule, requestedProtos []int) bool {
-	if len(rule.IPProto) == 0 {
+	if len(rule.protos) == 0 {
 		return true
 	}
 
-	for _, ruleProto := range rule.IPProto {
+	for _, ruleProto := range rule.protos {
 		if slices.Contains(requestedProtos, ruleProto) {
 			return true
 		}
@@ -597,15 +693,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 {
+func ruleAllowsAnyDest(
-	for _, dp := range rule.DstPorts {
+	rule testRule,
-		if !destEntryMatchesPrefixes(dp, dstPrefixes) {
+	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
 		}
 	}
@@ -613,17 +713,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,15 @@ package v2
 
 import (
 	"fmt"
+	"net/netip"
+	"slices"
 	"strings"
 	"testing"
 
 	"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 +291,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 +379,66 @@ 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()))
+}
+
 // 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.