From b92e2ebd24a5c33859ef2ba85cc554cc64a22a53 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Wed, 23 Jun 2021 22:09:31 -0700
Subject: [PATCH] wgengine/netstack: add Impl.DialContextUDP

Unused so far, but eventually we'll want this for SOCKS5 UDP binds (we
currently only do TCP with SOCKS5), and also for #2102 for forwarding
MagicDNS upstream to Tailscale IPs over netstack.

Updates #2102

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 wgengine/netstack/netstack.go | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/wgengine/netstack/netstack.go b/wgengine/netstack/netstack.go
index f587755ff..47103b9ce 100644
--- a/wgengine/netstack/netstack.go
+++ b/wgengine/netstack/netstack.go
@@ -362,6 +362,30 @@ func (ns *Impl) DialContextTCP(ctx context.Context, addr string) (*gonet.TCPConn
 	return gonet.DialContextTCP(ctx, ns.ipstack, remoteAddress, ipType)
 }
 
+func (ns *Impl) DialContextUDP(ctx context.Context, addr string) (*gonet.UDPConn, error) {
+	ns.mu.Lock()
+	dnsMap := ns.dns
+	ns.mu.Unlock()
+
+	remoteIPPort, err := dnsMap.Resolve(ctx, addr)
+	if err != nil {
+		return nil, err
+	}
+	remoteAddress := &tcpip.FullAddress{
+		NIC:  nicID,
+		Addr: tcpip.Address(remoteIPPort.IP().IPAddr().IP),
+		Port: remoteIPPort.Port(),
+	}
+	var ipType tcpip.NetworkProtocolNumber
+	if remoteIPPort.IP().Is4() {
+		ipType = ipv4.ProtocolNumber
+	} else {
+		ipType = ipv6.ProtocolNumber
+	}
+
+	return gonet.DialUDP(ns.ipstack, nil, remoteAddress, ipType)
+}
+
 func (ns *Impl) injectOutbound() {
 	for {
 		packetInfo, ok := ns.linkEP.ReadContext(context.Background())