wgengine/netstack: allow UDP listeners to receive traffic on Service VIP addresses (#18972)

Fixes UDP listeners on VIP Service addresses not receiving inbound traffic.

- Modified shouldProcessInbound to check for registered UDP transport endpoints when processing packets to service VIPs
- Uses FindTransportEndpoint to determine if a UDP listener exists for the destination VIP/port
- Supports both IPv4 and IPv6

The aim was to mirror the existing TCP logic, providing feature parity for UDP-based services on VIP Services.

Fixes #18971

Signed-off-by: chaosinthecrd <tom@tmlabs.co.uk>

wgengine/netstack/netstack.go

@@ -1231,6 +1231,34 @@ func (ns *Impl) shouldProcessInbound(p *packet.Parsed, t *tstun.Wrapper) bool {
 				return true
 			}
 		}
+		// check if there's a registered UDP endpoint for this service VIP
+		// This allows userspace UDP listeners (e.g., via tsnet.ListenPacket) to
+		// receive traffic on service VIP addresses.
+		if p.IPProto == ipproto.UDP {
+			var netProto tcpip.NetworkProtocolNumber
+			var id stack.TransportEndpointID
+			if p.Dst.Addr().Is4() {
+				netProto = ipv4.ProtocolNumber
+				id = stack.TransportEndpointID{
+					LocalAddress:  tcpip.AddrFrom4(p.Dst.Addr().As4()),
+					LocalPort:     p.Dst.Port(),
+					RemoteAddress: tcpip.AddrFrom4(p.Src.Addr().As4()),
+					RemotePort:    p.Src.Port(),
+				}
+			} else {
+				netProto = ipv6.ProtocolNumber
+				id = stack.TransportEndpointID{
+					LocalAddress:  tcpip.AddrFrom16(p.Dst.Addr().As16()),
+					LocalPort:     p.Dst.Port(),
+					RemoteAddress: tcpip.AddrFrom16(p.Src.Addr().As16()),
+					RemotePort:    p.Src.Port(),
+				}
+			}
+			ep := ns.ipstack.FindTransportEndpoint(netProto, udp.ProtocolNumber, id, nicID)
+			if ep != nil {
+				return true
+			}
+		}
 		return false
 	}
 	if p.IPVersion == 6 && !isLocal && viaRange.Contains(dstIP) {

wgengine/netstack/netstack_test.go

@@ -453,6 +453,194 @@ func TestShouldProcessInbound(t *testing.T) {
 			},
 			want: false,
 		},
