From 1d0488a3d7da6b6ed79202519f30e7a286e0d4e6 Mon Sep 17 00:00:00 2001
From: Jordan Whited <jordan@tailscale.com>
Date: Wed, 16 Jul 2025 08:43:20 -0700
Subject: [PATCH] conn,device: eval conn.PeerAwareEndpoint per-packet

Peer.SetEndpointFromPacket is not called per-packet. It is
guaranteed to be called at least once per packet batch.

Updates tailscale/corp#30042
Updates tailscale/corp#20732

Signed-off-by: Jordan Whited <jordan@tailscale.com>
---
 conn/conn.go      | 5 +++--
 device/peer.go    | 4 ----
 device/receive.go | 3 +++
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/conn/conn.go b/conn/conn.go
index b949641..f178161 100644
--- a/conn/conn.go
+++ b/conn/conn.go
@@ -109,8 +109,9 @@ type InitiationAwareEndpoint interface {
 // to learn the identification WireGuard can derive from the session
 // or handshake.
 //
-// wireguard-go never installs a [PeerAwareEndpoint] as the [Endpoint] for a
-// [Peer].
+// A [PeerAwareEndpoint] may be installed as the [conn.Endpoint] following
+// successful decryption unless endpoint roaming has been disabled for
+// the peer.
 type PeerAwareEndpoint interface {
 	// FromPeer is called at least once per successfully Cryptokey Routing ID'd
 	// [ReceiveFunc] packets batch for a given node key. wireguard-go will
diff --git a/device/peer.go b/device/peer.go
index c188c31..064feb2 100644
--- a/device/peer.go
+++ b/device/peer.go
@@ -282,10 +282,6 @@ func (peer *Peer) Stop() {
 func (peer *Peer) SetEndpointFromPacket(endpoint conn.Endpoint) {
 	peer.endpoint.Lock()
 	defer peer.endpoint.Unlock()
-	if ep, ok := endpoint.(conn.PeerAwareEndpoint); ok {
-		ep.FromPeer(peer.handshake.remoteStatic)
-		return
-	}
 	if peer.endpoint.disableRoaming {
 		return
 	}
diff --git a/device/receive.go b/device/receive.go
index e74de1a..02c8f21 100644
--- a/device/receive.go
+++ b/device/receive.go
@@ -460,6 +460,9 @@ func (peer *Peer) RoutineSequentialReceiver(maxBatchSize int) {
 				peer.timersHandshakeComplete()
 				peer.SendStagedPackets()
 			}
+			if ep, ok := elem.endpoint.(conn.PeerAwareEndpoint); ok {
+				ep.FromPeer(peer.handshake.remoteStatic)
+			}
 			rxBytesLen += uint64(len(elem.packet) + MinMessageSize)
 
 			if len(elem.packet) == 0 {