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