From badd0c4f93acd14fb58701fc6ab5dfb40f9aead7 Mon Sep 17 00:00:00 2001
From: Jordan Whited <jordan@tailscale.com>
Date: Tue, 23 Jun 2026 10:10:33 -0700
Subject: [PATCH] wgengine/magicsock: consider VNI as part of peer relay
 handshake suppression

Otherwise we may never handshake a new peer relay server endpoint
around remote client restarts and/or disco key rotation.

Updates #20215

Signed-off-by: Jordan Whited <jordan@tailscale.com>
---
 wgengine/magicsock/relaymanager.go      |  6 ++-
 wgengine/magicsock/relaymanager_test.go | 54 ++++++++++++++++++++++++-
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/wgengine/magicsock/relaymanager.go b/wgengine/magicsock/relaymanager.go
index 8ea15bce3..3be23b6ed 100644
--- a/wgengine/magicsock/relaymanager.go
+++ b/wgengine/magicsock/relaymanager.go
@@ -790,8 +790,10 @@ func (r *relayManager) handleNewServerEndpointRunLoop(newServerEndpoint newRelay
 		go r.sendCallMeMaybeVia(newServerEndpoint.wlb.ep, newServerEndpoint.se)
 	}
 
-	lastBestMatchingServer := newServerEndpoint.se.ServerDisco.Compare(newServerEndpoint.wlb.lastBest.relayServerDisco) == 0
-	if lastBestMatchingServer && newServerEndpoint.wlb.lastBestIsTrusted {
+	lastBestMatchingServerVNI := newServerEndpoint.wlb.lastBest.vni.IsSet() &&
+		newServerEndpoint.se.ServerDisco.Compare(newServerEndpoint.wlb.lastBest.relayServerDisco) == 0 &&
+		newServerEndpoint.se.VNI == newServerEndpoint.wlb.lastBest.vni.Get()
+	if lastBestMatchingServerVNI && newServerEndpoint.wlb.lastBestIsTrusted {
 		// This relay server endpoint is the same as [endpoint]'s bestAddr at
 		// the time UDP relay path discovery was started, and it was also a
 		// trusted path (see endpoint.trustBestAddrUntil), so return early.
diff --git a/wgengine/magicsock/relaymanager_test.go b/wgengine/magicsock/relaymanager_test.go
index 47d935404..cf07e4210 100644
--- a/wgengine/magicsock/relaymanager_test.go
+++ b/wgengine/magicsock/relaymanager_test.go
@@ -109,12 +109,36 @@ func TestRelayManager_handleNewServerEndpointRunLoop(t *testing.T) {
 	serverDiscoA := key.NewDisco().Public()
 	serverDiscoB := key.NewDisco().Public()
 
+	// getAddrQuality returns an addrQuality with serverDisco with the provided
+	// VNI set. A negative vni leaves the VNI unset.
+	getAddrQuality := func(serverDisco key.DiscoPublic, vni int) addrQuality {
+		aq := addrQuality{relayServerDisco: serverDisco}
+		if vni >= 0 {
+			aq.epAddr.vni.Set(uint32(vni))
+		}
+		return aq
+	}
+
 	serverAendpointALamport1VNI1 := newRelayServerEndpointEvent{
 		wlb: endpointWithLastBest{ep: epA},
 		se:  udprelay.ServerEndpoint{ServerDisco: serverDiscoA, LamportID: 1, VNI: 1},
 	}
+	// lastBest matches the new server endpoint on both server disco and VNI, and
+	// is trusted: suppression should fire.
 	serverAendpointALamport1VNI1LastBestMatching := newRelayServerEndpointEvent{
-		wlb: endpointWithLastBest{ep: epA, lastBestIsTrusted: true, lastBest: addrQuality{relayServerDisco: serverDiscoA}},
+		wlb: endpointWithLastBest{ep: epA, lastBestIsTrusted: true, lastBest: getAddrQuality(serverDiscoA, 1)},
+		se:  udprelay.ServerEndpoint{ServerDisco: serverDiscoA, LamportID: 1, VNI: 1},
+	}
+	// lastBest matches the new server endpoint on server disco but NOT VNI (1 vs
+	// 2), and is trusted: suppression should NOT fire.
+	serverAendpointALamport1VNI1LastBestMatchingServerNeqVNI := newRelayServerEndpointEvent{
+		wlb: endpointWithLastBest{ep: epA, lastBestIsTrusted: true, lastBest: getAddrQuality(serverDiscoA, 2)},
+		se:  udprelay.ServerEndpoint{ServerDisco: serverDiscoA, LamportID: 1, VNI: 1},
+	}
+	// lastBest matches the new server endpoint on server disco, is trusted, but
+	// has no VNI set: suppression should NOT fire.
+	serverAendpointALamport1VNI1LastBestMatchingServerUnsetVNI := newRelayServerEndpointEvent{
+		wlb: endpointWithLastBest{ep: epA, lastBestIsTrusted: true, lastBest: getAddrQuality(serverDiscoA, -1)},
 		se:  udprelay.ServerEndpoint{ServerDisco: serverDiscoA, LamportID: 1, VNI: 1},
 	}
 	serverAendpointALamport2VNI1 := newRelayServerEndpointEvent{
@@ -212,12 +236,38 @@ func TestRelayManager_handleNewServerEndpointRunLoop(t *testing.T) {
 			},
 		},
 		{
-			name: "trusted-last-best-with-matching-server",
+			// Trusted lastBest matching on both server disco and VNI suppresses
+			// the new handshake.
+			name: "trusted-last-best-matching-server-and-vni",
 			events: []newRelayServerEndpointEvent{
 				serverAendpointALamport1VNI1LastBestMatching,
 			},
 			want: []newRelayServerEndpointEvent{},
 		},
+		{
+			// Trusted lastBest matching on server disco but NOT VNI must not
+			// suppress the new handshake, otherwise we may never handshake a new
+			// peer relay server endpoint around remote client restarts and/or
+			// disco key rotation (#20215).
+			name: "trusted-last-best-matching-server-neq-vni",
+			events: []newRelayServerEndpointEvent{
+				serverAendpointALamport1VNI1LastBestMatchingServerNeqVNI,
+			},
+			want: []newRelayServerEndpointEvent{
+				serverAendpointALamport1VNI1LastBestMatchingServerNeqVNI,
+			},
+		},
+		{
+			// Trusted lastBest matching on server disco with an unset VNI must
+			// not suppress the new handshake.
+			name: "trusted-last-best-matching-server-unset-vni",
+			events: []newRelayServerEndpointEvent{
+				serverAendpointALamport1VNI1LastBestMatchingServerUnsetVNI,
+			},
+			want: []newRelayServerEndpointEvent{
+				serverAendpointALamport1VNI1LastBestMatchingServerUnsetVNI,
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {