initial commit

crates/rshc/src/cmd/connect.rs

@@ -0,0 +1,163 @@
+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),
+            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
+}
+
+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;
+            }
+        }
+    }
+}
+
+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;
+                    }
+                }
+            }
+        }
+    });
+}