Merge pull request #75 from block65/fix/relay-mode-and-peer-bugs

maxholman · web-flow · commit 6a7c9c653d6b · 2026-03-16T20:34:05.000+07:00
fix(relay): correct data plane wiring and peer management bugs
diff --git a/crates/core/src/client/client.rs b/crates/core/src/client/client.rs
@@ -44,6 +44,8 @@ pub struct ConnectResult<T: wallhack_transport::Transport + ?Sized> {
     control_tx: mpsc::Sender<ControlMessage>,
     /// Receiver for the server's `Handshake` (delivered via the control loop).
     peer_handshake_rx: Option<oneshot::Receiver<Handshake>>,
+    /// Pong-derived latency measurements from the control loop (milliseconds).
+    latency_rx: Option<mpsc::Receiver<f64>>,
 }
 
 impl<T: wallhack_transport::Transport + ?Sized> ConnectResult<T> {
@@ -55,6 +57,7 @@ impl<T: wallhack_transport::Transport + ?Sized> ConnectResult<T> {
         tasks: ConnectionTasks,
         control_tx: mpsc::Sender<ControlMessage>,
         peer_handshake_rx: Option<oneshot::Receiver<Handshake>>,
+        latency_rx: Option<mpsc::Receiver<f64>>,
     ) -> Self {
         Self {
             channels,
@@ -63,6 +66,7 @@ impl<T: wallhack_transport::Transport + ?Sized> ConnectResult<T> {
             transport,
             control_tx,
             peer_handshake_rx,
+            latency_rx,
         }
     }
 
@@ -106,6 +110,8 @@ pub struct ErasedConnectResult {
     pub control_tx: mpsc::Sender<ControlMessage>,
     pub peer_handshake_rx: Option<oneshot::Receiver<Handshake>>,
     pub peer_addr: String,
+    /// Pong-derived latency measurements from the control loop (milliseconds).
+    pub latency_rx: Option<mpsc::Receiver<f64>>,
 }
 
 impl<T> ConnectResult<T>
@@ -129,6 +135,7 @@ where
             channels: self.channels,
             tasks: self.tasks,
             control_tx: self.control_tx,
+            latency_rx: self.latency_rx,
         }
     }
 }
