wgengine/magicsock: set conn field in relayHandshakeDiscoMsgEvent (#16348)

Updates tailscale/corp#27502

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

wgengine/magicsock/endpoint.go

@@ -1601,7 +1601,7 @@ func (de *endpoint) handlePongConnLocked(m *disco.Pong, di *discoInfo, src epAdd
 	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(m, di, src)
+		de.c.relayManager.handleGeneveEncapDiscoMsgNotBestAddr(de.c, m, di, src)
 		return false
 	}
 

wgengine/magicsock/magicsock.go

@@ -2053,7 +2053,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(challenge, di, src)
+		c.relayManager.handleGeneveEncapDiscoMsgNotBestAddr(c, challenge, di, src)
 		return
 	}
 
@@ -2075,7 +2075,7 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 			return true
 		})
 		if !knownTxID && src.vni.isSet() {
-			c.relayManager.handleGeneveEncapDiscoMsgNotBestAddr(dm, di, src)
+			c.relayManager.handleGeneveEncapDiscoMsgNotBestAddr(c, dm, di, src)
 		}
 	case *disco.CallMeMaybe, *disco.CallMeMaybeVia:
 		var via *disco.CallMeMaybeVia
@@ -2221,7 +2221,7 @@ func (c *Conn) handlePingLocked(dm *disco.Ping, src epAddr, di *discoInfo, derpN
 			// using it as a bestAddr. [relayManager] might be in the middle of
 			// probing it or attempting to set it as best via
 			// [endpoint.relayEndpointReady()]. Make [relayManager] aware.
-			c.relayManager.handleGeneveEncapDiscoMsgNotBestAddr(dm, di, src)
+			c.relayManager.handleGeneveEncapDiscoMsgNotBestAddr(c, dm, di, src)
 			return
 		}
 	default: // no VNI

wgengine/magicsock/relaymanager.go

@@ -319,8 +319,8 @@ func (r *relayManager) handleCallMeMaybeVia(ep *endpoint, dm *disco.CallMeMaybeV
 // handleGeneveEncapDiscoMsgNotBestAddr handles reception of Geneve-encapsulated
 // disco messages if they are not associated with any known
 // [*endpoint.bestAddr].
-func (r *relayManager) handleGeneveEncapDiscoMsgNotBestAddr(dm disco.Message, di *discoInfo, src epAddr) {
-	relayManagerInputEvent(r, nil, &r.rxHandshakeDiscoMsgCh, relayHandshakeDiscoMsgEvent{msg: dm, disco: di.discoKey, from: src.ap, vni: src.vni.get(), at: time.Now()})
+func (r *relayManager) handleGeneveEncapDiscoMsgNotBestAddr(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()})
 }
 
 // handleRelayServersSet handles an update of the complete relay server set.

wgengine/magicsock/relaymanager_test.go

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