Skip to main content

shape_vm/
resource_limits.rs

1//! Runtime resource limits and sandboxing for the Shape VM.
2//!
3//! Provides configurable limits on instruction count, memory usage,
4//! wall-clock time, and output bytes. The VM checks these limits
5//! during execution and halts with an error when exceeded.
6
7use std::time::{Duration, Instant};
8
9/// Configurable resource limits for VM execution.
10#[derive(Debug, Clone)]
11pub struct ResourceLimits {
12    /// Maximum number of instructions before halting.
13    pub max_instructions: Option<u64>,
14    /// Maximum memory bytes the VM may allocate.
15    pub max_memory_bytes: Option<u64>,
16    /// Maximum wall-clock time for execution.
17    pub max_wall_time: Option<Duration>,
18    /// Maximum output bytes (stdout/stderr combined).
19    pub max_output_bytes: Option<u64>,
20}
21
22impl Default for ResourceLimits {
23    fn default() -> Self {
24        Self {
25            max_instructions: None,
26            max_memory_bytes: None,
27            max_wall_time: None,
28            max_output_bytes: None,
29        }
30    }
31}
32
33impl ResourceLimits {
34    /// Create limits with no restrictions.
35    pub fn unlimited() -> Self {
36        Self::default()
37    }
38
39    /// Create limits suitable for untrusted code execution.
40    pub fn sandboxed() -> Self {
41        Self {
42            max_instructions: Some(10_000_000),
43            max_memory_bytes: Some(256 * 1024 * 1024), // 256 MB
44            max_wall_time: Some(Duration::from_secs(30)),
45            max_output_bytes: Some(1024 * 1024), // 1 MB
46        }
47    }
48}
49
50/// Tracks resource usage during VM execution.
51#[derive(Debug)]
52pub struct ResourceUsage {
53    pub instructions_executed: u64,
54    pub memory_bytes_allocated: u64,
55    pub output_bytes_written: u64,
56    start_time: Option<Instant>,
57    limits: ResourceLimits,
58    /// Check wall time every N instructions (amortized cost).
59    wall_time_check_interval: u64,
60    instructions_since_time_check: u64,
61}
62
63/// Error returned when a resource limit is exceeded.
64#[derive(Debug, Clone)]
65pub enum ResourceLimitExceeded {
66    InstructionLimit { limit: u64, executed: u64 },
67    MemoryLimit { limit: u64, allocated: u64 },
68    WallTimeLimit { limit: Duration, elapsed: Duration },
69    OutputLimit { limit: u64, written: u64 },
70}
71
72impl std::fmt::Display for ResourceLimitExceeded {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        match self {
75            Self::InstructionLimit { limit, executed } => {
76                write!(f, "Instruction limit exceeded: {executed} >= {limit}")
77            }
78            Self::MemoryLimit { limit, allocated } => {
79                write!(
80                    f,
81                    "Memory limit exceeded: {allocated} bytes >= {limit} bytes"
82                )
83            }
84            Self::WallTimeLimit { limit, elapsed } => {
85                write!(f, "Wall time limit exceeded: {elapsed:?} >= {limit:?}")
86            }
87            Self::OutputLimit { limit, written } => {
88                write!(f, "Output limit exceeded: {written} bytes >= {limit} bytes")
89            }
90        }
91    }
92}
93
94impl ResourceUsage {
95    pub fn new(limits: ResourceLimits) -> Self {
96        Self {
97            instructions_executed: 0,
98            memory_bytes_allocated: 0,
99            output_bytes_written: 0,
100            start_time: None,
101            limits,
102            wall_time_check_interval: 1024,
103            instructions_since_time_check: 0,
104        }
105    }
106
107    /// Start tracking wall-clock time.
108    pub fn start(&mut self) {
109        self.start_time = Some(Instant::now());
110    }
111
112    /// Check instruction count limit and amortized wall-time check.
113    /// Called once per instruction in the dispatch loop.
114    #[inline]
115    pub fn tick_instruction(&mut self) -> Result<(), ResourceLimitExceeded> {
116        self.instructions_executed += 1;
117
118        if let Some(limit) = self.limits.max_instructions {
119            if self.instructions_executed >= limit {
120                return Err(ResourceLimitExceeded::InstructionLimit {
121                    limit,
122                    executed: self.instructions_executed,
123                });
124            }
125        }
126
127        // Amortized wall-time check every N instructions.
128        self.instructions_since_time_check += 1;
129        if self.instructions_since_time_check >= self.wall_time_check_interval {
130            self.instructions_since_time_check = 0;
131            self.check_wall_time()?;
132        }
133
134        Ok(())
135    }
136
137    /// Record memory allocation.
138    pub fn record_allocation(&mut self, bytes: u64) -> Result<(), ResourceLimitExceeded> {
139        self.memory_bytes_allocated += bytes;
140        if let Some(limit) = self.limits.max_memory_bytes {
141            if self.memory_bytes_allocated >= limit {
142                return Err(ResourceLimitExceeded::MemoryLimit {
143                    limit,
144                    allocated: self.memory_bytes_allocated,
145                });
146            }
147        }
148        Ok(())
149    }
150
151    /// Record output bytes written.
152    pub fn record_output(&mut self, bytes: u64) -> Result<(), ResourceLimitExceeded> {
153        self.output_bytes_written += bytes;
154        if let Some(limit) = self.limits.max_output_bytes {
155            if self.output_bytes_written >= limit {
156                return Err(ResourceLimitExceeded::OutputLimit {
157                    limit,
158                    written: self.output_bytes_written,
159                });
160            }
161        }
162        Ok(())
163    }
164
165    fn check_wall_time(&self) -> Result<(), ResourceLimitExceeded> {
166        if let (Some(limit), Some(start)) = (self.limits.max_wall_time, self.start_time) {
167            let elapsed = start.elapsed();
168            if elapsed >= limit {
169                return Err(ResourceLimitExceeded::WallTimeLimit { limit, elapsed });
170            }
171        }
172        Ok(())
173    }
174
175    /// Return current limits.
176    pub fn limits(&self) -> &ResourceLimits {
177        &self.limits
178    }
179
180    /// Elapsed wall time since start.
181    pub fn elapsed(&self) -> Option<Duration> {
182        self.start_time.map(|s| s.elapsed())
183    }
184}