8 files changed, 418 insertions, 0 deletions

diff --git a/crates/rebel-runner/src/util/checkable.rs b/crates/rebel-runner/src/util/checkable.rs
new file mode 100644
index 0000000..8528d29
--- /dev/null
+++ b/crates/rebel-runner/src/util/checkable.rs
@@ -0,0 +1,37 @@
+use std::{
+	io::{Error, ErrorKind, Result},
+	process::ExitStatus,
+};
+
+use nix::sys::wait;
+
+pub trait Checkable {
+	fn check(&self) -> Result<()>;
+}
+
+impl Checkable for ExitStatus {
+	fn check(&self) -> Result<()> {
+		if self.success() {
+			Ok(())
+		} else {
+			Err(Error::new(
+				ErrorKind::Other,
+				format!("Process exited with {}", self),
+			))
+		}
+	}
+}
+
+impl Checkable for wait::WaitStatus {
+	fn check(&self) -> Result<()> {
+		let message = match self {
+			wait::WaitStatus::Exited(_, 0) => return Ok(()),
+			wait::WaitStatus::Exited(_, code) => format!("Process exited with exit code: {}", code),
+			wait::WaitStatus::Signaled(_, signal, _) => {
+				format!("Process exited with signal: {}", signal)
+			}
+			_ => format!("Process in unexpected status: {:?}", self),
+		};
+		Err(Error::new(ErrorKind::Other, message))
+	}
+}
diff --git a/crates/rebel-runner/src/util/cjson.rs b/crates/rebel-runner/src/util/cjson.rs
new file mode 100644
index 0000000..e3840ce
--- /dev/null
+++ b/crates/rebel-runner/src/util/cjson.rs
@@ -0,0 +1,37 @@
+use std::{
+	fs::File,
+	io::{self, Write},
+	path::Path,
+};
+
+use digest::{self, Digest};
+use olpc_cjson::CanonicalFormatter;
+use serde::Serialize;
+use serde_json::error::Result;
+
+pub fn new_serializer<W: Write>(writer: W) -> serde_json::Serializer<W, CanonicalFormatter> {
+	serde_json::Serializer::with_formatter(writer, CanonicalFormatter::new())
+}
+
+pub fn to_writer<W: Write, T: ?Sized + Serialize>(writer: W, value: &T) -> Result<()> {
+	let mut ser = new_serializer(writer);
+	value.serialize(&mut ser)
+}
+
+pub fn to_file<P: AsRef<Path>, T: ?Sized + Serialize>(path: P, value: &T) -> io::Result<()> {
+	let file = File::create(path)?;
+	to_writer(&file, value)?;
+	file.sync_all()
+}
+
+pub fn to_string<T: ?Sized + Serialize>(value: &T) -> Result<String> {
+	let mut ret = Vec::new();
+	to_writer(&mut ret, value)?;
+	Ok(String::from_utf8(ret).unwrap())
+}
+
+pub fn digest<D: Digest + Write, T: ?Sized + Serialize>(value: &T) -> Result<digest::Output<D>> {
+	let mut digest = <D as Digest>::new();
+	to_writer(&mut digest, value)?;
+	Ok(digest.finalize())
+}
diff --git a/crates/rebel-runner/src/util/clone.rs b/crates/rebel-runner/src/util/clone.rs
new file mode 100644
index 0000000..51a31c3
--- /dev/null
+++ b/crates/rebel-runner/src/util/clone.rs
@@ -0,0 +1,59 @@
+use std::{mem, process};
+
+use nix::{
+	errno, sched,
+	unistd::{self, Pid},
+};
+
+#[repr(C)]
+#[derive(Debug, Default)]
+struct CloneArgs {
+	flags: u64,
+	pidfd: u64,
+	child_tid: u64,
+	parent_tid: u64,
+	exit_signal: u64,
+	stack: u64,
+	stack_size: u64,
+	tls: u64,
+}
+
+pub unsafe fn clone(flags: sched::CloneFlags) -> nix::Result<unistd::ForkResult> {
+	let mut args = CloneArgs {
+		flags: flags.bits() as u64,
+		exit_signal: libc::SIGCHLD as u64,
+		..CloneArgs::default()
+	};
+	let size = mem::size_of_val(&args) as libc::size_t;
+
+	let pid = libc::syscall(libc::SYS_clone3, &mut args, size);
+
+	#[allow(clippy::comparison_chain)]
+	if pid < 0 {
+		Err(errno::Errno::last())
+	} else if pid == 0 {
+		Ok(unistd::ForkResult::Child)
+	} else {
+		Ok(unistd::ForkResult::Parent {
+			child: Pid::from_raw(pid as libc::pid_t),
+		})
+	}
+}
+
+pub unsafe fn spawn<F>(flags: Option<sched::CloneFlags>, f: F) -> nix::Result<Pid>
+where
+	F: FnOnce(),
+{
+	let res = if let Some(flags) = flags {
+		clone(flags)
+	} else {
+		unistd::fork()
+	};
+	match res? {
+		unistd::ForkResult::Parent { child } => Ok(child),
+		unistd::ForkResult::Child => {
+			f();
+			process::exit(0)
+		}
+	}
+}
diff --git a/crates/rebel-runner/src/util/fs.rs b/crates/rebel-runner/src/util/fs.rs
new file mode 100644
index 0000000..9e33eb7
--- /dev/null
+++ b/crates/rebel-runner/src/util/fs.rs
@@ -0,0 +1,127 @@
+use std::{
+	fs::{self, File},
+	io,
+	os::unix::prelude::*,
+	path::{Path, PathBuf},
+};
+
+use nix::{
+	fcntl::OFlag,
+	mount::{self, MsFlags},
+	unistd,
+};
+
+use rebel_common::error::*;
+
+pub fn open<P: AsRef<Path>>(path: P) -> Result<fs::File> {
+	fs::File::open(path.as_ref())
+		.with_context(|| format!("Failed to open file {:?} for reading", path.as_ref()))
+}
+
+pub fn create<P: AsRef<Path>>(path: P) -> Result<fs::File> {
+	fs::File::create(path.as_ref())
+		.with_context(|| format!("Failed to open file {:?} for writing", path.as_ref()))
+}
+
+pub fn rename<P1: AsRef<Path>, P2: AsRef<Path>>(from: P1, to: P2) -> Result<()> {
+	fs::rename(from.as_ref(), to.as_ref())
+		.with_context(|| format!("Failed to rename {:?} to {:?}", from.as_ref(), to.as_ref()))
+}
+
+// Unlike fs::copy, this doesn't preserve file flags
+pub fn copy<P1: AsRef<Path>, P2: AsRef<Path>>(from: P1, to: P2) -> Result<()> {
+	(|| -> Result<()> {
+		let mut src = open(from.as_ref())?;
+		let mut dest = create(to.as_ref())?;
+		io::copy(&mut src, &mut dest)?;
+		dest.sync_all()?;
+		Ok(())
+	})()
+	.with_context(|| format!("Failed to copy {:?} to {:?}", from.as_ref(), to.as_ref()))
+}
+
+pub fn mkdir<P: AsRef<Path>>(path: P) -> Result<()> {
+	let mut builder = fs::DirBuilder::new();
+	builder.recursive(true);
+	builder
+		.create(path.as_ref())
+		.with_context(|| format!("Failed to create directory {:?}", path.as_ref()))
+}
+
+pub fn ensure_removed<P: AsRef<Path>>(path: P) -> Result<()> {
+	let result = if path.as_ref().is_dir() {
+		fs::remove_dir_all(path.as_ref())
+	} else {
+		fs::remove_file(path.as_ref())
+	};
+	result
+		.or_else(|err| match err.kind() {
+			io::ErrorKind::NotFound => Ok(()),
+			_ => Err(err),
+		})
+		.with_context(|| format!("Failed to delete {:?}", path.as_ref()))
+}
+
+pub fn is_dir_empty<P: AsRef<Path>>(path: P) -> Result<bool> {
+	Ok(fs::read_dir(path)?.next().is_none())
+}
+
+/// Fixes up weirdness of set-group-ID or bsdgroups
+pub fn fixup_permissions<P: AsRef<Path>>(path: P) -> Result<()> {
+	let path = path.as_ref();
+	let gid = unistd::getegid();
+
+	let metadata = path
+		.metadata()
+		.with_context(|| format!("Failed to get metadata of {:?}", path))?;
+
+	if metadata.gid() != gid.as_raw() {
+		unistd::chown(path, None, Some(gid))
+			.with_context(|| format!("Failed to set group of {:?}", path))?;
+	}
+
+	let mut perms = metadata.permissions();
+	let mode = perms.mode();
+	if (mode & 0o777) != mode {
+		perms.set_mode(mode & 0o777);
+		std::fs::set_permissions(path, perms)
+			.with_context(|| format!("Failed to set mode of {:?}", path))?;
+	}
+
+	Ok(())
+}
+
+#[must_use]
+pub struct Mount(PathBuf);
+
+impl Drop for Mount {
+	fn drop(&mut self) {
+		mount::umount(&self.0)
+			.with_context(|| format!("Failed to unmount {:?}", self.0))
+			.unwrap();
+	}
+}
+
+pub fn mount<P1: AsRef<Path>, P2: AsRef<Path>>(
+	source: P1,
+	target: P2,
+	fstype: Option<&str>,
+	flags: MsFlags,
+	data: Option<&str>,
+) -> Result<Mount> {
+	mkdir(target.as_ref()).with_context(|| format!("Failed to create {:?}", target.as_ref()))?;
+
+	let canon_target = target
+		.as_ref()
+		.canonicalize()
+		.with_context(|| format!("Failed to get absolute path for {:?}", target.as_ref()))?;
+	mount::mount(Some(source.as_ref()), &canon_target, fstype, flags, data)
+		.with_context(|| format!("Failed to mount {:?}", canon_target))?;
+	Ok(Mount(canon_target))
+}
+
+pub fn pipe() -> Result<(File, File)> {
+	unistd::pipe2(OFlag::O_CLOEXEC)
+		.context("pipe2()")
+		.map(|(piper, pipew)| (File::from(piper), File::from(pipew)))
+}
diff --git a/crates/rebel-runner/src/util/mod.rs b/crates/rebel-runner/src/util/mod.rs
new file mode 100644
index 0000000..0fbe3b5
--- /dev/null
+++ b/crates/rebel-runner/src/util/mod.rs
@@ -0,0 +1,7 @@
+pub mod checkable;
+pub mod cjson;
+pub mod clone;
+pub mod fs;
+pub mod stack;
+pub mod steal;
+pub mod unix;
diff --git a/crates/rebel-runner/src/util/stack.rs b/crates/rebel-runner/src/util/stack.rs
new file mode 100644
index 0000000..15d5daf
--- /dev/null
+++ b/crates/rebel-runner/src/util/stack.rs
@@ -0,0 +1,25 @@
+use std::mem;
+
+/// Simple inefficient datastructure with guaranteed drop order
+#[derive(Debug)]
+pub enum Stack<T> {
+	Cons(Box<(T, Stack<T>)>),
+	Empty,
+}
+
+impl<T> Default for Stack<T> {
+	fn default() -> Self {
+		Self::Empty
+	}
+}
+
+impl<T> Stack<T> {
+	pub fn new() -> Self {
+		Self::Empty
+	}
+
+	pub fn push(&mut self, value: T) {
+		let tmp = mem::take(self);
+		*self = Stack::Cons(Box::new((value, tmp)));
+	}
+}
diff --git a/crates/rebel-runner/src/util/steal.rs b/crates/rebel-runner/src/util/steal.rs
new file mode 100644
index 0000000..91b2cdf
--- /dev/null
+++ b/crates/rebel-runner/src/util/steal.rs
@@ -0,0 +1,40 @@
+use std::ops::{Deref, DerefMut};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub struct Steal<T>(pub Option<T>);
+
+impl<T> Steal<T> {
+	pub fn new(value: T) -> Steal<T> {
+		Steal(Some(value))
+	}
+
+	pub fn steal(&mut self) -> T {
+		self.0
+			.take()
+			.expect("Attempted to steal already stoken value")
+	}
+}
+
+impl<T> From<T> for Steal<T> {
+	fn from(value: T) -> Self {
+		Steal::new(value)
+	}
+}
+
+impl<T> Deref for Steal<T> {
+	type Target = T;
+
+	fn deref(&self) -> &Self::Target {
+		self.0
+			.as_ref()
+			.expect("Attempted to dereference stolen value")
+	}
+}
+
+impl<T> DerefMut for Steal<T> {
+	fn deref_mut(&mut self) -> &mut Self::Target {
+		self.0
+			.as_mut()
+			.expect("Attempted to dereference stolen value")
+	}
+}
diff --git a/crates/rebel-runner/src/util/unix.rs b/crates/rebel-runner/src/util/unix.rs
new file mode 100644
index 0000000..a97b1db
--- /dev/null
+++ b/crates/rebel-runner/src/util/unix.rs
@@ -0,0 +1,86 @@
+use std::{fs::File, os::unix::prelude::*, path::Path};
+
+use nix::{
+	fcntl::{self, FcntlArg, FdFlag, Flock, OFlag},
+	sched,
+	unistd::Pid,
+};
+
+use rebel_common::error::*;
+
+use super::fs;
+
+pub fn set_blocking<Fd: AsFd>(fd: &Fd, blocking: bool) -> Result<()> {
+	let raw_fd = fd.as_fd().as_raw_fd();
+
+	let flags =
+		OFlag::from_bits_retain(fcntl::fcntl(raw_fd, FcntlArg::F_GETFL).context("fcntl(F_GETFL)")?);
+
+	let new_flags = if blocking {
+		flags & !OFlag::O_NONBLOCK
+	} else {
+		flags | OFlag::O_NONBLOCK
+	};
+
+	if new_flags != flags {
+		fcntl::fcntl(raw_fd, FcntlArg::F_SETFL(new_flags)).context("fcntl(F_SETFL)")?;
+	}
+
+	Ok(())
+}
+
+pub fn set_cloexec<Fd: AsFd>(fd: &Fd, cloexec: bool) -> Result<()> {
+	let raw_fd = fd.as_fd().as_raw_fd();
+
+	let flags = FdFlag::from_bits_retain(
+		fcntl::fcntl(raw_fd, FcntlArg::F_GETFD).context("fcntl(F_GETFD)")?,
+	);
+
+	let new_flags = if cloexec {
+		flags | FdFlag::FD_CLOEXEC
+	} else {
+		flags & !FdFlag::FD_CLOEXEC
+	};
+
+	if new_flags != flags {
+		fcntl::fcntl(raw_fd, FcntlArg::F_SETFD(new_flags)).context("fcntl(F_SETFD)")?;
+	}
+
+	Ok(())
+}
+
+pub fn nproc() -> Result<usize> {
+	const MAXCPU: usize = sched::CpuSet::count();
+
+	let affinity = sched::sched_getaffinity(Pid::from_raw(0)).context("sched_getaffinity()")?;
+
+	let mut count = 0;
+
+	for cpu in 0..MAXCPU {
+		if affinity.is_set(cpu).unwrap() {
+			count += 1;
+		}
+	}
+
+	Ok(count)
+}
+
+pub fn lock<P: AsRef<Path>>(path: P, exclusive: bool, blocking: bool) -> Result<Flock<File>> {
+	use fcntl::FlockArg::*;
+
+	if let Some(parent) = path.as_ref().parent() {
+		fs::mkdir(parent)?;
+	}
+
+	let arg = match (exclusive, blocking) {
+		(true, true) => LockExclusive,
+		(true, false) => LockExclusiveNonblock,
+		(false, true) => LockShared,
+		(false, false) => LockSharedNonblock,
+	};
+
+	let file = fs::create(path.as_ref())?;
+	fcntl::Flock::lock(file, arg)
+		.map_err(|(_, errno)| errno)
+		.with_context(|| format!("flock failed on {:?}", path.as_ref()))
+}