use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
use std::time::Duration;
use crate::error::{SandboxError, SandlockError};
use crate::policy::Policy;
use crate::result::{ExitStatus, RunResult};
use crate::sandbox::Sandbox;
pub struct Stage {
pub policy: Policy,
pub args: Vec<String>,
}
impl Stage {
pub fn new(policy: &Policy, args: &[&str]) -> Self {
Self {
policy: policy.clone(),
args: args.iter().map(|s| s.to_string()).collect(),
}
}
pub async fn run(self, timeout: Option<Duration>) -> Result<RunResult, SandlockError> {
let cmd_refs: Vec<&str> = self.args.iter().map(|s| s.as_str()).collect();
if let Some(dur) = timeout {
match tokio::time::timeout(dur, Sandbox::run_interactive(&self.policy, &cmd_refs)).await
{
Ok(result) => result,
Err(_) => Ok(RunResult {
exit_status: ExitStatus::Timeout,
stdout: None,
stderr: None,
}),
}
} else {
Sandbox::run_interactive(&self.policy, &cmd_refs).await
}
}
}
impl std::ops::BitOr<Stage> for Stage {
type Output = Pipeline;
fn bitor(self, rhs: Stage) -> Pipeline {
Pipeline {
stages: vec![self, rhs],
}
}
}
pub struct Pipeline {
pub stages: Vec<Stage>,
}
impl Pipeline {
pub fn new(stages: Vec<Stage>) -> Result<Self, SandlockError> {
if stages.len() < 2 {
return Err(SandlockError::Sandbox(SandboxError::Child(
"Pipeline requires at least 2 stages".into(),
)));
}
Ok(Self { stages })
}
pub async fn run(self, timeout: Option<Duration>) -> Result<RunResult, SandlockError> {
if let Some(dur) = timeout {
match tokio::time::timeout(dur, run_pipeline(self.stages)).await {
Ok(result) => result,
Err(_) => Ok(RunResult {
exit_status: ExitStatus::Timeout,
stdout: None,
stderr: None,
}),
}
} else {
run_pipeline(self.stages).await
}
}
}
impl std::ops::BitOr<Stage> for Pipeline {
type Output = Pipeline;
fn bitor(mut self, rhs: Stage) -> Pipeline {
self.stages.push(rhs);
self
}
}
fn make_pipe() -> std::io::Result<(OwnedFd, OwnedFd)> {
let mut fds = [0i32; 2];
if unsafe { libc::pipe2(fds.as_mut_ptr(), libc::O_CLOEXEC) } < 0 {
return Err(std::io::Error::last_os_error());
}
Ok(unsafe {
(
OwnedFd::from_raw_fd(fds[0]),
OwnedFd::from_raw_fd(fds[1]),
)
})
}
async fn run_pipeline(stages: Vec<Stage>) -> Result<RunResult, SandlockError> {
let n = stages.len();
let mut inter_pipes: Vec<(OwnedFd, OwnedFd)> = Vec::with_capacity(n - 1);
for _ in 0..n - 1 {
inter_pipes.push(make_pipe().map_err(SandboxError::Io)?);
}
let (cap_stdout_r, cap_stdout_w) = make_pipe().map_err(SandboxError::Io)?;
let (cap_stderr_r, cap_stderr_w) = make_pipe().map_err(SandboxError::Io)?;
let mut sandboxes: Vec<Sandbox> = Vec::with_capacity(n);
for (i, stage) in stages.into_iter().enumerate() {
let mut sb = Sandbox::new(&stage.policy)?;
let stdin_fd: Option<RawFd> = if i == 0 {
None } else {
Some(inter_pipes[i - 1].0.as_raw_fd()) };
let stdout_fd: Option<RawFd> = if i == n - 1 {
Some(cap_stdout_w.as_raw_fd()) } else {
Some(inter_pipes[i].1.as_raw_fd()) };
let stderr_fd: Option<RawFd> = if i == n - 1 {
Some(cap_stderr_w.as_raw_fd()) } else {
None };
let cmd_refs: Vec<&str> = stage.args.iter().map(|s| s.as_str()).collect();
sb.spawn_with_io(&cmd_refs, stdin_fd, stdout_fd, stderr_fd)
.await?;
sandboxes.push(sb);
}
drop(inter_pipes);
drop(cap_stdout_w);
drop(cap_stderr_w);
let mut last_result = RunResult {
exit_status: ExitStatus::Killed,
stdout: None,
stderr: None,
};
for (i, mut sb) in sandboxes.into_iter().enumerate() {
let result = sb.wait().await?;
if i == n - 1 {
last_result.exit_status = result.exit_status;
}
}
last_result.stdout = Some(read_fd_to_end(cap_stdout_r));
last_result.stderr = Some(read_fd_to_end(cap_stderr_r));
Ok(last_result)
}
fn read_fd_to_end(fd: OwnedFd) -> Vec<u8> {
use std::io::Read;
let mut file = unsafe { std::fs::File::from_raw_fd(fd.into_raw_fd()) };
let mut buf = Vec::new();
let _ = file.read_to_end(&mut buf);
buf
}