wgengine/magicsock: simplify Geneve-encapsulated disco.Ping handling (#16448)

Just make [relayManager] always handle it, there's no benefit to
checking bestAddr's.

Also, remove passing of disco.Pong to [relayManager] in
endpoint.handlePongConnLocked(), which is redundant with the callsite in
Conn.handleDiscoMessage(). Conn.handleDiscoMessage() already passes to
[relayManager] if the txID us not known to any [*endpoint].

Updates tailscale/corp#27502

Signed-off-by: Jordan Whited <jordan@tailscale.com>

wgengine/magicsock/endpoint.go

@@ -1656,13 +1656,6 @@ func (de *endpoint) handlePongConnLocked(m *disco.Pong, di *discoInfo, src epAdd
 	de.mu.Lock()
 	defer de.mu.Unlock()
 
-	if src.vni.isSet() && src != de.bestAddr.epAddr {
-		// "src" is not our bestAddr, but [relayManager] might be in the
-		// middle of probing it, awaiting pong reception. Make it aware.
-		de.c.relayManager.handleGeneveEncapDiscoMsgNotBestAddr(de.c, m, di, src)
-		return false
-	}
-
 	isDerp := src.ap.Addr() == tailcfg.DerpMagicIPAddr
 
 	sp, ok := de.sentPing[m.TxID]

wgengine/magicsock/magicsock.go

@@ -2103,7 +2103,7 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 			c.logf("[unexpected] %T packets should not come from a relay server with Geneve control bit set", dm)
 			return
 		}
-		c.relayManager.handleGeneveEncapDiscoMsgNotBestAddr(c, challenge, di, src)
+		c.relayManager.handleGeneveEncapDiscoMsg(c, challenge, di, src)
 		return
 	}
 
@@ -2125,7 +2125,10 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 			return true
 		})
 		if !knownTxID && src.vni.isSet() {
-			c.relayManager.handleGeneveEncapDiscoMsgNotBestAddr(c, dm, di, src)
+			// If it's an unknown TxID, and it's Geneve-encapsulated, then
+			// make [relayManager] aware. It might be in the middle of probing
+			// src.
+			c.relayManager.handleGeneveEncapDiscoMsg(c, dm, di, src)
 		}
 	case *disco.CallMeMaybe, *disco.CallMeMaybeVia:
 		var via *disco.CallMeMaybeVia
