From 04fd94acd6f764ba7b8e96b9aa9cae5e28ae3b85 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Fri, 15 Oct 2021 21:55:59 -0700
Subject: [PATCH] wgengine/magicsock: remove endpointForDiscoKey call from
 handleDiscoMessage

A DiscoKey maps 1:n to endpoints. When we get a disco pong, we don't
necessarily know which endpoint sent it to us. Ask them all. There
will only usually be 1 (and in rare circumstances 2). So it's easier
to ask all two rather than building new maps from the random ping TxID
to its endpoint.

Updates #3088

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 wgengine/magicsock/magicsock.go | 34 ++++++++++++++++-----------------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 6450bff9c..baef28a41 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -1842,25 +1842,19 @@ func (c *Conn) handleDiscoMessage(msg []byte, src netaddr.IPPort, derpNodeSrc ta
 		return
 	}
 
-	// TODO(bradfitz): remove this endpointForDiscoKey lookup once handlePingLocked
-	// and handlePongConnLocked are updated to look up the endpoint on their own
-	// different ways (not by DiscoKey).
-	ep, ok := c.peerMap.endpointForDiscoKey(sender)
-	if !ok {
-		// Shouldn't be possible if anyEndpointForDiscoKey above passed.
-		return
-	}
-	if !ep.canP2P() {
-		// This endpoint allegedly sent us a disco packet, but we know
-		// they can't speak disco. Drop.
-		return
-	}
-
 	switch dm := dm.(type) {
 	case *disco.Ping:
 		c.handlePingLocked(dm, src, di, derpNodeSrc)
 	case *disco.Pong:
-		ep.handlePongConnLocked(dm, src)
+		// There might be multiple nodes for the sender's DiscoKey.
+		// Ask each to handle it, stopping once one reports that
+		// the Pong's TxID was theirs.
+		handled := false
+		c.peerMap.forEachEndpointWithDiscoKey(sender, func(ep *endpoint) {
+			if !handled && ep.handlePongConnLocked(dm, src) {
+				handled = true
+			}
+		})
 	case *disco.CallMeMaybe:
 		if src.IP() != derpMagicIPAddr || derpNodeSrc.IsZero() {
 			// CallMeMaybe messages should only come via DERP.
@@ -3610,7 +3604,9 @@ func (de *endpoint) noteConnectivityChange() {
 
 // handlePongConnLocked handles a Pong message (a reply to an earlier ping).
 // It should be called with the Conn.mu held.
-func (de *endpoint) handlePongConnLocked(m *disco.Pong, src netaddr.IPPort) {
+//
+// It reports whether m.TxID corresponds to a ping that this endpoint sent.
+func (de *endpoint) handlePongConnLocked(m *disco.Pong, src netaddr.IPPort) (knownTxID bool) {
 	de.mu.Lock()
 	defer de.mu.Unlock()
 
@@ -3618,9 +3614,10 @@ func (de *endpoint) handlePongConnLocked(m *disco.Pong, src netaddr.IPPort) {
 
 	sp, ok := de.sentPing[m.TxID]
 	if !ok {
-		// This is not a pong for a ping we sent. Ignore.
-		return
+		// This is not a pong for a ping we sent.
+		return false
 	}
+	knownTxID = true // for naked returns below
 	de.removeSentPingLocked(m.TxID, sp)
 
 	now := mono.Now()
@@ -3671,6 +3668,7 @@ func (de *endpoint) handlePongConnLocked(m *disco.Pong, src netaddr.IPPort) {
 			de.trustBestAddrUntil = now.Add(trustUDPAddrDuration)
 		}
 	}
+	return
 }
 
 // addrLatency is an IPPort with an associated latency.