4 files changed, 66 insertions, 3 deletions
diff --git a/wgengine/netstack/netstack.go b/wgengine/netstack/netstack.go
index 61b1cdb9b..1131490ec 100644
--- a/wgengine/netstack/netstack.go
+++ b/wgengine/netstack/netstack.go
@@ -367,6 +367,11 @@ func (ns *Impl) forwardTCP(client *gonet.TCPConn, wq *waiter.Queue, port uint16)
 		return
 	}
 	defer server.Close()
+	backendLocalAddr := server.LocalAddr().(*net.TCPAddr)
+	backendLocalIPPort, _ := netaddr.FromStdAddr(backendLocalAddr.IP, backendLocalAddr.Port, backendLocalAddr.Zone)
+	clientRemoteIP, _ := netaddr.FromStdIP(client.RemoteAddr().(*net.TCPAddr).IP)
+	ns.e.RegisterIPPortIdentity(backendLocalIPPort, clientRemoteIP)
+	defer ns.e.UnregisterIPPortIdentity(backendLocalIPPort)
 	connClosed := make(chan error, 2)
 	go func() {
 		_, err := io.Copy(server, client)
@@ -406,19 +411,28 @@ func (ns *Impl) acceptUDP(r *udp.ForwarderRequest) {
 func (ns *Impl) forwardUDP(client *gonet.UDPConn, wq *waiter.Queue, clientLocalAddr, clientRemoteAddr tcpip.FullAddress) {
 	port := clientLocalAddr.Port
 	ns.logf("[v2] netstack: forwarding incoming UDP connection on port %v", port)
-	backendLocalAddr := &net.UDPAddr{Port: int(clientRemoteAddr.Port)}
+	backendListenAddr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: int(clientRemoteAddr.Port)}
 	backendRemoteAddr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: int(port)}
-	backendConn, err := net.ListenUDP("udp4", backendLocalAddr)
+	backendConn, err := net.ListenUDP("udp4", backendListenAddr)
 	if err != nil {
 		ns.logf("netstack: could not bind local port %v: %v, trying again with random port", clientRemoteAddr.Port, err)
-		backendConn, err = net.ListenUDP("udp4", nil)
+		backendListenAddr.Port = 0
+		backendConn, err = net.ListenUDP("udp4", backendListenAddr)
 		if err != nil {
 			ns.logf("netstack: could not connect to local UDP server on port %v: %v", port, err)
 			return
 		}
 	}
+	backendLocalAddr := backendConn.LocalAddr().(*net.UDPAddr)
+	backendLocalIPPort, ok := netaddr.FromStdAddr(backendListenAddr.IP, backendLocalAddr.Port, backendLocalAddr.Zone)
+	if !ok {
+		ns.logf("could not get backend local IP:port from %v:%v", backendLocalAddr.IP, backendLocalAddr.Port)
+	}
+	clientRemoteIP, _ := netaddr.FromStdIP(net.ParseIP(clientRemoteAddr.Addr.String()))
+	ns.e.RegisterIPPortIdentity(backendLocalIPPort, clientRemoteIP)
 	ctx, cancel := context.WithCancel(context.Background())
 	timer := time.AfterFunc(2*time.Minute, func() {
+		ns.e.UnregisterIPPortIdentity(backendLocalIPPort)
 		ns.logf("netstack: UDP session between %s and %s timed out", clientRemoteAddr, backendRemoteAddr)
 		cancel()
 		client.Close()
diff --git a/wgengine/userspace.go b/wgengine/userspace.go
index 1baf7f040..fccca8c8b 100644
--- a/wgengine/userspace.go
+++ b/wgengine/userspace.go
@@ -125,6 +125,7 @@ type userspaceEngine struct {
 	pingers             map[wgkey.Key]*pinger                // legacy pingers for pre-discovery peers
 	pendOpen            map[flowtrack.Tuple]*pendingOpenFlow // see pendopen.go
 	networkMapCallbacks map[*someHandle]NetworkMapCallback
+	tsIPByIPPort        map[netaddr.IPPort]netaddr.IP // allows registration of IP:ports as belonging to a certain Tailscale IP for whois lookups
 
 	// Lock ordering: magicsock.Conn.mu, wgLock, then mu.
 }
@@ -1341,6 +1342,31 @@ func (e *userspaceEngine) Ping(ip netaddr.IP, cb func(*ipnstate.PingResult)) {
 	e.magicConn.Ping(ip, cb)
 }
 
+func (e *userspaceEngine) RegisterIPPortIdentity(ipport netaddr.IPPort, tsIP netaddr.IP) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	if e.tsIPByIPPort == nil {
+		e.tsIPByIPPort = make(map[netaddr.IPPort]netaddr.IP)
+	}
+	e.tsIPByIPPort[ipport] = tsIP
+}
+
+func (e *userspaceEngine) UnregisterIPPortIdentity(ipport netaddr.IPPort) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	if e.tsIPByIPPort == nil {
+		return
+	}
+	delete(e.tsIPByIPPort, ipport)
+}
+
+func (e *userspaceEngine) WhoIsIPPort(ipport netaddr.IPPort) (tsIP netaddr.IP, ok bool) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	tsIP, ok = e.tsIPByIPPort[ipport]
+	return tsIP, ok
+}
+
 // diagnoseTUNFailure is called if tun.CreateTUN fails, to poke around
 // the system and log some diagnostic info that might help debug why
 // TUN failed. Because TUN's already failed and things the program's
