Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
8.7 kB · 298 lines
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299use anyhow::{Context, Result};use nix::sys::signal::{Signal, kill};use nix::unistd::{Gid, Pid, Uid, User, getgrouplist, setgid, setgroups, setuid};use std::ffi::{CString, OsStr, OsString};use std::io;use std::os::unix::process::ExitStatusExt;use std::path::PathBuf;use std::process::Stdio;use std::time::Duration;use tokio::io::{AsyncRead, AsyncReadExt};use tokio::process::{Child, Command};use tokio::sync::mpsc::{self, Receiver, Sender};use tokio::task::JoinHandle;use tracing::warn;
#[derive(Clone, Debug)]pub struct Spec { pub program: OsString, pub args: Vec<OsString>, pub env: Vec<(OsString, OsString)>, pub cwd: Option<PathBuf>, pub timeout: Option<Duration>, pub uid: Option<u32>, pub gid: Option<u32>,}
impl Spec { pub fn new(program: impl Into<OsString>) -> Self { Self { program: program.into(), args: Vec::new(), env: Vec::new(), cwd: None, timeout: None, uid: None, gid: None, } }
pub fn arg(mut self, arg: impl Into<OsString>) -> Self { self.args.push(arg.into()); self }
pub fn args<I, S>(mut self, args: I) -> Self where I: IntoIterator<Item = S>, S: Into<OsString>, { self.args.extend(args.into_iter().map(Into::into)); self }
pub fn envs<I, K, V>(mut self, env: I) -> Self where I: IntoIterator<Item = (K, V)>, K: Into<OsString>, V: Into<OsString>, { self.env.extend( env.into_iter() .map(|(key, value)| (key.into(), value.into())), ); self }
pub fn cwd(mut self, cwd: impl Into<PathBuf>) -> Self { self.cwd = Some(cwd.into()); self }
pub fn timeout(mut self, timeout: Duration) -> Self { self.timeout = Some(timeout); self }
pub fn run_as(mut self, uid: u32, gid: u32) -> Self { self.uid = Some(uid); self.gid = Some(gid); self }}
#[derive(Clone, Debug)]pub struct ExitResult { pub exit_code: i32, pub error: Option<String>, pub timed_out: bool,}
#[derive(Clone, Debug)]pub struct CaptureOutput { pub exit: ExitResult, pub stdout: Vec<u8>, pub stderr: Vec<u8>,}
impl CaptureOutput { pub fn success(&self) -> bool { self.exit.exit_code == 0 && self.exit.error.is_none() }
pub fn combined_lossy(&self) -> String { let mut data = self.stdout.clone(); data.extend_from_slice(&self.stderr); String::from_utf8_lossy(&data).trim().to_owned() }}
#[derive(Clone, Copy, Debug)]pub enum OutKind { Stdout, Stderr,}
#[derive(Clone, Debug)]pub struct OutData { pub data: Vec<u8>, pub kind: OutKind,}
pub struct StreamingCommand { events: Receiver<OutData>, exit: JoinHandle<Result<ExitResult>>,}
impl StreamingCommand { pub fn into_parts(self) -> (Receiver<OutData>, JoinHandle<Result<ExitResult>>) { (self.events, self.exit) }}
pub async fn run_capture(spec: Spec) -> Result<CaptureOutput> { let running = spawn_streaming(spec)?; let mut stdout = Vec::new(); let mut stderr = Vec::new(); let (mut events, exit_task) = running.into_parts();
while let Some(event) = events.recv().await { match event.kind { OutKind::Stdout => stdout.extend_from_slice(&event.data), OutKind::Stderr => stderr.extend_from_slice(&event.data), } }
let exit = exit_task .await .unwrap_or_else(|error| Err(anyhow::anyhow!("command supervisor failed: {error}")))?; Ok(CaptureOutput { exit, stdout, stderr, })}
pub fn spawn_streaming(mut spec: Spec) -> Result<StreamingCommand> { let mut child = spawn(&mut spec)?; let stdout = child.stdout.take().context("stdout pipe missing")?; let stderr = child.stderr.take().context("stderr pipe missing")?;
let (events_tx, events_rx) = mpsc::channel(64); let stdout_thread = spawn_reader(stdout, events_tx.clone(), OutKind::Stdout); let stderr_thread = spawn_reader(stderr, events_tx.clone(), OutKind::Stderr); drop(events_tx);
let exit = tokio::spawn(async move { let exit = wait_child(&mut child, spec.timeout).await;
// ensure all output is observed before exiting // this assumes children dont daemonize and hold onto the stdout/err stdout_thread.await.context("stdout reader task failed")?; stderr_thread.await.context("stderr reader task failed")?;
Ok(exit) });
Ok(StreamingCommand { events: events_rx, exit, })}
fn spawn(spec: &mut Spec) -> Result<Child> { let mut cmd = Command::new(&spec.program); cmd.args(&spec.args) .envs(spec.env.iter().map(|(key, value)| (key, value))) .stdout(Stdio::piped()) .stderr(Stdio::piped());
if let Some(cwd) = &spec.cwd { cmd.current_dir(cwd); }
// don't use rust's .uid() / .gid() methods here because they clear // supplemantary groups, which means for example adding a user to "docker" // group won't actually let it access the sock. // https://github.com/rust-lang/rust/issues/90747 if let (Some(uid), Some(gid)) = (spec.uid, spec.gid) { let username = User::from_uid(Uid::from_raw(uid)) .ok() .flatten() .map(|u| u.name) .with_context(|| format!("lookup passwd entry for uid {uid}"))?; let cname = CString::new(username) .with_context(|| format!("username for uid {uid} contained a null byte"))?; // resolve groups beforehand so we don't have to read /etc/group in the pre_exec let groups = getgrouplist(&cname, Gid::from_raw(gid)).context("resolve supplementary groups")?; // SAFETY: pre_exec runs between fork and execve in the child. // we only call async-signal-safe syscalls and we don't touch any // shared state, no allocator, no mutexes, no globals. unsafe { cmd.pre_exec(move || { setgroups(&groups).map_err(io::Error::from)?; setgid(Gid::from_raw(gid)).map_err(io::Error::from)?; setuid(Uid::from_raw(uid)).map_err(io::Error::from)?; Ok(()) }); } }
// allow us to kill this whole process tree on deadline cmd.process_group(0);
cmd.spawn() .with_context(|| format!("spawn {}", display_os(&spec.program)))}
async fn wait_child(child: &mut Child, timeout: Option<Duration>) -> ExitResult { let wait = child.wait(); let status = match timeout { Some(timeout) => match tokio::time::timeout(timeout, wait).await { Ok(status) => status, Err(_) => { if let Some(pid) = child.id() && let Err(error) = kill(Pid::from_raw(-(pid as i32)), Signal::SIGKILL) { warn!(pid, %error, "failed to kill process group"); } let _ = child.wait().await; return ExitResult { exit_code: 124, error: Some("command timed out".to_owned()), timed_out: true, }; } }, None => wait.await, };
match status { Ok(status) => ExitResult { exit_code: status .code() .or_else(|| status.signal().map(|signal| 128 + signal)) .unwrap_or(1), error: None, timed_out: false, }, Err(error) => ExitResult { exit_code: 1, error: Some(error.to_string()), timed_out: false, }, }}
fn spawn_reader( mut reader: impl AsyncRead + Unpin + Send + 'static, events: Sender<OutData>, kind: OutKind,) -> JoinHandle<()> { tokio::spawn(async move { let mut buf = [0_u8; 32 * 1024]; loop { match reader.read(&mut buf).await { Ok(0) => return, Ok(n) => { let event = OutData { data: buf[..n].to_vec(), kind, }; if events.send(event).await.is_err() { return; } } Err(error) => { warn!(%error, "failed to read command stream"); return; } } } })}
fn display_os(value: &OsStr) -> String { value.to_string_lossy().into_owned()}