From 263e01c47bdb34ab3839d700d1d9fa022ec7f8db Mon Sep 17 00:00:00 2001
From: Andrew Lytvynov <awly@tailscale.com>
Date: Sat, 2 Dec 2023 16:30:33 -0600
Subject: [PATCH] wgengine/filter: add protocol-agnostic packet checker
 (#10446)

For use in ACL tests, we need a way to check whether a packet is allowed
not just with TCP, but any protocol.

Updates #3561

Signed-off-by: Andrew Lytvynov <awly@tailscale.com>
---
 wgengine/filter/filter.go | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/wgengine/filter/filter.go b/wgengine/filter/filter.go
index b5ed82a54..082d8a0f5 100644
--- a/wgengine/filter/filter.go
+++ b/wgengine/filter/filter.go
@@ -300,9 +300,9 @@ var dummyPacket = []byte{
 	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
 }
 
-// CheckTCP determines whether TCP traffic from srcIP to dstIP:dstPort
-// is allowed.
-func (f *Filter) CheckTCP(srcIP, dstIP netip.Addr, dstPort uint16) Response {
+// Check determines whether traffic from srcIP to dstIP:dstPort is allowed
+// using protocol proto.
+func (f *Filter) Check(srcIP, dstIP netip.Addr, dstPort uint16, proto ipproto.Proto) Response {
 	pkt := &packet.Parsed{}
 	pkt.Decode(dummyPacket) // initialize private fields
 	switch {
@@ -319,12 +319,20 @@ func (f *Filter) CheckTCP(srcIP, dstIP netip.Addr, dstPort uint16) Response {
 	}
 	pkt.Src = netip.AddrPortFrom(srcIP, 0)
 	pkt.Dst = netip.AddrPortFrom(dstIP, dstPort)
-	pkt.IPProto = ipproto.TCP
-	pkt.TCPFlags = packet.TCPSyn
+	pkt.IPProto = proto
+	if proto == ipproto.TCP {
+		pkt.TCPFlags = packet.TCPSyn
+	}
 
 	return f.RunIn(pkt, 0)
 }
 
+// CheckTCP determines whether TCP traffic from srcIP to dstIP:dstPort
+// is allowed.
+func (f *Filter) CheckTCP(srcIP, dstIP netip.Addr, dstPort uint16) Response {
+	return f.Check(srcIP, dstIP, dstPort, ipproto.TCP)
+}
+
 // CapsWithValues appends to base the capabilities that srcIP has talking
 // to dstIP.
 func (f *Filter) CapsWithValues(srcIP, dstIP netip.Addr) tailcfg.PeerCapMap {