policy/matcher: binary-search ranges in overlap checks

hscontrol/policy/matcher/export_test.go

@@ -6,8 +6,5 @@ package matcher
 // wider. Callers that need strict validation should pre-validate
 // their inputs via [util.ParseIPSet].
 func MatchFromStrings(sources, destinations []string) Match {
-	return Match{
-		srcs:  buildIPSet(sources),
-		dests: buildIPSet(destinations),
-	}
+	return newMatch(buildIPSet(sources), buildIPSet(destinations))
 }

hscontrol/policy/matcher/matcher.go

@@ -14,6 +14,18 @@ import (
 type Match struct {
 	srcs  *netipx.IPSet
 	dests *netipx.IPSet
+
+	srcRanges  []netipx.IPRange
+	destRanges []netipx.IPRange
+}
+
+func newMatch(srcs, dests *netipx.IPSet) Match {
+	return Match{
+		srcs:       srcs,
+		dests:      dests,
+		srcRanges:  srcs.Ranges(),
+		destRanges: dests.Ranges(),
+	}
 }
 
 func (m *Match) DebugString() string {
@@ -67,10 +79,7 @@ func MatchFromFilterRule(rule tailcfg.FilterRule) Match {
 
 	destsSet, _ := dests.IPSet()
 
-	return Match{
-		srcs:  buildIPSet(rule.SrcIPs),
-		dests: destsSet,
-	}
+	return newMatch(buildIPSet(rule.SrcIPs), destsSet)
 }
 
 // buildIPSet parses each string via [util.ParseIPSet] and unions the
@@ -99,11 +108,28 @@ func (m *Match) DestsContainsIP(ips ...netip.Addr) bool {
 }
 
 func (m *Match) SrcsOverlapsPrefixes(prefixes ...netip.Prefix) bool {
-	return slices.ContainsFunc(prefixes, m.srcs.OverlapsPrefix)
+	return slices.ContainsFunc(prefixes, func(prefix netip.Prefix) bool {
+		return overlaps(m.srcRanges, netipx.RangeOfPrefix(prefix))
+	})
 }
 
 func (m *Match) DestsOverlapsPrefixes(prefixes ...netip.Prefix) bool {
-	return slices.ContainsFunc(prefixes, m.dests.OverlapsPrefix)
+	return slices.ContainsFunc(prefixes, func(prefix netip.Prefix) bool {
+		return overlaps(m.destRanges, netipx.RangeOfPrefix(prefix))
+	})
+}
+
+func overlaps(ranges []netipx.IPRange, r netipx.IPRange) bool {
+	if !r.IsValid() {
+		return false
+	}
+
+	i, _ := slices.BinarySearchFunc(ranges, r.From(),
+		func(x netipx.IPRange, a netip.Addr) int {
+			return x.To().Compare(a)
+		})
+
+	return i < len(ranges) && ranges[i].Overlaps(r)
 }
 
 // DestsIsTheInternet reports whether the destination covers "the

hscontrol/policy/matcher/matcher_test.go

@@ -2,15 +2,87 @@ package matcher
 
 import (
 	"net/netip"
+	"slices"
 	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"pgregory.net/rapid"
 	"tailscale.com/tailcfg"
 	"tailscale.com/tailcfg/peercap"
 )
 
+func TestOverlapsPrefixesMatchesIPSet(t *testing.T) {
+	t.Parallel()
+
+	prefixGen := rapid.Custom(func(rt *rapid.T) netip.Prefix {
+		if rapid.Bool().Draw(rt, "invalid") {
+			return netip.Prefix{}
+		}
+
+		n := rapid.Uint16().Draw(rt, "address")
+		if rapid.Bool().Draw(rt, "ipv6") {
+			addr := netip.AddrFrom16([16]byte{
+				0x20, 0x01, 0x0d, 0xb8, byte(n >> 8), byte(n & 0xff),
+			})
+
+			return netip.PrefixFrom(addr, rapid.IntRange(0, 128).Draw(rt, "bits"))
+		}
+
+		addr := netip.AddrFrom4([4]byte{10, byte(n >> 8), byte(n & 0xff), 0})
+
+		return netip.PrefixFrom(addr, rapid.IntRange(0, 32).Draw(rt, "bits"))
+	})
+	stringsGen := rapid.Map(prefixGen, netip.Prefix.String)
+
+	rapid.Check(t, func(rt *rapid.T) {
+		sources := rapid.SliceOfN(stringsGen, 0, 32).Draw(rt, "sources")
+		destinations := rapid.SliceOfN(stringsGen, 0, 32).Draw(rt, "destinations")
+		prefixes := rapid.SliceOfN(prefixGen, 0, 16).Draw(rt, "prefixes")
+		m := MatchFromStrings(sources, destinations)
+		wantSrc := slices.ContainsFunc(prefixes, m.srcs.OverlapsPrefix)
+		wantDst := slices.ContainsFunc(prefixes, m.dests.OverlapsPrefix)
+
+		require.Equal(rt, wantSrc, m.SrcsOverlapsPrefixes(prefixes...))
+		require.Equal(rt, wantDst, m.DestsOverlapsPrefixes(prefixes...))
+
+		for _, prefix := range prefixes {
+			require.Equal(rt, m.srcs.OverlapsPrefix(prefix),
+				m.SrcsOverlapsPrefixes(prefix))
+			require.Equal(rt, m.dests.OverlapsPrefix(prefix),
+				m.DestsOverlapsPrefixes(prefix))
+		}
+	})
+}
+
+func BenchmarkOverlapsPrefixesManyRanges(b *testing.B) {
+	addresses := make([]string, 16384)
+	for i := range addresses {
+		addresses[i] = netip.AddrFrom4([4]byte{
+			10, byte(i >> 7), byte(i << 1), 0,
+		}).String()
+	}
+
+	m := MatchFromStrings(addresses, addresses)
+	queries := []netip.Prefix{
+		netip.MustParsePrefix("10.255.255.255/32"),
+		netip.MustParsePrefix("10.127.254.0/32"),
+	}
+
+	for _, query := range queries {
+		b.Run(query.String(), func(b *testing.B) {
+			b.ReportAllocs()
+
+			for b.Loop() {
+				if m.SrcsOverlapsPrefixes(query) != m.DestsOverlapsPrefixes(query) {
+					b.Fatal("identical sets must give identical overlap results")
+				}
+			}
+		})
+	}
+}
+
 func TestMatchFromStrings(t *testing.T) {
 	t.Parallel()
 

hscontrol/policy/v2/policy_test.go

@@ -9,11 +9,13 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/juanfont/headscale/hscontrol/policy/matcher"
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/puzpuzpuz/xsync/v4"
 	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/require"
+	"go4.org/netipx"
 	"pgregory.net/rapid"
 	"tailscale.com/net/tsaddr"
 	"tailscale.com/tailcfg"
@@ -82,6 +84,7 @@ func TestPolicyManager(t *testing.T) {
 				tt.wantMatchers,
 				matchers,
 				cmp.AllowUnexported(matcher.Match{}),
+				cmpopts.EquateComparable(netipx.IPRange{}),
 			); diff != "" {
 				t.Errorf("Filter() matchers mismatch (-want +got):\n%s", diff)
 			}