diff --git a/crates/maestro-daemon/src/local_api.rs b/crates/maestro-daemon/src/local_api.rs index efb2459..9a1a985 100644 --- a/crates/maestro-daemon/src/local_api.rs +++ b/crates/maestro-daemon/src/local_api.rs @@ -359,28 +359,39 @@ async fn broker_attach_ws( } }; + // Dial the peer BEFORE completing the upgrade. A failure here — peer down, + // perimeter-side rejection, or the peer's verifier refusing the pair JWT + // (all of which the peer reports pre-upgrade) — then surfaces as a + // structured `502` the client can act on, instead of an opaque + // post-upgrade close indistinguishable from a clean session end (NAP-265). + let upstream = match tokio_tungstenite::connect_async(&dial_url).await { + Ok((up, _)) => up, + Err(e) => { + tracing::warn!(%id, url = %dial_url, error = %e, "T2 dial to peer failed"); + return Err(error_response( + StatusCode::BAD_GATEWAY, + "peer_unreachable", + &format!("could not open the data-plane connection to the session's daemon: {e}"), + )); + } + }; + Ok(ws.on_upgrade(move |socket| async move { - let _ = pump_broker_ws(socket, dial_url).await; + let _ = pump_broker_ws(socket, upstream).await; })) } -/// Bidirectional pump between the client WS and the peer daemon's WS. +/// Bidirectional pump between the client WS and an already-connected peer WS. async fn pump_broker_ws( client: axum::extract::ws::WebSocket, - upstream_url: String, + upstream: tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >, ) -> Result<(), axum::Error> { use axum::extract::ws::Message as AxMessage; use futures_util::{SinkExt, StreamExt}; use tokio_tungstenite::tungstenite::Message as WsMessage; - let (upstream, _) = match tokio_tungstenite::connect_async(&upstream_url).await { - Ok(pair) => pair, - Err(e) => { - tracing::warn!(url = %upstream_url, error = %e, "T2 dial failed"); - return Ok(()); - } - }; - let (mut client_sink, mut client_stream) = client.split(); let (mut up_sink, mut up_stream) = upstream.split(); diff --git a/crates/maestro-e2e/tests/e2e.rs b/crates/maestro-e2e/tests/e2e.rs index 1755c55..310f9a0 100644 --- a/crates/maestro-e2e/tests/e2e.rs +++ b/crates/maestro-e2e/tests/e2e.rs @@ -613,29 +613,24 @@ async fn broker_rejects_when_target_doesnt_allow_inbound() { .await; let broker_url = format!("ws://{}/api/attach/{}", alpha.bind_addr, sid); - // The handshake itself succeeds (alpha's broker accepts and starts - // proxying), but beta will reject the pair JWT — the upstream WS closes - // immediately, and alpha's proxy closes the client WS in turn. - let (mut ws, _) = tokio_tungstenite::connect_async(broker_url) + // NAP-265: beta rejects the pair JWT (empty allow_inbound) pre-upgrade. + // alpha's broker now dials the peer *before* completing the upgrade, so the + // client's handshake fails with a structured 502 rather than succeeding and + // then closing opaquely. + let err = tokio_tungstenite::connect_async(broker_url) .await - .expect("alpha handshake"); - // We should get a close shortly without any binary frames. - let deadline = std::time::Instant::now() + Duration::from_secs(3); - let mut got_binary = false; - while std::time::Instant::now() < deadline { - tokio::select! { - msg = ws.next() => match msg { - Some(Ok(WsMessage::Binary(_))) => got_binary = true, - None | Some(Err(_)) => break, - Some(Ok(_)) => {} - }, - _ = tokio::time::sleep(Duration::from_millis(50)) => {} + .expect_err("handshake must fail when the peer rejects the attach"); + match err { + tokio_tungstenite::tungstenite::Error::Http(resp) => { + assert_eq!( + resp.status(), + 502, + "expected BAD_GATEWAY from alpha's broker, got {}", + resp.status() + ); } + other => panic!("expected an HTTP 502 response, got {other:?}"), } - assert!( - !got_binary, - "beta rejected the JWT but alpha somehow forwarded PTY bytes" - ); client.kill_session(sid).await.expect("kill"); drop(alpha);