ts_derp: factor out transport behavior

The baseline derp client doesn't need to know about translating NodeKeys
to transport `PeerId`s, that's a higher level concern now provided by
`ts_derp::Transport`, which wraps the underlying client (abstracted
behind `NodekeyTransport`) with a `PeerLookup` in order to provide
`UnderlayTransport`. This keeps the examples and other functionality
dependent on bare derp working without needing to bring in a dummy
peer id lookup layer.

Signed-off-by: Nathan Perry <nathan@tailscale.com>
Change-Id: Ib965f787e92880ac3d74c364760acc546a6a6964

Author: npry

Cargo.lock

@@ -3,7 +3,6 @@
 //! Intended to test ping/pong/keepalive.
 
 use ts_keys::NodeKeyPair;
-use ts_transport::UnderlayTransport;
 
 mod common;
 
@@ -16,20 +15,16 @@ async fn main() -> ts_cli_util::Result<()> {
 
     let keypair = NodeKeyPair::new();
 
-    let client =
-        ts_derp::Client::connect(region, &keypair, ts_derp::DummyStaticLookup::default()).await?;
+    let client = ts_derp::Client::connect(region, &keypair).await?;
     tracing::info!("derp handshake done");
 
     loop {
-        for result in client.recv().await {
-            match result {
-                Ok((peer, pkts)) => {
-                    let pkts = pkts.into_iter().collect::<Vec<_>>();
-                    tracing::info!(?peer, ?pkts);
-                }
-                Err(e) => {
-                    tracing::error!(err = %e, "recv");
-                }
+        match client.recv_one().await {
+            Ok((peer, pkt)) => {
+                tracing::info!(?peer, ?pkt);
+            }
+            Err(e) => {
+                tracing::error!(err = %e, "recv");
             }
         }
     }

ts_derp/examples/listen.rs

@@ -3,9 +3,7 @@
 use std::{sync::Arc, time::Duration};
 
 use tokio::task::JoinSet;
-use ts_derp::PeerLookup;
 use ts_keys::NodeKeyPair;
-use ts_transport::UnderlayTransport;
 
 mod common;
 
@@ -18,10 +16,7 @@ async fn main() -> ts_cli_util::Result<()> {
 
     let keypair = NodeKeyPair::new();
 
-    let peer_map = &*Box::leak(Box::new(ts_derp::DummyStaticLookup::default()));
-    let self_id = peer_map.key_to_id(&keypair.public).unwrap();
-
-    let client = ts_derp::Client::connect(region, &keypair, peer_map).await?;
+    let client = ts_derp::Client::connect(region, &keypair).await?;
     tracing::info!("derp handshake done");
 
     let client = Arc::new(client);
@@ -33,7 +28,7 @@ async fn main() -> ts_cli_util::Result<()> {
         let mut ticker = tokio::time::interval(Duration::from_secs(1));
 
         loop {
-            if let Err(e) = pinger.send([(self_id, vec![vec![1].into()])]).await {
+            if let Err(e) = pinger.send_one(keypair.public, &[1]).await {
                 tracing::error!(err = %e, "ping");
             } else {
                 tracing::info!("ping");
@@ -47,10 +42,8 @@ async fn main() -> ts_cli_util::Result<()> {
     js.spawn(async move {
         loop {
             match recv.recv_one().await {
-                Ok((peer_id, pkt)) => {
-                    let peer_key = peer_map.id_to_key(peer_id);
-
-                    tracing::info!(?pkt, %peer_id, ?peer_key, "pong");
+                Ok((peer_key, pkt)) => {
+                    tracing::info!(?pkt, %peer_key, "pong");
                 }
                 Err(e) => {
                     tracing::error!(err = %e, "recv");

ts_derp/examples/ping.rs

@@ -8,27 +8,25 @@ use tokio::{
 };
 use tokio_util::codec::{FramedRead, FramedWrite};
 use ts_http_util::Client as _;
-use ts_keys::NodeKeyPair;
+use ts_keys::{NodeKeyPair, NodePublicKey};
 use ts_packet::PacketMut;
-use ts_transport::{PeerId, UnderlayTransport};
 use url::Url;
 
 use crate::{
-    Error, ServerConnInfo, frame,
+    Error, ServerConnInfo, Transport, frame,
     frame::{ClientInfo, FrameType, PeerGone, Ping, RawFrame, ServerInfo, ServerKey},
-    peer_lookup::PeerLookup,
+    transport::NodekeyTransport,
 };
 
 type DefaultIo = ts_http_util::Upgraded;
 
 /// Type alias for the default derp client over upgraded HTTP on a tokio executor.
-pub type DefaultClient<Lookup> = Client<DefaultIo, Lookup>;
+pub type DefaultClient = Client<DefaultIo>;
 
-/// Asynchronous DERP transport for a single DERP region.
-pub struct Client<Io, Lookup> {
+/// Single-region DERP client.
+pub struct Client<Io> {
     read_conn: Mutex<FramedRead<ReadHalf<Io>, frame::Codec>>,
     write_conn: Mutex<FramedWrite<WriteHalf<Io>, frame::Codec>>,
-    peer_lookup: Lookup,
 }
 
 /// Establish and upgrade a http connection to the derp region.
@@ -56,18 +54,19 @@ pub async fn connect<'c>(
     Ok(Some(upgraded))
 }
 
-impl<Io, Lookup> Client<Io, Lookup>
+impl<Io> Client<Io>
 where
     Io: AsyncRead + AsyncWrite,
-    Lookup: PeerLookup,
 {
+    /// Convert this derp client into a [`ts_transport::UnderlayTransport`] given the
+    /// specified `Lookup`.
+    pub fn into_transport<Lookup>(self, lookup: Lookup) -> Transport<Self, Lookup> {
+        Transport::new(self, lookup)
+    }
+
     /// Perform a derp handshake over the given transport and return a [`Client`].
     #[tracing::instrument(skip_all)]
-    pub async fn handshake(
-        conn: Io,
-        node_keypair: &NodeKeyPair,
-        peer_lookup: Lookup,
-    ) -> Result<Self, Error> {
+    pub async fn handshake(conn: Io, node_keypair: &NodeKeyPair) -> Result<Self, Error> {
         let (read_conn, write_conn) = tokio::io::split(conn);
 
         let mut fw = FramedWrite::new(write_conn, frame::Codec);
@@ -121,10 +120,15 @@ where
         Ok(Self {
             read_conn: Mutex::new(fr),
             write_conn: Mutex::new(fw),
-            peer_lookup,
         })
     }
 
+    /// Send a message to a nodekey on the derp server.
+    pub async fn send_one(&self, node_key: NodePublicKey, msg: &[u8]) -> Result<(), Error> {
+        self.send_frame_with_extra(&frame::SendPacket { dest: node_key }, msg)
+            .await
+    }
+
     /// Send a frame to the derp server.
     pub async fn send_frame(
         &self,
@@ -151,7 +155,7 @@ where
 
     /// Waits for a single data packet from a peer to arrive via this DERP server and returns it.
     /// DERP control messages (KeepAlive, Ping, etc) are handled inline and are not returned.
-    pub async fn recv_one(&self) -> Result<(PeerId, PacketMut), Error> {
+    pub async fn recv_one(&self) -> Result<(NodePublicKey, PacketMut), Error> {
         // DERP exchanges control messages (KeepAlives, Pings, etc) in-band with data messages
         // (SendPacket, RecvPacket, etc). The caller only cares about the payloads of data
         // messages, so we recv_one_raw() in a loop to handle any control messages while waiting
@@ -199,12 +203,8 @@ where
                 }
                 FrameType::RecvPacket => {
                     let (recv, payload) = frame.as_type::<frame::RecvPacket>().unwrap();
-                    let Some(id) = self.peer_lookup.key_to_id(&recv.src) else {
-                        tracing::trace!(src = %recv.src, "no known peer for node key");
-                        continue;
-                    };
 
-                    return Ok((id, payload.into()));
+                    return Ok((recv.src, payload.into()));
                 }
                 t => {
                     return Err(Error::UnexpectedRecvFrameType(t));
@@ -214,29 +214,25 @@ where
     }
 }
 
-impl<Lookup> Client<DefaultIo, Lookup>
-where
-    Lookup: PeerLookup,
-{
+impl Client<DefaultIo> {
     /// Connect to and handshake with the derp server with the given URL over HTTP.
     pub async fn connect<'c>(
         region: impl IntoIterator<Item = &'c ServerConnInfo>,
         node_keypair: &NodeKeyPair,
-        lookup: Lookup,
     ) -> Result<Self, Error> {
         let conn = connect(region).await?.unwrap();
 
-        Client::handshake(conn, node_keypair, lookup).await
+        Client::handshake(conn, node_keypair).await
     }
 }
 
-impl<Io, Lookup> fmt::Debug for Client<Io, Lookup> {
+impl<Io> fmt::Debug for Client<Io> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt::Display::fmt(self, f)
     }
 }
 
-impl<Io, Lookup> fmt::Display for Client<Io, Lookup> {
+impl<Io> fmt::Display for Client<Io> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_tuple("Client").finish()
     }
@@ -291,41 +287,15 @@ fn decrypt_server_info(
     Ok(sip)
 }
 
-impl<Io, Lookup> UnderlayTransport for Client<Io, Lookup>
+impl<Io> NodekeyTransport for Client<Io>
 where
     Io: AsyncRead + AsyncWrite + Send,
-    Lookup: PeerLookup,
 {
-    type Error = Error;
-
-    #[tracing::instrument(fields(%self))]
-    async fn recv(
-        &self,
-    ) -> impl IntoIterator<Item = Result<(PeerId, impl IntoIterator<Item = PacketMut>), Self::Error>>
-    {
-        [self.recv_one().await.map(|(k, pkt)| (k, [pkt]))]
+    async fn send_one(&self, node_key: NodePublicKey, body: &[u8]) -> Result<(), Error> {
+        self.send_one(node_key, body).await
     }
 
-    /// Send a batch of packets to a peer via this DERP server.
-    async fn send<BatchIter, PacketIter>(&self, peer_packets: BatchIter) -> Result<(), Self::Error>
-    where
-        BatchIter: IntoIterator<Item = (PeerId, PacketIter)> + Send,
-        BatchIter::IntoIter: Send,
-        PacketIter: IntoIterator<Item = PacketMut> + Send,
-        PacketIter::IntoIter: Send,
-    {
-        for (peer, packets) in peer_packets {
-            let Some(node_key) = self.peer_lookup.id_to_key(peer) else {
-                tracing::warn!(peer_id = %peer, "no node key known for peer");
-                continue;
-            };
-
-            for packet in packets {
-                self.send_frame_with_extra(&frame::SendPacket { dest: node_key }, packet.as_ref())
-                    .await?;
-            }
-        }
-
-        Ok(())
+    async fn recv_one(&self) -> Result<(NodePublicKey, PacketMut), Error> {
+        self.recv_one().await
     }
 }

ts_derp/src/async_tokio.rs ts_derp/src/client.rsts_derp/src/async_tokio.rs renamed to ts_derp/src/client.rs

@@ -8,15 +8,17 @@ use core::{
 
 use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};
 
-mod async_tokio;
+mod client;
 pub mod dial;
 mod error;
 pub mod frame;
 mod peer_lookup;
+mod transport;
 
-pub use async_tokio::{Client, DefaultClient};
+pub use client::{Client, DefaultClient};
 pub use error::Error;
 pub use peer_lookup::{DummyStaticLookup, PeerLookup};
+pub use transport::Transport;
 
 /// A 24-byte nonce for symmetric encryption with ChaCha20Poly1305.
 #[repr(C)]

ts_derp/src/lib.rs

@@ -3,8 +3,7 @@ use std::sync::Mutex;
 /// Trait providing conversion between [`ts_keys::NodePublicKey`] (required to send and
 /// receive derp messages) and [`ts_transport::PeerId`] (tailscale-rs internal type).
 pub trait PeerLookup: Send + Sync {
-    /// Convert `key` to a [`ts_transport::PeerId`], allocating a new id if the peer doesn't
-    /// exist yet.
+    /// Convert `key` to a [`ts_transport::PeerId`].
     fn key_to_id(&self, key: &ts_keys::NodePublicKey) -> Option<ts_transport::PeerId>;
 
     /// Convert the `id` to a [`ts_keys::NodePublicKey`].

ts_derp/src/peer_lookup.rs

@@ -0,0 +1,86 @@
+use ts_keys::NodePublicKey;
+use ts_packet::PacketMut;
+use ts_transport::{PeerId, UnderlayTransport};
+
+use crate::{Error, PeerLookup};
+
+pub trait NodekeyTransport: Send + Sync {
+    /// Send a message addressed to the given node key.
+    fn send_one(
+        &self,
+        node_key: NodePublicKey,
+        body: &[u8],
+    ) -> impl Future<Output = Result<(), Error>> + Send;
+
+    /// Receive a frame from a particular node key.
+    fn recv_one(&self) -> impl Future<Output = Result<(NodePublicKey, PacketMut), Error>> + Send;
+}
+
+/// An implementation of [`UnderlayTransport`] wrapping a `NodekeyTransport` with a
+/// [`PeerLookup`].
+pub struct Transport<NkT, Lookup> {
+    inner: NkT,
+    peer_lookup: Lookup,
+}
+
+impl<NkT, Lookup> Transport<NkT, Lookup> {
+    /// Construct a new [`Transport`] with the given `NodekeyTransport` and [`PeerLookup`].
+    pub fn new(client: NkT, lookup: Lookup) -> Self {
+        Self {
+            inner: client,
+            peer_lookup: lookup,
+        }
+    }
+
+    /// Destruct this [`Transport`] into its constituent parts.
+    pub fn into_parts(self) -> (NkT, Lookup) {
+        (self.inner, self.peer_lookup)
+    }
+}
+
+impl<NkT, Lookup> UnderlayTransport for Transport<NkT, Lookup>
+where
+    NkT: NodekeyTransport,
+    Lookup: PeerLookup,
+{
+    type Error = Error;
+
+    async fn recv(
+        &self,
+    ) -> impl IntoIterator<Item = Result<(PeerId, impl IntoIterator<Item = PacketMut>), Self::Error>>
+    {
+        match self.inner.recv_one().await {
+            Ok((k, pkt)) => {
+                let Some(id) = self.peer_lookup.key_to_id(&k) else {
+                    tracing::warn!(node_key = %k, lookup_ty = std::any::type_name::<Lookup>(), "no peer id known for node key");
+                    return None;
+                };
+
+                Some(Ok((id, [pkt])))
+            }
+            Err(e) => Some(Err(e)),
+        }
+    }
+
+    /// Send a batch of packets to a peer via this DERP server.
+    async fn send<BatchIter, PacketIter>(&self, peer_packets: BatchIter) -> Result<(), Self::Error>
+    where
+        BatchIter: IntoIterator<Item = (PeerId, PacketIter)> + Send,
+        BatchIter::IntoIter: Send,
+        PacketIter: IntoIterator<Item = PacketMut> + Send,
+        PacketIter::IntoIter: Send,
+    {
+        for (peer, packets) in peer_packets {
+            let Some(node_key) = self.peer_lookup.id_to_key(peer) else {
+                tracing::warn!(peer_id = %peer, "no node key known for peer");
+                continue;
+            };
+
+            for packet in packets {
+                self.inner.send_one(node_key, packet.as_ref()).await?;
+            }
+        }
+
+        Ok(())
+    }
+}