wgengine/filter: add experiment to use ART for filter packet checksbradfitz/ipx_set_contains

Updates #10049

Change-Id: Iae56d44d2cf5f6f57d47237aa62876c45b130a43
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

4 files changed, 48 insertions, 22 deletions

diff --git a/go.mod b/go.mod
index 3118bc3af..7b251bb00 100644
--- a/go.mod
+++ b/go.mod
@@ -30,7 +30,7 @@ require (
 	github.com/godbus/dbus/v5 v5.1.1-0.20230522191255-76236955d466
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da
 	github.com/golangci/golangci-lint v1.52.2
-	github.com/google/go-cmp v0.5.9
+	github.com/google/go-cmp v0.6.0
 	github.com/google/go-containerregistry v0.16.1
 	github.com/google/nftables v0.1.1-0.20230115205135-9aa6fdf5a28c
 	github.com/google/uuid v1.3.1
@@ -57,6 +57,7 @@ require (
 	github.com/prometheus/client_golang v1.17.0
 	github.com/prometheus/common v0.44.0
 	github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e
+	github.com/tailscale/art v0.0.0-20231101034115-3827a3c782e9
 	github.com/tailscale/certstore v0.1.1-0.20231020161753-77811a65f4ff
 	github.com/tailscale/depaware v0.0.0-20210622194025-720c4b409502
 	github.com/tailscale/goexpect v0.0.0-20210902213824-6e8c725cea41
diff --git a/go.sum b/go.sum
index 6cc857bec..852d7f514 100644
--- a/go.sum
+++ b/go.sum
@@ -440,8 +440,9 @@ github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
-github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-containerregistry v0.16.1 h1:rUEt426sR6nyrL3gt+18ibRcvYpKYdpsa5ZW7MA08dQ=
 github.com/google/go-containerregistry v0.16.1/go.mod h1:u0qB2l7mvtWVR5kNcbFIhFY1hLbf8eeGapA+vbFDCtQ=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
@@ -864,6 +865,8 @@ github.com/subosito/gotenv v1.4.2 h1:X1TuBLAMDFbaTAChgCBLu3DU3UPyELpnF2jjJ2cz/S8
 github.com/subosito/gotenv v1.4.2/go.mod h1:ayKnFf/c6rvx/2iiLrJUk1e6plDbT3edrFNGqEflhK0=
 github.com/t-yuki/gocover-cobertura v0.0.0-20180217150009-aaee18c8195c h1:+aPplBwWcHBo6q9xrfWdMrT9o4kltkmmvpemgIjep/8=
 github.com/t-yuki/gocover-cobertura v0.0.0-20180217150009-aaee18c8195c/go.mod h1:SbErYREK7xXdsRiigaQiQkI9McGRzYMvlKYaP3Nimdk=
+github.com/tailscale/art v0.0.0-20231101034115-3827a3c782e9 h1:0jZYJEIxi+gRFUVhTEugH5GWC3muM3At60Pbvgx6EDU=
+github.com/tailscale/art v0.0.0-20231101034115-3827a3c782e9/go.mod h1:vqKZQktbl+AhzNtqxc7rSm6t+lvuHd729sHSw/ElykI=
 github.com/tailscale/certstore v0.1.1-0.20231020161753-77811a65f4ff h1:vnxdYZUJbsSRcIcduDW3DcQqfqaiv4FUgy25q8X+vfI=
 github.com/tailscale/certstore v0.1.1-0.20231020161753-77811a65f4ff/go.mod h1:XrBNfAFN+pwoWuksbFS9Ccxnopa15zJGgXRFN90l3K4=
 github.com/tailscale/depaware v0.0.0-20210622194025-720c4b409502 h1:34icjjmqJ2HPjrSuJYEkdZ+0ItmGQAQ75cRHIiftIyE=
diff --git a/wgengine/filter/filter.go b/wgengine/filter/filter.go
index b5ed82a54..86d7feb30 100644
--- a/wgengine/filter/filter.go
+++ b/wgengine/filter/filter.go
@@ -11,6 +11,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/tailscale/art"
 	"go4.org/netipx"
 	"tailscale.com/envknob"
 	"tailscale.com/net/flowtrack"
@@ -26,16 +27,16 @@ import (
 // Filter is a stateful packet filter.
 type Filter struct {
 	logf logger.Logf
-	// local is the set of IPs prefixes that we know to be "local" to
-	// this node. All packets coming in over tailscale must have a
-	// destination within local, regardless of the policy filter
+	// localContains tests where an IP is in the set of IPs prefixes that we
+	// know to be "local" to this node. All packets coming in over tailscale
+	// must have a destination within local, regardless of the policy filter
 	// below.
-	local *netipx.IPSet
+	localContains func(netip.Addr) bool
 
-	// logIPs is the set of IPs that are allowed to appear in flow
-	// logs. If a packet is to or from an IP not in logIPs, it will
-	// never be logged.
-	logIPs *netipx.IPSet
+	// logIPsContains tests whether an IP is in the set of IPs that are allowed
+	// to appear in flow logs. If a packet is to or from an IP not in logIPs, it
+	// will never be logged.
+	logIPsContains func(netip.Addr) bool
 
 	// matches4 and matches6 are lists of match->action rules
 	// applied to all packets arriving over tailscale
@@ -167,6 +168,8 @@ func NewShieldsUpFilter(localNets *netipx.IPSet, logIPs *netipx.IPSet, shareStat
 	return f
 }
 
+var useART = envknob.RegisterBool("TS_DEBUG_FILTER_USE_ART")
+
 // New creates a new packet filter. The filter enforces that incoming
 // packets must be destined to an IP in localNets, and must be allowed
 // by matches. If shareStateWith is non-nil, the returned filter
@@ -182,18 +185,36 @@ func New(matches []Match, localNets *netipx.IPSet, logIPs *netipx.IPSet, shareSt
 		}
 	}
 	f := &Filter{
-		logf:     logf,
-		matches4: matchesFamily(matches, netip.Addr.Is4),
-		matches6: matchesFamily(matches, netip.Addr.Is6),
-		cap4:     capMatchesFunc(matches, netip.Addr.Is4),
-		cap6:     capMatchesFunc(matches, netip.Addr.Is6),
-		local:    localNets,
-		logIPs:   logIPs,
-		state:    state,
+		logf:           logf,
+		matches4:       matchesFamily(matches, netip.Addr.Is4),
+		matches6:       matchesFamily(matches, netip.Addr.Is6),
+		cap4:           capMatchesFunc(matches, netip.Addr.Is4),
+		cap6:           capMatchesFunc(matches, netip.Addr.Is6),
+		localContains:  localNets.Contains,
+		logIPsContains: logIPs.Contains,
+		state:          state,
+	}
+	if useART() {
+		f.localContains = containsFuncForART(localNets)
+		f.logIPsContains = containsFuncForART(logIPs)
 	}
 	return f
 }
 
+func containsFuncForART(s *netipx.IPSet) func(netip.Addr) bool {
+	if s == nil {
+		return func(netip.Addr) bool { return false }
+	}
+	t := &art.Table[struct{}]{}
+	for _, p := range s.Prefixes() {
+		t.Insert(p, struct{}{})
+	}
+	return func(ip netip.Addr) bool {
+		_, ok := t.Get(ip)
+		return ok
+	}
+}
+
 // matchesFamily returns the subset of ms for which keep(srcNet.IP)
 // and keep(dstNet.IP) are both true.
 func matchesFamily(ms matches, keep func(netip.Addr) bool) matches {
@@ -410,7 +431,7 @@ func (f *Filter) runIn4(q *packet.Parsed) (r Response, why string) {
 	// A compromised peer could try to send us packets for
 	// destinations we didn't explicitly advertise. This check is to
 	// prevent that.
-	if !f.local.Contains(q.Dst.Addr()) {
+	if !f.localContains(q.Dst.Addr()) {
 		return Drop, "destination not allowed"
 	}
 
@@ -470,7 +491,7 @@ func (f *Filter) runIn6(q *packet.Parsed) (r Response, why string) {
 	// A compromised peer could try to send us packets for
 	// destinations we didn't explicitly advertise. This check is to
 	// prevent that.
-	if !f.local.Contains(q.Dst.Addr()) {
+	if !f.localContains(q.Dst.Addr()) {
 		return Drop, "destination not allowed"
 	}
 
@@ -596,7 +617,7 @@ func (f *Filter) pre(q *packet.Parsed, rf RunFlags, dir direction) Response {
 
 // loggingAllowed reports whether p can appear in logs at all.
 func (f *Filter) loggingAllowed(p *packet.Parsed) bool {
-	return f.logIPs.Contains(p.Src.Addr()) && f.logIPs.Contains(p.Dst.Addr())
+	return f.logIPsContains(p.Src.Addr()) && f.logIPsContains(p.Dst.Addr())
 }
 
 // omitDropLogging reports whether packet p, which has already been
diff --git a/wgengine/filter/filter_test.go b/wgengine/filter/filter_test.go
index 1d7521821..fe944402a 100644
--- a/wgengine/filter/filter_test.go
+++ b/wgengine/filter/filter_test.go
@@ -436,7 +436,8 @@ func TestLoggingPrivacy(t *testing.T) {
 	logB.AddPrefix(netip.MustParsePrefix("100.64.0.0/10"))
 	logB.AddPrefix(tsaddr.TailscaleULARange())
 	f := newFilter(logf)
-	f.logIPs, _ = logB.IPSet()
+	logIPsSet, _ := logB.IPSet()
+	f.logIPsContains = logIPsSet.Contains
 
 	var (
 		ts4       = netip.AddrPortFrom(tsaddr.CGNATRange().Addr().Next(), 1234)