@@ -2233,50 +2236,22 @@ func (c *Conn) handlePingLocked(dm *disco.Ping, src epAddr, di *discoInfo, derpN
 	di.lastPingTime = time.Now()
 	isDerp := src.ap.Addr() == tailcfg.DerpMagicIPAddr
 
-	// numNodes tracks how many nodes (node keys) are associated with the disco
-	// key tied to this inbound ping. Multiple nodes may share the same disco
-	// key in the case of node sharing and users switching accounts.
-	var numNodes int
-	// If we got a ping over DERP, then derpNodeSrc is non-zero and we reply
-	// over DERP (in which case ipDst is also a DERP address).
-	// But if the ping was over UDP (ipDst is not a DERP address), then dstKey
-	// will be zero here, but that's fine: sendDiscoMessage only requires
-	// a dstKey if the dst ip:port is DERP.
-	dstKey := derpNodeSrc
-
-	switch {
-	case src.vni.isSet():
+	if src.vni.isSet() {
 		if isDerp {
 			c.logf("[unexpected] got Geneve-encapsulated disco ping from %v/%v over DERP", src, derpNodeSrc)
 			return
 		}
 
-		var bestEpAddr epAddr
-		var discoKey key.DiscoPublic
-		ep, ok := c.peerMap.endpointForEpAddr(src)
-		if ok {
-			ep.mu.Lock()
-			bestEpAddr = ep.bestAddr.epAddr
-			ep.mu.Unlock()
-			disco := ep.disco.Load()
-			if disco != nil {
-				discoKey = disco.key
-			}
-		}
-
-		if src == bestEpAddr && discoKey == di.discoKey {
-			// We have an associated endpoint with src as its bestAddr. Set
-			// numNodes so we TX a pong further down.
-			numNodes = 1
-		} else {
-			// We have no [endpoint] in the [peerMap] for this relay [epAddr]
-			// using it as a bestAddr. [relayManager] might be in the middle of
-			// probing it or attempting to set it as best via
-			// [endpoint.udpRelayEndpointReady()]. Make [relayManager] aware.
-			c.relayManager.handleGeneveEncapDiscoMsgNotBestAddr(c, dm, di, src)
+		// [relayManager] is always responsible for handling (replying) to
+		// Geneve-encapsulated [disco.Ping] messages in the interest of
+		// simplicity. It might be in the middle of probing src, so it must be
+		// made aware.
+		c.relayManager.handleGeneveEncapDiscoMsg(c, dm, di, src)
 		return
 	}
-	default: // no VNI
+
+	// This is a naked [disco.Ping] without a VNI.
+
 	// If we can figure out with certainty which node key this disco
 	// message is for, eagerly update our [epAddr]<>node and disco<>node
 	// mappings to make p2p path discovery faster in simple
@@ -2290,6 +2265,17 @@ func (c *Conn) handlePingLocked(dm *disco.Ping, src epAddr, di *discoInfo, derpN
 		}
 	}
 
+	// numNodes tracks how many nodes (node keys) are associated with the disco
+	// key tied to this inbound ping. Multiple nodes may share the same disco
+	// key in the case of node sharing and users switching accounts.
+	var numNodes int
+	// If we got a ping over DERP, then derpNodeSrc is non-zero and we reply
+	// over DERP (in which case ipDst is also a DERP address).
+	// But if the ping was over UDP (ipDst is not a DERP address), then dstKey
+	// will be zero here, but that's fine: sendDiscoMessage only requires
+	// a dstKey if the dst ip:port is DERP.
+	dstKey := derpNodeSrc
+
 	// Remember this route if not present.
 	var dup bool
 	if isDerp {
@@ -2320,7 +2306,6 @@ func (c *Conn) handlePingLocked(dm *disco.Ping, src epAddr, di *discoInfo, derpN
 			dstKey = key.NodePublic{}
 		}
 	}
-	}
 
 	if numNodes == 0 {
 		c.logf("[unexpected] got disco ping from %v/%v for node not in peers", src, derpNodeSrc)

wgengine/magicsock/relaymanager.go

@@ -325,10 +325,9 @@ func (r *relayManager) handleCallMeMaybeVia(ep *endpoint, lastBest addrQuality,
 	})
 }
 
-// handleGeneveEncapDiscoMsgNotBestAddr handles reception of Geneve-encapsulated
-// disco messages if they are not associated with any known
-// [*endpoint.bestAddr].
-func (r *relayManager) handleGeneveEncapDiscoMsgNotBestAddr(conn *Conn, dm disco.Message, di *discoInfo, src epAddr) {
+// handleGeneveEncapDiscoMsg handles reception of Geneve-encapsulated disco
+// messages.
+func (r *relayManager) handleGeneveEncapDiscoMsg(conn *Conn, dm disco.Message, di *discoInfo, src epAddr) {
 	relayManagerInputEvent(r, nil, &r.rxHandshakeDiscoMsgCh, relayHandshakeDiscoMsgEvent{conn: conn, msg: dm, disco: di.discoKey, from: src.ap, vni: src.vni.get(), at: time.Now()})
 }
 

wgengine/magicsock/relaymanager_test.go

@@ -26,7 +26,7 @@ func TestRelayManagerInitAndIdle(t *testing.T) {
 	<-rm.runLoopStoppedCh
 
 	rm = relayManager{}
-	rm.handleGeneveEncapDiscoMsgNotBestAddr(&Conn{discoPrivate: key.NewDisco()}, &disco.BindUDPRelayEndpointChallenge{}, &discoInfo{}, epAddr{})
+	rm.handleGeneveEncapDiscoMsg(&Conn{discoPrivate: key.NewDisco()}, &disco.BindUDPRelayEndpointChallenge{}, &discoInfo{}, epAddr{})
 	<-rm.runLoopStoppedCh
 
 	rm = relayManager{}