use super::*;
use eyre::eyre;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::io::AsyncRead;
impl CmdLineRunner<'_> {
pub(crate) async fn read_isolated(mut self, limit: usize) -> Result<String> {
let _read_lock = RAW_LOCK.read().await;
let timeout = self.timeout.unwrap_or(Duration::from_secs(5));
self.cmd.kill_on_drop(true);
self.cmd
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
#[cfg(unix)]
{
self.cmd.env(TASK_PGID_MANAGED_ENV, "1");
self.cmd.process_group(0);
}
let mut child = self.spawn_async_with_etxtbsy_retry().await?;
let _running = RunningPidGuard::new(child.id());
let tree = ChildTree::new(&mut child)?;
let stdout = child.stdout.take().expect("piped stdout");
let stderr = child.stderr.take().expect("piped stderr");
let budget = AtomicUsize::new(limit);
let result = tokio::time::timeout(timeout, async {
tokio::try_join!(
child.wait(),
capture(stdout, &budget, limit),
capture(stderr, &budget, limit)
)
})
.await;
let (status, stdout, _stderr) = match result {
Ok(Ok(output)) => output,
Ok(Err(err)) => {
end(&mut child, tree).await;
return Err(err.into());
}
Err(_) => {
end(&mut child, tree).await;
bail!("timed out after {timeout:?}");
}
};
if !status.success() {
bail!("command exited with non-zero status: {status}");
}
Ok(String::from_utf8(stdout)?.trim_end().to_string())
}
}
const REAP: Duration = Duration::from_secs(1);
async fn end(child: &mut tokio::process::Child, tree: ChildTree) {
drop(tree);
let _ = child.start_kill();
if tokio::time::timeout(REAP, child.wait()).await.is_err() {
debug!("gave up reaping a killed command after {REAP:?}");
}
}
async fn capture(
mut stream: impl AsyncRead + Unpin,
budget: &AtomicUsize,
limit: usize,
) -> std::io::Result<Vec<u8>> {
let mut output = Vec::new();
let mut buf = [0u8; 8192];
loop {
let n = stream.read(&mut buf).await?;
if n == 0 {
return Ok(output);
}
let spend = |left: usize| left.checked_sub(n);
if budget
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, spend)
.is_err()
{
return Err(std::io::Error::other(format!(
"command output exceeded {limit} bytes"
)));
}
output.extend_from_slice(&buf[..n]);
}
}
#[cfg(unix)]
struct ChildTree(nix::unistd::Pid);
#[cfg(unix)]
impl ChildTree {
fn new(child: &mut tokio::process::Child) -> Result<Self> {
Ok(Self(nix::unistd::Pid::from_raw(
child.id().ok_or_else(|| eyre!("child has no pid"))? as i32,
)))
}
}
#[cfg(unix)]
impl Drop for ChildTree {
fn drop(&mut self) {
let _ = nix::sys::signal::killpg(self.0, nix::sys::signal::Signal::SIGKILL);
}
}
#[cfg(windows)]
struct ChildTree(Option<usize>);
#[cfg(windows)]
impl ChildTree {
fn new(child: &mut tokio::process::Child) -> Result<Self> {
use windows_sys::Win32::System::JobObjects::*;
unsafe {
let handle = CreateJobObjectW(std::ptr::null(), std::ptr::null());
if handle.is_null() {
return Err(std::io::Error::last_os_error().into());
}
let mut job = Self(Some(handle as usize));
let mut info: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed();
info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
if SetInformationJobObject(
handle,
JobObjectExtendedLimitInformation,
&info as *const _ as *const _,
std::mem::size_of_val(&info) as u32,
) == 0
{
return Err(std::io::Error::last_os_error().into());
}
let process = child
.raw_handle()
.ok_or_else(|| eyre!("child has no handle"))?;
if AssignProcessToJobObject(handle, process as _) == 0 {
let err = std::io::Error::last_os_error();
return match child.try_wait() {
Ok(Some(_)) => {
debug!("child exited before it could join a job object: {err}");
job.close();
Ok(job)
}
_ => Err(err.into()),
};
}
Ok(job)
}
}
fn close(&mut self) {
if let Some(handle) = self.0.take() {
unsafe {
windows_sys::Win32::Foundation::CloseHandle(handle as _);
}
}
}
}
#[cfg(windows)]
impl Drop for ChildTree {
fn drop(&mut self) {
self.close();
}
}
#[cfg(all(test, unix))]
mod tests {
use super::*;
#[tokio::test]
async fn bounds_output_and_inherited_pipes() {
let started = Instant::now();
for script in ["yes x", "sleep 30 & exit 0", "sleep 30 & wait"] {
let err = CmdLineRunner::new("/bin/sh")
.args(["-c", script])
.with_timeout(Duration::from_millis(150))
.read_isolated(1024)
.await
.unwrap_err();
assert!(
err.to_string().contains("exceeded") || err.to_string().contains("timed out"),
"{err}"
);
}
assert!(started.elapsed() < Duration::from_secs(3));
}
#[tokio::test]
async fn the_two_pipes_share_one_budget() {
let half = "x".repeat(700);
let script = "printf %s \"$H\"; printf %s \"$H\" >&2; sleep 30";
let err = CmdLineRunner::new("/bin/sh")
.args(["-c", script])
.env("H", &half)
.with_timeout(Duration::from_secs(5))
.read_isolated(1024)
.await
.unwrap_err();
assert!(
err.to_string().contains("exceeded"),
"the overrun is caught while the command still runs: {err}"
);
}
#[tokio::test]
async fn cleans_up_descendants_when_parent_exits() {
let dir = tempfile::tempdir().unwrap();
let pidfile = dir.path().join("child.pid");
let script = "sleep 30 >/dev/null 2>&1 & echo $! > \"$PIDFILE\"; printf ok";
let result = CmdLineRunner::new("/bin/sh")
.args(["-c", script])
.env("PIDFILE", &pidfile)
.read_isolated(1024)
.await
.unwrap();
assert_eq!(result, "ok");
let pid: i32 = std::fs::read_to_string(pidfile)
.unwrap()
.trim()
.parse()
.unwrap();
let status = CmdLineRunner::new("/bin/ps")
.args(["-o", "stat=", "-p", &pid.to_string()])
.read()
.await;
assert!(status.is_err() || status.unwrap().trim().starts_with('Z'));
}
}