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) 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")
}
#[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>>>;
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) -> io::Result<()> {
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 {
let n = self.read_src(&mut buf)?;
if n == 0 {
return Ok(());
}
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 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::*;
#[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);
}
}