wgengine/magicsock: trigger TSMP disco key advert on wg handshake resp

Both inbound via DERP, and outbound via any connection type. Gating to
handshake response reduces the chances we end up in a closed packet
scheduling loop between wireguard-go and magicsock. We trigger in both
directions since initiator is one-sided for the lifetime of the
WireGuard "session".

Updates #20081

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

wgengine/magicsock/derp.go

@@ -769,6 +769,12 @@ func (c *Conn) processDERPReadResult(dm derpReadResult, b []byte) (n int, ep *en
 		update(0, netip.AddrPortFrom(ep.nodeAddr, 0), srcAddr.ap, 1, dm.n, true)
 	}
 
+	if looksLikeHandshakeResponse(b[:n]) {
+		ep.mu.Lock()
+		ep.maybeSendTSMPDiscoAdvertLocked()
+		ep.mu.Unlock()
+	}
+
 	c.metrics.inboundPacketsDERPTotal.Add(1)
 	c.metrics.inboundBytesDERPTotal.Add(int64(n))
 	return n, ep

wgengine/magicsock/endpoint.go

@@ -20,9 +20,12 @@
 	"sync/atomic"
 	"time"
 
+	"github.com/tailscale/wireguard-go/device"
 	"golang.org/x/net/ipv4"
 	"golang.org/x/net/ipv6"
 	"tailscale.com/disco"
+	"tailscale.com/envknob"
+	"tailscale.com/feature/buildfeatures"
 	"tailscale.com/ipn/ipnstate"
 	"tailscale.com/net/packet"
 	"tailscale.com/net/stun"
@@ -40,10 +43,9 @@
 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
+// discoKeyAdvertisementInterval gates how frequently a disco update via TSMP
+// can happen. The update is triggered via outbound WireGuard handshake response.
+const discoKeyAdvertisementInterval = device.RekeyAfterTime*2 + sessionActiveTimeout
 
 func init() {
 	for _, m := range tstun.WireMTUsToProbe {
@@ -1056,6 +1058,15 @@ func (de *endpoint) send(buffs [][]byte, offset int) error {
 		return errExpired
 	}
 
+	if len(buffs) == 1 && looksLikeHandshakeResponse(buffs[0][offset:]) {
+		// We hook TSMP disco advert around outbound (and inbound via DERP)
+		// WireGuard handshake response as it's synonymous with handshake
+		// completion. TSMP sits on the other side of wireguard-go, so we must
+		// be careful as to not close a packet scheduling loop between
+		// wireguard-go and magicsock.
+		de.maybeSendTSMPDiscoAdvertLocked()
+	}
+
 	now := mono.Now()
 	udpAddr, derpAddr, startWGPing := de.addrForSendLocked(now)
 
@@ -2106,3 +2117,31 @@ func (de *endpoint) setDERPHome(regionID uint16) {
 		de.c.relayManager.handleDERPHomeChange(de.publicKey, regionID)
 	}
 }
+
+// maybeSendTSMPDiscoAdvertLocked conditionally emits an event indicating that we
+// should send our DiscoKey to the first node address of the [endpoint].
+//
+// The event is suppressed if we are communicating with de over a direct
+// connection, or it has been less than [discoKeyAdvertisementInterval] since
+// the previous emission.
+func (de *endpoint) maybeSendTSMPDiscoAdvertLocked() {
+	if !buildfeatures.HasCacheNetMap || !envknob.BoolDefaultTrue("TS_USE_CACHED_NETMAP") {
+		return
+	}
+
+	if !de.nodeAddr.IsValid() {
+		return
+	}
+
+	now := mono.Now()
+	if now.Sub(de.lastDiscoKeyAdvertisement) <= discoKeyAdvertisementInterval ||
+		(!de.lastDiscoKeyAdvertisement.IsZero() && de.bestAddr.isDirect()) {
+		return
+	}
+
+	de.lastDiscoKeyAdvertisement = now
+	de.c.tsmpDiscoKeyAvailablePub.Publish(NewDiscoKeyAvailable{
+		NodeFirstAddr: de.nodeAddr,
+		NodeID:        de.nodeID,
+	})
+}

wgengine/magicsock/magicsock.go

@@ -1815,6 +1815,11 @@ func (c *Conn) mkReceiveFunc(ruc *RebindingUDPConn, healthItem *health.ReceiveFu
 	}
 }
 
+func looksLikeHandshakeResponse(b []byte) bool {
+	return len(b) == device.MessageResponseSize &&
+		binary.LittleEndian.Uint32(b) == device.MessageResponseType
+}
+
 // looksLikeInitiationMsg returns true if b looks like a WireGuard initiation
 // message, otherwise it returns false.
 func looksLikeInitiationMsg(b []byte) bool {
@@ -2675,8 +2680,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)
@@ -4523,34 +4526,3 @@ type NewDiscoKeyAvailable struct {
 	NodeFirstAddr netip.Addr
 	NodeID        tailcfg.NodeID
 }
-
-// maybeSendTSMPDiscoAdvert conditionally emits an event indicating that we
-// should send our DiscoKey to the first node address of the magicksock endpoint.
-// The event is only emitted if we are not already communicating directly and
-// more than 60 seconds has passed since the last DiscoKey was sent.
-//
-// 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
-	}
-
-	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.isDirect()) {
-		return
-	}
-
-	de.lastDiscoKeyAdvertisement = now
-	c.tsmpDiscoKeyAvailablePub.Publish(NewDiscoKeyAvailable{
-		NodeFirstAddr: de.nodeAddr,
-		NodeID:        de.nodeID,
-	})
-}

wgengine/magicsock/magicsock_test.go

@@ -4544,8 +4544,10 @@ func TestSendingTSMPDiscoTimer(t *testing.T) {
 	}
 
 	// Only one gets through, second is rate limited.
-	conn.maybeSendTSMPDiscoAdvert(ep)
-	conn.maybeSendTSMPDiscoAdvert(ep)
+	ep.mu.Lock()
+	ep.maybeSendTSMPDiscoAdvertLocked()
+	ep.maybeSendTSMPDiscoAdvertLocked()
+	ep.mu.Unlock()
 	if err := eventbustest.ExpectExactly(tw, eventbustest.Type[NewDiscoKeyAvailable]()); err != nil {
 		t.Errorf("expected only one event, got: %s", err)
 	}
@@ -4553,8 +4555,8 @@ func TestSendingTSMPDiscoTimer(t *testing.T) {
 	// Reset to get the event firing again.
 	ep.mu.Lock()
 	ep.lastDiscoKeyAdvertisement = 0
+	ep.maybeSendTSMPDiscoAdvertLocked()
 	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)
 	}
@@ -4564,14 +4566,14 @@ func TestSendingTSMPDiscoTimer(t *testing.T) {
 	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.maybeSendTSMPDiscoAdvertLocked()
 	ep.mu.Unlock()
-	conn.maybeSendTSMPDiscoAdvert(ep)
 
 	// Simulating restart should send an advert.
 	ep.mu.Lock()
 	ep.lastDiscoKeyAdvertisement = 0
+	ep.maybeSendTSMPDiscoAdvertLocked()
 	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)
 	}