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

@@ -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
+// A [PeerAwareEndpoint] may be installed as the [conn.Endpoint] following
-// [Peer].
+// 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

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
 	}

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 {