tinysandbox 0.2.0

use std::collections::{BTreeMap, BTreeSet};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use tokio::io::{AsyncRead, AsyncWrite};

use super::fs::Fs;

pub type CommandFuture = Pin<Box<dyn Future<Output = CommandResult> + Send>>;
pub type BoxAsyncRead = Pin<Box<dyn AsyncRead + Send>>;
pub type BoxAsyncWrite = Pin<Box<dyn AsyncWrite + Send>>;

pub trait Command: Send + Sync {
    fn run(&self, ctx: CommandContext) -> CommandFuture;
}

impl<F, Fut> Command for F
where
    F: Fn(CommandContext) -> Fut + Send + Sync,
    Fut: Future<Output = CommandResult> + Send + 'static,
{
    fn run(&self, ctx: CommandContext) -> CommandFuture {
        Box::pin(self(ctx))
    }
}

pub struct CommandContext {
    pub args: Vec<String>,
    pub env: BTreeMap<String, String>,
    pub cwd: String,
    pub stdin: BoxAsyncRead,
    pub stdout: BoxAsyncWrite,
    pub stderr: BoxAsyncWrite,
    pub fs: Fs,
    pub limits: Limits,
    pub commands: Arc<BTreeSet<String>>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CommandResult {
    pub exit_code: i32,
    pub peak_wasm_memory_bytes: Option<usize>,
}

impl CommandResult {
    pub const fn new(exit_code: i32) -> Self {
        Self {
            exit_code,
            peak_wasm_memory_bytes: None,
        }
    }

    pub const fn with_peak_wasm_memory(mut self, bytes: usize) -> Self {
        self.peak_wasm_memory_bytes = Some(bytes);
        self
    }

    pub const fn success() -> Self {
        Self::new(0)
    }

    pub const fn failure() -> Self {
        Self::new(1)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Limits {
    pub wall_time: Duration,
    pub stdout_bytes: usize,
    pub stderr_bytes: usize,
    pub max_commands: usize,
    pub sort_input_bytes: usize,
    pub wasm_memory_bytes: usize,
}

impl Default for Limits {
    fn default() -> Self {
        Self {
            wall_time: Duration::from_secs(30),
            stdout_bytes: 1024 * 1024,
            stderr_bytes: 1024 * 1024,
            max_commands: 1024,
            sort_input_bytes: 8 * 1024 * 1024,
            wasm_memory_bytes: 64 * 1024 * 1024,
        }
    }
}