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

@@ -1772,18 +1772,12 @@ func (c *Conn) handleDiscoMessage(msg []byte, src netaddr.IPPort, derpNodeSrc ke
 		return
 	}
 
-	ep, ok := c.peerMap.endpointForDiscoKey(sender)
+	if !c.peerMap.anyEndpointForDiscoKey(sender) {
-	if !ok {
 		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)