rust: local shell client on portable-pty
This commit is contained in:
Andy Dixon
@@ -20,3 +20,4 @@ tokio = { version = "1.53.2", features = ["full"] }
|
||||
sha2 = "0.11.0"
|
||||
hmac = "0.13.0"
|
||||
sha1 = "0.11.0"
|
||||
portable-pty = "0.9.0"
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
//! Local shell in a PTY (forkpty on Unix, ConPTY on Windows), a port of internal/localshell.
|
||||
|
||||
use super::Callbacks;
|
||||
use crate::Result;
|
||||
use portable_pty::{CommandBuilder, MasterPty, PtySize, native_pty_system};
|
||||
use std::io::{Read, Write};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
pub fn default_shell() -> String {
|
||||
if cfg!(windows) {
|
||||
let ps = std::env::var_os("SystemRoot")
|
||||
.map(|r| PathBuf::from(r).join(r"System32\WindowsPowerShell\v1.0\powershell.exe"))
|
||||
.filter(|p| p.exists());
|
||||
return ps.map_or_else(|| "cmd.exe".into(), |p| p.to_string_lossy().into());
|
||||
}
|
||||
std::env::var("SHELL").ok().filter(|s| !s.is_empty()).unwrap_or_else(|| "/bin/sh".into())
|
||||
}
|
||||
|
||||
fn is_utf8_locale(v: &str) -> bool {
|
||||
let v = v.to_lowercase();
|
||||
v.contains("utf-8") || v.contains("utf8")
|
||||
}
|
||||
|
||||
/// TERM=xterm-256color and a UTF-8 character locale, keeping the user's environment otherwise.
|
||||
pub fn local_shell_env(env: &mut Vec<(String, String)>) {
|
||||
let set = |env: &mut Vec<(String, String)>, k: &str, v: &str| {
|
||||
env.retain(|(ek, _)| ek != k);
|
||||
env.push((k.into(), v.into()));
|
||||
};
|
||||
set(env, "TERM", "xterm-256color");
|
||||
let get = |k: &str| env.iter().find(|(ek, _)| ek == k).map(|(_, v)| v.clone());
|
||||
let effective = get("LC_ALL").or_else(|| get("LC_CTYPE")).or_else(|| get("LANG"));
|
||||
if effective.as_deref().is_some_and(is_utf8_locale) {
|
||||
return;
|
||||
}
|
||||
let locale = if cfg!(target_os = "macos") { "en_US.UTF-8" } else { "C.UTF-8" };
|
||||
let key = if get("LC_ALL").is_some() {
|
||||
"LC_ALL"
|
||||
} else if get("LC_CTYPE").is_some() || get("LANG").is_some() {
|
||||
"LC_CTYPE"
|
||||
} else {
|
||||
"LANG"
|
||||
};
|
||||
set(env, key, locale);
|
||||
}
|
||||
|
||||
struct Running {
|
||||
master: Box<dyn MasterPty + Send>,
|
||||
writer: Box<dyn Write + Send>,
|
||||
killer: Box<dyn portable_pty::ChildKiller + Send + Sync>,
|
||||
pid: Option<u32>,
|
||||
}
|
||||
|
||||
pub struct LocalShell {
|
||||
shell: String,
|
||||
cb: Callbacks,
|
||||
state: Mutex<(Option<Running>, PathBuf, u16, u16, bool)>, // running, cwd, cols, rows, closed
|
||||
}
|
||||
|
||||
impl LocalShell {
|
||||
pub fn new(shell: &str, start_dir: Option<PathBuf>, cols: u16, rows: u16, cb: Callbacks) -> Arc<LocalShell> {
|
||||
let shell = if shell.is_empty() { default_shell() } else { shell.to_string() };
|
||||
let dir = start_dir.filter(|d| d.is_dir()).or_else(home_dir).unwrap_or_else(|| PathBuf::from("/"));
|
||||
Arc::new(LocalShell { shell, cb, state: Mutex::new((None, dir, cols.max(1), rows.max(1), false)) })
|
||||
}
|
||||
|
||||
pub fn start(self: &Arc<Self>) -> Result<()> {
|
||||
let mut st = self.state.lock().unwrap();
|
||||
if st.4 {
|
||||
return Err("shell already closed".into());
|
||||
}
|
||||
if st.0.is_some() {
|
||||
return Err("shell already running".into());
|
||||
}
|
||||
let pair = native_pty_system()
|
||||
.openpty(PtySize { rows: st.3, cols: st.2, pixel_width: 0, pixel_height: 0 })
|
||||
.map_err(|e| format!("failed to start shell {:?}: {e}", self.shell))?;
|
||||
let mut cmd = CommandBuilder::new(&self.shell);
|
||||
let mut env: Vec<(String, String)> = std::env::vars().collect();
|
||||
local_shell_env(&mut env);
|
||||
cmd.env_clear();
|
||||
for (k, v) in env {
|
||||
cmd.env(k, v);
|
||||
}
|
||||
cmd.cwd(&st.1);
|
||||
let mut child = pair.slave.spawn_command(cmd).map_err(|e| format!("failed to start shell {:?}: {e}", self.shell))?;
|
||||
drop(pair.slave);
|
||||
let mut reader = pair.master.try_clone_reader()?;
|
||||
let writer = pair.master.take_writer()?;
|
||||
let pid = child.process_id();
|
||||
st.0 = Some(Running { master: pair.master, writer, killer: child.clone_killer(), pid });
|
||||
drop(st);
|
||||
|
||||
let me = self.clone();
|
||||
std::thread::Builder::new().name("nexus-local-read".into()).spawn(move || {
|
||||
let mut buf = vec![0u8; 32 * 1024];
|
||||
loop {
|
||||
match reader.read(&mut buf) {
|
||||
Ok(0) | Err(_) => break,
|
||||
Ok(n) => (me.cb.on_data)(&buf[..n]),
|
||||
}
|
||||
}
|
||||
})?;
|
||||
let me = self.clone();
|
||||
std::thread::Builder::new().name("nexus-local-wait".into()).spawn(move || {
|
||||
let _ = child.wait();
|
||||
let was_closed = {
|
||||
let mut st = me.state.lock().unwrap();
|
||||
let was = st.4;
|
||||
st.4 = true;
|
||||
st.0 = None;
|
||||
was
|
||||
};
|
||||
if !was_closed {
|
||||
(me.cb.on_close)();
|
||||
}
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn write(&self, data: &[u8]) -> Result<()> {
|
||||
if data.iter().any(|b| matches!(b, b'\r' | b'\n' | 4)) {
|
||||
self.capture_cwd();
|
||||
}
|
||||
let mut st = self.state.lock().unwrap();
|
||||
let r = st.0.as_mut().ok_or("shell not running")?;
|
||||
r.writer.write_all(data)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn resize(&self, cols: u16, rows: u16) -> Result<()> {
|
||||
let mut st = self.state.lock().unwrap();
|
||||
if cols > 0 {
|
||||
st.2 = cols;
|
||||
}
|
||||
if rows > 0 {
|
||||
st.3 = rows;
|
||||
}
|
||||
let size = PtySize { rows: st.3, cols: st.2, pixel_width: 0, pixel_height: 0 };
|
||||
if let Some(r) = &st.0 {
|
||||
r.master.resize(size)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Idempotent; does not fire on_close (matches Go: an explicit close is not an exit).
|
||||
pub fn close(&self) {
|
||||
let mut st = self.state.lock().unwrap();
|
||||
if st.4 {
|
||||
return;
|
||||
}
|
||||
st.4 = true;
|
||||
if let Some(mut r) = st.0.take() {
|
||||
let _ = r.killer.kill();
|
||||
}
|
||||
}
|
||||
|
||||
/// The last directory the shell was seen in (used to restart in the same place).
|
||||
pub fn working_directory(&self) -> PathBuf {
|
||||
self.state.lock().unwrap().1.clone()
|
||||
}
|
||||
|
||||
fn capture_cwd(&self) {
|
||||
let pid = self.state.lock().unwrap().0.as_ref().and_then(|r| r.pid);
|
||||
if let Some(dir) = pid.and_then(process_cwd).filter(|d| d.is_dir()) {
|
||||
self.state.lock().unwrap().1 = dir;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn home_dir() -> Option<PathBuf> {
|
||||
std::env::var_os(if cfg!(windows) { "USERPROFILE" } else { "HOME" }).map(PathBuf::from)
|
||||
}
|
||||
|
||||
fn process_cwd(pid: u32) -> Option<PathBuf> {
|
||||
if cfg!(target_os = "linux") {
|
||||
return std::fs::read_link(format!("/proc/{pid}/cwd")).ok();
|
||||
}
|
||||
if cfg!(target_os = "macos") {
|
||||
let out = std::process::Command::new("/usr/sbin/lsof").args(["-a", "-p", &pid.to_string(), "-d", "cwd", "-Fn"]).output().ok()?;
|
||||
return String::from_utf8_lossy(&out.stdout).lines().find_map(|l| l.strip_prefix('n').map(PathBuf::from));
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(all(test, unix))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
#[test]
|
||||
fn runs_echoes_and_reports_exit() {
|
||||
let seen: Arc<Mutex<Vec<u8>>> = Arc::default();
|
||||
let closed = Arc::new(Mutex::new(0));
|
||||
let (s2, c2) = (seen.clone(), closed.clone());
|
||||
let cb = Callbacks {
|
||||
on_data: Arc::new(move |d| s2.lock().unwrap().extend_from_slice(d)),
|
||||
on_error: Arc::new(|_| {}),
|
||||
on_close: Arc::new(move || *c2.lock().unwrap() += 1),
|
||||
};
|
||||
let sh = LocalShell::new("/bin/sh", Some("/tmp".into()), 80, 24, cb);
|
||||
sh.start().unwrap();
|
||||
sh.write(b"cd / && echo local-$((2*21))\n").unwrap();
|
||||
let t = Instant::now();
|
||||
while !String::from_utf8_lossy(&seen.lock().unwrap()).contains("local-42") {
|
||||
assert!(t.elapsed() < Duration::from_secs(5));
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
sh.resize(100, 30).unwrap();
|
||||
sh.write(b"exit\n").unwrap();
|
||||
while *closed.lock().unwrap() == 0 {
|
||||
assert!(t.elapsed() < Duration::from_secs(5));
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
assert_eq!(sh.working_directory(), PathBuf::from("/"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn env_gets_term_and_utf8() {
|
||||
let mut e = vec![("LANG".to_string(), "C".to_string())];
|
||||
local_shell_env(&mut e);
|
||||
assert!(e.contains(&("TERM".into(), "xterm-256color".into())));
|
||||
assert!(e.iter().any(|(k, v)| k == "LC_CTYPE" && v.ends_with("UTF-8")));
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,7 @@
|
||||
//! - `write`/`resize`/`close` are callable from any thread; `close` is idempotent and must not
|
||||
//! fire `on_close` re-entrantly while holding locks.
|
||||
|
||||
pub mod local;
|
||||
pub mod sftp;
|
||||
pub mod ssh;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user