From 633d48604e8d1e9adb742181aff14603be931ca9 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Sat, 17 Oct 2026 10:16:36 +0000 Subject: [PATCH] policy/matcher: binary-search ranges in overlap checks --- hscontrol/policy/matcher/export_test.go | 5 +- hscontrol/policy/matcher/matcher.go | 38 +++++++++++-- hscontrol/policy/matcher/matcher_test.go | 72 ++++++++++++++++++++++++ hscontrol/policy/v2/policy_test.go | 3 + 4 files changed, 108 insertions(+), 10 deletions(-) diff --git a/hscontrol/policy/matcher/export_test.go b/hscontrol/policy/matcher/export_test.go index 63ec9f09d..2c271cc50 100644 --- a/hscontrol/policy/matcher/export_test.go +++ b/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)) } diff --git a/hscontrol/policy/matcher/matcher.go b/hscontrol/policy/matcher/matcher.go index acfd35b4a..1061cae01 100644 --- a/hscontrol/policy/matcher/matcher.go +++ b/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 diff --git a/hscontrol/policy/matcher/matcher_test.go b/hscontrol/policy/matcher/matcher_test.go index bfd528cbd..e0e0f7e37 100644 --- a/hscontrol/policy/matcher/matcher_test.go +++ b/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() diff --git a/hscontrol/policy/v2/policy_test.go b/hscontrol/policy/v2/policy_test.go index 3d72a5173..613fc4d26 100644 --- a/hscontrol/policy/v2/policy_test.go +++ b/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) }