1use crate::fd_table::FdTableManager;
2use crate::pipe_manager::PipeManager;
3use crate::process_table::{ProcessStatus, ProcessTable};
4use crate::pty::PtyManager;
5use crate::socket_table::{SocketState, SocketTable};
6use crate::vfs::{VfsResult, VirtualFileSystem};
7use std::collections::{BTreeMap, BTreeSet};
8use std::error::Error;
9use std::fmt;
10
11pub const DEFAULT_MAX_FILESYSTEM_BYTES: u64 = 64 * 1024 * 1024;
12pub const DEFAULT_MAX_INODE_COUNT: usize = 16_384;
13pub const DEFAULT_MAX_PROCESSES: usize = 256;
14pub const DEFAULT_MAX_OPEN_FDS: usize = 256;
15pub const DEFAULT_MAX_PIPES: usize = 128;
16pub const DEFAULT_MAX_PTYS: usize = 128;
17pub const DEFAULT_MAX_SOCKETS: usize = 256;
18pub const DEFAULT_MAX_CONNECTIONS: usize = 256;
19pub const DEFAULT_BLOCKING_READ_TIMEOUT_MS: u64 = 5_000;
20pub const DEFAULT_MAX_PREAD_BYTES: usize = 64 * 1024 * 1024;
21pub const DEFAULT_MAX_FD_WRITE_BYTES: usize = 64 * 1024 * 1024;
22pub const DEFAULT_MAX_PROCESS_ARGV_BYTES: usize = 1024 * 1024;
23pub const DEFAULT_MAX_PROCESS_ENV_BYTES: usize = 1024 * 1024;
24pub const DEFAULT_MAX_READDIR_ENTRIES: usize = 4_096;
25pub const DEFAULT_VIRTUAL_CPU_COUNT: usize = 1;
26
27#[derive(Debug, Clone, PartialEq, Eq, Default)]
28pub struct ResourceSnapshot {
29 pub running_processes: usize,
30 pub exited_processes: usize,
31 pub fd_tables: usize,
32 pub open_fds: usize,
33 pub pipes: usize,
34 pub pipe_buffered_bytes: usize,
35 pub ptys: usize,
36 pub pty_buffered_input_bytes: usize,
37 pub pty_buffered_output_bytes: usize,
38 pub sockets: usize,
39 pub socket_listeners: usize,
40 pub socket_connections: usize,
41}
42
43#[derive(Debug, Clone, PartialEq, Eq)]
44pub struct ResourceLimits {
45 pub virtual_cpu_count: Option<usize>,
46 pub max_processes: Option<usize>,
47 pub max_open_fds: Option<usize>,
48 pub max_pipes: Option<usize>,
49 pub max_ptys: Option<usize>,
50 pub max_sockets: Option<usize>,
51 pub max_connections: Option<usize>,
52 pub max_filesystem_bytes: Option<u64>,
53 pub max_inode_count: Option<usize>,
54 pub max_blocking_read_ms: Option<u64>,
55 pub max_pread_bytes: Option<usize>,
56 pub max_fd_write_bytes: Option<usize>,
57 pub max_process_argv_bytes: Option<usize>,
58 pub max_process_env_bytes: Option<usize>,
59 pub max_readdir_entries: Option<usize>,
60 pub max_wasm_fuel: Option<u64>,
61 pub max_wasm_memory_bytes: Option<u64>,
62 pub max_wasm_stack_bytes: Option<usize>,
63}
64
65impl Default for ResourceLimits {
66 fn default() -> Self {
67 Self {
68 virtual_cpu_count: Some(DEFAULT_VIRTUAL_CPU_COUNT),
69 max_processes: Some(DEFAULT_MAX_PROCESSES),
70 max_open_fds: Some(DEFAULT_MAX_OPEN_FDS),
71 max_pipes: Some(DEFAULT_MAX_PIPES),
72 max_ptys: Some(DEFAULT_MAX_PTYS),
73 max_sockets: Some(DEFAULT_MAX_SOCKETS),
74 max_connections: Some(DEFAULT_MAX_CONNECTIONS),
75 max_filesystem_bytes: Some(DEFAULT_MAX_FILESYSTEM_BYTES),
76 max_inode_count: Some(DEFAULT_MAX_INODE_COUNT),
77 max_blocking_read_ms: Some(DEFAULT_BLOCKING_READ_TIMEOUT_MS),
78 max_pread_bytes: Some(DEFAULT_MAX_PREAD_BYTES),
79 max_fd_write_bytes: Some(DEFAULT_MAX_FD_WRITE_BYTES),
80 max_process_argv_bytes: Some(DEFAULT_MAX_PROCESS_ARGV_BYTES),
81 max_process_env_bytes: Some(DEFAULT_MAX_PROCESS_ENV_BYTES),
82 max_readdir_entries: Some(DEFAULT_MAX_READDIR_ENTRIES),
83 max_wasm_fuel: None,
84 max_wasm_memory_bytes: None,
85 max_wasm_stack_bytes: None,
86 }
87 }
88}
89
90#[derive(Debug, Clone, PartialEq, Eq, Default)]
91pub struct FileSystemUsage {
92 pub total_bytes: u64,
93 pub inode_count: usize,
94}
95
96#[derive(Debug, Clone, PartialEq, Eq)]
97pub struct ResourceError {
98 code: &'static str,
99 message: String,
100}
101
102impl ResourceError {
103 pub fn code(&self) -> &'static str {
104 self.code
105 }
106
107 fn exhausted(message: impl Into<String>) -> Self {
108 Self {
109 code: "EAGAIN",
110 message: message.into(),
111 }
112 }
113
114 fn file_table_full(message: impl Into<String>) -> Self {
115 Self {
116 code: "ENFILE",
117 message: message.into(),
118 }
119 }
120
121 fn filesystem_full(message: impl Into<String>) -> Self {
122 Self {
123 code: "ENOSPC",
124 message: message.into(),
125 }
126 }
127
128 fn invalid_input(message: impl Into<String>) -> Self {
129 Self {
130 code: "EINVAL",
131 message: message.into(),
132 }
133 }
134
135 fn out_of_memory(message: impl Into<String>) -> Self {
136 Self {
137 code: "ENOMEM",
138 message: message.into(),
139 }
140 }
141}
142
143impl fmt::Display for ResourceError {
144 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
145 write!(f, "{}: {}", self.code, self.message)
146 }
147}
148
149impl Error for ResourceError {}
150
151#[derive(Debug, Clone, Default)]
152pub struct ResourceAccountant {
153 limits: ResourceLimits,
154}
155
156impl ResourceAccountant {
157 pub fn new(limits: ResourceLimits) -> Self {
158 Self { limits }
159 }
160
161 pub fn limits(&self) -> &ResourceLimits {
162 &self.limits
163 }
164
165 pub fn snapshot(
166 &self,
167 processes: &ProcessTable,
168 fd_tables: &FdTableManager,
169 pipes: &PipeManager,
170 ptys: &PtyManager,
171 sockets: &SocketTable,
172 ) -> ResourceSnapshot {
173 let process_list = processes.list_processes();
174 let running_processes = process_list
175 .values()
176 .filter(|process| process.status == ProcessStatus::Running)
177 .count();
178 let exited_processes = process_list
179 .values()
180 .filter(|process| process.status == ProcessStatus::Exited)
181 .count();
182 let socket_snapshot = sockets.snapshot();
183
184 ResourceSnapshot {
185 running_processes,
186 exited_processes,
187 fd_tables: fd_tables.len(),
188 open_fds: fd_tables.total_open_fds(),
189 pipes: pipes.pipe_count(),
190 pipe_buffered_bytes: pipes.buffered_bytes(),
191 ptys: ptys.pty_count(),
192 pty_buffered_input_bytes: ptys.buffered_input_bytes(),
193 pty_buffered_output_bytes: ptys.buffered_output_bytes(),
194 sockets: socket_snapshot.sockets,
195 socket_listeners: socket_snapshot.listeners,
196 socket_connections: socket_snapshot.connections,
197 }
198 }
199
200 pub fn check_process_spawn(
201 &self,
202 snapshot: &ResourceSnapshot,
203 additional_fds: usize,
204 ) -> Result<(), ResourceError> {
205 if let Some(limit) = self.limits.max_processes {
206 if snapshot.running_processes + snapshot.exited_processes >= limit {
207 return Err(ResourceError::exhausted("maximum process limit reached"));
208 }
209 }
210
211 self.check_open_fds(snapshot, additional_fds)
212 }
213
214 pub fn check_process_argv_bytes(
215 &self,
216 command: &str,
217 args: &[String],
218 ) -> Result<(), ResourceError> {
219 if let Some(limit) = self.limits.max_process_argv_bytes {
220 let total = argv_payload_bytes(command, args);
221 if total > limit {
222 return Err(ResourceError::invalid_input(format!(
223 "process argv payload {total} bytes exceeds configured limit {limit}"
224 )));
225 }
226 }
227
228 Ok(())
229 }
230
231 pub fn check_process_env_bytes(
232 &self,
233 inherited_env: &BTreeMap<String, String>,
234 overrides: &BTreeMap<String, String>,
235 ) -> Result<(), ResourceError> {
236 if let Some(limit) = self.limits.max_process_env_bytes {
237 let total = merged_env_payload_bytes(inherited_env, overrides);
238 if total > limit {
239 return Err(ResourceError::invalid_input(format!(
240 "process environment payload {total} bytes exceeds configured limit {limit}"
241 )));
242 }
243 }
244
245 Ok(())
246 }
247
248 pub fn check_pipe_allocation(&self, snapshot: &ResourceSnapshot) -> Result<(), ResourceError> {
249 if let Some(limit) = self.limits.max_pipes {
250 if snapshot.pipes >= limit {
251 return Err(ResourceError::exhausted("maximum pipe count reached"));
252 }
253 }
254
255 self.check_open_fds(snapshot, 2)
256 }
257
258 pub fn check_pty_allocation(&self, snapshot: &ResourceSnapshot) -> Result<(), ResourceError> {
259 if let Some(limit) = self.limits.max_ptys {
260 if snapshot.ptys >= limit {
261 return Err(ResourceError::exhausted("maximum PTY count reached"));
262 }
263 }
264
265 self.check_open_fds(snapshot, 2)
266 }
267
268 pub fn check_socket_allocation(
269 &self,
270 snapshot: &ResourceSnapshot,
271 ) -> Result<(), ResourceError> {
272 if let Some(limit) = self.limits.max_sockets {
273 if snapshot.sockets >= limit {
274 return Err(ResourceError::exhausted("maximum socket count reached"));
275 }
276 }
277
278 Ok(())
279 }
280
281 pub fn check_socket_state_transition(
282 &self,
283 snapshot: &ResourceSnapshot,
284 current: SocketState,
285 next: SocketState,
286 ) -> Result<(), ResourceError> {
287 if !current.counts_as_connection() && next.counts_as_connection() {
288 if let Some(limit) = self.limits.max_connections {
289 if snapshot.socket_connections >= limit {
290 return Err(ResourceError::exhausted("maximum connection count reached"));
291 }
292 }
293 }
294
295 Ok(())
296 }
297
298 pub fn check_pread_length(&self, length: usize) -> Result<(), ResourceError> {
299 if let Some(limit) = self.limits.max_pread_bytes {
300 if length > limit {
301 return Err(ResourceError::invalid_input(format!(
302 "pread length {length} exceeds configured limit {limit}"
303 )));
304 }
305 }
306
307 Ok(())
308 }
309
310 pub fn check_fd_write_size(&self, size: usize) -> Result<(), ResourceError> {
311 if let Some(limit) = self.limits.max_fd_write_bytes {
312 if size > limit {
313 return Err(ResourceError::invalid_input(format!(
314 "write size {size} exceeds configured limit {limit}"
315 )));
316 }
317 }
318
319 Ok(())
320 }
321
322 pub fn check_fd_allocation(
323 &self,
324 snapshot: &ResourceSnapshot,
325 additional_fds: usize,
326 ) -> Result<(), ResourceError> {
327 self.check_open_fds(snapshot, additional_fds)
328 }
329
330 pub fn max_readdir_entries(&self) -> Option<usize> {
331 self.limits.max_readdir_entries
332 }
333
334 pub fn check_readdir_entries(&self, entries: usize) -> Result<(), ResourceError> {
335 if let Some(limit) = self.limits.max_readdir_entries {
336 if entries > limit {
337 return Err(ResourceError::out_of_memory(format!(
338 "directory listing with {entries} entries exceeds configured limit {limit}"
339 )));
340 }
341 }
342
343 Ok(())
344 }
345
346 fn check_open_fds(
347 &self,
348 snapshot: &ResourceSnapshot,
349 additional_fds: usize,
350 ) -> Result<(), ResourceError> {
351 if let Some(limit) = self.limits.max_open_fds {
352 if snapshot.open_fds.saturating_add(additional_fds) > limit {
353 return Err(ResourceError::file_table_full(
354 "maximum open file descriptor limit reached",
355 ));
356 }
357 }
358
359 Ok(())
360 }
361
362 pub fn check_filesystem_usage(
363 &self,
364 _usage: &FileSystemUsage,
365 resulting_bytes: u64,
366 resulting_inodes: usize,
367 ) -> Result<(), ResourceError> {
368 if let Some(limit) = self.limits.max_filesystem_bytes {
369 if resulting_bytes > limit {
370 return Err(ResourceError::filesystem_full(
371 "maximum filesystem size limit reached",
372 ));
373 }
374 }
375
376 if let Some(limit) = self.limits.max_inode_count {
377 if resulting_inodes > limit {
378 return Err(ResourceError::filesystem_full(
379 "maximum inode count limit reached",
380 ));
381 }
382 }
383 Ok(())
384 }
385}
386
387fn argv_payload_bytes(command: &str, args: &[String]) -> usize {
388 let command_bytes = command.len().saturating_add(1);
389 command_bytes.saturating_add(
390 args.iter()
391 .map(|arg| arg.len().saturating_add(1))
392 .sum::<usize>(),
393 )
394}
395
396fn env_entry_payload_bytes(key: &str, value: &str) -> usize {
397 key.len()
398 .saturating_add(1)
399 .saturating_add(value.len())
400 .saturating_add(1)
401}
402
403fn merged_env_payload_bytes(
404 inherited_env: &BTreeMap<String, String>,
405 overrides: &BTreeMap<String, String>,
406) -> usize {
407 let mut total = inherited_env
408 .iter()
409 .map(|(key, value)| env_entry_payload_bytes(key, value))
410 .sum::<usize>();
411
412 for (key, value) in overrides {
413 if let Some(previous) = inherited_env.get(key) {
414 total = total.saturating_sub(env_entry_payload_bytes(key, previous));
415 }
416 total = total.saturating_add(env_entry_payload_bytes(key, value));
417 }
418
419 total
420}
421
422pub fn measure_filesystem_usage<F: VirtualFileSystem>(
423 filesystem: &mut F,
424) -> VfsResult<FileSystemUsage> {
425 let mut visited = BTreeSet::new();
426 measure_path_usage(filesystem, "/", &mut visited)
427}
428
429fn measure_path_usage<F: VirtualFileSystem>(
430 filesystem: &mut F,
431 path: &str,
432 visited: &mut BTreeSet<u64>,
433) -> VfsResult<FileSystemUsage> {
434 let stat = filesystem.lstat(path)?;
435 let mut usage = FileSystemUsage::default();
436
437 if visited.insert(stat.ino) {
438 usage.inode_count += 1;
439 if !stat.is_directory {
440 usage.total_bytes = usage.total_bytes.saturating_add(stat.size);
441 }
442 }
443
444 if !stat.is_directory || stat.is_symbolic_link {
445 return Ok(usage);
446 }
447
448 for entry in filesystem.read_dir_with_types(path)? {
449 if matches!(entry.name.as_str(), "." | "..") {
450 continue;
451 }
452
453 let child_path = if path == "/" {
454 format!("/{}", entry.name)
455 } else {
456 format!("{path}/{}", entry.name)
457 };
458 let child_usage = measure_path_usage(filesystem, &child_path, visited)?;
459 usage.total_bytes = usage.total_bytes.saturating_add(child_usage.total_bytes);
460 usage.inode_count = usage.inode_count.saturating_add(child_usage.inode_count);
461 }
462
463 Ok(usage)
464}