+		{
+			name: "udp-on-service-vip-with-listener-ipv4",
+			pkt: &packet.Parsed{
+				IPVersion: 4,
+				IPProto:   ipproto.UDP,
+				Src:       netip.MustParseAddrPort("100.101.102.103:1234"),
+				Dst:       netip.MustParseAddrPort("100.100.100.100:8080"),
+			},
+			beforeStart: func(i *Impl) {
+				i.ProcessLocalIPs = false
+				i.ProcessSubnets = false
+			},
+			afterStart: func(i *Impl) {
+				IPServiceMap := netmap.IPServiceMappings{
+					serviceIP: "svc:test-service",
+				}
+				i.lb.SetIPServiceMappingsForTest(IPServiceMap)
+
+				i.atomicIsVIPServiceIPFunc.Store(func(addr netip.Addr) bool {
+					return addr == serviceIP
+				})
+
+				// Register the service VIP address on the NIC so gVisor can route to it
+				protocolAddr := tcpip.ProtocolAddress{
+					Protocol:          header.IPv4ProtocolNumber,
+					AddressWithPrefix: tcpip.AddrFrom4(serviceIP.As4()).WithPrefix(),
+				}
+
+				if err := i.ipstack.AddProtocolAddress(nicID, protocolAddr, stack.AddressProperties{}); err != nil {
+					t.Fatalf("AddProtocolAddress: %v", err)
+				}
+
+				// Create a UDP listener on the service VIP
+				pc, err := gonet.DialUDP(i.ipstack, &tcpip.FullAddress{
+					NIC:  nicID,
+					Addr: tcpip.AddrFrom4(serviceIP.As4()),
+					Port: 8080,
+				}, nil, header.IPv4ProtocolNumber)
+				if err != nil {
+					t.Fatalf("DialUDP: %v", err)
+				}
+				t.Cleanup(func() { pc.Close() })
+
+				i.atomicIsLocalIPFunc.Store(looksLikeATailscaleSelfAddress)
+			},
+			want: true,
+		},
+		{
+			name: "udp-on-service-vip-no-listener-ipv4",
+			pkt: &packet.Parsed{
+				IPVersion: 4,
+				IPProto:   ipproto.UDP,
+				Src:       netip.MustParseAddrPort("100.101.102.103:1234"),
+				Dst:       netip.MustParseAddrPort("100.100.100.100:9999"),
+			},
+			beforeStart: func(i *Impl) {
+				i.ProcessLocalIPs = false
+				i.ProcessSubnets = false
+			},
+			afterStart: func(i *Impl) {
+				IPServiceMap := netmap.IPServiceMappings{
+					serviceIP: "svc:test-service",
+				}
+				i.lb.SetIPServiceMappingsForTest(IPServiceMap)
+
+				i.atomicIsVIPServiceIPFunc.Store(func(addr netip.Addr) bool {
+					return addr == serviceIP
+				})
+
+				i.atomicIsLocalIPFunc.Store(looksLikeATailscaleSelfAddress)
+			},
+			want: false,
+		},
+		{
+			name: "udp-on-service-vip-with-listener-ipv6",
+			pkt: &packet.Parsed{
+				IPVersion: 6,
+				IPProto:   ipproto.UDP,
+				Src:       netip.MustParseAddrPort("[fd7a:115c:a1e0::1]:1234"),
+				Dst:       netip.MustParseAddrPort("[fd7a:115c:a1e0::53]:8080"),
+			},
+			beforeStart: func(i *Impl) {
+				i.ProcessLocalIPs = false
+				i.ProcessSubnets = false
+			},
+			afterStart: func(i *Impl) {
+				IPServiceMap := netmap.IPServiceMappings{
+					serviceIPv6: "svc:test-service",
+				}
+				i.lb.SetIPServiceMappingsForTest(IPServiceMap)
+
+				i.atomicIsVIPServiceIPFunc.Store(func(addr netip.Addr) bool {
+					return addr == serviceIPv6
+				})
+
+				protocolAddr := tcpip.ProtocolAddress{
+					Protocol:          header.IPv6ProtocolNumber,
+					AddressWithPrefix: tcpip.AddrFrom16(serviceIPv6.As16()).WithPrefix(),
+				}
+				if err := i.ipstack.AddProtocolAddress(nicID, protocolAddr, stack.AddressProperties{}); err != nil {
+					t.Fatalf("AddProtocolAddress: %v", err)
+				}
+
+				pc, err := gonet.DialUDP(i.ipstack, &tcpip.FullAddress{
+					NIC:  nicID,
+					Addr: tcpip.AddrFrom16(serviceIPv6.As16()),
+					Port: 8080,
+				}, nil, header.IPv6ProtocolNumber)
+				if err != nil {
+					t.Fatalf("DialUDP: %v", err)
+				}
+				t.Cleanup(func() { pc.Close() })
+
+				i.atomicIsLocalIPFunc.Store(looksLikeATailscaleSelfAddress)
+			},
+			want: true,
+		},
+		{
+			name: "udp-on-service-vip-no-listener-ipv6",
+			pkt: &packet.Parsed{
+				IPVersion: 6,
+				IPProto:   ipproto.UDP,
+				Src:       netip.MustParseAddrPort("[fd7a:115c:a1e0::1]:1234"),
+				Dst:       netip.AddrPortFrom(serviceIPv6, 9999),
+			},
+			beforeStart: func(i *Impl) {
+				i.ProcessLocalIPs = false
+				i.ProcessSubnets = false
+			},
+			afterStart: func(i *Impl) {
+				IPServiceMap := netmap.IPServiceMappings{
+					serviceIPv6: "svc:test-service",
+				}
+				i.lb.SetIPServiceMappingsForTest(IPServiceMap)
+
+				i.atomicIsVIPServiceIPFunc.Store(func(addr netip.Addr) bool {
+					return addr == serviceIPv6
+				})
+
+				i.atomicIsLocalIPFunc.Store(looksLikeATailscaleSelfAddress)
+			},
+			want: false,
+		},
+		{
+			name: "tcp-on-service-vip-with-udp-listener",
+			pkt: &packet.Parsed{
+				IPVersion: 4,
+				IPProto:   ipproto.TCP,
+				Src:       netip.MustParseAddrPort("100.101.102.103:1234"),
+				Dst:       netip.MustParseAddrPort("100.100.100.100:8080"), // serviceIP
+				TCPFlags:  packet.TCPSyn,
+			},
+			beforeStart: func(i *Impl) {
+				i.ProcessLocalIPs = false
+				i.ProcessSubnets = false
+			},
+			afterStart: func(i *Impl) {
+				IPServiceMap := netmap.IPServiceMappings{
+					serviceIP: "svc:test-service",
+				}
+				i.lb.SetIPServiceMappingsForTest(IPServiceMap)
+
+				i.atomicIsVIPServiceIPFunc.Store(func(addr netip.Addr) bool {
+					return addr == serviceIP
+				})
+
+				protocolAddr := tcpip.ProtocolAddress{
+					Protocol:          header.IPv4ProtocolNumber,
+					AddressWithPrefix: tcpip.AddrFrom4(serviceIP.As4()).WithPrefix(),
+				}
+				if err := i.ipstack.AddProtocolAddress(nicID, protocolAddr, stack.AddressProperties{}); err != nil {
+					t.Fatalf("AddProtocolAddress: %v", err)
+				}
+
+				pc, err := gonet.DialUDP(i.ipstack, &tcpip.FullAddress{
+					NIC:  nicID,
+					Addr: tcpip.AddrFrom4(serviceIP.As4()),
+					Port: 8080,
+				}, nil, header.IPv4ProtocolNumber)
+				if err != nil {
+					t.Fatalf("DialUDP: %v", err)
+				}
+				t.Cleanup(func() { pc.Close() })
+
+				i.atomicIsLocalIPFunc.Store(looksLikeATailscaleSelfAddress)
+			},
+			want: false,
+		},
 
 		// TODO(andrew): test PeerAPI
 		// TODO(andrew): test TCP packets without the SYN flag set