use std::io::Read;
use std::process::{Child, Command, Stdio};
use std::sync::mpsc;
use std::time::{Duration, Instant};
use super::reasons::RefusalReason;
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(300);
pub const DEFAULT_OUTPUT_CAP: usize = 256 * 1024 * 1024;
enum ReaderMsg {
Done(Vec<u8>),
CapExceeded,
ReadError,
}
fn spawn_reader(mut stdout: impl Read + Send + 'static, cap: usize) -> mpsc::Receiver<ReaderMsg> {
let (tx, rx) = mpsc::channel();
std::thread::spawn(move || {
let mut buf = Vec::new();
let mut chunk = [0u8; 64 * 1024];
loop {
match stdout.read(&mut chunk) {
Ok(0) => {
let _ = tx.send(ReaderMsg::Done(buf));
return;
}
Ok(n) => {
buf.extend_from_slice(&chunk[..n]);
if buf.len() > cap {
let _ = tx.send(ReaderMsg::CapExceeded);
return;
}
}
Err(_) => {
let _ = tx.send(ReaderMsg::ReadError);
return;
}
}
}
});
rx
}
fn kill_and_wait(child: &mut Child) {
let _ = child.kill();
let _ = child.wait();
}
pub fn run_engine(
mut command: Command,
timeout: Duration,
output_cap: usize,
) -> Result<Vec<u8>, RefusalReason> {
command
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::null());
let mut child = command.spawn().map_err(|_| RefusalReason::EngineKilled)?;
let stdout = child.stdout.take().expect("stdout was piped");
let rx = spawn_reader(stdout, output_cap);
let deadline = Instant::now() + timeout;
loop {
match rx.recv_timeout(Duration::from_millis(25)) {
Ok(ReaderMsg::Done(bytes)) => {
let status = match child.wait() {
Ok(s) => s,
Err(_) => return Err(RefusalReason::EngineKilled),
};
return if status.success() {
Ok(bytes)
} else {
Err(RefusalReason::EngineKilled)
};
}
Ok(ReaderMsg::CapExceeded) => {
kill_and_wait(&mut child);
return Err(RefusalReason::OutputBounded);
}
Ok(ReaderMsg::ReadError) => {
kill_and_wait(&mut child);
return Err(RefusalReason::EngineKilled);
}
Err(mpsc::RecvTimeoutError::Timeout) => {
if Instant::now() >= deadline {
kill_and_wait(&mut child);
return Err(RefusalReason::EngineTimeout);
}
}
Err(mpsc::RecvTimeoutError::Disconnected) => {
return Err(RefusalReason::EngineKilled);
}
}
}
}