wgengine/netstack: add support for custom UDP flow handlers

To be used by tsnet and sniproxy later.

Updates #5871
Updates #1748

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

tsnet/tsnet.go

@@ -41,6 +41,7 @@ import (
 	"tailscale.com/net/tsdial"
 	"tailscale.com/smallzstd"
 	"tailscale.com/types/logger"
+	"tailscale.com/types/nettype"
 	"tailscale.com/util/mak"
 	"tailscale.com/wgengine"
 	"tailscale.com/wgengine/monitor"
@@ -440,6 +441,7 @@ func (s *Server) start() (reterr error) {
 	}
 	ns.ProcessLocalIPs = true
 	ns.ForwardTCPIn = s.forwardTCP
+	ns.GetUDPHandlerForFlow = s.getUDPHandlerForFlow
 	s.netstack = ns
 	s.dialer.UseNetstackForIP = func(ip netip.Addr) bool {
 		_, ok := eng.PeerForIP(ip)
@@ -579,6 +581,12 @@ func (s *Server) forwardTCP(c net.Conn, port uint16) {
 	}
 }
 
+func (s *Server) getUDPHandlerForFlow(src, dst netip.AddrPort) (handler func(nettype.ConnPacketConn), intercept bool) {
+	s.logf("rejecting incoming UDP flow: (%v, %v)", src, dst)
+	// TODO(bradfitz): hook up to Listen("udp", dst) so users of tsnet can hook into this.
+	return nil, true
+}
+
 // getTSNetDir usually just returns filepath.Join(confDir, "tsnet-"+prog)
 // with no error.
 //

types/nettype/nettype.go

@@ -48,3 +48,9 @@ func (a packetListenerAdapter) ListenPacket(ctx context.Context, network, addres
 	}
 	return pc.(PacketConn), nil
 }
+
+// ConnPacketConn is the interface that's a superset of net.Conn and net.PacketConn.
+type ConnPacketConn interface {
+	net.Conn
+	net.PacketConn
+}

wgengine/netstack/netstack.go

@@ -47,6 +47,7 @@ import (
 	"tailscale.com/types/ipproto"
 	"tailscale.com/types/logger"
 	"tailscale.com/types/netmap"
+	"tailscale.com/types/nettype"
 	"tailscale.com/version/distro"
 	"tailscale.com/wgengine"
 	"tailscale.com/wgengine/filter"
@@ -78,12 +79,26 @@ func init() {
 // and implements wgengine.FakeImpl to act as a userspace network
 // stack when Tailscale is running in fake mode.
 type Impl struct {
-	// ForwardTCPIn, if non-nil, handles forwarding an inbound TCP
-	// connection.
-	// TODO(bradfitz): provide mechanism for tsnet to reject a
-	// port other than accepting it and closing it.
+	// ForwardTCPIn, if non-nil, handles forwarding an inbound TCP connection.
+	//
+	// TODO(bradfitz): convert this to the GetUDPHandlerForFlow pattern below to
+	// provide mechanism for tsnet to reject a port other than accepting it and
+	// closing it.
 	ForwardTCPIn func(c net.Conn, port uint16)
 
+	// GetUDPHandlerForFlow conditionally handles an incoming UDP flow for the
+	// provided (src/port, dst/port) 4-tuple.
+	//
+	// A nil value is equivalent to a func returning (nil, false).
+	//
+	// If func returns intercept=false, the default forwarding behavior (if
+	// ProcessLocalIPs and/or ProcesssSubnetIPs) takes place.
+	//
+	// When intercept=true, the behavior depends on whether the returned handler
+	// is non-nil: if nil, the connection is rejected. If non-nil, handler takes
+	// over the UDP flow.
+	GetUDPHandlerForFlow func(src, dst netip.AddrPort) (handler func(nettype.ConnPacketConn), intercept bool)
+
 	// ProcessLocalIPs is whether netstack should handle incoming
 	// traffic directed at the Node.Addresses (local IPs).
 	// It can only be set before calling Start.
@@ -1020,8 +1035,20 @@ func (ns *Impl) acceptUDP(r *udp.ForwarderRequest) {
 		return
 	}
 
+	if get := ns.GetUDPHandlerForFlow; get != nil {
+		h, intercept := get(srcAddr, dstAddr)
+		if intercept {
+			if h == nil {
+				ep.Close()
+				return
+			}
+			go h(gonet.NewUDPConn(ns.ipstack, &wq, ep))
+			return
+		}
+	}
+
 	c := gonet.NewUDPConn(ns.ipstack, &wq, ep)
-	go ns.forwardUDP(c, &wq, srcAddr, dstAddr)
+	go ns.forwardUDP(c, srcAddr, dstAddr)
 }
 
 func (ns *Impl) handleMagicDNSUDP(srcAddr netip.AddrPort, c *gonet.UDPConn) {
@@ -1065,7 +1092,7 @@ func (ns *Impl) handleMagicDNSUDP(srcAddr netip.AddrPort, c *gonet.UDPConn) {
 // dstAddr may be either a local Tailscale IP, in which we case we proxy to
 // 127.0.0.1, or any other IP (from an advertised subnet), in which case we
 // proxy to it directly.
-func (ns *Impl) forwardUDP(client *gonet.UDPConn, wq *waiter.Queue, clientAddr, dstAddr netip.AddrPort) {
+func (ns *Impl) forwardUDP(client *gonet.UDPConn, clientAddr, dstAddr netip.AddrPort) {
 	port, srcPort := dstAddr.Port(), clientAddr.Port()
 	if debugNetstack() {
 		ns.logf("[v2] netstack: forwarding incoming UDP connection on port %v", port)