use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::Duration;
#[derive(Debug, Clone)]
pub struct SandboxConfig {
pub max_memory_bytes: usize,
pub max_cpu_time_secs: u64,
pub max_call_duration: Duration,
pub allow_network: bool,
pub allow_filesystem: bool,
pub allowed_syscalls: Vec<u32>,
}
impl Default for SandboxConfig {
fn default() -> Self {
Self {
max_memory_bytes: 128 * 1024 * 1024, max_cpu_time_secs: 60,
max_call_duration: Duration::from_secs(30),
allow_network: false,
allow_filesystem: true,
allowed_syscalls: Vec::new(),
}
}
}
#[derive(Debug, Clone)]
pub struct ResourceUsage {
pub memory_bytes: usize,
pub peak_memory_bytes: usize,
pub cpu_time_ms: u64,
pub total_calls: u64,
pub total_syscalls: u64,
}
#[derive(Debug)]
pub struct ResourceTracker {
memory_bytes: AtomicUsize,
peak_memory_bytes: AtomicUsize,
cpu_time_ms: AtomicU64,
total_calls: AtomicU64,
total_syscalls: AtomicU64,
}
impl ResourceTracker {
pub fn new() -> Self {
Self {
memory_bytes: AtomicUsize::new(0),
peak_memory_bytes: AtomicUsize::new(0),
cpu_time_ms: AtomicU64::new(0),
total_calls: AtomicU64::new(0),
total_syscalls: AtomicU64::new(0),
}
}
pub fn update_memory(&self, bytes: usize) {
self.memory_bytes.store(bytes, Ordering::Release);
let mut peak = self.peak_memory_bytes.load(Ordering::Acquire);
while bytes > peak {
match self.peak_memory_bytes.compare_exchange_weak(
peak,
bytes,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => break,
Err(x) => peak = x,
}
}
}
pub fn record_call(&self) {
self.total_calls.fetch_add(1, Ordering::Relaxed);
}
pub fn record_cpu_time(&self, ms: u64) {
self.cpu_time_ms.fetch_add(ms, Ordering::Relaxed);
}
pub fn record_syscall(&self) {
self.total_syscalls.fetch_add(1, Ordering::Relaxed);
}
pub fn usage(&self) -> ResourceUsage {
ResourceUsage {
memory_bytes: self.memory_bytes.load(Ordering::Acquire),
peak_memory_bytes: self.peak_memory_bytes.load(Ordering::Acquire),
cpu_time_ms: self.cpu_time_ms.load(Ordering::Acquire),
total_calls: self.total_calls.load(Ordering::Acquire),
total_syscalls: self.total_syscalls.load(Ordering::Acquire),
}
}
pub fn exceeds_limits(&self, config: &SandboxConfig) -> bool {
self.memory_bytes.load(Ordering::Acquire) > config.max_memory_bytes
}
pub fn reset(&self) {
self.memory_bytes.store(0, Ordering::Release);
self.cpu_time_ms.store(0, Ordering::Relaxed);
self.total_calls.store(0, Ordering::Relaxed);
self.total_syscalls.store(0, Ordering::Relaxed);
}
}
impl Default for ResourceTracker {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Copy)]
pub struct Permissions {
pub can_allocate: bool,
pub can_network: bool,
pub can_filesystem: bool,
pub can_thread: bool,
pub can_env: bool,
}
impl Default for Permissions {
fn default() -> Self {
Self {
can_allocate: true,
can_network: false,
can_filesystem: true,
can_thread: true,
can_env: false,
}
}
}
impl Permissions {
pub fn restrictive() -> Self {
Self {
can_allocate: true,
can_network: false,
can_filesystem: false,
can_thread: false,
can_env: false,
}
}
pub fn permissive() -> Self {
Self {
can_allocate: true,
can_network: true,
can_filesystem: true,
can_thread: true,
can_env: true,
}
}
pub fn check(&self, permission: Permission) -> bool {
match permission {
Permission::Allocate => self.can_allocate,
Permission::Network => self.can_network,
Permission::Filesystem => self.can_filesystem,
Permission::Thread => self.can_thread,
Permission::Env => self.can_env,
}
}
}
#[derive(Debug, Clone, Copy)]
pub enum Permission {
Allocate,
Network,
Filesystem,
Thread,
Env,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_sandbox_config() {
let config = SandboxConfig::default();
assert_eq!(config.max_memory_bytes, 128 * 1024 * 1024);
assert!(!config.allow_network);
}
#[test]
fn test_resource_tracker() {
let tracker = ResourceTracker::new();
tracker.update_memory(1000);
tracker.record_call();
let usage = tracker.usage();
assert_eq!(usage.memory_bytes, 1000);
assert_eq!(usage.total_calls, 1);
}
#[test]
fn test_permissions() {
let perms = Permissions::restrictive();
assert!(perms.can_allocate);
assert!(!perms.can_network);
assert!(!perms.can_filesystem);
}
}