Something went wrong. Try again.
TRAMP-style remote editing for nushell: an ssh agent and a nu plugin, in Rust
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828//! Process operations for the tramp-agent RPC server.//!//! Implements the following RPC methods://!//! | Method | Description |//! |-----------------|-------------------------------------------------------|//! | `process.run` | Run a command synchronously, collect stdout/stderr |//! | `process.start` | Start a long-running process, return a handle |//! | `process.read` | Read buffered stdout/stderr from a started process |//! | `process.write` | Write data to a started process's stdin |//! | `process.kill` | Kill a started process by handle |
use std::collections::HashMap;use std::process::Stdio;use std::sync::atomic::{AtomicU64, Ordering};
use rmpv::Value;use tokio::io::{AsyncReadExt, AsyncWriteExt};use tokio::process::{Child, Command};use tokio::sync::Mutex;
use crate::rpc::{Response, error_code};
// ---------------------------------------------------------------------------// Helpers// ---------------------------------------------------------------------------
/// Extract a required string parameter from a MsgPack map by key.fn get_str_param<'a>(params: &'a Value, key: &str) -> Result<&'a str, Response> { params .as_map() .and_then(|m| { m.iter() .find(|(k, _)| k.as_str() == Some(key)) .and_then(|(_, v)| v.as_str()) }) .ok_or_else(|| { Response::err( 0, error_code::INVALID_PARAMS, format!("missing or invalid parameter: {key}"), ) })}
/// Extract an optional array of string parameters from a MsgPack map by key.fn get_str_array_param<'a>(params: &'a Value, key: &str) -> Option<Vec<&'a str>> { params.as_map().and_then(|m| { m.iter() .find(|(k, _)| k.as_str() == Some(key)) .and_then(|(_, v)| v.as_array()) .map(|arr| arr.iter().filter_map(|v| v.as_str()).collect()) })}
/// Extract an optional binary parameter from a MsgPack map by key.fn get_bin_param<'a>(params: &'a Value, key: &str) -> Option<&'a [u8]> { params.as_map().and_then(|m| { m.iter() .find(|(k, _)| k.as_str() == Some(key)) .and_then(|(_, v)| v.as_slice()) })}
/// Extract a required u64 parameter from a MsgPack map by key.fn get_u64_param(params: &Value, key: &str) -> Result<u64, Response> { params .as_map() .and_then(|m| { m.iter() .find(|(k, _)| k.as_str() == Some(key)) .and_then(|(_, v)| v.as_u64()) }) .ok_or_else(|| { Response::err( 0, error_code::INVALID_PARAMS, format!("missing or invalid parameter: {key}"), ) })}
/// Extract an optional string parameter from a MsgPack map by key.fn get_optional_str_param<'a>(params: &'a Value, key: &str) -> Option<&'a str> { params.as_map().and_then(|m| { m.iter() .find(|(k, _)| k.as_str() == Some(key)) .and_then(|(_, v)| v.as_str()) })}
/// Extract an optional map of string→string environment variables from a/// MsgPack map by key.fn get_env_param(params: &Value, key: &str) -> Option<Vec<(String, String)>> { params.as_map().and_then(|m| { m.iter() .find(|(k, _)| k.as_str() == Some(key)) .and_then(|(_, v)| v.as_map()) .map(|pairs| { pairs .iter() .filter_map(|(k, v)| Some((k.as_str()?.to_owned(), v.as_str()?.to_owned()))) .collect() }) })}
// ---------------------------------------------------------------------------// Process table// ---------------------------------------------------------------------------
/// Global counter for process handles.static NEXT_HANDLE: AtomicU64 = AtomicU64::new(1);
/// A managed child process with buffered I/O.struct ManagedProcess { child: Child, /// Accumulated stdout that hasn't been read yet. stdout_buf: Vec<u8>, /// Accumulated stderr that hasn't been read yet. stderr_buf: Vec<u8>,}
/// Shared process table. Maps handle IDs to running child processes.////// The table is wrapped in a `Mutex` because process reads/writes must be/// serialised (we're dealing with streaming I/O on a single child).pub struct ProcessTable { processes: Mutex<HashMap<u64, ManagedProcess>>,}
impl ProcessTable { /// Create an empty process table. pub fn new() -> Self { Self { processes: Mutex::new(HashMap::new()), } }}
impl Default for ProcessTable { fn default() -> Self { Self::new() }}
// ---------------------------------------------------------------------------// RPC method handlers// ---------------------------------------------------------------------------
/// Maximum number of bytes to read from a process's stdout/stderr in a/// single `process.read` call.const READ_BUF_SIZE: usize = 64 * 1024;
/// `process.run` — run a command synchronously and collect its output.////// Params:/// - `program`: the program to execute (string, required)/// - `args`: arguments (array of strings, optional, default `[]`)/// - `cwd`: working directory (string, optional)/// - `env`: environment variables to set (map of string→string, optional)/// - `stdin_data`: data to pipe to stdin (binary, optional)////// Result: `{ stdout: <binary>, stderr: <binary>, exit_code: <integer> }`pub async fn run(id: u64, params: &Value) -> Response { let program = match get_str_param(params, "program") { Ok(p) => p, Err(mut e) => { e.id = id; return e; } };
let args = get_str_array_param(params, "args").unwrap_or_default(); let cwd = get_optional_str_param(params, "cwd"); let env_vars = get_env_param(params, "env"); let stdin_data = get_bin_param(params, "stdin_data");
let has_stdin = stdin_data.is_some();
let mut cmd = Command::new(program); cmd.args(&args) .stdout(Stdio::piped()) .stderr(Stdio::piped());
if has_stdin { cmd.stdin(Stdio::piped()); } else { cmd.stdin(Stdio::null()); }
if let Some(dir) = cwd { cmd.current_dir(dir); }
if let Some(vars) = &env_vars { for (k, v) in vars { cmd.env(k, v); } }
let result = if let Some(data) = stdin_data { // Spawn and write to stdin before waiting. match cmd.spawn() { Ok(mut child) => { if let Some(mut stdin_handle) = child.stdin.take() { // Ignore write errors — the child may have already exited. let _ = stdin_handle.write_all(data).await; drop(stdin_handle); } child.wait_with_output().await } Err(e) => { return Response::err( id, error_code::IO_ERROR, format!("failed to spawn `{program}`: {e}"), ); } } } else { cmd.output().await };
match result { Ok(output) => Response::ok( id, Value::Map(vec![ (Value::String("stdout".into()), Value::Binary(output.stdout)), (Value::String("stderr".into()), Value::Binary(output.stderr)), ( Value::String("exit_code".into()), Value::Integer(output.status.code().unwrap_or(-1).into()), ), ]), ), Err(e) => Response::err( id, error_code::IO_ERROR, format!("failed to execute `{program}`: {e}"), ), }}
/// `process.start` — start a long-running process and return a handle.////// Params:/// - `program`: the program to execute (string, required)/// - `args`: arguments (array of strings, optional, default `[]`)/// - `cwd`: working directory (string, optional)/// - `env`: environment variables to set (map of string→string, optional)////// Result: `{ handle: <u64> }`pub async fn start(id: u64, params: &Value, table: &ProcessTable) -> Response { let program = match get_str_param(params, "program") { Ok(p) => p, Err(mut e) => { e.id = id; return e; } };
let args = get_str_array_param(params, "args").unwrap_or_default(); let cwd = get_optional_str_param(params, "cwd"); let env_vars = get_env_param(params, "env");
let mut cmd = Command::new(program); cmd.args(&args) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::piped());
if let Some(dir) = cwd { cmd.current_dir(dir); }
if let Some(vars) = &env_vars { for (k, v) in vars { cmd.env(k, v); } }
match cmd.spawn() { Ok(child) => { let handle = NEXT_HANDLE.fetch_add(1, Ordering::Relaxed); let managed = ManagedProcess { child, stdout_buf: Vec::new(), stderr_buf: Vec::new(), }; table.processes.lock().await.insert(handle, managed);
Response::ok( id, Value::Map(vec![( Value::String("handle".into()), Value::Integer(handle.into()), )]), ) } Err(e) => Response::err( id, error_code::IO_ERROR, format!("failed to spawn `{program}`: {e}"), ), }}
/// `process.read` — read buffered output from a started process.////// This performs a non-blocking read: it returns whatever data is currently/// available (up to 64 KiB from each stream), plus whether the process has/// exited.////// Params:/// - `handle`: process handle from `process.start` (u64, required)////// Result: `{ stdout: <binary>, stderr: <binary>, running: <bool>,/// exit_code: <integer|null> }`pub async fn read(id: u64, params: &Value, table: &ProcessTable) -> Response { let handle = match get_u64_param(params, "handle") { Ok(h) => h, Err(mut e) => { e.id = id; return e; } };
let mut processes = table.processes.lock().await;
let Some(managed) = processes.get_mut(&handle) else { return Response::err( id, error_code::NOT_FOUND, format!("no process with handle {handle}"), ); };
// Try to read available data from stdout. let mut stdout_chunk = vec![0u8; READ_BUF_SIZE]; if let Some(ref mut stdout) = managed.child.stdout { match tokio::time::timeout( std::time::Duration::from_millis(10), stdout.read(&mut stdout_chunk), ) .await { Ok(Ok(n)) if n > 0 => { managed.stdout_buf.extend_from_slice(&stdout_chunk[..n]); } _ => {} } }
// Try to read available data from stderr. let mut stderr_chunk = vec![0u8; READ_BUF_SIZE]; if let Some(ref mut stderr) = managed.child.stderr { match tokio::time::timeout( std::time::Duration::from_millis(10), stderr.read(&mut stderr_chunk), ) .await { Ok(Ok(n)) if n > 0 => { managed.stderr_buf.extend_from_slice(&stderr_chunk[..n]); } _ => {} } }
// Drain the accumulated buffers. let stdout_data = std::mem::take(&mut managed.stdout_buf); let stderr_data = std::mem::take(&mut managed.stderr_buf);
// Check if the process has exited. let (running, exit_code) = match managed.child.try_wait() { Ok(Some(status)) => (false, Some(status.code().unwrap_or(-1))), Ok(None) => (true, None), Err(_) => (false, Some(-1)), };
let exit_code_val = match exit_code { Some(code) => Value::Integer(code.into()), None => Value::Nil, };
// If the process has exited, remove it from the table. if !running { processes.remove(&handle); }
Response::ok( id, Value::Map(vec![ (Value::String("stdout".into()), Value::Binary(stdout_data)), (Value::String("stderr".into()), Value::Binary(stderr_data)), (Value::String("running".into()), Value::Boolean(running)), (Value::String("exit_code".into()), exit_code_val), ]), )}
/// `process.write` — write data to a started process's stdin.////// Params:/// - `handle`: process handle from `process.start` (u64, required)/// - `data`: binary data to write (binary, required)/// - `close`: if `true`, close stdin after writing (boolean, optional, default `false`)////// Result: `{}` (empty map on success).pub async fn write(id: u64, params: &Value, table: &ProcessTable) -> Response { let handle = match get_u64_param(params, "handle") { Ok(h) => h, Err(mut e) => { e.id = id; return e; } };
let data = match params.as_map().and_then(|m| { m.iter() .find(|(k, _)| k.as_str() == Some("data")) .and_then(|(_, v)| v.as_slice()) }) { Some(d) => d, None => { return Response::err( id, error_code::INVALID_PARAMS, "missing or invalid binary parameter: data", ); } };
let close = params .as_map() .and_then(|m| { m.iter() .find(|(k, _)| k.as_str() == Some("close")) .and_then(|(_, v)| v.as_bool()) }) .unwrap_or(false);
let mut processes = table.processes.lock().await;
let Some(managed) = processes.get_mut(&handle) else { return Response::err( id, error_code::NOT_FOUND, format!("no process with handle {handle}"), ); };
if let Some(ref mut stdin_handle) = managed.child.stdin { if let Err(e) = stdin_handle.write_all(data).await { return Response::err( id, error_code::IO_ERROR, format!("failed to write to process stdin: {e}"), ); } if let Err(e) = stdin_handle.flush().await { return Response::err( id, error_code::IO_ERROR, format!("failed to flush process stdin: {e}"), ); } } else { return Response::err( id, error_code::IO_ERROR, "process stdin is not available (already closed?)", ); }
// Close stdin if requested (signals EOF to the child). if close { managed.child.stdin.take(); }
Response::ok(id, Value::Map(vec![]))}
/// `process.kill` — kill a started process.////// Params:/// - `handle`: process handle from `process.start` (u64, required)////// Result: `{}` (empty map on success).pub async fn kill(id: u64, params: &Value, table: &ProcessTable) -> Response { let handle = match get_u64_param(params, "handle") { Ok(h) => h, Err(mut e) => { e.id = id; return e; } };
let mut processes = table.processes.lock().await;
let Some(mut managed) = processes.remove(&handle) else { return Response::err( id, error_code::NOT_FOUND, format!("no process with handle {handle}"), ); };
match managed.child.kill().await { Ok(()) => Response::ok(id, Value::Map(vec![])), Err(e) => Response::err( id, error_code::IO_ERROR, format!("failed to kill process: {e}"), ), }}
// ---------------------------------------------------------------------------// Tests// ---------------------------------------------------------------------------
#[cfg(test)]mod tests { use super::*; use rmpv::Value;
fn make_params(pairs: Vec<(&str, Value)>) -> Value { Value::Map( pairs .into_iter() .map(|(k, v)| (Value::String(k.into()), v)) .collect(), ) }
#[tokio::test] async fn run_simple_command() { let params = make_params(vec![ ("program", Value::String("echo".into())), ( "args", Value::Array(vec![Value::String("hello world".into())]), ), ]); let resp = run(1, ¶ms).await; assert!(resp.error.is_none(), "expected ok, got: {:?}", resp.error);
let result = resp.result.unwrap(); let map = result.as_map().unwrap();
let stdout = map .iter() .find(|(k, _)| k.as_str() == Some("stdout")) .unwrap() .1 .as_slice() .unwrap(); assert_eq!(String::from_utf8_lossy(stdout).trim(), "hello world");
let exit_code = map .iter() .find(|(k, _)| k.as_str() == Some("exit_code")) .unwrap() .1 .as_i64() .unwrap(); assert_eq!(exit_code, 0); }
#[tokio::test] async fn run_failing_command() { let params = make_params(vec![("program", Value::String("false".into()))]); let resp = run(2, ¶ms).await; assert!(resp.error.is_none()); // The RPC itself succeeds
let result = resp.result.unwrap(); let map = result.as_map().unwrap(); let exit_code = map .iter() .find(|(k, _)| k.as_str() == Some("exit_code")) .unwrap() .1 .as_i64() .unwrap(); assert_ne!(exit_code, 0); }
#[tokio::test] async fn run_nonexistent_command() { let params = make_params(vec![( "program", Value::String("__tramp_agent_no_such_command_99999__".into()), )]); let resp = run(3, ¶ms).await; assert!(resp.error.is_some()); assert_eq!(resp.error.unwrap().code, error_code::IO_ERROR); }
#[tokio::test] async fn run_with_stdin_data() { let params = make_params(vec![ ("program", Value::String("cat".into())), ("stdin_data", Value::Binary(b"piped input".to_vec())), ]); let resp = run(4, ¶ms).await; assert!(resp.error.is_none(), "expected ok, got: {:?}", resp.error);
let result = resp.result.unwrap(); let map = result.as_map().unwrap(); let stdout = map .iter() .find(|(k, _)| k.as_str() == Some("stdout")) .unwrap() .1 .as_slice() .unwrap(); assert_eq!(stdout, b"piped input"); }
#[tokio::test] async fn run_with_cwd() { let params = make_params(vec![ ("program", Value::String("pwd".into())), ("cwd", Value::String("/tmp".into())), ]); let resp = run(5, ¶ms).await; assert!(resp.error.is_none());
let result = resp.result.unwrap(); let map = result.as_map().unwrap(); let stdout = map .iter() .find(|(k, _)| k.as_str() == Some("stdout")) .unwrap() .1 .as_slice() .unwrap(); // /tmp may be a symlink (e.g. to /private/tmp on macOS), so // we just check it contains "tmp". let output = String::from_utf8_lossy(stdout); assert!(output.contains("tmp"), "unexpected pwd output: {output}"); }
#[tokio::test] async fn run_with_env() { let params = make_params(vec![ ("program", Value::String("sh".into())), ( "args", Value::Array(vec![ Value::String("-c".into()), Value::String("echo $TRAMP_TEST_VAR".into()), ]), ), ( "env", Value::Map(vec![( Value::String("TRAMP_TEST_VAR".into()), Value::String("agent_value".into()), )]), ), ]); let resp = run(6, ¶ms).await; assert!(resp.error.is_none());
let result = resp.result.unwrap(); let map = result.as_map().unwrap(); let stdout = map .iter() .find(|(k, _)| k.as_str() == Some("stdout")) .unwrap() .1 .as_slice() .unwrap(); assert_eq!(String::from_utf8_lossy(stdout).trim(), "agent_value"); }
#[tokio::test] async fn run_missing_program_param() { let params = Value::Map(vec![]); let resp = run(99, ¶ms).await; assert!(resp.error.is_some()); assert_eq!(resp.error.unwrap().code, error_code::INVALID_PARAMS); }
#[tokio::test] async fn start_read_kill_lifecycle() { let table = ProcessTable::new();
// Start a long-running process (sleep). let start_params = make_params(vec![ ("program", Value::String("sleep".into())), ("args", Value::Array(vec![Value::String("60".into())])), ]); let resp = start(10, &start_params, &table).await; assert!(resp.error.is_none(), "start failed: {:?}", resp.error);
let result = resp.result.unwrap(); let map = result.as_map().unwrap(); let handle = map .iter() .find(|(k, _)| k.as_str() == Some("handle")) .unwrap() .1 .as_u64() .unwrap();
// Read — process should be running with no output yet. let read_params = make_params(vec![("handle", Value::Integer(handle.into()))]); let resp = read(11, &read_params, &table).await; assert!(resp.error.is_none());
let result = resp.result.unwrap(); let map = result.as_map().unwrap(); let running = map .iter() .find(|(k, _)| k.as_str() == Some("running")) .unwrap() .1 .as_bool() .unwrap(); assert!(running);
// Kill the process. let kill_params = make_params(vec![("handle", Value::Integer(handle.into()))]); let resp = kill(12, &kill_params, &table).await; assert!(resp.error.is_none(), "kill failed: {:?}", resp.error);
// Reading from the killed handle should fail (removed from table). let resp = read(13, &read_params, &table).await; assert!(resp.error.is_some()); assert_eq!(resp.error.unwrap().code, error_code::NOT_FOUND); }
#[tokio::test] async fn start_and_write_stdin() { let table = ProcessTable::new();
// Start `cat` which reads stdin and echoes to stdout. let start_params = make_params(vec![("program", Value::String("cat".into()))]); let resp = start(20, &start_params, &table).await; assert!(resp.error.is_none());
let result = resp.result.unwrap(); let map = result.as_map().unwrap(); let handle = map .iter() .find(|(k, _)| k.as_str() == Some("handle")) .unwrap() .1 .as_u64() .unwrap();
// Write data and close stdin. let write_params = make_params(vec![ ("handle", Value::Integer(handle.into())), ("data", Value::Binary(b"hello from agent\n".to_vec())), ("close", Value::Boolean(true)), ]); let resp = write(21, &write_params, &table).await; assert!(resp.error.is_none(), "write failed: {:?}", resp.error);
// Give cat a moment to process and exit. tokio::time::sleep(std::time::Duration::from_millis(200)).await;
// Read output — cat should have exited. let read_params = make_params(vec![("handle", Value::Integer(handle.into()))]); let resp = read(22, &read_params, &table).await; assert!(resp.error.is_none(), "read failed: {:?}", resp.error);
let result = resp.result.unwrap(); let map = result.as_map().unwrap();
let stdout = map .iter() .find(|(k, _)| k.as_str() == Some("stdout")) .unwrap() .1 .as_slice() .unwrap(); assert_eq!(stdout, b"hello from agent\n");
let running = map .iter() .find(|(k, _)| k.as_str() == Some("running")) .unwrap() .1 .as_bool() .unwrap(); assert!(!running); }
#[tokio::test] async fn kill_nonexistent_handle() { let table = ProcessTable::new();
let params = make_params(vec![("handle", Value::Integer(999999.into()))]); let resp = kill(30, ¶ms, &table).await; assert!(resp.error.is_some()); assert_eq!(resp.error.unwrap().code, error_code::NOT_FOUND); }
#[tokio::test] async fn write_to_nonexistent_handle() { let table = ProcessTable::new();
let params = make_params(vec![ ("handle", Value::Integer(888888.into())), ("data", Value::Binary(b"data".to_vec())), ]); let resp = write(31, ¶ms, &table).await; assert!(resp.error.is_some()); assert_eq!(resp.error.unwrap().code, error_code::NOT_FOUND); }
#[tokio::test] async fn read_nonexistent_handle() { let table = ProcessTable::new();
let params = make_params(vec![("handle", Value::Integer(777777.into()))]); let resp = read(32, ¶ms, &table).await; assert!(resp.error.is_some()); assert_eq!(resp.error.unwrap().code, error_code::NOT_FOUND); }}