From e5779f019e0be46f3977c41f9056372a78f77bb8 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Fri, 15 Oct 2021 19:22:30 -0700
Subject: [PATCH] wgengine/magicsock: move temporary endpoint lookup later, add
 TODO to remove

Updates #3088

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

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 0dcd3280a..05756fb92 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -1772,18 +1772,12 @@ func (c *Conn) handleDiscoMessage(msg []byte, src netaddr.IPPort, derpNodeSrc ke
 		return
 	}
 
-	ep, ok := c.peerMap.endpointForDiscoKey(sender)
-	if !ok {
+	if !c.peerMap.anyEndpointForDiscoKey(sender) {
 		if debugDisco {
 			c.logf("magicsock: disco: ignoring disco-looking frame, don't know endpoint for %v", sender.ShortString())
 		}
 		return
 	}
-	if !ep.canP2P() {
-		// This endpoint allegedly sent us a disco packet, but we know
-		// they can't speak disco. Drop.
-		return
-	}
 
 	// We're now reasonably sure we're expecting communication from
 	// this peer, do the heavy crypto lifting to see what they want.
@@ -1824,6 +1818,20 @@ func (c *Conn) handleDiscoMessage(msg []byte, src netaddr.IPPort, derpNodeSrc ke
 		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, ep, src, sender)