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

//! `maestro attach` — raw-TTY proxy over the daemon's per-session WS endpoint.
//!
//! Wire-format mirrors `crates/maestro-daemon/src/attach.rs`:
//! - PTY bytes ride in **binary** frames both directions.
//! - `ClientControl::{Resize, Ping, Pong}` ride in **text** frames.

use std::io::{self, IsTerminal};
use std::os::fd::AsFd;
use std::sync::Arc;

use futures_util::{SinkExt, StreamExt};
use maestro_protocol::ClientControl;
use nix::sys::termios::{cfmakeraw, tcgetattr, tcsetattr, SetArg, Termios};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio_tungstenite::tungstenite::Message;

#[derive(Debug, thiserror::Error)]
#[allow(clippy::result_large_err)]
pub enum AttachError {
    #[error("ws: {0}")]
    WebSocket(#[from] tokio_tungstenite::tungstenite::Error),
    #[error("daemon returned {status} ({code}): {message}")]
    Rejected {
        status: String,
        code: String,
        message: String,
    },
    #[error("io: {0}")]
    Io(#[from] io::Error),
    #[error("termios: {0}")]
    Termios(#[from] nix::Error),
    #[error(transparent)]
    Url(#[from] maestro_protocol::routes::UrlError),
    #[error("invalid detach key: {0}")]
    InvalidDetachKey(String),
    #[error("stdin must be a TTY for attach")]
    NotATty,
}

impl AttachError {
    fn from_websocket(error: tokio_tungstenite::tungstenite::Error) -> Self {
        if let tokio_tungstenite::tungstenite::Error::Http(response) = &error {
            if let Some(body) = response.body().as_deref() {
                if let Ok(error_body) = serde_json::from_slice::<maestro_protocol::ApiError>(body) {
                    return Self::Rejected {
                        status: response.status().to_string(),
                        code: error_body.error.code,
                        message: error_body.error.message,
                    };
                }
            }
        }
        Self::WebSocket(error)
    }

    /// True when the failure was an inability to reach the WS endpoint at the
    /// transport level (connection refused), as opposed to an HTTP-level
    /// rejection from a daemon that did answer.
    pub fn is_unreachable(&self) -> bool {
        use std::io::ErrorKind;
        use tokio_tungstenite::tungstenite::Error as WsError;
        match self {
            AttachError::WebSocket(WsError::Io(e)) => matches!(
                e.kind(),
                ErrorKind::ConnectionRefused
                    | ErrorKind::ConnectionReset
                    | ErrorKind::ConnectionAborted
            ),
            _ => false,
        }
    }
}

/// Parse a detach key spec. Accepts `^X` for Ctrl-X (case-insensitive) or a
/// single literal character. Returns the resulting byte.
#[allow(clippy::result_large_err)]
pub fn parse_detach_key(spec: &str) -> Result<u8, AttachError> {
    let bytes = spec.as_bytes();
    match bytes {
        [b'^', c] if c.is_ascii_alphabetic() => Ok(c.to_ascii_uppercase() - b'@'),
        [c] => Ok(*c),
        _ => Err(AttachError::InvalidDetachKey(spec.to_string())),
    }
}

/// RAII guard that puts the terminal into raw mode and restores on drop.
#[derive(Debug)]
struct RawTty {
    fd: std::os::fd::OwnedFd,
    original: Termios,
}

impl RawTty {
    #[allow(clippy::result_large_err)]
    fn enable() -> Result<Self, AttachError> {
        let stdin = io::stdin();
        if !stdin.is_terminal() {
            return Err(AttachError::NotATty);
        }
        // Dup stdin so we hold an OwnedFd whose lifetime we control.
        let fd = stdin
            .as_fd()
            .try_clone_to_owned()
            .map_err(AttachError::Io)?;
        let original = tcgetattr(&fd)?;
        let mut raw = original.clone();
        cfmakeraw(&mut raw);
        tcsetattr(&fd, SetArg::TCSANOW, &raw)?;
        Ok(Self { fd, original })
    }
}

impl Drop for RawTty {
    fn drop(&mut self) {
        let _ = tcsetattr(&self.fd, SetArg::TCSANOW, &self.original);
    }
}

fn current_winsize() -> Option<(u16, u16)> {
    use std::mem::MaybeUninit;
    let mut ws = MaybeUninit::<libc::winsize>::uninit();
    let r = unsafe { libc::ioctl(libc::STDOUT_FILENO, libc::TIOCGWINSZ, ws.as_mut_ptr()) };
    if r != 0 {
        return None;
    }
    let ws = unsafe { ws.assume_init() };
    Some((ws.ws_col, ws.ws_row))
}

/// Drive the attach loop against an already-built WS `url` ([`attach_url`])
/// until the client detaches or the remote closes. `detach_byte` is the single
/// byte that exits cleanly without killing the session.
#[allow(clippy::result_large_err)]
pub async fn run_attach_url(url: &str, detach_byte: u8) -> Result<(), AttachError> {
    let (ws, _resp) = tokio_tungstenite::connect_async(url)
        .await
        .map_err(AttachError::from_websocket)?;
    let (mut sink, mut stream) = ws.split();

    let raw = RawTty::enable()?;

    // Send initial window size so the remote PTY matches our terminal.
    if let Some((cols, rows)) = current_winsize() {
        let frame =
            serde_json::to_string(&ClientControl::Resize { cols, rows }).expect("serialize");
        sink.send(Message::Text(frame)).await?;
    }

    let mut stdin = tokio::io::stdin();
    let mut stdout = tokio::io::stdout();

    // Resize signal stream.
    let mut sigwinch =
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::window_change())?;

    let mut in_buf = [0u8; 4096];

    let detach_msg = Arc::new(format!(
        "\r\n[maestro] detached (key 0x{detach_byte:02x})\r\n"
    ));

    let mut session_ended = false;

    loop {
        tokio::select! {
            // Server → local stdout
            msg = stream.next() => {
                match msg {
                    None => {
                        session_ended = true;
                        break;
                    }
                    Some(Err(e)) => return Err(AttachError::WebSocket(e)),
                    Some(Ok(Message::Binary(bytes))) => {
                        stdout.write_all(&bytes).await?;
                        stdout.flush().await?;
                    }
                    Some(Ok(Message::Text(text))) => {
                        if let Ok(ClientControl::Ping) = serde_json::from_str::<ClientControl>(&text) {
                            let pong = serde_json::to_string(&ClientControl::Pong)
                                .expect("serialize pong");
                            sink.send(Message::Text(pong)).await?;
                        }
                    }
                    Some(Ok(Message::Ping(p))) => {
                        sink.send(Message::Pong(p)).await?;
                    }
                    Some(Ok(Message::Close(_))) => {
                        session_ended = true;
                        break;
                    }
                    Some(Ok(_)) => {}
                }
            }

            // Local stdin → server
            n = stdin.read(&mut in_buf) => {
                let n = n?;
                if n == 0 {
                    break;
                }
                let chunk = &in_buf[..n];
                if let Some(pos) = chunk.iter().position(|b| *b == detach_byte) {
                    if pos > 0 {
                        sink.send(Message::Binary(chunk[..pos].to_vec())).await.ok();
                    }
                    let _ = sink.send(Message::Close(None)).await;
                    // best-effort goodbye on the local terminal
                    stdout.write_all(detach_msg.as_bytes()).await.ok();
                    break;
                }
                sink.send(Message::Binary(chunk.to_vec())).await?;
            }

            // Window resize
            _ = sigwinch.recv() => {
                if let Some((cols, rows)) = current_winsize() {
                    let frame = serde_json::to_string(&ClientControl::Resize { cols, rows })
                        .expect("serialize");
                    sink.send(Message::Text(frame)).await?;
                }
            }
        }
    }

    // Restore the TTY before printing so the message lands in cooked mode.
    drop(raw);
    if session_ended {
        let _ = stdout.write_all(b"[maestro] session ended\n").await;
        let _ = stdout.flush().await;
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_detach_key_handles_caret_form() {
        assert_eq!(parse_detach_key("^q").unwrap(), 0x11);
        assert_eq!(parse_detach_key("^Q").unwrap(), 0x11);
        assert_eq!(parse_detach_key("^c").unwrap(), 0x03);
    }

    #[test]
    fn parse_detach_key_handles_literal() {
        assert_eq!(parse_detach_key("x").unwrap(), b'x');
    }

    #[test]
    fn parse_detach_key_rejects_multi_char() {
        assert!(matches!(
            parse_detach_key("ab"),
            Err(AttachError::InvalidDetachKey(_))
        ));
    }

    #[test]
    fn attach_surfaces_structured_handshake_errors() {
        let body = serde_json::to_vec(&maestro_protocol::ApiError::new(
            "not_terminal",
            "session is not a terminal session",
        ))
        .unwrap();
        let response = tokio_tungstenite::tungstenite::http::Response::builder()
            .status(409)
            .body(Some(body))
            .unwrap();

        let error =
            AttachError::from_websocket(tokio_tungstenite::tungstenite::Error::Http(response));

        assert_eq!(
            error.to_string(),
            "daemon returned 409 Conflict (not_terminal): session is not a terminal session"
        );
    }
}