shape_vm/
resource_limits.rs1use std::time::{Duration, Instant};
8
9#[derive(Debug, Clone)]
11pub struct ResourceLimits {
12 pub max_instructions: Option<u64>,
14 pub max_memory_bytes: Option<u64>,
16 pub max_wall_time: Option<Duration>,
18 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 pub fn unlimited() -> Self {
36 Self::default()
37 }
38
39 pub fn sandboxed() -> Self {
41 Self {
42 max_instructions: Some(10_000_000),
43 max_memory_bytes: Some(256 * 1024 * 1024), max_wall_time: Some(Duration::from_secs(30)),
45 max_output_bytes: Some(1024 * 1024), }
47 }
48}
49
50#[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 wall_time_check_interval: u64,
60 instructions_since_time_check: u64,
61}
62
63#[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 pub fn start(&mut self) {
109 self.start_time = Some(Instant::now());
110 }
111
112 #[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 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 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 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 pub fn limits(&self) -> &ResourceLimits {
177 &self.limits
178 }
179
180 pub fn elapsed(&self) -> Option<Duration> {
182 self.start_time.map(|s| s.elapsed())
183 }
184}