rsh/crates/rshc/src/cmd/connect.rs

use crate::auth::AuthedClient;
use crate::cmd::connection;
use crate::ui;
use anyhow::{anyhow, Result};
use crossterm::terminal;
use rsh_types::{AttachIOFrame, OpReq, OpResp};
use std::io::Write;
use tokio::io::AsyncReadExt;

pub async fn run(
    client: &AuthedClient,
    session: String,
    connection_id: Option<u64>,
    no_pty: bool,
) -> Result<()> {
    let conns = connection::fetch(client, Some(session.clone())).await?;
    if conns.is_empty() {
        return Err(anyhow!("no connections for session '{session}'"));
    }
    let target = match connection_id {
        Some(id) => {
            conns
                .iter()
                .find(|c| c.connection_id == id)
                .ok_or_else(|| anyhow!("no connection {id} in session '{session}'"))?
                .clone()
        }
        None => {
            if conns.len() == 1 {
                conns[0].clone()
            } else {
                let labels: Vec<String> = conns
                    .iter()
                    .map(|c| format!("#{} {}@{} ({})", c.connection_id, c.info.user, c.info.hostname, ui::fmt_time(c.connected_at)))
                    .collect();
                let pick = inquire::Select::new("select connection:", labels.clone())
                    .prompt()
                    .map_err(|e| anyhow!("prompt: {e}"))?;
                let idx = labels.iter().position(|l| l == &pick).unwrap();
                conns[idx].clone()
            }
        }
    };

    let (cols, rows) = terminal::size().unwrap_or((80, 24));
    let pty = !no_pty;
    let (attach_id, mut resps) = client
        .req_stream(OpReq::Attach {
            session: session.clone(),
            connection_id: Some(target.connection_id),
            shell_id: None,
            pty,
            cols,
            rows,
        })
        .await?;

    let ready = resps.recv().await.ok_or_else(|| anyhow!("attach: no response"))?;
    match ready {
        OpResp::AttachReady { .. } => {}
        OpResp::Err(e) => {
            client.drop_stream(attach_id).await;
            return Err(anyhow!(e));
        }
        other => {
            client.drop_stream(attach_id).await;
            return Err(anyhow!("unexpected: {other:?}"));
        }
    }
    ui::print_info(&format!(
        "attached to #{} ({}@{}) — Ctrl-] to detach",
        target.connection_id, target.info.user, target.info.hostname,
    ));

    if pty {
        terminal::enable_raw_mode().ok();
    }
    let result = pump(client, &mut resps, pty).await;
    if pty {
        terminal::disable_raw_mode().ok();
    }
    client.drop_stream(attach_id).await;
    println!();
    result
}

pub async fn pump(
    client: &AuthedClient,
    resps: &mut tokio::sync::mpsc::Receiver<OpResp>,
    pty: bool,
) -> Result<()> {
    let mut stdin = tokio::io::stdin();
    let mut buf = [0u8; 4096];

    let mut resize_rx: Option<tokio::sync::mpsc::Receiver<(u16, u16)>> = if pty {
        let (tx, rx) = tokio::sync::mpsc::channel(8);
        spawn_resize_watch(tx);
        Some(rx)
    } else {
        None
    };

    loop {
        tokio::select! {
            r = stdin.read(&mut buf) => {
                let n = r?;
                if n == 0 {
                    let _ = client.send_attach_io(AttachIOFrame::Eof).await;
                    continue;
                }
                if pty && buf[..n].iter().any(|&b| b == 0x1d) {
                    let _ = client.send_attach_io(AttachIOFrame::Kill).await;
                    return Ok(());
                }
                let _ = client.send_attach_io(AttachIOFrame::Stdin(buf[..n].to_vec())).await;
            }
            Some(resp) = resps.recv() => {
                match resp {
                    OpResp::Stdout(b) => {
                        let mut out = std::io::stdout();
                        out.write_all(&b).ok();
                        out.flush().ok();
                    }
                    OpResp::Stderr(b) => {
                        let mut err = std::io::stderr();
                        err.write_all(&b).ok();
                        err.flush().ok();
                    }
                    OpResp::Exited { code } => {
                        ui::print_info(&format!("remote exited (code {:?})", code));
                        return Ok(());
                    }
                    OpResp::Err(e) => return Err(anyhow!(e)),
                    _ => {}
                }
            }
            Some((cols, rows)) = async {
                match resize_rx.as_mut() {
                    Some(r) => r.recv().await,
                    None => std::future::pending().await,
                }
            } => {
                let _ = client.send_attach_io(AttachIOFrame::Resize { cols, rows }).await;
            }
        }
    }
}

pub fn spawn_resize_watch(tx: tokio::sync::mpsc::Sender<(u16, u16)>) {
    tokio::spawn(async move {
        let mut last = terminal::size().unwrap_or((80, 24));
        loop {
            tokio::time::sleep(std::time::Duration::from_millis(500)).await;
            if let Ok(sz) = terminal::size() {
                if sz != last {
                    last = sz;
                    if tx.send(sz).await.is_err() {
                        break;
                    }
                }
            }
        }
    });
}