dataplane: plumb disco and stun to separate streams

npry · npry · commit 5799dc41be1d · 2026-06-13T06:34:42.000-04:00
Signed-off-by: Nathan Perry &lt;nathan@tailscale.com&gt;
Change-Id: Ibf7d20ff1a3ab10e1277e78179285eb46a6a6964
diff --git a/ts_dataplane/src/async_tokio.rs b/ts_dataplane/src/async_tokio.rs
@@ -9,19 +9,35 @@ use ts_tunnel::NodeKeyPair;
 
 use crate::{EventResult, InboundResult, OutboundResult};
 
-/// Queue for packets leaving the data plane "up" into an overlay transport.
-pub type DataplaneToOverlay = mpsc::UnboundedSender<Vec<PacketMut>>;
-
-/// Queue for packets entering the data plane "down" from an overlay transport.
-pub type DataplaneFromOverlay = mpsc::UnboundedReceiver<Vec<PacketMut>>;
-
-/// Queue for packets leaving the data plane "down" into an underlay transport.
-pub type DataplaneToUnderlay = mpsc::UnboundedSender<(PeerId, Vec<PacketMut>)>;
-
-/// Queue for packets entering the data plane "up" from an underlay transport.
-pub type DataplaneFromUnderlay = mpsc::UnboundedReceiver<(PeerId, Vec<PacketMut>)>;
-
-// TODO: wire in overlay/underlay transport traits
+// NOTE(npry): this used to have unique types for each queue, but the names got confusing due to
+// having to think about the cartesian product of PacketType x QueueDirection x Network
+// (this is a "sender handle for receive packets on the overlay", vs. "receive handle for sender
+// packets on the overlay", etc.). It wasn't always clear to distinguish what referred to the
+// channel direction (sender/receiver) and what referred to the actual traffic kind (packets
+// received _from_ the underlay are different from packets sent _to_ the underlay). So now we name
+// the packet type here (to/from overlay/underlay) and use separate helper types to make the channel
+// directions easier to name.
+
+/// Packet batches sent to an overlay.
+pub type ToOverlay = Vec<PacketMut>;
+/// Packet batches received from an overlay.
+pub type FromOverlay = Vec<PacketMut>;
+
+/// Packet batches sent to an underlay.
+pub type ToUnderlay = (PeerId, Vec<PacketMut>);
+/// Packet batches received from an underlay.
+pub type FromUnderlay = Vec<PacketMut>;
+
+/// A batch of disco packets received from an underlay transport.
+pub type DiscoBatch = Vec<PacketMut>;
+/// A batch of stun packets received from an underlay transport.
+pub type StunBatch = Vec<PacketMut>;
+
+/// Shorthand for a sender channel.
+pub type Tx<T> = mpsc::UnboundedSender<T>;
+
+/// Shorthand for a receiver channel.
+pub type Rx<T> = mpsc::UnboundedReceiver<T>;
 
 /// Transforms packets to make tailscale happen.
 pub struct DataPlane {
@@ -30,8 +46,10 @@ pub struct DataPlane {
 
     transports_changed: tokio::sync::Notify,
 
-    underlay_down: DataplaneToUnderlay,
-    overlay_up: DataplaneToOverlay,
+    // These are the senders handed out in new_*_transport, just held here so we can clone them, we
+    // never send to them.
+    underlay_down: Tx<FromUnderlay>,
+    overlay_up: Tx<FromOverlay>,
 
     next_underlay_transport: AtomicU32,
     next_overlay_transport: AtomicU32,
@@ -42,31 +60,42 @@ struct CoreState {
     sync: crate::DataPlane,
 
     /// Queues to write packets to overlay transports.
-    overlay_transports: HashMap<OverlayTransportId, DataplaneToOverlay>,
+    overlay_transports: HashMap<OverlayTransportId, Tx<ToOverlay>>,
     /// Queues to write packets to underlay transports.
-    underlay_transports: HashMap<UnderlayTransportId, DataplaneToUnderlay>,
+    underlay_transports: HashMap<UnderlayTransportId, Tx<ToUnderlay>>,
+
+    /// Send handle for disco packets received from underlays.
+    disco_out: Tx<DiscoBatch>,
+    /// Send handle for stun packets received from underlays.
+    stun_out: Tx<DiscoBatch>,
 }
 
 /// State that must be held during async polling.
 struct PollState {
     /// Queue for packets entering the data plane ("coming down") from overlay transports.
-    from_overlay: DataplaneFromOverlay,
+    from_overlay: Rx<FromOverlay>,
     /// Queue for packets entering the data plane ("coming up") from underlay transports.
-    from_underlay: DataplaneFromUnderlay,
+    from_underlay: Rx<FromUnderlay>,
 }
 
 impl DataPlane {
     /// Create a new data plane for a wireguard node key.
     ///
     /// The caller must configure overlay/underlay output queues for the data plane to be useful,
     /// otherwise all it can do is drop packets.
-    pub fn new(my_key: NodeKeyPair) -> Self {
+    ///
+    /// The second and third elements of the return tuple are output queues for disco and
+    /// STUN messages, respectively.
+    pub fn new(my_key: NodeKeyPair) -> (Self, Rx<DiscoBatch>, Rx<StunBatch>) {
         let (overlay_up, overlay_down) = mpsc::unbounded_channel();
         let (underlay_down, underlay_up) = mpsc::unbounded_channel();
 
+        let (disco_tx, disco_rx) = mpsc::unbounded_channel();
+        let (stun_tx, stun_rx) = mpsc::unbounded_channel();
+
         let sync = crate::DataPlane::new(my_key);
 
-        Self {
+        let dp = Self {
             underlay_down,
             overlay_up,
 
@@ -77,6 +106,8 @@ impl DataPlane {
 
             core_state: Mutex::new(CoreState {
                 sync,
+                stun_out: stun_tx,
+                disco_out: disco_tx,
                 overlay_transports: Default::default(),
                 underlay_transports: Default::default(),
             }),
@@ -85,17 +116,15 @@ impl DataPlane {
                 from_overlay: overlay_down,
                 from_underlay: underlay_up,
             }),
-        }
+        };
+
+        (dp, disco_rx, stun_rx)
     }
 
     /// Allocate a new underlay transport.
     pub async fn new_underlay_transport(
         &self,
-    ) -> (
-        UnderlayTransportId,
-        DataplaneFromUnderlay,
-        DataplaneToUnderlay,
-    ) {
+    ) -> (UnderlayTransportId, Rx<ToUnderlay>, Tx<FromUnderlay>) {
         let id = self
             .next_underlay_transport
             .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
@@ -116,7 +145,7 @@ impl DataPlane {
     /// Allocate a new overlay transport.
     pub async fn new_overlay_transport(
         &self,
-    ) -> (OverlayTransportId, DataplaneToOverlay, DataplaneFromOverlay) {
+    ) -> (OverlayTransportId, Tx<FromOverlay>, Rx<ToOverlay>) {
         let id = self
             .next_overlay_transport
             .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
@@ -146,7 +175,7 @@ impl DataPlane {
     pub async fn step(&self) {
         enum SelectResult {
             OverlayDown(Vec<PacketMut>),
-            UnderlayUp(PeerId, Vec<PacketMut>),
+            UnderlayUp(Vec<PacketMut>),
             TransportsChanged,
             Event,
         }
@@ -185,10 +214,10 @@ impl DataPlane {
                 }
 
                 underlay_pkts = underlay_up.recv() => {
-                    let (peer_id, underlay_pkts) = underlay_pkts.unwrap();
-                    tracing::trace!(%peer_id, n_underlay_pkts = underlay_pkts.len());
+                    let underlay_pkts = underlay_pkts.unwrap();
+                    tracing::trace!(n_underlay_pkts = underlay_pkts.len());
 
-                    SelectResult::UnderlayUp(peer_id, underlay_pkts)
+                    SelectResult::UnderlayUp(underlay_pkts)
                 }
 
                 _ = self.transports_changed.notified() => {
@@ -214,8 +243,21 @@ impl DataPlane {
 
                 (Some(to_peers), Some(loopback))
             }
-            SelectResult::UnderlayUp(_peer_id, underlay_up) => {
-                let InboundResult { to_local, to_peers } = core.sync.process_inbound(underlay_up);
+            SelectResult::UnderlayUp(underlay_up) => {
+                let InboundResult {
+                    to_local,
+                    to_peers,
+                    disco,
+                    stun,
+                } = core.sync.process_inbound(underlay_up);
+
+                if !disco.is_empty() && core.disco_out.send(disco).is_err() {
+                    tracing::warn!("disco packets dropped: no receiver");
+                }
+
+                if !stun.is_empty() && core.stun_out.send(stun).is_err() {
+                    tracing::warn!("stun packets dropped: no receiver");
+                }
 
                 (Some(to_peers), Some(to_local))
             }
diff --git a/ts_dataplane/src/lib.rs b/ts_dataplane/src/lib.rs
@@ -101,7 +101,26 @@ impl DataPlane {
         &mut self,
         packets: impl IntoIterator<Item = PacketMut>,
     ) -> InboundResult {
-        let ts_tunnel::RecvResult { to_local, to_peers } = self.wireguard.recv(packets);
+        let (wireguard, disco, stun) = packets.into_iter().fold(
+            (vec![], vec![], vec![]),
+            |(mut wg, mut disco, mut stun), pkt| {
+                let ident = PacketIdent::identify(pkt.as_ref());
+
+                match ident.ty {
+                    PacketType::Disco => {
+                        disco.push(pkt);
+                    }
+                    PacketType::StunBinding => {
+                        stun.push(pkt);
+                    }
+                    PacketType::Wireguard | PacketType::Unknown => wg.push(pkt),
+                }
+
+                (wg, disco, stun)
+            },
+        );
+
+        let ts_tunnel::RecvResult { to_local, to_peers } = self.wireguard.recv(wireguard);
 
         let to_local = to_local
             .into_iter()
@@ -200,7 +219,12 @@ impl DataPlane {
             prev.cancel();
         }
 
-        InboundResult { to_local, to_peers }
+        InboundResult {
+            to_local,
+            to_peers,
+            disco,
+            stun,
+        }
     }
 
     /// Return the next time at which [`DataPlane::process_events`] must be called.
@@ -260,6 +284,12 @@ pub struct InboundResult {
     pub to_local: HashMap<OverlayTransportId, Vec<PacketMut>>,
     /// Encrypted packets to be sent to wireguard peers by the underlay.
     pub to_peers: HashMap<(UnderlayTransportId, PeerId), Vec<PacketMut>>,
+
+    /// Encrypted disco packets to be handled externally.
+    pub disco: Vec<PacketMut>,
+
+    /// STUN packets to be handled externally.
+    pub stun: Vec<PacketMut>,
 }
 
 /// The result of processing an event.
diff --git a/ts_runtime/src/dataplane.rs b/ts_runtime/src/dataplane.rs
@@ -4,9 +4,8 @@ use kameo::{
     actor::ActorRef,
     message::{Context, Message},
 };
-use tokio::sync::mpsc;
-use ts_packet::PacketMut;
-use ts_transport::{OverlayTransportId, PeerId, UnderlayTransportId};
+use ts_dataplane::async_tokio::{FromOverlay, FromUnderlay, Rx, ToOverlay, ToUnderlay, Tx};
+use ts_transport::{OverlayTransportId, UnderlayTransportId};
 
 use crate::{
     Error,
@@ -17,18 +16,6 @@ use crate::{
     src_filter::SourceFilterState,
 };
 
-/// Queue for packets sent from the overlay to the dataplane.
-pub type OverlayToDataplane = mpsc::UnboundedSender<Vec<PacketMut>>;
-
-/// Queue for packets entering the overlay from the dataplane.
-pub type OverlayFromDataplane = mpsc::UnboundedReceiver<Vec<PacketMut>>;
-
-/// Queue for packets leaving the underlay to the dataplane.
-pub type UnderlayToDataplane = mpsc::UnboundedSender<(PeerId, Vec<PacketMut>)>;
-
-/// Queue for packets entering an underlay from the dataplane.
-pub type UnderlayFromDataplane = mpsc::UnboundedReceiver<(PeerId, Vec<PacketMut>)>;
-
 pub struct DataplaneActor {
     dataplane: Arc<ts_dataplane::async_tokio::DataPlane>,
     task: tokio::task::JoinHandle<()>,
@@ -45,18 +32,14 @@ impl DataplaneActor {
     #[message]
     pub async fn new_overlay_transport(
         &self,
-    ) -> (OverlayTransportId, OverlayToDataplane, OverlayFromDataplane) {
+    ) -> (OverlayTransportId, Tx<FromOverlay>, Rx<ToOverlay>) {
         self.dataplane.new_overlay_transport().await
     }
 
     #[message]
     pub async fn new_underlay_transport(
         &self,
-    ) -> (
-        UnderlayTransportId,
-        UnderlayFromDataplane,
-        UnderlayToDataplane,
-    ) {
+    ) -> (UnderlayTransportId, Rx<ToUnderlay>, Tx<FromUnderlay>) {
         self.dataplane.new_underlay_transport().await
     }
 }
@@ -66,9 +49,8 @@ impl kameo::Actor for DataplaneActor {
     type Error = Error;
 
     async fn on_start(env: Self::Args, slf: ActorRef<Self>) -> Result<Self, Self::Error> {
-        let dataplane = Arc::new(ts_dataplane::async_tokio::DataPlane::new(
-            env.keys.node_keys,
-        ));
+        let (dataplane, ..) = ts_dataplane::async_tokio::DataPlane::new(env.keys.node_keys);
+        let dataplane = Arc::new(dataplane);
 
         env.subscribe::<PeerRouteUpdate>(&slf).await?;
         env.subscribe::<SelfRouteUpdate>(&slf).await?;
diff --git a/ts_runtime/src/multiderp.rs b/ts_runtime/src/multiderp.rs
@@ -11,6 +11,7 @@ use kameo::{
 };
 use tokio::{sync::watch, task::JoinSet};
 use ts_control::DerpRegion;
+use ts_dataplane::async_tokio::{FromUnderlay, Rx, ToUnderlay, Tx};
 use ts_derp::RegionId;
 use ts_keys::{NodeKeyPair, NodePublicKey};
 use ts_transport::{
@@ -19,7 +20,7 @@ use ts_transport::{
 
 use crate::{
     Env, Error,
-    dataplane::{DataplaneActor, NewUnderlayTransport, UnderlayFromDataplane, UnderlayToDataplane},
+    dataplane::{DataplaneActor, NewUnderlayTransport},
     derp_latency::DerpLatencyMeasurement,
     peer_tracker::{PeerDb, PeerState},
 };
@@ -156,8 +157,8 @@ async fn run_derp_once(
     id: RegionId,
     region: &DerpRegion,
     keys: NodeKeyPair,
-    to_dataplane: &UnderlayToDataplane,
-    from_dataplane: &mut UnderlayFromDataplane,
+    to_dataplane: &Tx<FromUnderlay>,
+    from_dataplane: &mut Rx<ToUnderlay>,
     home_derp_rx: &mut watch::Receiver<bool>,
     peer_db: &RwLock<Option<Arc<PeerDb>>>,
 ) -> Result<(), ts_derp::Error> {
@@ -210,7 +211,7 @@ async fn run_derp_once(
 
                         tracing::trace!(parent: &span, %peer_id, len = pkts.len(), "packet from derp server");
 
-                        let Ok(()) = to_dataplane.send((peer_id, pkts)) else {
+                        let Ok(()) = to_dataplane.send(pkts) else {
                             tracing::error!(parent: &span, "underlay receive channel closed");
                             break;
                         };
diff --git a/ts_runtime/src/netstack_actor.rs b/ts_runtime/src/netstack_actor.rs
@@ -9,13 +9,10 @@ use netstack::{
     netcore::{Channel, NetstackControl},
 };
 use tokio::task::JoinSet;
+use ts_dataplane::async_tokio::{FromOverlay, Rx, ToOverlay, Tx};
 use ts_packet::PacketMut;
 
-use crate::{
-    Error,
-    dataplane::{OverlayFromDataplane, OverlayToDataplane},
-    env::Env,
-};
+use crate::{Error, env::Env};
 
 pub struct NetstackActor {
     _joinset: JoinSet<()>,
@@ -26,8 +23,8 @@ impl kameo::Actor for NetstackActor {
     type Args = (
         Env,
         netstack::netcore::Config,
-        OverlayToDataplane,
-        OverlayFromDataplane,
+        Tx<FromOverlay>,
+        Rx<ToOverlay>,
     );
     type Error = Error;
 
