From 2f5d148bcfe13a65d0f85dca8698200543c226da Mon Sep 17 00:00:00 2001
From: James Tucker <james@tailscale.com>
Date: Fri, 7 Jun 2024 16:57:40 -0700
Subject: [PATCH] conn,device: enable cryptorouting via PeerAwareEndpoint

Introduce an optional extension point for Endpoint that enables a path
for WireGuard to inform an integration about the peer public key that is
associated with an Endpoint.

The API is expected to return either the same or a new Endpoint in
response to this function. A future version of this patch could
potentially remove the returned Endpoint, but would require larger
integrator changes downstream.

This adds a small per-packet cost that could later be removed with a
larger refactor of the wireguard-go interface and Tailscale magicsock
code, as well as introducing a generic bound for Endpoint in a device &
bind instance.

Updates tailscale/corp#20732
---
 conn/conn.go             | 14 ++++++++++++++
 device/noise-protocol.go |  2 +-
 device/peer.go           |  3 +++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/conn/conn.go b/conn/conn.go
index a1f57d2..8df5aaa 100644
--- a/conn/conn.go
+++ b/conn/conn.go
@@ -84,6 +84,20 @@ type Endpoint interface {
 	SrcIP() netip.Addr
 }
 
+// PeerAwareEndpoint is an optional Endpoint specialization for
+// integrations that want to know about the outcome of cryptorouting
+// identification.
+//
+// If they receive a packet from a source they had not pre-identified,
+// to learn the identification WireGuard can derive from the session
+// or handshake.
+//
+// If GetPeerEndpoint returns nil, WireGuard will be unable to respond
+// to the peer until a new endpoint is written by a later packet.
+type PeerAwareEndpoint interface {
+	GetPeerEndpoint(peerPublicKey [32]byte) Endpoint
+}
+
 var (
 	ErrBindAlreadyOpen   = errors.New("bind is already open")
 	ErrWrongEndpointType = errors.New("endpoint type does not correspond with bind type")
diff --git a/device/noise-protocol.go b/device/noise-protocol.go
index 9f2ba50..2d8f984 100644
--- a/device/noise-protocol.go
+++ b/device/noise-protocol.go
@@ -124,7 +124,7 @@ type Handshake struct {
 	localEphemeral            NoisePrivateKey          // ephemeral secret key
 	localIndex                uint32                   // used to clear hash-table
 	remoteIndex               uint32                   // index for sending
-	remoteStatic              NoisePublicKey           // long term key
+	remoteStatic              NoisePublicKey           // long term key, never changes, can be accessed without mutex
 	remoteEphemeral           NoisePublicKey           // ephemeral public key
 	precomputedStaticStatic   [NoisePublicKeySize]byte // precomputed shared secret
 	lastTimestamp             tai64n.Timestamp
diff --git a/device/peer.go b/device/peer.go
index 89b719b..876e5da 100644
--- a/device/peer.go
+++ b/device/peer.go
@@ -283,6 +283,9 @@ func (peer *Peer) SetEndpointFromPacket(endpoint conn.Endpoint) {
 		return
 	}
 	peer.endpoint.clearSrcOnTx = false
+	if ep, ok := endpoint.(conn.PeerAwareEndpoint); ok {
+		endpoint = ep.GetPeerEndpoint(peer.handshake.remoteStatic)
+	}
 	peer.endpoint.val = endpoint
 }