diff --git a/wgengine/watchdog.go b/wgengine/watchdog.go
index 3b96f2e8e..f4f7d3085 100644
--- a/wgengine/watchdog.go
+++ b/wgengine/watchdog.go
@@ -120,6 +120,16 @@ func (e *watchdogEngine) DiscoPublicKey() (k tailcfg.DiscoKey) {
 func (e *watchdogEngine) Ping(ip netaddr.IP, cb func(*ipnstate.PingResult)) {
 	e.watchdog("Ping", func() { e.wrap.Ping(ip, cb) })
 }
+func (e *watchdogEngine) RegisterIPPortIdentity(ipp netaddr.IPPort, tsIP netaddr.IP) {
+	e.watchdog("RegisterIPPortIdentity", func() { e.wrap.RegisterIPPortIdentity(ipp, tsIP) })
+}
+func (e *watchdogEngine) UnregisterIPPortIdentity(ipp netaddr.IPPort) {
+	e.watchdog("UnregisterIPPortIdentity", func() { e.wrap.UnregisterIPPortIdentity(ipp) })
+}
+func (e *watchdogEngine) WhoIsIPPort(ipp netaddr.IPPort) (tsIP netaddr.IP, ok bool) {
+	e.watchdog("UnregisterIPPortIdentity", func() { tsIP, ok = e.wrap.WhoIsIPPort(ipp) })
+	return tsIP, ok
+}
 func (e *watchdogEngine) Close() {
 	e.watchdog("Close", e.wrap.Close)
 }
diff --git a/wgengine/wgengine.go b/wgengine/wgengine.go
index 57bc7cb77..c8e7963db 100644
--- a/wgengine/wgengine.go
+++ b/wgengine/wgengine.go
@@ -137,4 +137,17 @@ type Engine interface {
 	// Ping is a request to start a discovery ping with the peer handling
 	// the given IP and then call cb with its ping latency & method.
 	Ping(ip netaddr.IP, cb func(*ipnstate.PingResult))
+
+	// RegisterIPPortIdentity registers a given node (identified by its
+	// Tailscale IP) as temporarily having the given IP:port for whois lookups.
+	// The IP:port is generally a localhost IP and an ephemeral port, used
+	// while proxying connections to localhost.
+	RegisterIPPortIdentity(netaddr.IPPort, netaddr.IP)
+
+	// UnregisterIPPortIdentity removes a temporary IP:port registration.
+	UnregisterIPPortIdentity(netaddr.IPPort)
+
+	// WhoIsIPPort looks up an IP:port in the temporary registrations,
+	// and returns a matching Tailscale IP, if it exists.
+	WhoIsIPPort(netaddr.IPPort) (netaddr.IP, bool)
 }