use std::collections::BTreeMap;
use std::ffi::OsString;
use std::fs::{self, File, Metadata};
use std::io::{self, BufRead, BufWriter, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
pub(crate) const CHUNK: usize = 5000;
pub(crate) const SANDBOX_OUTPUT_CAP: u64 = 8 * 1024 * 1024;
pub(crate) fn reset_sigpipe() {
#[cfg(unix)]
unsafe {
libc::signal(libc::SIGPIPE, libc::SIG_DFL);
}
}
pub(crate) type Signals = &'static [std::os::raw::c_int];
#[cfg(unix)]
pub(crate) const XZ_SIGNALS: Signals = &[
libc::SIGINT,
libc::SIGTERM,
libc::SIGHUP,
libc::SIGPIPE,
libc::SIGXCPU,
libc::SIGXFSZ,
];
#[cfg(not(unix))]
pub(crate) const XZ_SIGNALS: Signals = &[];
pub(crate) const GNU_GZIP_SIGNALS: Signals = XZ_SIGNALS;
#[cfg(unix)]
pub(crate) const APPLE_GZIP_SIGNALS: Signals = &[libc::SIGINT];
#[cfg(not(unix))]
pub(crate) const APPLE_GZIP_SIGNALS: Signals = &[];
#[cfg(unix)]
pub(crate) const BZIP2_SIGNALS: Signals = &[libc::SIGINT, libc::SIGTERM, libc::SIGHUP];
#[cfg(not(unix))]
pub(crate) const BZIP2_SIGNALS: Signals = &[];
#[cfg(unix)]
const SLOT_LEN: usize = libc::PATH_MAX as usize;
#[cfg(unix)]
struct PathSlot(std::cell::UnsafeCell<[u8; SLOT_LEN]>);
#[cfg(unix)]
unsafe impl Sync for PathSlot {}
#[cfg(unix)]
static SLOT: PathSlot = PathSlot(std::cell::UnsafeCell::new([0; SLOT_LEN]));
#[cfg(unix)]
static ARMED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
#[must_use = "dropping the guard disarms the cleanup at once"]
pub(crate) struct CleanupGuard {
#[cfg(unix)]
armed: bool,
#[cfg(unix)]
kind: GuardKind,
}
#[cfg(unix)]
#[derive(Clone, Copy, PartialEq, Eq)]
enum GuardKind {
Reraise,
Exit,
}
impl Drop for CleanupGuard {
fn drop(&mut self) {
#[cfg(unix)]
if self.armed {
if self.kind == GuardKind::Exit {
unsafe { libc::signal(libc::SIGINT, libc::SIG_DFL) };
}
ARMED.store(false, Ordering::Release);
}
}
}
#[cfg(unix)]
fn fill_slot(bytes: &[u8], signals: Signals) {
debug_assert!(bytes.len() < SLOT_LEN && !bytes.contains(&0));
unsafe {
let mut block: libc::sigset_t = std::mem::zeroed();
let mut old: libc::sigset_t = std::mem::zeroed();
libc::sigemptyset(&mut block);
for &sig in signals {
libc::sigaddset(&mut block, sig);
}
libc::pthread_sigmask(libc::SIG_BLOCK, &block, &mut old);
ARMED.store(false, Ordering::SeqCst);
let slot = SLOT.0.get().cast::<u8>();
std::ptr::copy_nonoverlapping(bytes.as_ptr(), slot, bytes.len());
slot.add(bytes.len()).write(0);
ARMED.store(true, Ordering::Release);
libc::pthread_sigmask(libc::SIG_SETMASK, &old, std::ptr::null_mut());
}
}
pub(crate) fn arm_cleanup(host: &Host, path: &Path, signals: Signals) -> CleanupGuard {
#[cfg(unix)]
{
use std::os::unix::ffi::OsStrExt;
let bytes = path.as_os_str().as_bytes();
let kind = GuardKind::Reraise;
if matches!(host, Host::Sandbox(_)) || bytes.len() >= SLOT_LEN || bytes.contains(&0) {
return CleanupGuard { armed: false, kind };
}
debug_assert!(
!ARMED.load(Ordering::Acquire),
"arm_cleanup: another output is still armed"
);
install_handlers(signals);
fill_slot(bytes, signals);
CleanupGuard { armed: true, kind }
}
#[cfg(not(unix))]
{
let _ = (host, path, signals);
CleanupGuard {}
}
}
#[cfg(unix)]
fn install_handlers(signals: Signals) {
let handler = on_cleanup_signal as extern "C" fn(libc::c_int) as libc::sighandler_t;
unsafe {
let mut action: libc::sigaction = std::mem::zeroed();
libc::sigemptyset(&mut action.sa_mask);
for &sig in signals {
libc::sigaddset(&mut action.sa_mask, sig);
}
action.sa_flags = 0;
action.sa_sigaction = handler;
for &sig in signals {
let mut old: libc::sigaction = std::mem::zeroed();
if libc::sigaction(sig, std::ptr::null(), &mut old) == 0
&& (old.sa_sigaction == libc::SIG_IGN || old.sa_sigaction == handler)
{
continue;
}
libc::sigaction(sig, &action, std::ptr::null_mut());
}
}
}
#[cfg(unix)]
extern "C" fn on_cleanup_signal(sig: libc::c_int) {
if ARMED.load(Ordering::Acquire) {
unsafe { libc::unlink(SLOT.0.get().cast::<libc::c_char>()) };
}
unsafe {
libc::signal(sig, libc::SIG_DFL);
libc::raise(sig);
}
}
pub(crate) fn arm_cleanup_exit(host: &Host, path: &Path) -> CleanupGuard {
#[cfg(unix)]
{
use std::os::unix::ffi::OsStrExt;
let bytes = path.as_os_str().as_bytes();
let unarmed = CleanupGuard {
armed: false,
kind: GuardKind::Exit,
};
if matches!(host, Host::Sandbox(_)) || bytes.len() >= SLOT_LEN || bytes.contains(&0) {
return unarmed;
}
if !fs::metadata(path).is_ok_and(|m| m.is_file()) {
return unarmed;
}
debug_assert!(
!ARMED.load(Ordering::Acquire),
"arm_cleanup_exit: another output is still armed"
);
const INT: Signals = &[libc::SIGINT];
fill_slot(bytes, INT);
let handler = on_interrupt_exit as extern "C" fn(libc::c_int) as libc::sighandler_t;
unsafe {
let mut action: libc::sigaction = std::mem::zeroed();
libc::sigemptyset(&mut action.sa_mask);
action.sa_flags = libc::SA_RESTART;
action.sa_sigaction = handler;
libc::sigaction(libc::SIGINT, &action, std::ptr::null_mut());
}
CleanupGuard {
armed: true,
kind: GuardKind::Exit,
}
}
#[cfg(not(unix))]
{
let _ = (host, path);
CleanupGuard {}
}
}
#[cfg(unix)]
extern "C" fn on_interrupt_exit(_sig: libc::c_int) {
unsafe { libc::signal(libc::SIGINT, libc::SIG_IGN) };
if ARMED.load(Ordering::Acquire) {
unsafe { libc::unlink(SLOT.0.get().cast::<libc::c_char>()) };
}
unsafe {
libc::write(2, b"\n".as_ptr().cast(), 1);
libc::_exit(2);
}
}
#[cfg(unix)]
static INTERRUPTED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
#[must_use = "dropping the guard disarms the flag at once"]
pub(crate) struct InterruptGuard {
#[cfg(unix)]
old: Vec<(libc::c_int, libc::sigaction)>,
}
impl Drop for InterruptGuard {
fn drop(&mut self) {
#[cfg(unix)]
for (sig, old) in &self.old {
unsafe { libc::sigaction(*sig, old, std::ptr::null_mut()) };
}
}
}
pub(crate) fn arm_interrupt_flag(host: &Host, signals: Signals) -> InterruptGuard {
#[cfg(unix)]
{
let mut old = Vec::new();
if matches!(host, Host::Sandbox(_)) {
return InterruptGuard { old };
}
INTERRUPTED.store(false, Ordering::SeqCst);
let handler = on_interrupt_signal as extern "C" fn(libc::c_int) as libc::sighandler_t;
unsafe {
let mut action: libc::sigaction = std::mem::zeroed();
libc::sigemptyset(&mut action.sa_mask);
action.sa_flags = libc::SA_RESTART;
action.sa_sigaction = handler;
for &sig in signals {
let mut prev: libc::sigaction = std::mem::zeroed();
if libc::sigaction(sig, &action, &mut prev) == 0 {
old.push((sig, prev));
}
}
}
InterruptGuard { old }
}
#[cfg(not(unix))]
{
let _ = (host, signals);
InterruptGuard {}
}
}
#[cfg(unix)]
extern "C" fn on_interrupt_signal(_sig: libc::c_int) {
INTERRUPTED.store(true, Ordering::SeqCst);
}
pub(crate) fn take_interrupt(host: &Host) -> bool {
#[cfg(unix)]
{
!matches!(host, Host::Sandbox(_)) && INTERRUPTED.swap(false, Ordering::SeqCst)
}
#[cfg(not(unix))]
{
let _ = host;
false
}
}
#[cfg(unix)]
const PROGRESS_SIGNALS: &[libc::c_int] = &[
libc::SIGALRM,
#[cfg(any(
target_vendor = "apple",
target_os = "freebsd",
target_os = "dragonfly",
target_os = "netbsd",
target_os = "openbsd"
))]
libc::SIGINFO,
libc::SIGUSR1,
];
pub(crate) static PROGRESS_REQUESTED: std::sync::atomic::AtomicBool =
std::sync::atomic::AtomicBool::new(false);
pub(crate) fn install_progress_signals(host: &Host) {
#[cfg(unix)]
if matches!(host, Host::Real) {
let handler = on_progress_signal as extern "C" fn(libc::c_int) as libc::sighandler_t;
unsafe {
let mut action: libc::sigaction = std::mem::zeroed();
libc::sigemptyset(&mut action.sa_mask);
action.sa_flags = libc::SA_RESTART;
action.sa_sigaction = handler;
for &sig in PROGRESS_SIGNALS {
libc::sigaction(sig, &action, std::ptr::null_mut());
}
}
}
#[cfg(not(unix))]
let _ = host;
}
#[cfg(unix)]
extern "C" fn on_progress_signal(_sig: libc::c_int) {
PROGRESS_REQUESTED.store(true, Ordering::Relaxed);
}
static STOPPED_MS: AtomicU64 = AtomicU64::new(0);
fn monotonic_ms() -> u64 {
#[cfg(unix)]
{
let mut ts: libc::timespec = unsafe { std::mem::zeroed() };
unsafe { libc::clock_gettime(libc::CLOCK_MONOTONIC, &mut ts) };
(ts.tv_sec as u64).wrapping_mul(1000) + (ts.tv_nsec as u64) / 1_000_000
}
#[cfg(not(unix))]
{
static BASE: std::sync::OnceLock<std::time::Instant> = std::sync::OnceLock::new();
BASE.get_or_init(std::time::Instant::now)
.elapsed()
.as_millis() as u64
}
}
pub(crate) fn running_ms() -> u64 {
monotonic_ms().wrapping_sub(STOPPED_MS.load(Ordering::Relaxed))
}
pub(crate) fn install_stop_handler(host: &Host) {
#[cfg(unix)]
if matches!(host, Host::Real) {
let handler = on_stop_signal as extern "C" fn(libc::c_int) as libc::sighandler_t;
unsafe {
let mut old: libc::sigaction = std::mem::zeroed();
if libc::sigaction(libc::SIGTSTP, std::ptr::null(), &mut old) == 0
&& old.sa_sigaction == libc::SIG_IGN
{
return;
}
let mut action: libc::sigaction = std::mem::zeroed();
libc::sigemptyset(&mut action.sa_mask);
for &sig in XZ_SIGNALS.iter().chain(PROGRESS_SIGNALS) {
libc::sigaddset(&mut action.sa_mask, sig);
}
libc::sigaddset(&mut action.sa_mask, libc::SIGTSTP);
action.sa_flags = 0;
action.sa_sigaction = handler;
libc::sigaction(libc::SIGTSTP, &action, std::ptr::null_mut());
}
}
#[cfg(not(unix))]
let _ = host;
}
#[cfg(unix)]
extern "C" fn on_stop_signal(_sig: libc::c_int) {
let t = monotonic_ms();
unsafe { libc::raise(libc::SIGSTOP) };
STOPPED_MS.fetch_add(monotonic_ms().wrapping_sub(t), Ordering::Relaxed);
}
pub(crate) fn create_temp(dir: &Path, family: &str) -> io::Result<(PathBuf, File)> {
let mut opts = File::options();
opts.write(true).create_new(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
opts.mode(0o600);
}
for n in 0u32.. {
let tmp = dir.join(format!(".stuffr-{family}-{}-{n}.tmp", std::process::id()));
match opts.open(&tmp) {
Ok(f) => return Ok((tmp, f)),
Err(e) if e.kind() == io::ErrorKind::AlreadyExists && n < 1000 => {}
Err(e) => return Err(e),
}
}
unreachable!("the loop returns")
}
pub(crate) fn create_temp_armed(
host: &Host,
dir: &Path,
family: &str,
signals: Signals,
) -> io::Result<(PathBuf, File, CleanupGuard)> {
held_back(signals, || {
create_temp(dir, family).map(|(p, f)| {
let g = arm_cleanup(host, &p, signals);
(p, f, g)
})
})
}
pub(crate) fn create_temp_armed_exit(
dir: &Path,
family: &str,
host: &Host,
) -> io::Result<(PathBuf, File, CleanupGuard)> {
#[cfg(unix)]
const INT: Signals = &[libc::SIGINT];
#[cfg(not(unix))]
const INT: Signals = &[];
held_back(INT, || {
create_temp(dir, family).map(|(p, f)| {
let g = arm_cleanup_exit(host, &p);
(p, f, g)
})
})
}
fn held_back<T>(signals: Signals, f: impl FnOnce() -> T) -> T {
#[cfg(unix)]
{
let mut block: libc::sigset_t = unsafe { std::mem::zeroed() };
let mut old: libc::sigset_t = unsafe { std::mem::zeroed() };
unsafe {
libc::sigemptyset(&mut block);
for &sig in signals {
libc::sigaddset(&mut block, sig);
}
libc::pthread_sigmask(libc::SIG_BLOCK, &block, &mut old);
}
let made = f();
unsafe { libc::pthread_sigmask(libc::SIG_SETMASK, &old, std::ptr::null_mut()) };
made
}
#[cfg(not(unix))]
{
let _ = signals;
f()
}
}
#[cfg(all(unix, debug_assertions))]
pub(crate) fn signal_test(args: &[OsString]) -> std::process::ExitCode {
let [mode, file] = args else {
return std::process::ExitCode::from(64);
};
if mode != "exit" {
return std::process::ExitCode::from(64);
}
let file = Path::new(file);
let _guard = arm_cleanup_exit(&Host::Real, file);
let mut ready = file.as_os_str().to_owned();
ready.push(".ready");
if fs::write(&ready, b"").is_err() {
return std::process::ExitCode::from(64);
}
std::thread::sleep(std::time::Duration::from_secs(60));
std::process::ExitCode::SUCCESS
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Published {
Done,
Exists,
}
pub(crate) fn publish(tmp: &Path, out: &Path, force: bool) -> io::Result<Published> {
if force {
fs::rename(tmp, out)?;
return Ok(Published::Done);
}
match fs::hard_link(tmp, out) {
Ok(()) => fs::remove_file(tmp)?,
Err(e) if e.kind() == io::ErrorKind::AlreadyExists => {
let _ = fs::remove_file(tmp);
return Ok(Published::Exists);
}
Err(_) => fs::rename(tmp, out)?,
}
Ok(Published::Done)
}
pub(crate) fn path(b: &[u8]) -> PathBuf {
#[cfg(unix)]
{
use std::os::unix::ffi::OsStringExt;
PathBuf::from(OsString::from_vec(b.to_vec()))
}
#[cfg(not(unix))]
{
PathBuf::from(String::from_utf8_lossy(b).into_owned())
}
}
pub(crate) fn strerror(e: &io::Error) -> Vec<u8> {
match e.raw_os_error() {
Some(n) => errno_text(n),
None => e.to_string().into_bytes(),
}
}
pub(crate) fn errno_text(n: i32) -> Vec<u8> {
let s = io::Error::from_raw_os_error(n).to_string();
let suffix = format!(" (os error {n})");
s.strip_suffix(&suffix).unwrap_or(&s).as_bytes().to_vec()
}
pub(crate) fn not_a_standard_file(p: &Path) -> bool {
fs::symlink_metadata(p).map_or(true, |m| !m.file_type().is_file())
}
pub(crate) fn same_file(a: &Metadata, b: &Metadata) -> Option<bool> {
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
Some(a.dev() == b.dev() && a.ino() == b.ino())
}
#[cfg(not(unix))]
{
let _ = (a, b);
None
}
}
pub(crate) fn count_hard_links(p: &Path) -> u64 {
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
fs::symlink_metadata(p).map_or(0, |m| m.nlink().saturating_sub(1))
}
#[cfg(not(unix))]
{
let _ = p;
0
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Kind {
Regular,
Device,
Pipe,
}
pub(crate) fn kind_of(f: &File) -> Kind {
#[cfg(unix)]
{
use std::os::unix::fs::FileTypeExt;
match f.metadata().map(|m| m.file_type()) {
Ok(t) if t.is_fifo() || t.is_socket() => Kind::Pipe,
Ok(t) if t.is_char_device() || t.is_block_device() => Kind::Device,
_ => Kind::Regular,
}
}
#[cfg(not(unix))]
{
let _ = f;
Kind::Regular
}
}
pub(crate) fn stdin_file() -> io::Result<File> {
#[cfg(unix)]
{
use std::os::fd::AsFd;
Ok(File::from(io::stdin().as_fd().try_clone_to_owned()?))
}
#[cfg(not(unix))]
{
Err(io::Error::from(io::ErrorKind::Unsupported))
}
}
pub(crate) fn stdout_writer() -> Box<dyn Write + Send> {
#[cfg(unix)]
{
use std::os::fd::AsFd;
if let Ok(fd) = io::stdout().as_fd().try_clone_to_owned() {
return Box::new(File::from(fd));
}
}
Box::new(io::stdout())
}
#[doc(hidden)]
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct InputReads {
pub opens: u64,
pub bytes: u64,
pub len: u64,
}
pub(crate) type ReadLog = Arc<Mutex<BTreeMap<PathBuf, InputReads>>>;
#[derive(Debug)]
pub(crate) struct InFile {
file: File,
log: Option<(ReadLog, PathBuf)>,
}
impl InFile {
pub(crate) fn try_clone(&self) -> io::Result<InFile> {
Ok(InFile {
file: self.file.try_clone()?,
log: self.log.clone(),
})
}
}
impl From<File> for InFile {
fn from(file: File) -> Self {
InFile { file, log: None }
}
}
impl std::ops::Deref for InFile {
type Target = File;
fn deref(&self) -> &File {
&self.file
}
}
impl Read for InFile {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let n = self.file.read(buf)?;
if let Some((log, path)) = &self.log {
let mut log = log.lock().unwrap_or_else(|p| p.into_inner());
log.entry(path.clone()).or_default().bytes += n as u64;
}
Ok(n)
}
}
impl Seek for InFile {
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
self.file.seek(pos)
}
}
pub(crate) struct Input {
src: InputSrc,
buf: Box<[u8; CHUNK]>,
pos: usize,
len: usize,
eof: bool,
pub(crate) kind: Kind,
}
pub(crate) enum InputSrc {
File(InFile),
Stdin(io::Stdin),
Bytes(io::Cursor<Vec<u8>>),
}
impl Read for InputSrc {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self {
InputSrc::File(f) => f.read(buf),
InputSrc::Stdin(s) => s.read(buf),
InputSrc::Bytes(c) => c.read(buf),
}
}
}
impl Input {
pub(crate) fn file(f: impl Into<InFile>) -> Self {
let f = f.into();
let kind = kind_of(&f);
Self::new(InputSrc::File(f), kind)
}
pub(crate) fn stdin() -> Self {
match stdin_file() {
Ok(f) => Self::file(f),
Err(_) => Self::new(InputSrc::Stdin(io::stdin()), Kind::Pipe),
}
}
pub(crate) fn bytes(data: Vec<u8>) -> Self {
Self::new(InputSrc::Bytes(io::Cursor::new(data)), Kind::Pipe)
}
fn new(src: InputSrc, kind: Kind) -> Self {
Self {
src,
buf: Box::new([0; CHUNK]),
pos: 0,
len: 0,
eof: false,
kind,
}
}
fn read_src(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if self.eof {
return Ok(0);
}
loop {
match self.src.read(buf) {
Ok(0) => {
self.eof = true;
return Ok(0);
}
Ok(n) => return Ok(n),
Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
Err(e) => return Err(e),
}
}
}
pub(crate) fn copy_raw(
&mut self,
out: &mut dyn Write,
stop: &mut dyn FnMut() -> bool,
) -> io::Result<bool> {
if let InputSrc::File(f) = &mut self.src
&& f.seek(SeekFrom::Start(0)).is_ok()
{
self.eof = false;
}
self.pos = 0;
self.len = 0;
let mut buf = [0u8; CHUNK];
loop {
if stop() {
return Ok(true);
}
let n = self.read_src(&mut buf)?;
if n == 0 {
return Ok(false);
}
out.write_all(&buf[..n])?;
}
}
}
impl Read for Input {
fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
let avail = self.fill_buf()?;
let n = avail.len().min(out.len());
out[..n].copy_from_slice(&avail[..n]);
self.consume(n);
Ok(n)
}
}
impl BufRead for Input {
fn fill_buf(&mut self) -> io::Result<&[u8]> {
if self.pos == self.len {
let mut n = 0;
let mut chunk = [0u8; CHUNK];
while n < CHUNK {
match self.read_src(&mut chunk[n..])? {
0 => break,
k => n += k,
}
}
self.buf[..n].copy_from_slice(&chunk[..n]);
self.pos = 0;
self.len = n;
}
Ok(&self.buf[self.pos..self.len])
}
fn consume(&mut self, n: usize) {
self.pos = (self.pos + n).min(self.len);
}
}
pub(crate) struct Counting<W> {
pub(crate) inner: W,
pub(crate) n: Arc<AtomicU64>,
}
impl<W: Write> Write for Counting<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let k = self.inner.write(buf)?;
self.n.fetch_add(k as u64, Ordering::Relaxed);
Ok(k)
}
fn flush(&mut self) -> io::Result<()> {
self.inner.flush()
}
}
pub(crate) enum Host {
Real,
Sandbox(Sandbox),
}
pub(crate) struct Sandbox {
pub(crate) dir: PathBuf,
pub(crate) stdin: Option<Vec<u8>>,
pub(crate) stdout: Arc<Mutex<Vec<u8>>>,
pub(crate) budget: Arc<AtomicU64>,
pub(crate) reads: ReadLog,
}
impl Host {
pub(crate) fn path(&self, b: &[u8]) -> PathBuf {
match self {
Host::Real => path(b),
Host::Sandbox(s) => s.dir.join(path(b)),
}
}
pub(crate) fn track(&self, path: &Path, f: File) -> InFile {
match self {
Host::Real => InFile::from(f),
Host::Sandbox(s) => {
let len = f.metadata().map_or(0, |m| m.len());
let mut log = s.reads.lock().unwrap_or_else(|p| p.into_inner());
let entry = log.entry(path.to_path_buf()).or_default();
entry.opens += 1;
entry.len = entry.len.max(len);
InFile {
file: f,
log: Some((s.reads.clone(), path.to_path_buf())),
}
}
}
}
pub(crate) fn stdin_is_tty(&self) -> bool {
matches!(self, Host::Real) && crate::compat::stdin_is_tty()
}
pub(crate) fn stdout_is_tty(&self) -> bool {
matches!(self, Host::Real) && crate::compat::stdout_is_tty()
}
pub(crate) fn stderr_is_tty(&self) -> bool {
use std::io::IsTerminal;
matches!(self, Host::Real) && io::stderr().is_terminal()
}
pub(crate) fn stdin(&mut self) -> Input {
match self {
Host::Real => Input::stdin(),
Host::Sandbox(s) => Input::bytes(s.stdin.take().unwrap_or_default()),
}
}
pub(crate) fn charge(&self, n: u64) -> io::Result<()> {
match self {
Host::Real => Ok(()),
Host::Sandbox(s) => take(&s.budget, n),
}
}
pub(crate) fn out_writer(&self, f: Option<File>) -> BufWriter<Box<dyn Write + Send>> {
let w: Box<dyn Write + Send> = match (self, f) {
(Host::Real, Some(f)) => Box::new(f),
(Host::Real, None) => stdout_writer(),
(Host::Sandbox(s), f) => {
let inner: Box<dyn Write + Send> = match f {
Some(f) => Box::new(f),
None => Box::new(Captured(s.stdout.clone())),
};
Box::new(Capped {
inner,
budget: s.budget.clone(),
})
}
};
BufWriter::with_capacity(64 * 1024, w)
}
}
struct Captured(Arc<Mutex<Vec<u8>>>);
impl Write for Captured {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let mut v = self.0.lock().unwrap_or_else(|p| p.into_inner());
v.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
struct Capped {
inner: Box<dyn Write + Send>,
budget: Arc<AtomicU64>,
}
impl Write for Capped {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let want = buf.len() as u64;
take(&self.budget, want)?;
let n = match self.inner.write(buf) {
Ok(n) => n,
Err(e) => {
self.budget.fetch_add(want, Ordering::Relaxed);
return Err(e);
}
};
self.budget.fetch_add(want - n as u64, Ordering::Relaxed);
Ok(n)
}
fn flush(&mut self) -> io::Result<()> {
self.inner.flush()
}
}
fn take(budget: &AtomicU64, n: u64) -> io::Result<()> {
let mut left = budget.load(Ordering::Relaxed);
loop {
let rest = left.checked_sub(n).ok_or_else(no_space)?;
match budget.compare_exchange_weak(left, rest, Ordering::Relaxed, Ordering::Relaxed) {
Ok(_) => return Ok(()),
Err(now) => left = now,
}
}
}
fn no_space() -> io::Error {
#[cfg(unix)]
{
io::Error::from_raw_os_error(libc::ENOSPC)
}
#[cfg(not(unix))]
{
io::Error::from(io::ErrorKind::StorageFull)
}
}
pub(crate) fn printf(w: &mut dyn Write, fmt: &str, args: &[&[u8]]) {
let mut out = Vec::with_capacity(fmt.len() + 64);
let mut args = args.iter();
let mut bytes = fmt.as_bytes().iter().copied();
while let Some(b) = bytes.next() {
if b != b'%' {
out.push(b);
continue;
}
match bytes.next() {
Some(b'%') => out.push(b'%'),
Some(b's' | b'd') => out.extend_from_slice(args.next().copied().unwrap_or_default()),
Some(other) => out.extend_from_slice(&[b'%', other]),
None => out.push(b'%'),
}
}
let _ = w.write_all(&out);
}
pub(crate) fn set_times(meta: &Metadata, to: &Path) -> io::Result<()> {
let mut times = fs::FileTimes::new();
if let Ok(t) = meta.accessed() {
times = times.set_accessed(t);
}
if let Ok(t) = meta.modified() {
times = times.set_modified(t);
}
File::options().write(true).open(to)?.set_times(times)
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(unix)]
#[test]
fn signal_sets_are_the_measured_ones() {
assert_eq!(
XZ_SIGNALS,
&[
libc::SIGINT,
libc::SIGTERM,
libc::SIGHUP,
libc::SIGPIPE,
libc::SIGXCPU,
libc::SIGXFSZ
]
);
assert_eq!(GNU_GZIP_SIGNALS, XZ_SIGNALS);
assert_eq!(APPLE_GZIP_SIGNALS, &[libc::SIGINT]);
assert_eq!(BZIP2_SIGNALS, &[libc::SIGINT, libc::SIGTERM, libc::SIGHUP]);
}
#[cfg(unix)]
#[test]
fn the_flag_records_a_signal_once_and_restores_on_drop() {
let host = Host::Real;
let guard = arm_interrupt_flag(&host, &[libc::SIGUSR2]);
assert!(!take_interrupt(&host));
unsafe { libc::raise(libc::SIGUSR2) };
assert!(take_interrupt(&host));
assert!(!take_interrupt(&host));
drop(guard);
let mut now: libc::sigaction = unsafe { std::mem::zeroed() };
unsafe { libc::sigaction(libc::SIGUSR2, std::ptr::null(), &mut now) };
assert_eq!(now.sa_sigaction, libc::SIG_DFL);
}
#[test]
fn a_sandbox_never_sees_an_interrupt() {
let host = Host::Sandbox(Sandbox {
dir: PathBuf::new(),
stdin: None,
stdout: Arc::default(),
budget: Arc::new(AtomicU64::new(0)),
reads: ReadLog::default(),
});
let _g = arm_interrupt_flag(&host, BZIP2_SIGNALS);
assert!(!take_interrupt(&host));
}
#[test]
fn the_capped_writer_never_wraps_its_budget() {
let budget = Arc::new(AtomicU64::new(10));
let mut w = Capped {
inner: Box::new(io::sink()),
budget: budget.clone(),
};
assert_eq!(w.write(b"1234567").unwrap(), 7);
assert_eq!(budget.load(Ordering::Relaxed), 3);
let e = w.write(b"1234567").unwrap_err();
assert_eq!(e.kind(), no_space().kind());
assert_eq!(budget.load(Ordering::Relaxed), 3);
assert_eq!(w.write(b"123").unwrap(), 3);
assert_eq!(budget.load(Ordering::Relaxed), 0);
assert!(w.write(b"x").is_err());
assert_eq!(w.write(b"").unwrap(), 0);
}
fn scratch(tag: &str) -> PathBuf {
let dir = std::env::temp_dir().join(format!("stuffr-host-{tag}-{}", std::process::id()));
let _ = fs::remove_dir_all(&dir);
fs::create_dir_all(&dir).unwrap();
dir
}
#[test]
fn publish_never_replaces_an_output_that_appeared_meanwhile() {
let dir = scratch("exists");
let (tmp, out) = (dir.join(".tmp"), dir.join("a.out"));
fs::write(&tmp, b"new").unwrap();
fs::write(&out, b"appeared").unwrap();
assert_eq!(publish(&tmp, &out, false).unwrap(), Published::Exists);
assert!(!tmp.exists(), "the temporary is left");
assert_eq!(fs::read(&out).unwrap(), b"appeared");
fs::write(&tmp, b"new").unwrap();
fs::remove_file(&out).unwrap();
assert_eq!(publish(&tmp, &out, false).unwrap(), Published::Done);
assert!(!tmp.exists());
assert_eq!(fs::read(&out).unwrap(), b"new");
let _ = fs::remove_dir_all(&dir);
}
#[test]
fn publish_with_force_replaces() {
let dir = scratch("force");
let (tmp, out) = (dir.join(".tmp"), dir.join("a.out"));
fs::write(&tmp, b"new").unwrap();
fs::write(&out, b"old").unwrap();
assert_eq!(publish(&tmp, &out, true).unwrap(), Published::Done);
assert!(!tmp.exists());
assert_eq!(fs::read(&out).unwrap(), b"new");
let _ = fs::remove_dir_all(&dir);
}
#[cfg(unix)]
#[test]
fn create_temp_names_the_family_and_is_0600() {
use std::os::unix::fs::PermissionsExt;
let dir = scratch("temp");
let (a, _fa) = create_temp(&dir, "gzip").unwrap();
let (b, _fb) = create_temp(&dir, "gzip").unwrap();
assert_ne!(a, b);
let name = a.file_name().unwrap().to_str().unwrap().to_owned();
assert!(
name.starts_with(&format!(".stuffr-gzip-{}-", std::process::id()))
&& name.ends_with(".tmp"),
"{name}"
);
assert_eq!(
fs::metadata(&a).unwrap().permissions().mode() & 0o777,
0o600
);
let _ = fs::remove_dir_all(&dir);
}
#[cfg(unix)]
#[test]
fn arm_then_drop_disarms() {
use std::os::unix::ffi::OsStrExt;
let dir = scratch("arm");
let file = dir.join("partial.xz");
fs::write(&file, b"partial").unwrap();
let bytes = file.as_os_str().as_bytes();
{
let _guard = arm_cleanup(&Host::Real, &file, XZ_SIGNALS);
assert!(ARMED.load(Ordering::Acquire), "not armed");
let slot = unsafe { &*SLOT.0.get() };
assert_eq!(&slot[..bytes.len()], bytes);
assert_eq!(slot[bytes.len()], 0);
}
assert!(!ARMED.load(Ordering::Acquire), "still armed after drop");
let long = dir.join("x".repeat(SLOT_LEN));
let guard = arm_cleanup(&Host::Real, &long, XZ_SIGNALS);
assert!(!ARMED.load(Ordering::Acquire), "an over-long path armed");
drop(guard);
let _ = fs::remove_dir_all(&dir);
}
#[cfg(unix)]
fn in_child(child: impl FnOnce()) -> libc::c_int {
unsafe {
let pid = libc::fork();
assert!(pid >= 0, "fork: {}", io::Error::last_os_error());
if pid == 0 {
child();
libc::_exit(99);
}
let mut status = 0;
while libc::waitpid(pid, &mut status, 0) < 0 {
assert_eq!(io::Error::last_os_error().raw_os_error(), Some(libc::EINTR));
}
status
}
}
#[cfg(unix)]
fn default_dispositions() {
for &sig in XZ_SIGNALS {
unsafe { libc::signal(sig, libc::SIG_DFL) };
}
}
#[cfg(unix)]
#[test]
fn an_interrupt_removes_the_armed_file() {
let dir = scratch("interrupt");
for &sig in XZ_SIGNALS {
let file = dir.join(format!("partial-{sig}.xz"));
fs::write(&file, b"partial").unwrap();
let status = in_child(|| {
default_dispositions();
ARMED.store(false, Ordering::SeqCst);
let _guard = arm_cleanup(&Host::Real, &file, XZ_SIGNALS);
unsafe { libc::raise(sig) };
});
assert!(
libc::WIFSIGNALED(status),
"signal {sig}: the child exited (status {status:#x})"
);
assert_eq!(libc::WTERMSIG(status), sig);
assert!(!file.exists(), "signal {sig}: the partial file survived");
}
let _ = fs::remove_dir_all(&dir);
}
#[cfg(unix)]
fn current_mask() -> libc::sigset_t {
unsafe {
let mut cur: libc::sigset_t = std::mem::zeroed();
libc::pthread_sigmask(libc::SIG_SETMASK, std::ptr::null(), &mut cur);
cur
}
}
#[cfg(unix)]
fn same_mask(a: &libc::sigset_t, b: &libc::sigset_t) -> bool {
(1..32).all(|s| unsafe { libc::sigismember(a, s) == libc::sigismember(b, s) })
}
#[cfg(unix)]
#[test]
fn create_temp_armed_arms_the_new_temporary() {
use std::os::unix::ffi::OsStrExt;
let dir = scratch("temp-armed");
let status = in_child(|| {
default_dispositions();
ARMED.store(false, Ordering::SeqCst);
let code = match create_temp_armed(&Host::Real, &dir, "xz", XZ_SIGNALS) {
Err(_) => 1,
Ok((path, _file, guard)) => {
let bytes = path.as_os_str().as_bytes();
let slot = unsafe { &*SLOT.0.get() };
let armed = ARMED.load(Ordering::Acquire);
let held = &slot[..bytes.len()] == bytes && slot[bytes.len()] == 0;
drop(guard);
let released = !ARMED.load(Ordering::Acquire);
if !armed {
2
} else if !held {
3
} else if !released {
4
} else {
0
}
}
};
unsafe { libc::_exit(code) };
});
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"status {status:#x}"
);
let _ = fs::remove_dir_all(&dir);
}
#[cfg(unix)]
#[test]
fn create_temp_armed_failure_leaves_nothing_armed_and_the_mask_restored() {
let dir = scratch("temp-armed-fail").join("missing");
let status = in_child(|| {
default_dispositions();
ARMED.store(false, Ordering::SeqCst);
let before = current_mask();
let failed = create_temp_armed(&Host::Real, &dir, "xz", XZ_SIGNALS).is_err();
let after = current_mask();
let code = if !failed {
1
} else if ARMED.load(Ordering::Acquire) {
2
} else if !same_mask(&before, &after) {
3
} else {
0
};
unsafe { libc::_exit(code) };
});
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"status {status:#x}"
);
let _ = fs::remove_dir_all(dir.parent().unwrap());
}
#[cfg(unix)]
#[test]
fn create_temp_armed_in_a_sandbox_arms_nothing() {
let dir = scratch("temp-armed-sandbox");
let host = Host::Sandbox(Sandbox {
dir: dir.clone(),
stdin: None,
stdout: Arc::default(),
budget: Arc::new(AtomicU64::new(0)),
reads: ReadLog::default(),
});
let status = in_child(|| {
default_dispositions();
ARMED.store(false, Ordering::SeqCst);
let code = match create_temp_armed(&host, &dir, "xz", XZ_SIGNALS) {
Err(_) => 1,
Ok((path, _file, _guard)) => {
if !path.exists() {
2
} else if ARMED.load(Ordering::Acquire) {
3
} else {
0
}
}
};
unsafe { libc::_exit(code) };
});
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"status {status:#x}"
);
let _ = fs::remove_dir_all(&dir);
}
#[cfg(unix)]
#[test]
fn sandbox_never_installs_handlers() {
let dir = scratch("sandbox-signals");
let file = dir.join("partial.xz");
fs::write(&file, b"partial").unwrap();
let host = Host::Sandbox(Sandbox {
dir: dir.clone(),
stdin: None,
stdout: Arc::default(),
budget: Arc::new(AtomicU64::new(0)),
reads: ReadLog::default(),
});
let status = in_child(|| {
default_dispositions();
ARMED.store(false, Ordering::SeqCst);
let _guard = arm_cleanup(&host, &file, XZ_SIGNALS);
let mut all_default = !ARMED.load(Ordering::Acquire);
for &sig in XZ_SIGNALS {
unsafe {
let mut old: libc::sigaction = std::mem::zeroed();
let queried = libc::sigaction(sig, std::ptr::null(), &mut old) == 0;
all_default &= queried && old.sa_sigaction == libc::SIG_DFL;
}
}
unsafe { libc::_exit(if all_default { 0 } else { 1 }) };
});
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"a sandboxed arm installed a handler (status {status:#x})"
);
assert!(file.exists());
let _ = fs::remove_dir_all(&dir);
}
#[cfg(unix)]
fn sigint_action() -> libc::sighandler_t {
unsafe {
let mut old: libc::sigaction = std::mem::zeroed();
libc::sigaction(libc::SIGINT, std::ptr::null(), &mut old);
old.sa_sigaction
}
}
#[cfg(unix)]
#[test]
fn exit_kind_restores_sig_dfl_on_drop() {
let dir = scratch("exit-kind");
let file = dir.join("partial.zst");
fs::write(&file, b"partial").unwrap();
let status = in_child(|| {
default_dispositions();
unsafe { libc::signal(libc::SIGINT, libc::SIG_IGN) };
ARMED.store(false, Ordering::SeqCst);
let handler = on_interrupt_exit as extern "C" fn(libc::c_int) as libc::sighandler_t;
let guard = arm_cleanup_exit(&Host::Real, &file);
let armed = ARMED.load(Ordering::Acquire) && sigint_action() == handler;
let term_default = {
let mut old: libc::sigaction = unsafe { std::mem::zeroed() };
unsafe { libc::sigaction(libc::SIGTERM, std::ptr::null(), &mut old) };
old.sa_sigaction == libc::SIG_DFL
};
drop(guard);
let restored = !ARMED.load(Ordering::Acquire) && sigint_action() == libc::SIG_DFL;
unsafe { libc::signal(libc::SIGINT, libc::SIG_IGN) };
let g = arm_cleanup_exit(&Host::Real, &dir);
let skipped = !ARMED.load(Ordering::Acquire) && sigint_action() == libc::SIG_IGN;
drop(g);
let untouched = sigint_action() == libc::SIG_IGN;
let code = match (armed, term_default, restored, skipped, untouched) {
(false, ..) => 1,
(_, false, ..) => 2,
(_, _, false, ..) => 3,
(_, _, _, false, _) => 4,
(.., false) => 5,
_ => 0,
};
unsafe { libc::_exit(code) };
});
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"status {status:#x}"
);
assert!(file.exists());
let _ = fs::remove_dir_all(&dir);
}
#[cfg(unix)]
#[test]
fn sandbox_installs_nothing() {
let dir = scratch("exit-sandbox");
let file = dir.join("partial.zst");
fs::write(&file, b"partial").unwrap();
let host = Host::Sandbox(Sandbox {
dir: dir.clone(),
stdin: None,
stdout: Arc::default(),
budget: Arc::new(AtomicU64::new(0)),
reads: ReadLog::default(),
});
let status = in_child(|| {
default_dispositions();
ARMED.store(false, Ordering::SeqCst);
let g = arm_cleanup_exit(&host, &file);
let direct = !ARMED.load(Ordering::Acquire) && sigint_action() == libc::SIG_DFL;
drop(g);
let made = match create_temp_armed_exit(&dir, "zstd", &host) {
Ok((path, _file, _g)) => {
path.exists()
&& !ARMED.load(Ordering::Acquire)
&& sigint_action() == libc::SIG_DFL
}
Err(_) => false,
};
let after = sigint_action() == libc::SIG_DFL;
let code = match (direct, made, after) {
(false, ..) => 1,
(_, false, _) => 2,
(.., false) => 3,
_ => 0,
};
unsafe { libc::_exit(code) };
});
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"status {status:#x}"
);
assert!(file.exists());
let _ = fs::remove_dir_all(&dir);
}
#[cfg(unix)]
#[test]
fn create_temp_armed_exit_arms_the_new_temporary() {
let dir = scratch("exit-temp");
let status = in_child(|| {
default_dispositions();
ARMED.store(false, Ordering::SeqCst);
unsafe {
let null = libc::open(c"/dev/null".as_ptr(), libc::O_WRONLY);
libc::dup2(null, 2);
}
if let Ok((_path, _file, _guard)) = create_temp_armed_exit(&dir, "zstd", &Host::Real) {
unsafe { libc::raise(libc::SIGINT) };
}
unsafe { libc::_exit(1) };
});
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 2,
"status {status:#x}"
);
assert!(
fs::read_dir(&dir).unwrap().next().is_none(),
"the temporary survived"
);
let _ = fs::remove_dir_all(&dir);
}
#[cfg(unix)]
#[test]
fn progress_signals_only_set_the_flag() {
for &sig in PROGRESS_SIGNALS {
let status = in_child(|| {
unsafe { libc::signal(sig, libc::SIG_IGN) };
PROGRESS_REQUESTED.store(false, Ordering::SeqCst);
install_progress_signals(&Host::Real);
unsafe { libc::raise(sig) };
let set = PROGRESS_REQUESTED.load(Ordering::SeqCst);
unsafe { libc::_exit(if set { 0 } else { 1 }) };
});
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"signal {sig}: status {status:#x}"
);
}
assert!(PROGRESS_SIGNALS.contains(&libc::SIGUSR1));
#[cfg(target_vendor = "apple")]
assert!(PROGRESS_SIGNALS.contains(&libc::SIGINFO));
let host = Host::Sandbox(Sandbox {
dir: std::env::temp_dir(),
stdin: None,
stdout: Arc::default(),
budget: Arc::new(AtomicU64::new(0)),
reads: ReadLog::default(),
});
let status = in_child(|| {
let mut all_default = true;
for &sig in PROGRESS_SIGNALS.iter().chain(&[libc::SIGTSTP]) {
unsafe { libc::signal(sig, libc::SIG_DFL) };
}
install_progress_signals(&host);
install_stop_handler(&host);
for &sig in PROGRESS_SIGNALS.iter().chain(&[libc::SIGTSTP]) {
unsafe {
let mut old: libc::sigaction = std::mem::zeroed();
let queried = libc::sigaction(sig, std::ptr::null(), &mut old) == 0;
all_default &= queried && old.sa_sigaction == libc::SIG_DFL;
}
}
unsafe { libc::_exit(if all_default { 0 } else { 1 }) };
});
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"a sandbox installed a progress handler (status {status:#x})"
);
}
#[cfg(unix)]
#[test]
fn a_stop_is_left_out_of_the_running_clock() {
unsafe {
let pid = libc::fork();
assert!(pid >= 0, "fork: {}", io::Error::last_os_error());
if pid == 0 {
libc::signal(libc::SIGTSTP, libc::SIG_DFL);
install_stop_handler(&Host::Real);
let before = STOPPED_MS.load(Ordering::SeqCst);
libc::raise(libc::SIGTSTP);
let stopped = STOPPED_MS.load(Ordering::SeqCst) - before;
libc::_exit(if stopped >= 150 { 0 } else { 1 });
}
let mut status = 0;
while libc::waitpid(pid, &mut status, libc::WUNTRACED) < 0 {}
assert!(libc::WIFSTOPPED(status), "status {status:#x}");
assert_eq!(libc::WSTOPSIG(status), libc::SIGSTOP);
std::thread::sleep(std::time::Duration::from_millis(250));
libc::kill(pid, libc::SIGCONT);
while libc::waitpid(pid, &mut status, 0) < 0 {}
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"the stop was not counted (status {status:#x})"
);
}
let status = in_child(|| {
unsafe {
libc::signal(libc::SIGTSTP, libc::SIG_IGN);
install_stop_handler(&Host::Real);
let mut old: libc::sigaction = std::mem::zeroed();
libc::sigaction(libc::SIGTSTP, std::ptr::null(), &mut old);
libc::_exit(if old.sa_sigaction == libc::SIG_IGN {
0
} else {
1
});
}
});
assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"an ignored SIGTSTP was hooked (status {status:#x})"
);
}
}