From 34ac18d20c13a78914d447fee83204811a27b1e4 Mon Sep 17 00:00:00 2001 From: Matthias Schiffer Date: Mon, 25 Oct 2021 00:19:45 +0200 Subject: Move runner into separate crate --- crates/runner/src/util/cjson.rs | 31 ++++++++++ crates/runner/src/util/clone.rs | 58 +++++++++++++++++++ crates/runner/src/util/fs.rs | 122 ++++++++++++++++++++++++++++++++++++++++ crates/runner/src/util/mod.rs | 42 ++++++++++++++ crates/runner/src/util/unix.rs | 61 ++++++++++++++++++++ 5 files changed, 314 insertions(+) create mode 100644 crates/runner/src/util/cjson.rs create mode 100644 crates/runner/src/util/clone.rs create mode 100644 crates/runner/src/util/fs.rs create mode 100644 crates/runner/src/util/mod.rs create mode 100644 crates/runner/src/util/unix.rs (limited to 'crates/runner/src/util') diff --git a/crates/runner/src/util/cjson.rs b/crates/runner/src/util/cjson.rs new file mode 100644 index 0000000..5990943 --- /dev/null +++ b/crates/runner/src/util/cjson.rs @@ -0,0 +1,31 @@ +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(writer: W) -> serde_json::Serializer { + serde_json::Serializer::with_formatter(writer, CanonicalFormatter::new()) +} + +pub fn to_writer(writer: W, value: &T) -> Result<()> { + let mut ser = new_serializer(writer); + value.serialize(&mut ser) +} + +pub fn to_file, T: ?Sized + Serialize>(path: P, value: &T) -> io::Result<()> { + let file = File::create(path)?; + to_writer(&file, value)?; + file.sync_all() +} + +pub fn digest(value: &T) -> Result> { + let mut digest = ::new(); + to_writer(&mut digest, value)?; + Ok(digest.finalize()) +} diff --git a/crates/runner/src/util/clone.rs b/crates/runner/src/util/clone.rs new file mode 100644 index 0000000..4835b53 --- /dev/null +++ b/crates/runner/src/util/clone.rs @@ -0,0 +1,58 @@ +use std::{mem, process}; + +use nix::{errno, sched, unistd}; + +#[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 { + 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); + if pid < 0 { + Err(errno::Errno::last()) + } else if pid == 0 { + Ok(unistd::ForkResult::Child) + } else { + Ok(unistd::ForkResult::Parent { + child: unistd::Pid::from_raw(pid as libc::pid_t), + }) + } +} + +pub unsafe fn spawn( + flags: Option, + arg: T, + f: F, +) -> nix::Result<(unistd::Pid, T)> +where + F: FnOnce(T), +{ + let res = if let Some(flags) = flags { + clone(flags) + } else { + unistd::fork() + }; + match res? { + unistd::ForkResult::Parent { child } => Ok((child, arg)), + unistd::ForkResult::Child => { + f(arg); + process::exit(0) + } + } +} diff --git a/crates/runner/src/util/fs.rs b/crates/runner/src/util/fs.rs new file mode 100644 index 0000000..099a339 --- /dev/null +++ b/crates/runner/src/util/fs.rs @@ -0,0 +1,122 @@ +use std::{ + fs::{self, File}, + io, + os::unix::prelude::*, + path::{Path, PathBuf}, +}; + +use nix::{ + fcntl::OFlag, + mount::{self, MsFlags}, + unistd, +}; + +use common::error::*; + +pub fn open>(path: P) -> Result { + fs::File::open(path.as_ref()) + .with_context(|| format!("Failed to open file {:?} for reading", path.as_ref())) +} + +pub fn create>(path: P) -> Result { + fs::File::create(path.as_ref()) + .with_context(|| format!("Failed to open file {:?} for writing", path.as_ref())) +} + +pub fn rename, P2: AsRef>(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, P2: AsRef>(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>(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>(path: P) -> Result<()> { + fs::remove_dir_all(path.as_ref()) + .or_else(|err| match err.kind() { + io::ErrorKind::NotFound => Ok(()), + _ => Err(err), + }) + .with_context(|| format!("Failed to delete directory {:?}", path.as_ref())) +} + +pub fn is_dir_empty>(path: P) -> Result { + Ok(fs::read_dir(path)?.next().is_none()) +} + +/// Fixes up weirdness of set-group-ID or bsdgroups +pub fn fixup_permissions>(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, P2: AsRef>( + source: P1, + target: P2, + fstype: Option<&str>, + flags: MsFlags, + data: Option<&str>, +) -> Result { + 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)| unsafe { (File::from_raw_fd(piper), File::from_raw_fd(pipew)) }) +} diff --git a/crates/runner/src/util/mod.rs b/crates/runner/src/util/mod.rs new file mode 100644 index 0000000..eff589d --- /dev/null +++ b/crates/runner/src/util/mod.rs @@ -0,0 +1,42 @@ +pub mod cjson; +pub mod clone; +pub mod fs; +pub mod unix; + +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/runner/src/util/unix.rs b/crates/runner/src/util/unix.rs new file mode 100644 index 0000000..48db764 --- /dev/null +++ b/crates/runner/src/util/unix.rs @@ -0,0 +1,61 @@ +use std::os::unix::prelude::*; + +use nix::{ + fcntl::{self, FcntlArg, FdFlag, OFlag}, + sched, + unistd::Pid, +}; + +use common::error::*; + +pub fn set_blocking(fd: RawFd, blocking: bool) -> Result<()> { + let flags = unsafe { + OFlag::from_bits_unchecked(fcntl::fcntl(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(fd, FcntlArg::F_SETFL(new_flags)).context("fcntl(F_SETFL)")?; + } + + Ok(()) +} + +pub fn set_cloexec(fd: RawFd, cloexec: bool) -> Result<()> { + let flags = unsafe { + FdFlag::from_bits_unchecked(fcntl::fcntl(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(fd, FcntlArg::F_SETFD(new_flags)).context("fcntl(F_SETFD)")?; + } + + Ok(()) +} + +pub fn nproc() -> Result { + 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) +} -- cgit v1.2.3