ipn/ipnlocal,wgengine: implement wg-go SetPriorityMessageOnEstablishmentFunc (#20606)

Implements a way to send TSMPDiscoAdverts based on a trigger from
wireguard-go when a rekey happens. This lets us distribute disco keys
consistently, but also sets us up for a minimal message that can be
distributed to other clients.

A benchmark is implemented to make it easier to keep the call cheap and
to avoid locking up anything in wireguard-go.

Updates #20081

Signed-off-by: Claus Lensbøl <claus@tailscale.com>

control/controlclient/direct.go

@@ -426,7 +426,7 @@ func NewDirect(opts Options) (*Direct, error) {
 	c.controlTimePub = eventbus.Publish[ControlTime](c.busClient)
 	discoKeyPub := eventbus.Publish[events.PeerDiscoKeyUpdate](c.busClient)
 	eventbus.SubscribeFunc(c.busClient, func(update events.DiscoKeyAdvertisement) {
-		c.logf("controlclient direct: got TSMP disco key advertisement from %v via eventbus", update.Src)
+		c.logf("[v1] controlclient direct: got TSMP disco key advertisement from %v via eventbus", update.Src)
 		var peerID tailcfg.NodeID
 		var peerKey key.NodePublic
 		var ok bool
@@ -438,7 +438,7 @@ func NewDirect(opts Options) (*Direct, error) {
 		}
 
 		if sess != nil && ok {
-			c.logf("controlclient direct: updating discoKey for %v via mapSession", update.Src)
+			c.logf("[v1] controlclient direct: updating discoKey for %v via mapSession", update.Src)
 
 			// If we update without error, return. If the err indicates that the
 			// mapSession has gone away, we want to fall back to pushing the key
@@ -453,7 +453,7 @@ func NewDirect(opts Options) (*Direct, error) {
 		// We need to push the update further down the chain. Either because we do
 		// not have a mapSession (we are not connected to control) or because the
 		// mapSession queue has closed.
-		c.logf("controlclient direct: updating discoKey for %v via magicsock", update.Src)
+		c.logf("[v1] controlclient direct: updating discoKey for %v via magicsock", update.Src)
 		discoKeyPub.Publish(events.PeerDiscoKeyUpdate(update))
 	})
 

ipn/ipnlocal/local.go

@@ -617,6 +617,7 @@ func NewLocalBackend(logf logger.Logf, logID logid.PublicID, sys *tsd.System, lo
 	e.SetPeerForIPFunc(b.PeerForIP)
 	e.SetPeerSessionStateFunc(b.onPeerWireGuardState)
 	e.SetNetLogSource(netLogNodeSource{b})
+	e.SetPeerPriorityMessageOnEstablishmentFunc(b.MagicConn().PriorityMessageForPeer)
 	e.SetWGPeerLookup(b.lookupPeerWireGuardString)
 	b.dialer.SetResolveMagicDNS(b.resolveMagicDNS)
 	if buildfeatures.HasDNS {

ipn/ipnlocal/state_test.go

@@ -2015,6 +2015,8 @@ func (e *mockEngine) ResetDevicePeer(key.NodePublic) {}
 func (e *mockEngine) SetPeerSessionStateFunc(func(key.NodePublic, wgengine.PeerWireGuardState)) {
 }
 func (e *mockEngine) SetNetLogSource(wgengine.NetLogSource)                            {}
+func (e *mockEngine) SetPeerPriorityMessageOnEstablishmentFunc(fn func(key.NodePublic) (msg []byte)) {
+}
 func (e *mockEngine) SetWGPeerLookup(func(wgString string) (tsString string, ok bool)) {}
 func (e *mockEngine) ProbeLocks()                                                      {}
 

tstest/natlab/vmtest/vmtest_test.go

@@ -1314,6 +1314,7 @@ func TestDirectConnectionWithCachedNetmapOnTwoNodes(t *testing.T) {
 	checkClientMetrics(t, "Node A", env.ClientMetrics(a), map[string]int64{
 		"magicsock_cached_peer_contact_derp":   0,
 		"magicsock_cached_peer_contact_direct": 0,
+		"magicsock_tsmp_disco_key_advertisement_sent": 0,
 	})
 	checkInitialMetrics.End(nil)
 
@@ -1349,6 +1350,7 @@ func TestDirectConnectionWithCachedNetmapOnTwoNodes(t *testing.T) {
 	checkFinalMetrics.Begin()
 	checkClientMetrics(t, "Node A", env.ClientMetrics(a), map[string]int64{
 		"magicsock_cached_peer_contact_direct": 1,
+		"magicsock_tsmp_disco_key_advertisement_sent": 1,
 	})
 	checkFinalMetrics.End(nil)
 }

wgengine/magicsock/endpoint.go

@@ -40,11 +40,6 @@
 var mtuProbePingSizesV4 []int
 var mtuProbePingSizesV6 []int
 
-// discoKeyAdvertisementInterval tells how often a disco update via TSMP can
-// happen. The update is triggered via enqueueCallMeMaybe, and thus it will
-// only be sent if the magicsock is in a state to send out CallMeMaybe.
-const discoKeyAdvertisementInterval = time.Minute * 2
-
 func init() {
 	for _, m := range tstun.WireMTUsToProbe {
 		mtuProbePingSizesV4 = append(mtuProbePingSizesV4, pktLenToPingSize(m, false))
@@ -85,7 +80,6 @@ type endpoint struct {
 	lastSendAny               mono.Time      // last time there were outgoing packets sent this peer from any trigger, internal or external to magicsock
 	lastFullPing              mono.Time      // last time we pinged all disco or wireguard only endpoints
 	lastUDPRelayPathDiscovery mono.Time      // last time we ran UDP relay path discovery
-	lastDiscoKeyAdvertisement mono.Time      // last time we sent a TSMPDiscoAdvertisement or not to this endpoint
 	derpAddr                  netip.AddrPort // fallback/bootstrap path, if non-zero (non-zero for well-behaved clients)
 
 	bestAddr           addrQuality // best non-DERP path; zero if none; mutate via setBestAddrLocked()
@@ -102,7 +96,7 @@ type endpoint struct {
 	probeUDPLifetime  *probeUDPLifetime // UDP path lifetime probing; nil if disabled
 
 	expired         bool // whether the node has expired
-	isWireguardOnly bool // whether the endpoint is WireGuard only
+	isWireguardOnly bool // whether the endpoint is WireGuard only. Must not be changed after initializing the endpont.
 	relayCapable    bool // whether the node is capable of speaking via a [tailscale.com/net/udprelay.Server]
 }
 
@@ -1845,8 +1839,6 @@ type addrQuality struct {
 	wireMTU          tstun.WireMTU
 }
 
-func (a addrQuality) isZero() bool { return a == addrQuality{} }
-
 func (a addrQuality) String() string {
 	// TODO(jwhited): consider including relayServerDisco
 	return fmt.Sprintf("%v@%v+%v", a.epAddr, a.latency, a.wireMTU)

wgengine/magicsock/magicsock.go

@@ -181,10 +181,9 @@ type Conn struct {
 	connCtxCancel func()          // closes connCtx
 	donec         <-chan struct{} // connCtx.Done()'s to avoid context.cancelCtx.Done()'s mutex per call
 
-	allocRelayEndpointPub    *eventbus.Publisher[UDPRelayAllocReq]
-	portUpdatePub            *eventbus.Publisher[router.PortUpdate]
-	tsmpDiscoKeyAvailablePub *eventbus.Publisher[NewDiscoKeyAvailable]
-	homeDERPChangedPub       *eventbus.Publisher[HomeDERPChanged]
+	allocRelayEndpointPub *eventbus.Publisher[UDPRelayAllocReq]
+	portUpdatePub         *eventbus.Publisher[router.PortUpdate]
+	homeDERPChangedPub    *eventbus.Publisher[HomeDERPChanged]
 
 	// pconn4 and pconn6 are the underlying UDP sockets used to
 	// send/receive packets for wireguard and other magicsock
@@ -678,7 +677,6 @@ func NewConn(opts Options) (*Conn, error) {
 	c.eventClient = ec
 	c.allocRelayEndpointPub = eventbus.Publish[UDPRelayAllocReq](ec)
 	c.portUpdatePub = eventbus.Publish[router.PortUpdate](ec)
-	c.tsmpDiscoKeyAvailablePub = eventbus.Publish[NewDiscoKeyAvailable](ec)
 	c.homeDERPChangedPub = eventbus.Publish[HomeDERPChanged](ec)
 	eventbus.SubscribeFunc(ec, c.onPortMapChanged)
 	eventbus.SubscribeFunc(ec, c.onUDPRelayAllocResp)
@@ -1254,8 +1252,7 @@ func (c *Conn) DiscoPublicKey() key.DiscoPublic {
 
 // RotateDiscoKey generates a new discovery key pair and updates the connection
 // to use it. This invalidates all existing disco sessions and will cause peers
-// to re-establish discovery sessions with the new key. Addtionally, the
-// lastTSMPDiscoAdvertisement on all endpoints is reset to 0.
+// to re-establish discovery sessions with the new key.
 //
 // This is primarily for debugging and testing purposes, a future enhancement
 // should provide a mechanism for seamless rotation by supporting short term use
@@ -1269,11 +1266,6 @@ func (c *Conn) RotateDiscoKey() {
 	newShort := c.discoAtomic.Short()
 	c.discoInfo = make(map[key.DiscoPublic]*discoInfo)
 	connCtx := c.connCtx
-	for _, endpoint := range c.peerMap.byEpAddr {
-		endpoint.ep.mu.Lock()
-		endpoint.ep.lastDiscoKeyAdvertisement = 0
-		endpoint.ep.mu.Unlock()
-	}
 	c.mu.Unlock()
 
 	c.logf("magicsock: rotated disco key from %v to %v", oldShort, newShort)
@@ -2678,8 +2670,6 @@ func (c *Conn) enqueueCallMeMaybe(derpAddr netip.AddrPort, de *endpoint) {
 		return
 	}
 
-	c.maybeSendTSMPDiscoAdvert(de)
-
 	eps := make([]netip.AddrPort, 0, len(c.lastEndpoints))
 	for _, ep := range c.lastEndpoints {
 		eps = append(eps, ep.Addr)
@@ -4325,6 +4315,8 @@ type discoInfo struct {
 	metricTSMPDiscoKeyAdvertisementReceived  = clientmetric.NewCounter("magicsock_tsmp_disco_key_advertisement_received")
 	metricTSMPDiscoKeyAdvertisementApplied   = clientmetric.NewCounter("magicsock_tsmp_disco_key_advertisement_applied")
 	metricTSMPDiscoKeyAdvertisementUnchanged = clientmetric.NewCounter("magicsock_tsmp_disco_key_advertisement_unchanged")
+	metricTSMPDiscoKeyAdvertisementSent      = clientmetric.NewCounter("magicsock_tsmp_disco_key_advertisement_sent")
+	metricTSMPDiscoKeyAdvertisementError     = clientmetric.NewCounter("magicsock_tsmp_disco_key_advertisement_error")
 
 	// Counters for peer contacts established using cached network map data.
 	metricCachedPeerContactDERP   = clientmetric.NewCounter("magicsock_cached_peer_contact_derp")
@@ -4511,7 +4503,7 @@ func (c *Conn) HandleDiscoKeyAdvertisement(node tailcfg.NodeView, update packet.
 		return
 	}
 
-	c.logf("magicsock: received disco key update %v from %v", discoKey.ShortString(), node.StableID())
+	c.logf("[v1] magicsock: received disco key update %v from %v", discoKey.ShortString(), node.StableID())
 	metricTSMPDiscoKeyAdvertisementReceived.Add(1)
 
 	c.mu.Lock()
@@ -4531,7 +4523,7 @@ func (c *Conn) HandleDiscoKeyAdvertisement(node tailcfg.NodeView, update packet.
 	// If the key did not change, count it and return.
 	if oldDiscoKey.Compare(discoKey) == 0 {
 		metricTSMPDiscoKeyAdvertisementUnchanged.Add(1)
-		c.logf("magicsock: disco key did not change for node %v", nodeKey.ShortString())
+		c.logf("[v1] magicsock: disco key did not change for node %v", nodeKey.ShortString())
 		return
 	}
 	c.discoInfoForKnownPeerLocked(discoKey)
@@ -4544,67 +4536,69 @@ func (c *Conn) HandleDiscoKeyAdvertisement(node tailcfg.NodeView, update packet.
 	metricTSMPDiscoKeyAdvertisementApplied.Add(1)
 }
 
-// NewDiscoKeyAvailable is an eventbus topic that is emitted when we're sending
-// a packet to a node and observe we haven't told it our current DiscoKey before.
+// PriorityMessageForPeer is a [github.com/tailscale/wireguard-go/device.PeerPriorityMessageFunc]
+// that returns a marshaled plaintext [packet.TSMPDiscoKeyAdvertisement] if
+// nodeKey supports TSMP, otherwise it returns nil.
 //
-// The publisher is magicsock, when we're sending a packet.
-// The subscriber is userspaceEngine, which sends a TSMP packet, also via
-// magicsock. This doesn't recurse infinitely because we only publish it once per
-// DiscoKey.
-// In the common case, a DiscoKey is not rotated within a process generation
-// (as of 2026-01-21), except with debug commands to simulate process restarts.
-//
-// The address is the first node address (tailscale address) of the node. It
-// does not matter if the address is v4/v6, the receiver should handle either.
-//
-// Since we have not yet communicated with the node at the time we are
-// sending this event, the resulting TSMPDiscoKeyAdvertisement will with all
-// likelihood be transmitted via DERP.
-type NewDiscoKeyAvailable struct {
-	NodeFirstAddr netip.Addr
-	NodeID        tailcfg.NodeID
+// This callback must be cheap and must not call back into the
+// [github.com/tailscale/wireguard-go/device.Device]. The returned message must
+// not exceed [github.com/tailscale/wireguard-go/device.MaxPriorityMessageContentSize].
+func (c *Conn) PriorityMessageForPeer(nodeKey key.NodePublic) []byte {
+	disco := c.DiscoPublicKey()
+	if disco.IsZero() {
+		metricTSMPDiscoKeyAdvertisementError.Add(1)
+		return nil
+	}
+
+	c.mu.Lock()
+	self := c.self
+	ep, ok := c.peerMap.endpointForNodeKey(nodeKey)
+	c.mu.Unlock()
+	if !ok || !self.Valid() {
+		metricTSMPDiscoKeyAdvertisementError.Add(1)
+		return nil
+	}
+
+	// Do not send TSMP messages to peers that only speaks wireguard.
+	// The bool is only written once at creation of the endpoint so it is
+	// not necessary to hold the endpoint lock.
+	if ep.isWireguardOnly {
+		return nil
+	}
+
+	ep.mu.Lock()
+	dst := ep.nodeAddr
+	ep.mu.Unlock()
+
+	// Resolve our own Tailscale address in the same family as dst.
+	src := selfIPMatchingFamily(self, dst)
+	if !src.IsValid() {
+		metricTSMPDiscoKeyAdvertisementError.Add(1)
+		return nil
+	}
+
+	tdka := packet.TSMPDiscoKeyAdvertisement{Src: src, Dst: dst, Key: disco}
+	payload, err := tdka.Marshal()
+	if err != nil {
+		metricTSMPDiscoKeyAdvertisementError.Add(1)
+		return nil
+	}
+
+	// The metric is called sent, but since sending the payload is controlled by
+	// wireguard-go, we can only assume it to be sent. Thus this is an estimation
+	// of it being sent based on generation, not the actual time the message has
+	// been sent.
+	metricTSMPDiscoKeyAdvertisementSent.Add(1)
+	return payload
 }
 
-// maybeSendTSMPDiscoAdvert conditionally emits an event indicating that we
-// should send our DiscoKey to the first node address of the magicksock endpoint.
-//
-// The event is suppressed if we are communicating over a non-DERP path, or
-// less than [discoKeyAdvertisementInterval] has passed since the last DiscoKey
-// was sent, or netmap caching is disabled on this node.
-//
-// We do not need the Conn to be locked, but the endpoint should be.
-func (c *Conn) maybeSendTSMPDiscoAdvert(de *endpoint) {
-	if !buildfeatures.HasCacheNetMap || !envknob.BoolDefaultTrue("TS_USE_CACHED_NETMAP") {
-		return
+// selfIPMatchingFamily returns self's first single-IP Tailscale address whose
+// family matches want, or the zero Addr. self must be Valid.
+func selfIPMatchingFamily(self tailcfg.NodeView, want netip.Addr) netip.Addr {
+	for _, p := range self.Addresses().All() {
+		if p.IsSingleIP() && p.Addr().BitLen() == want.BitLen() {
+			return p.Addr()
+		}
 	}
-
-	// Disable TSMP disco advert by default, unless network map caching is
-	// enabled for the local node. Caching network maps on the remote node is
-	// what really matters in terms of handling a TSMP disco advert and applying
-	// it in a useful way, but the TSMP disco advert implementation as it exists
-	// here has pathological behaviors. Therefore, it should be disabled for
-	// almost all tailnets, and we lean on the network map caching control knob
-	// for this purpose. See #20081.
-	if c.controlKnobs == nil || !c.controlKnobs.CacheNetworkMaps.Load() {
-		return
-	}
-
-	de.mu.Lock()
-	defer de.mu.Unlock()
-
-	if !de.nodeAddr.IsValid() {
-		return
-	}
-
-	now := mono.Now()
-	if now.Sub(de.lastDiscoKeyAdvertisement) <= discoKeyAdvertisementInterval ||
-		(!de.lastDiscoKeyAdvertisement.IsZero() && !de.bestAddr.isZero()) {
-		return
-	}
-
-	de.lastDiscoKeyAdvertisement = now
-	c.tsmpDiscoKeyAvailablePub.Publish(NewDiscoKeyAvailable{
-		NodeFirstAddr: de.nodeAddr,
-		NodeID:        de.nodeID,
-	})
+	return netip.Addr{}
 }

wgengine/magicsock/magicsock_test.go

@@ -4709,172 +4709,163 @@ func TestReceiveTSMPDiscoKeyAdvertisement(t *testing.T) {
 	}
 }
 
-func TestSendingTSMPDiscoTimer(t *testing.T) {
-	conn := newTestConn(t)
-	tw := eventbustest.NewWatcher(t, conn.eventBus)
-	t.Cleanup(func() { conn.Close() })
+func TestPriorityMessageForPeer(t *testing.T) {
+	conn := &Conn{}
+	conn.discoAtomic.pair.Store(&discoKeyPair{})
 
-	// maybeSendTSMPDiscoAdvert only advertises when netmap caching is enabled.
-	conn.controlKnobs = new(controlknobs.Knobs)
-	conn.controlKnobs.CacheNetworkMaps.Store(true)
-
-	peerKey := key.NewNode().Public()
-	ep := &endpoint{
-		nodeID:    1,
-		publicKey: peerKey,
-		nodeAddr:  netip.MustParseAddr("100.64.0.1"),
+	// Test early return when self key is zero.
+	if res := conn.PriorityMessageForPeer(key.NewNode().Public()); res != nil {
+		t.Errorf("expected nil, got %v", res)
 	}
+
+	conn = newTestConn(t)
+	conn.SetPrivateKey(key.NewNode())
+
+	selfNode := (&tailcfg.Node{
+		ID: 0,
+		Addresses: []netip.Prefix{
+			netip.MustParsePrefix("fd7a:115c:a1e0::/128"),
+		},
+	}).View()
+	conn.mu.Lock()
+	conn.self = selfNode
+	conn.mu.Unlock()
+
+	nodeID := tailcfg.NodeID(1)
+
+	ip4 := netip.MustParseAddr("100.64.0.1")
+	ep := &endpoint{
+		nodeID:    nodeID,
+		publicKey: key.NewNode().Public(),
+		nodeAddr:  ip4,
+	}
+
 	discoKey := key.NewDisco().Public()
 	ep.disco.Store(&endpointDisco{
 		key:   discoKey,
 		short: discoKey.ShortString(),
 	})
+
 	ep.c = conn
+
+	// Test the EP missing from the peerMap.
+	if res := conn.PriorityMessageForPeer(ep.publicKey); res != nil {
+		t.Errorf("expected nil, got %v", res)
+	}
+
 	conn.mu.Lock()
-	nodeView := (&tailcfg.Node{
-		Key: ep.publicKey,
-		Addresses: []netip.Prefix{
-			netip.MustParsePrefix("100.64.0.1/32"),
-		},
-	}).View()
-	conn.peersByID = map[tailcfg.NodeID]tailcfg.NodeView{nodeView.ID(): nodeView}
+	conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
 	conn.mu.Unlock()
 
-	conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
+	// Test isWireguardOnly.
+	// It is OK for us to modify the endpoint unsynchronized here, because
+	// the callback is not running concurrently.
+	ep.isWireguardOnly = true
+	if res := conn.PriorityMessageForPeer(ep.publicKey); res != nil {
+		t.Errorf("expected nil, got %v", res)
+	}
+	ep.isWireguardOnly = false
 
-	if ep.discoShort() != discoKey.ShortString() {
-		t.Errorf("Original disco key %s, does not match %s", discoKey.ShortString(), ep.discoShort())
+	// Test address family mismatch.
+	if res := conn.PriorityMessageForPeer(ep.publicKey); res != nil {
+		t.Errorf("expected nil, got %v", res)
 	}
 
-	// Only one gets through, second is rate limited.
-	conn.maybeSendTSMPDiscoAdvert(ep)
-	conn.maybeSendTSMPDiscoAdvert(ep)
-	if err := eventbustest.ExpectExactly(tw, eventbustest.Type[NewDiscoKeyAvailable]()); err != nil {
-		t.Errorf("expected only one event, got: %s", err)
+	selfNode = (&tailcfg.Node{
+		ID: 0,
+		Addresses: []netip.Prefix{
+			netip.MustParsePrefix("100.64.0.0/32"),
+			netip.MustParsePrefix("fd7a:115c:a1e0::/128"),
+		},
+	}).View()
+	conn.mu.Lock()
+	conn.self = selfNode
+	conn.mu.Unlock()
+
+	// Test successful message.
+	expected, err := (&packet.TSMPDiscoKeyAdvertisement{
+		Src: netip.MustParseAddr("100.64.0.0"),
+		Dst: netip.MustParseAddr("100.64.0.1"),
+		Key: conn.DiscoPublicKey(),
+	}).Marshal()
+	if err != nil {
+		t.Fatalf("Failed to marshal expected packet: %v", err)
 	}
-
-	// Reset to get the event firing again.
-	ep.mu.Lock()
-	ep.lastDiscoKeyAdvertisement = 0
-	ep.mu.Unlock()
-	conn.maybeSendTSMPDiscoAdvert(ep)
-	if err := eventbustest.Expect(tw, eventbustest.Type[NewDiscoKeyAvailable]()); err != nil {
-		t.Errorf("expected only one event, got: %s", err)
-	}
-
-	// With a direct bestAddr and a non-zero lastDiscoKeyAdvertisement past the
-	// rate-limit interval. No advert should be sent due to the active bestAddr.
-	ep.mu.Lock()
-	ep.lastDiscoKeyAdvertisement = mono.Now().Add(-discoKeyAdvertisementInterval - time.Second)
-	ep.bestAddr = addrQuality{epAddr: epAddr{ap: netip.MustParseAddrPort("1.2.3.4:567")}}
-	ep.mu.Unlock()
-	conn.maybeSendTSMPDiscoAdvert(ep)
-
-	// Simulating restart should send an advert.
-	ep.mu.Lock()
-	ep.lastDiscoKeyAdvertisement = 0
-	ep.mu.Unlock()
-	conn.maybeSendTSMPDiscoAdvert(ep)
-	if err := eventbustest.ExpectExactly(tw, eventbustest.Type[NewDiscoKeyAvailable]()); err != nil {
-		t.Errorf("expected only one event, got: %s", err)
+	res := conn.PriorityMessageForPeer(ep.publicKey)
+	if !slices.Equal(res, expected) {
+		t.Errorf("expected \n%v, got \n%v", expected, res)
 	}
 }
 
-// TestSendingTSMPDiscoCachingDisabled verifies that maybeSendTSMPDiscoAdvert
-// early-returns (sends no advert) when netmap caching is not enabled via the
-// CacheNetworkMaps control knob, including when no knobs are present at all.
-func TestSendingTSMPDiscoCachingDisabled(t *testing.T) {
-	tests := []struct {
-		name  string
-		knobs *controlknobs.Knobs
-	}{
-		{name: "no-knobs", knobs: nil},
-		// Knobs present but CacheNetworkMaps left at its false default.
-		{name: "caching-disabled", knobs: new(controlknobs.Knobs)},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			conn := newTestConn(t)
-			t.Cleanup(func() { conn.Close() })
-			conn.controlKnobs = tt.knobs
+func BenchmarkPriorityMessageForPeer(b *testing.B) {
+	// Can test up to 2^16 nodes given the address generation.
+	nodeCount := []int{10, 10000}
 
-			ep := &endpoint{
-				nodeID:    1,
-				publicKey: key.NewNode().Public(),
-				nodeAddr:  netip.MustParseAddr("100.64.0.1"),
+	for _, tt := range nodeCount {
+		b.Run(fmt.Sprintf("%d_nodes", tt), func(b *testing.B) {
+			conn := newTestConn(b)
+			conn.SetPrivateKey(key.NewNode())
+			peersByID := make(map[tailcfg.NodeID]tailcfg.NodeView, tt)
+			var targetKey key.NodePublic
+
+			selfNode := (&tailcfg.Node{
+				ID: 0,
+				Addresses: []netip.Prefix{
+					netip.MustParsePrefix("100.64.0.0/32"),
+					netip.MustParsePrefix("fd7a:115c:a1e0::/128"),
+				},
+			}).View()
+			conn.mu.Lock()
+			conn.self = selfNode
+			conn.mu.Unlock()
+
+			for i := range tt {
+				nodeID := tailcfg.NodeID(i + 1)
+				nodeKey := key.NewNode().Public()
+				if i == 0 {
+					targetKey = nodeKey
+				}
+
+				addrIdx := i + 1
+				ip4 := netip.AddrFrom4([4]byte{100, 64, byte(addrIdx >> 8), byte(addrIdx)})
+				ip6 := netip.AddrFrom16([16]byte{
+					0xfd, 0x7a, 0x11, 0x5c, 0xa1, 0xe0,
+					0, 0, 0, 0, 0, 0, 0, 0, byte(addrIdx >> 8), byte(addrIdx),
+				})
+				ep := &endpoint{
+					nodeID:    nodeID,
+					publicKey: nodeKey,
+					nodeAddr:  ip4,
+				}
+
+				discoKey := key.NewDisco().Public()
+				ep.disco.Store(&endpointDisco{
+					key:   discoKey,
+					short: discoKey.ShortString(),
+				})
+
+				ep.c = conn
+				nodeView := (&tailcfg.Node{
+					ID:  1,
+					Key: ep.publicKey,
+					Addresses: []netip.Prefix{
+						netip.PrefixFrom(ip4, 32),
+						netip.PrefixFrom(ip6, 128),
+					},
+				}).View()
+				peersByID[nodeID] = nodeView
+				conn.mu.Lock()
+				conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
+				conn.mu.Unlock()
 			}
-			ep.c = conn
 
-			// A fresh endpoint with a zero lastDiscoKeyAdvertisement and no
-			// direct bestAddr would otherwise advertise; the only thing
-			// suppressing it here is the disabled caching knob. On early
-			// return the timestamp is left untouched (zero).
-			conn.maybeSendTSMPDiscoAdvert(ep)
+			conn.mu.Lock()
+			conn.peersByID = peersByID
+			conn.mu.Unlock()
 
-			ep.mu.Lock()
-			defer ep.mu.Unlock()
-			if !ep.lastDiscoKeyAdvertisement.IsZero() {
-				t.Errorf("lastDiscoKeyAdvertisement = %v; want zero (advert should have been suppressed)", ep.lastDiscoKeyAdvertisement)
+			for b.Loop() {
+				conn.PriorityMessageForPeer(targetKey)
 			}
 		})
 	}
 }
-
-// TestSendingTSMPDiscoPeerRelaySuppressed verifies that maybeSendTSMPDiscoAdvert
-// suppresses the advert when the bestAddr is a peer relay path (a non-zero
-// addrQuality whose epAddr has a VNI set), even though such a path is not
-// direct. Suppression is observed via lastDiscoKeyAdvertisement remaining
-// unchanged, since a fired advert would overwrite it with the current time.
-func TestSendingTSMPDiscoPeerRelaySuppressed(t *testing.T) {
-	conn := newTestConn(t)
-	t.Cleanup(func() { conn.Close() })
-
-	// maybeSendTSMPDiscoAdvert only advertises when netmap caching is enabled.
-	conn.controlKnobs = new(controlknobs.Knobs)
-	conn.controlKnobs.CacheNetworkMaps.Store(true)
-
-	peerKey := key.NewNode().Public()
-	ep := &endpoint{
-		nodeID:    1,
-		publicKey: peerKey,
-		nodeAddr:  netip.MustParseAddr("100.64.0.1"),
-	}
-	discoKey := key.NewDisco().Public()
-	ep.disco.Store(&endpointDisco{
-		key:   discoKey,
-		short: discoKey.ShortString(),
-	})
-	ep.c = conn
-	conn.mu.Lock()
-	nodeView := (&tailcfg.Node{
-		Key: ep.publicKey,
-		Addresses: []netip.Prefix{
-			netip.MustParsePrefix("100.64.0.1/32"),
-		},
-	}).View()
-	conn.peersByID = map[tailcfg.NodeID]tailcfg.NodeView{nodeView.ID(): nodeView}
-	conn.mu.Unlock()
-
-	conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
-
-	// A peer relay bestAddr: an epAddr with a VNI set. It is past the
-	// rate-limit interval with a non-zero lastDiscoKeyAdvertisement, so the
-	// only thing suppressing the advert is the active (non-zero) bestAddr.
-	var vni packet.VirtualNetworkID
-	vni.Set(7)
-	lastAdvert := mono.Now().Add(-discoKeyAdvertisementInterval - time.Second)
-	ep.mu.Lock()
-	ep.lastDiscoKeyAdvertisement = lastAdvert
-	ep.bestAddr = addrQuality{epAddr: epAddr{ap: netip.MustParseAddrPort("1.2.3.4:567"), vni: vni}}
-	ep.mu.Unlock()
-
-	conn.maybeSendTSMPDiscoAdvert(ep)
-
-	// A fired advert would have overwritten lastDiscoKeyAdvertisement with the
-	// current time; confirm it was left untouched, indicating suppression.
-	ep.mu.Lock()
-	defer ep.mu.Unlock()
-	if ep.lastDiscoKeyAdvertisement != lastAdvert {
-		t.Errorf("lastDiscoKeyAdvertisement = %v; want unchanged %v (advert should have been suppressed)", ep.lastDiscoKeyAdvertisement, lastAdvert)
-	}
-}

wgengine/userspace.go

@@ -52,7 +52,6 @@
 	"tailscale.com/util/eventbus"
 	"tailscale.com/util/execqueue"
 	"tailscale.com/util/mak"
-	"tailscale.com/util/singleflight"
 	"tailscale.com/util/testenv"
 	"tailscale.com/util/usermetric"
 	"tailscale.com/version"
@@ -597,7 +596,7 @@ func NewUserspaceEngine(logf logger.Logf, conf Config) (_ Engine, reterr error)
 		e.linkChangeQueue.Add(func() { e.linkChange(&cd) })
 	})
 	eventbus.SubscribeFunc(ec, func(update events.PeerDiscoKeyUpdate) {
-		e.logf("wgengine: got TSMP disco key advertisement from %v via eventbus", update.Src)
+		e.logf("[v1] wgengine: got TSMP disco key advertisement from %v via eventbus", update.Src)
 		if e.magicConn == nil {
 			e.logf("wgengine: no magicConn")
 			return
@@ -613,17 +612,6 @@ func NewUserspaceEngine(logf logger.Logf, conf Config) (_ Engine, reterr error)
 		}
 		e.magicConn.HandleDiscoKeyAdvertisement(peer.Node, pkt)
 	})
-	var tsmpRequestGroup singleflight.Group[netip.Addr, struct{}]
-	eventbus.SubscribeFunc(ec, func(req magicsock.NewDiscoKeyAvailable) {
-		if !req.NodeFirstAddr.IsValid() {
-			return
-		}
-		go tsmpRequestGroup.Do(req.NodeFirstAddr, func() (struct{}, error) {
-			e.sendTSMPDiscoAdvertisement(req.NodeFirstAddr)
-			e.logf("wgengine: sending TSMP disco key advertisement to %v", req.NodeFirstAddr)
-			return struct{}{}, nil
-		})
-	})
 	e.eventClient = ec
 	e.logf("Engine created.")
 	return e, nil
@@ -758,6 +746,23 @@ func (e *userspaceEngine) SetPeerSessionStateFunc(fn func(key.NodePublic, PeerWi
 	})
 }
 
+// SetPeerPriorityMessageOnEstablishmentFunc registers a callback with a
+// [github.com/tailscale/wireguard-go/device] to be sent on session establishement.
+// This establishment happens at every wireguard rekey event.
+//
+// This callback must be cheap and must not call back into the
+// [github.com/tailscale/wireguard-go/device.Device]. The returned message must
+// not exceed [github.com/tailscale/wireguard-go/device.MaxPriorityMessageContentSize].
+func (e *userspaceEngine) SetPeerPriorityMessageOnEstablishmentFunc(fn func(key.NodePublic) (msg []byte)) {
+	if fn != nil {
+		e.wgdev.SetPriorityMessageOnEstablishmentFunc(func(pk device.NoisePublicKey) (msg []byte) {
+			return fn(key.NodePublicFromRaw32(mem.B(pk[:])))
+		})
+	} else {
+		e.wgdev.SetPriorityMessageOnEstablishmentFunc(nil)
+	}
+}
+
 // SetNetLogSource installs the [NetLogSource] consulted by the engine's
 // network flow logger.
 func (e *userspaceEngine) SetNetLogSource(src NetLogSource) {
@@ -1274,7 +1279,6 @@ func (e *userspaceEngine) Ping(ip netip.Addr, pingType tailcfg.PingType, size in
 		e.magicConn.Ping(peer, res, size, cb)
 	case "TSMP":
 		e.sendTSMPPing(ip, peer, res, cb)
-		e.sendTSMPDiscoAdvertisement(ip)
 	case "ICMP":
 		e.sendICMPEchoRequest(ip, peer, res, cb)
 	}
@@ -1395,29 +1399,6 @@ func (e *userspaceEngine) sendTSMPPing(ip netip.Addr, peer tailcfg.NodeView, res
 	e.tundev.InjectOutbound(tsmpPing)
 }
 
-func (e *userspaceEngine) sendTSMPDiscoAdvertisement(ip netip.Addr) {
-	srcIP, err := e.mySelfIPMatchingFamily(ip)
-	if err != nil {
-		e.logf("getting matching node: %s", err)
-		return
-	}
-	tdka := packet.TSMPDiscoKeyAdvertisement{
-		Src: srcIP,
-		Dst: ip,
-		Key: e.magicConn.DiscoPublicKey(),
-	}
-	payload, err := tdka.Marshal()
-	if err != nil {
-		e.logf("error generating TSMP Advertisement: %s", err)
-		metricTSMPDiscoKeyAdvertisementError.Add(1)
-	} else if err := e.tundev.InjectOutbound(payload); err != nil {
-		e.logf("error sending TSMP Advertisement: %s", err)
-		metricTSMPDiscoKeyAdvertisementError.Add(1)
-	} else {
-		metricTSMPDiscoKeyAdvertisementSent.Add(1)
-	}
-}
-
 func (e *userspaceEngine) setTSMPPongCallback(data [8]byte, cb func(packet.TSMPPongReply)) {
 	e.mu.Lock()
 	defer e.mu.Unlock()
@@ -1554,9 +1535,6 @@ func (ls fwdDNSLinkSelector) PickLink(ip netip.Addr) (linkName string) {
 
 	metricNumMajorChanges = clientmetric.NewCounter("wgengine_major_changes")
 	metricNumMinorChanges = clientmetric.NewCounter("wgengine_minor_changes")
-
-	metricTSMPDiscoKeyAdvertisementSent  = clientmetric.NewCounter("magicsock_tsmp_disco_key_advertisement_sent")
-	metricTSMPDiscoKeyAdvertisementError = clientmetric.NewCounter("magicsock_tsmp_disco_key_advertisement_error")
 )
 
 func (e *userspaceEngine) InstallCaptureHook(cb packet.CaptureCallback) {

wgengine/userspace_test.go

@@ -334,57 +334,6 @@ func TestUserspaceEnginePeerMTUReconfig(t *testing.T) {
 	}
 }
 
-func TestTSMPKeyAdvertisement(t *testing.T) {
-	var knobs controlknobs.Knobs
-
-	bus := eventbustest.NewBus(t)
-	ht := health.NewTracker(bus)
-	reg := new(usermetric.Registry)
-	e, err := NewFakeUserspaceEngine(t.Logf, 0, &knobs, ht, reg, bus)
-	if err != nil {
-		t.Fatal(err)
-	}
-	t.Cleanup(e.Close)
-	ue := e.(*userspaceEngine)
-	routerCfg := &router.Config{}
-	nodeKey := nkFromHex("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb")
-	nm := &netmap.NetworkMap{
-		Peers: nodeViews([]*tailcfg.Node{
-			{
-				ID:  1,
-				Key: nodeKey,
-			},
-		}),
-		SelfNode: (&tailcfg.Node{
-			StableID:  "TESTCTRL00000001",
-			Name:      "test-node.test.ts.net",
-			Addresses: []netip.Prefix{netip.MustParsePrefix("100.64.0.1/32"), netip.MustParsePrefix("fd7a:115c:a1e0:ab12:4843:cd96:0:1/128")},
-		}).View(),
-	}
-	cfg := &wgcfg.Config{
-		Addresses: nm.SelfNode.Addresses().AsSlice(),
-	}
-
-	ue.SetSelfNode(nm.SelfNode)
-	err = ue.Reconfig(cfg, routerCfg, &dns.Config{})
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	addr := netip.MustParseAddr("100.100.99.1")
-	previousValue := metricTSMPDiscoKeyAdvertisementSent.Value()
-	ue.sendTSMPDiscoAdvertisement(addr)
-	if val := metricTSMPDiscoKeyAdvertisementSent.Value(); val <= previousValue {
-		errs := metricTSMPDiscoKeyAdvertisementError.Value()
-		t.Errorf("Expected 1 disco key advert, got %d, errors %d", val, errs)
-	}
-	// Remove config to have the engine shut down more consistently
-	err = ue.Reconfig(&wgcfg.Config{}, &router.Config{}, &dns.Config{})
-	if err != nil {
-		t.Fatal(err)
-	}
-}
-
 func nkFromHex(hex string) key.NodePublic {
 	if len(hex) != 64 {
 		panic(fmt.Sprintf("%q is len %d; want 64", hex, len(hex)))

wgengine/wgengine.go

@@ -255,6 +255,8 @@ type Engine interface {
 	// LocalBackend construction.
 	SetPeerSessionStateFunc(func(key.NodePublic, PeerWireGuardState))
 
+	SetPeerPriorityMessageOnEstablishmentFunc(fn func(key.NodePublic) (msg []byte))
+
 	// ProbeLocks acquires and releases the engine's internal locks so
 	// that [ipnlocal.LocalBackend]'s watchdog can detect deadlocks in
 	// the engine. It is otherwise a no-op.