dataplane: plumb disco and stun to separate streams

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

Author: npry

ts_dataplane/src/async_tokio.rs

@@ -28,6 +28,11 @@ 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>;
 
@@ -58,6 +63,11 @@ struct CoreState {
     overlay_transports: HashMap<OverlayTransportId, Tx<ToOverlay>>,
     /// Queues to write packets to underlay transports.
     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<StunBatch>,
 }
 
 /// State that must be held during async polling.
@@ -73,13 +83,19 @@ impl DataPlane {
     ///
     /// 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,
 
@@ -90,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(),
             }),
@@ -98,7 +116,9 @@ impl DataPlane {
                 from_overlay: overlay_down,
                 from_underlay: underlay_up,
             }),
-        }
+        };
+
+        (dp, disco_rx, stun_rx)
     }
 
     /// Allocate a new underlay transport.
@@ -230,7 +250,20 @@ impl DataPlane {
                 (Some(to_peers), Some(loopback))
             }
             SelectResult::UnderlayUp(underlay_up) => {
-                let InboundResult { to_local, to_peers } = core.sync.process_inbound(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))
             }

ts_dataplane/src/lib.rs

@@ -101,8 +101,27 @@ impl DataPlane {
         &mut self,
         packets: impl IntoIterator<Item = PacketMut>,
     ) -> InboundResult {
+        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(Instant::now(), packets);
+            self.wireguard.recv(Instant::now(), wireguard);
 
         let to_local = to_local
             .into_iter()
@@ -201,7 +220,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.
@@ -261,6 +285,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.

ts_runtime/src/dataplane.rs

@@ -42,9 +42,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.clone(),
-        ));
+        let (dataplane, ..) = ts_dataplane::async_tokio::DataPlane::new(env.keys.node_keys.clone());
+        let dataplane = Arc::new(dataplane);
 
         env.subscribe::<PeerRouteUpdate>(&slf).await?;
         env.subscribe::<SelfRouteUpdate>(&slf).await?;