Skip to main content

secure_exec_kernel/
resource_accounting.rs

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