diff --git a/crates/core/src/client/quic/mod.rs b/crates/core/src/client/quic/mod.rs
@@ -169,15 +169,16 @@ impl Client for QuicClient {
 
         // Create oneshot for receiving server's Handshake via the control loop.
         let (handshake_tx, handshake_rx) = tokio::sync::oneshot::channel::<Handshake>();
+        let (latency_tx, latency_rx) = tokio::sync::mpsc::channel::<f64>(4);
 
         // Spawn control stream task
         let transport_ctrl = Arc::clone(&transport);
         let control_handle = tokio::spawn(async move {
             let mut channels = protocol::ControlChannels {
                 outgoing_rx: control_rx,
                 handshake_tx: Some(handshake_tx), // receive server's Handshake
-                latency_tx: None,                 // pong handled inline
-                control_response_tx: None,        // no ControlResponse channel needed now
+                latency_tx: Some(latency_tx),
+                control_response_tx: None,
                 role_transition_tx: None,
             };
             match protocol::run_control_stream_initiator(
@@ -230,6 +231,7 @@ impl Client for QuicClient {
             tasks,
             control_tx,
             Some(handshake_rx),
+            Some(latency_rx),
         ))
     }
 
diff --git a/crates/core/src/client/ws/mod.rs b/crates/core/src/client/ws/mod.rs
@@ -386,14 +386,15 @@ impl WsClient {
 
         // Create oneshot for receiving server's Handshake via the control loop.
         let (handshake_tx, handshake_rx) = tokio::sync::oneshot::channel::<Handshake>();
+        let (latency_tx, latency_rx) = tokio::sync::mpsc::channel::<f64>(4);
 
         // Spawn control stream task
         let transport_ctrl = Arc::clone(&transport);
         let control_handle = tokio::spawn(async move {
             let mut channels = protocol::ControlChannels {
                 outgoing_rx: control_rx,
                 handshake_tx: Some(handshake_tx), // receive server's Handshake
-                latency_tx: None,
+                latency_tx: Some(latency_tx),
                 control_response_tx: None,
                 role_transition_tx: None,
             };
@@ -447,6 +448,7 @@ impl WsClient {
             tasks,
             control_tx,
             Some(handshake_rx),
+            Some(latency_rx),
         ))
     }
 }
diff --git a/crates/daemon/src/mode/auto.rs b/crates/daemon/src/mode/auto.rs
@@ -259,17 +259,8 @@ async fn run_auto_connector(
                         let r = Arc::clone(&routes);
                         let ru = ru.resubscribe();
                         async move {
-                            let peer_handshake_rx = e.peer_handshake_rx;
-                            let transport = e.transport;
-                            let channels = e.channels;
-                            let tasks = e.tasks;
-                            let control_tx = e.control_tx;
                             run_auto_connect_session_dispatch(
-                                peer_handshake_rx,
-                                transport,
-                                channels,
-                                tasks,
-                                control_tx,
+                                e,
                                 &lhs,
                                 &pa,
                                 metrics,
@@ -319,17 +310,8 @@ async fn run_auto_connector(
                         let r = Arc::clone(&routes);
                         let ru = ru.resubscribe();
                         async move {
-                            let peer_handshake_rx = e.peer_handshake_rx;
-                            let transport = e.transport;
-                            let channels = e.channels;
-                            let tasks = e.tasks;
-                            let control_tx = e.control_tx;
                             run_auto_connect_session_dispatch(
-                                peer_handshake_rx,
-                                transport,
-                                channels,
-                                tasks,
-                                control_tx,
+                                e,
                                 &lhs,
                                 &pa,
                                 metrics,
@@ -354,11 +336,7 @@ async fn run_auto_connector(
 /// Non-generic auto-connector dispatch: negotiates role and runs the session.
 #[allow(clippy::too_many_arguments, clippy::too_many_lines)]
 async fn run_auto_connect_session_dispatch(
-    peer_handshake_rx: Option<tokio::sync::oneshot::Receiver<Handshake>>,
-    transport: Arc<dyn ErasedTransport>,
-    channels: DataChannels,
-    tasks: wallhack_core::client::client::ConnectionTasks,
-    control_tx: tokio::sync::mpsc::Sender<wallhack_wire::control::ControlMessage>,
+    connect_result: wallhack_core::client::client::ErasedConnectResult,
     local_hs: &Handshake,
     peer_addr: &str,
     metrics: Arc<Metrics>,
@@ -369,6 +347,16 @@ async fn run_auto_connect_session_dispatch(
         tokio::sync::broadcast::Receiver<wallhack_core::control::routes::RouteUpdate>,
     >,
 ) -> Result<(), NodeError> {
+    let wallhack_core::client::client::ErasedConnectResult {
+        peer_handshake_rx,
+        transport,
+        channels,
+        tasks,
+        control_tx,
+        peer_addr: _,
+        latency_rx,
+    } = connect_result;
+
     let DataChannels {
         instructions_tx,
         instructions_rx,
@@ -435,6 +423,8 @@ async fn run_auto_connect_session_dispatch(
                 &metrics,
                 peer_addr,
                 Some(peer_name),
+                Some(Arc::clone(&peers)),
+                latency_rx,
                 routes,
                 route_updates,
             )
diff --git a/crates/daemon/src/mode/entry.rs b/crates/daemon/src/mode/entry.rs
@@ -316,6 +316,7 @@ pub(crate) async fn run_entry_connect(
     }
 
     let peer_addr = endpoint.to_string();
+    let peers = Arc::clone(&res.peers);
     let security = SecurityParams {
         psk: global.psk.clone(),
         accept_fingerprint: None,
@@ -352,11 +353,12 @@ pub(crate) async fn run_entry_connect(
                     move |connect_result| {
                         let e = connect_result.erase();
                         let m = Arc::clone(&metrics);
+                        let p = Arc::clone(&peers);
                         let r = Arc::clone(&routes);
                         let ru = r_updates.resubscribe();
                         let pa = peer_addr.clone();
                         async move {
-                            run_entry_connected_erased(e, &m, &pa, Some(r), Some(ru)).await
+                            run_entry_connected_erased(e, &m, &pa, Some(p), Some(r), Some(ru)).await
                         }
                     },
                     RECONNECT_DELAY,
@@ -388,11 +390,12 @@ pub(crate) async fn run_entry_connect(
                     move |connect_result| {
                         let e = connect_result.erase();
                         let m = Arc::clone(&metrics);
+                        let p = Arc::clone(&peers);
                         let r = Arc::clone(&routes);
                         let ru = r_updates.resubscribe();
                         let pa = peer_addr.clone();
                         async move {
-                            run_entry_connected_erased(e, &m, &pa, Some(r), Some(ru)).await
+                            run_entry_connected_erased(e, &m, &pa, Some(p), Some(r), Some(ru)).await
                         }
                     },
                     RECONNECT_DELAY,
@@ -429,6 +432,7 @@ pub(crate) async fn run_entry_connected_erased(
     connect_result: wallhack_core::client::client::ErasedConnectResult,
     metrics: &Arc<Metrics>,
     peer_addr: &str,
+    peers: Option<Arc<Registry>>,
     routes: Option<SharedRouteTable>,
     route_updates: Option<
         tokio::sync::broadcast::Receiver<wallhack_core::control::routes::RouteUpdate>,
@@ -443,6 +447,7 @@ pub(crate) async fn run_entry_connected_erased(
         tasks: _tasks,
         control_tx,
         peer_addr: _,
+        latency_rx,
     } = connect_result;
 
     // Wait for the server's handshake to get the peer name
@@ -478,6 +483,8 @@ pub(crate) async fn run_entry_connected_erased(
         metrics,
         peer_addr,
         peer_name.as_deref(),
+        peers,
+        latency_rx,
         routes,
         route_updates,
     )
@@ -495,6 +502,8 @@ pub(crate) async fn run_entry_connected_inner(
     metrics: &Arc<Metrics>,
     peer_addr: &str,
     peer_name: Option<&str>,
+    peers: Option<Arc<Registry>>,
+    latency_rx: Option<tokio::sync::mpsc::Receiver<f64>>,
     routes: Option<SharedRouteTable>,
     route_updates: Option<
         tokio::sync::broadcast::Receiver<wallhack_core::control::routes::RouteUpdate>,
@@ -574,6 +583,19 @@ pub(crate) async fn run_entry_connected_inner(
         tracing::debug!("Initial ping failed: {e}");
     }
 
+    // Consume latency measurements from the control loop and update the registry.
+    if let (Some(mut lat_rx), Some(reg)) = (latency_rx, &peers) {
+        let peer_id = peer_name.map(std::string::ToString::to_string);
+        let reg = Arc::clone(reg);
+        tokio::spawn(async move {
+            while let Some(ms) = lat_rx.recv().await {
+                if let Some(ref id) = peer_id {
+                    reg.update_latency(id, ms);
+                }
+            }
+        });
+    }
+
     let manager_handle = tokio::spawn(async move { manager.run().await });
 
     match manager_handle.await {
diff --git a/crates/daemon/src/mode/exit.rs b/crates/daemon/src/mode/exit.rs
@@ -416,9 +416,13 @@ async fn run_exit_loop_inner(
 
     tracing::info!("Connected to {peer_addr}");
 
-    // Resolve the peer's handshake name (delivered asynchronously via the
+    // Resolve the peer's handshake (delivered asynchronously via the
     // control loop). Fall back to the address if unavailable.
-    let peer_name = resolve_peer_name(peer_handshake_rx, peer_addr).await;
+    let peer_handshake = resolve_peer_handshake(peer_handshake_rx).await;
+    let peer_name = peer_handshake
+        .as_ref()
+        .filter(|h| !h.name.is_empty())
+        .map_or_else(|| peer_addr.to_string(), |h| h.name.clone());
 
     // We connected to the entry peer, so side=Connect.
     ctx.peers.register(
@@ -428,6 +432,14 @@ async fn run_exit_loop_inner(
         ConnectionSide::Connect,
     );
 
+    // Apply the peer's advertised capabilities from the handshake.
+    let peer_capabilities = peer_handshake
+        .as_ref()
+        .and_then(|h| h.capabilities)
+        .unwrap_or_default();
+    ctx.peers
+        .update_capabilities(&peer_name, &peer_capabilities);
+
     // Outgoing: open uni stream to entry peer, send responses.
     let transport_out = Arc::clone(&transport);
     tokio::spawn(async move {
@@ -475,20 +487,18 @@ async fn run_exit_loop_inner(
     Ok(())
 }
 
-/// Await the peer's handshake name from a oneshot receiver (with timeout).
+/// Await the peer's handshake from a oneshot receiver (with timeout).
 ///
-/// Returns the handshake name if non-empty, otherwise falls back to the
-/// peer address.
-async fn resolve_peer_name(
+/// Returns the received `Handshake` if available within the timeout, or
+/// `None` if the receiver is absent, the sender dropped, or the timeout
+/// expires.
+async fn resolve_peer_handshake(
     rx: Option<tokio::sync::oneshot::Receiver<wallhack_wire::data::Handshake>>,
-    peer_addr: &str,
-) -> String {
-    let Some(rx) = rx else {
-        return peer_addr.to_string();
-    };
+) -> Option<wallhack_wire::data::Handshake> {
+    let rx = rx?;
     match tokio::time::timeout(std::time::Duration::from_secs(10), rx).await {
-        Ok(Ok(hs)) if !hs.name.is_empty() => hs.name,
-        _ => peer_addr.to_string(),
+        Ok(Ok(hs)) => Some(hs),
+        _ => None,
     }
 }
 
diff --git a/crates/daemon/src/mode/relay.rs b/crates/daemon/src/mode/relay.rs
diff --git a/crates/daemon/src/transport.rs b/crates/daemon/src/transport.rs

Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,8 @@ pub struct ConnectResult<T: wallhack_transport::Transport + ?Sized> {`
`44`	`44`	`control_tx: mpsc::Sender<ControlMessage>,`
`45`	`45`	/// Receiver for the server's `Handshake` (delivered via the control loop).
`46`	`46`	`peer_handshake_rx: Option<oneshot::Receiver<Handshake>>,`
	`47`	`+ /// Pong-derived latency measurements from the control loop (milliseconds).`
	`48`	`+ latency_rx: Option<mpsc::Receiver<f64>>,`
`47`	`49`	`}`
`48`	`50`
`49`	`51`	`impl<T: wallhack_transport::Transport + ?Sized> ConnectResult<T> {`
`@@ -55,6 +57,7 @@ impl<T: wallhack_transport::Transport + ?Sized> ConnectResult<T> {`
`55`	`57`	`tasks: ConnectionTasks,`
`56`	`58`	`control_tx: mpsc::Sender<ControlMessage>,`
`57`	`59`	`peer_handshake_rx: Option<oneshot::Receiver<Handshake>>,`
	`60`	`+ latency_rx: Option<mpsc::Receiver<f64>>,`
`58`	`61`	`) -> Self {`
`59`	`62`	`Self {`
`60`	`63`	`channels,`
`@@ -63,6 +66,7 @@ impl<T: wallhack_transport::Transport + ?Sized> ConnectResult<T> {`
`63`	`66`	`transport,`
`64`	`67`	`control_tx,`
`65`	`68`	`peer_handshake_rx,`
	`69`	`+ latency_rx,`
`66`	`70`	`}`
`67`	`71`	`}`
`68`	`72`
`@@ -106,6 +110,8 @@ pub struct ErasedConnectResult {`
`106`	`110`	`pub control_tx: mpsc::Sender<ControlMessage>,`
`107`	`111`	`pub peer_handshake_rx: Option<oneshot::Receiver<Handshake>>,`
`108`	`112`	`pub peer_addr: String,`
	`113`	`+ /// Pong-derived latency measurements from the control loop (milliseconds).`
	`114`	`+ pub latency_rx: Option<mpsc::Receiver<f64>>,`
`109`	`115`	`}`
`110`	`116`
`111`	`117`	`impl<T> ConnectResult<T>`
`@@ -129,6 +135,7 @@ where`
`129`	`135`	`channels: self.channels,`
`130`	`136`	`tasks: self.tasks,`
`131`	`137`	`control_tx: self.control_tx,`
	`138`	`+ latency_rx: self.latency_rx,`
`132`	`139`	`}`
`133`	`140`	`}`
`134`	`141`	`}`