maestro/crates/maestro-daemon/src/attach.rs

//! Per-session attach endpoint: `WS /sessions/:id/stream?token=<jwt>`.
//!
//! - Binary frames carry raw PTY bytes both directions
//! - Text frames carry [`maestro_protocol::ClientControl`] (Resize/Ping/Pong)
//! - JWT is HS256, validated locally with no callback to the registry

use std::sync::Arc;

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Path as AxumPath, Query, State};
use axum::http::StatusCode;
use axum::response::Response;
use axum::routing::any;
use axum::Router;
use futures_util::{SinkExt, StreamExt};
use maestro_protocol::ClientControl;
use serde::Deserialize;
use tokio::sync::broadcast::error::RecvError;
use uuid::Uuid;

use crate::jwt::{JwtVerifier, PairAttachVerifier};
use crate::session::SessionManager;

#[derive(Debug, Clone)]
pub struct AttachState {
    pub manager: SessionManager,
    pub verifier: Arc<JwtVerifier>,
    /// Pair-scoped JWT verifier (NAP-247). When the legacy verifier fails
    /// (e.g. the token was minted as a pair-attach JWT for the broker path),
    /// we fall back to this one before rejecting the client outright. Local
    /// attaches always go through the legacy verifier.
    pub pair_verifier: Arc<PairAttachVerifier>,
}

#[derive(Debug, Deserialize)]
struct AttachParams {
    token: String,
}

pub fn router(
    manager: SessionManager,
    verifier: JwtVerifier,
    pair_verifier: PairAttachVerifier,
) -> Router {
    Router::new()
        .route("/sessions/:id/stream", any(attach_ws))
        .with_state(AttachState {
            manager,
            verifier: Arc::new(verifier),
            pair_verifier: Arc::new(pair_verifier),
        })
}

async fn attach_ws(
    AxumPath(id): AxumPath<Uuid>,
    Query(params): Query<AttachParams>,
    State(state): State<AttachState>,
    ws: WebSocketUpgrade,
) -> Result<Response, (StatusCode, String)> {
    // Try the legacy `JwtClaims` shape first (browser/CLI direct attach). On
    // any failure, try the pair-scoped shape used by the home-daemon broker
    // (NAP-247). One — and only one — must accept.
    let legacy_err = match state.verifier.verify(&params.token) {
        Ok(claims) => {
            if claims.session_id != id {
                return Err((
                    StatusCode::FORBIDDEN,
                    format!(
                        "token is for session {} but path requested {}",
                        claims.session_id, id
                    ),
                ));
            }
            None
        }
        Err(e) => Some(e),
    };
    if let Some(legacy_err) = legacy_err {
        match state.pair_verifier.verify(&params.token, id) {
            Ok(sub) => {
                tracing::debug!(home_daemon = %sub, %id, "pair-scoped attach accepted");
            }
            Err(pair_err) => {
                // Both rejected — return the legacy error since most clients
                // still mint that shape today.
                return Err((
                    StatusCode::UNAUTHORIZED,
                    format!("legacy: {legacy_err}; pair: {pair_err}"),
                ));
            }
        }
    }
    if !state.manager.has(id).await {
        return Err((StatusCode::NOT_FOUND, format!("no session {id}")));
    }
    let manager = state.manager.clone();
    Ok(ws.on_upgrade(move |socket| async move {
        let _ = serve_local(socket, manager, id).await;
    }))
}

/// Pump a WebSocket against a local session's PTY. Reused by both the legacy
/// `/sessions/:id/stream` endpoint (after JWT auth) and the new broker
/// `WS /api/attach/:id` endpoint (when the session is local-owned — NAP-243).
pub(crate) async fn serve_local(
    socket: WebSocket,
    manager: SessionManager,
    id: Uuid,
) -> Result<(), axum::Error> {
    let (mut sink, mut stream) = socket.split();

    let (replay, mut rx) = match manager.subscribe_pty_with_replay(id).await {
        Some(pair) => pair,
        None => {
            // session was killed between the auth check and now — let the
            // client see a clean close instead of a dropped socket.
            let _ = sink.send(Message::Close(None)).await;
            return Ok(());
        }
    };

    // NAP-229: send buffered scrollback up-front so the client sees the last
    // screenful immediately instead of a blank terminal. The replay lock the
    // bridge held during snapshot+subscribe guarantees no byte appears both
    // here and on the live receiver, and none is dropped at the seam.
    if !replay.is_empty() && sink.send(Message::Binary(replay)).await.is_err() {
        return Ok(());
    }

    // Set to true when the session itself terminates (broadcast closes); we
    // owe the client a Close frame in that case. Local breakouts (e.g. the
    // client sent its own Close, or the socket died) skip it.
    let mut session_ended = false;

    loop {
        tokio::select! {
            // PTY output → client
            received = rx.recv() => match received {
                Ok(bytes) => {
                    if sink.send(Message::Binary(bytes)).await.is_err() {
                        break;
                    }
                }
                Err(RecvError::Lagged(_)) => {
                    // dropped some frames under back-pressure; keep going
                    continue;
                }
                Err(RecvError::Closed) => {
                    session_ended = true;
                    break;
                }
            },

            // Client input → tmux
            from_client = stream.next() => match from_client {
                None => break,
                Some(Err(_)) => break,
                Some(Ok(Message::Close(_))) => break,
                Some(Ok(Message::Binary(bytes))) => {
                    let _ = manager.send_keys(id, &bytes).await;
                }
                Some(Ok(Message::Text(text))) => {
                    if let Ok(ctrl) = serde_json::from_str::<ClientControl>(&text) {
                        match ctrl {
                            ClientControl::Resize { cols, rows } => {
                                let _ = manager.resize(id, cols, rows).await;
                            }
                            ClientControl::Ping => {
                                let pong = serde_json::to_string(&ClientControl::Pong)
                                    .expect("serialize pong");
                                if sink.send(Message::Text(pong)).await.is_err() {
                                    break;
                                }
                            }
                            ClientControl::Pong | ClientControl::Error(_) => {}
                        }
                    }
                }
                Some(Ok(_)) => {}
            },
        }
    }

    if session_ended {
        let _ = sink.send(Message::Close(None)).await;
    }
    Ok(())
}