diff --git a/Cargo.lock b/Cargo.lock index 849e52446a..732c57b1cb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10702,6 +10702,7 @@ dependencies = [ "infra_utils", "mempool_test_utils", "papyrus_config", + "rlimit", "rstest", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index bfac7a76f6..c034bacb93 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -199,6 +199,7 @@ regex = "1.10.4" replace_with = "0.1.7" reqwest = "0.11" retry = "2.0.0" +rlimit = "0.10.2" rstest = "0.17.0" rustc-hex = "2.1.0" schemars = "0.8.12" diff --git a/crates/starknet_sierra_compile/Cargo.toml b/crates/starknet_sierra_compile/Cargo.toml index bd559d226e..e5eea12893 100644 --- a/crates/starknet_sierra_compile/Cargo.toml +++ b/crates/starknet_sierra_compile/Cargo.toml @@ -17,6 +17,7 @@ cairo-lang-starknet-classes.workspace = true cairo-lang-utils.workspace = true cairo-native = { workspace = true, optional = true } papyrus_config.workspace = true +rlimit.workspace = true serde.workspace = true serde_json.workspace = true starknet-types-core.workspace = true diff --git a/crates/starknet_sierra_compile/src/lib.rs b/crates/starknet_sierra_compile/src/lib.rs index 80a543d24b..11773f51f5 100644 --- a/crates/starknet_sierra_compile/src/lib.rs +++ b/crates/starknet_sierra_compile/src/lib.rs @@ -11,6 +11,7 @@ pub mod config; pub mod constants; pub mod errors; pub mod paths; +pub mod resource_limits; pub mod utils; #[cfg(test)] diff --git a/crates/starknet_sierra_compile/src/resource_limits.rs b/crates/starknet_sierra_compile/src/resource_limits.rs new file mode 100644 index 0000000000..248ad270ff --- /dev/null +++ b/crates/starknet_sierra_compile/src/resource_limits.rs @@ -0,0 +1,98 @@ +use std::io; +use std::os::unix::process::CommandExt; +use std::process::Command; + +use rlimit::{setrlimit, Resource}; + +#[cfg(test)] +#[path = "resource_limits_test.rs"] +pub mod test; + +#[allow(dead_code)] +struct ResourceLimits { + resource: Resource, + soft: u64, + hard: u64, + units: String, +} + +impl ResourceLimits { + #[allow(dead_code)] + fn set(&self) -> io::Result<()> { + // Use `println!` and not a logger because this method is called in an unsafe block, and we + // don't want to risk unexpected behavior. + println!("Setting {:?} limits to {} {}.", self.resource, self.soft, self.units); + setrlimit(self.resource, self.soft, self.hard) + } +} + +#[allow(dead_code)] +struct ResourcesLimits { + cpu_time: Option, + file_size: Option, + memory_size: Option, +} + +impl ResourcesLimits { + #[allow(dead_code)] + fn new( + cpu_time: Option, + file_size: Option, + memory_size: Option, + ) -> ResourcesLimits { + ResourcesLimits { + cpu_time: cpu_time.map(|x| ResourceLimits { + resource: Resource::CPU, + soft: x, + hard: x, + units: "seconds".to_string(), + }), + file_size: file_size.map(|x| ResourceLimits { + resource: Resource::FSIZE, + soft: x, + hard: x, + units: "bytes".to_string(), + }), + memory_size: memory_size.map(|x| ResourceLimits { + resource: Resource::AS, + soft: x, + hard: x, + units: "bytes".to_string(), + }), + } + } + + #[allow(dead_code)] + fn set(&self) -> io::Result<()> { + [self.cpu_time.as_ref(), self.file_size.as_ref(), self.memory_size.as_ref()] + .iter() + .flatten() + .try_for_each(|resource_limit| resource_limit.set()) + } + + #[allow(dead_code)] + fn apply(self, command: &mut Command) -> &mut Command { + #[cfg(unix)] + unsafe { + // The `pre_exec` method runs a given closure after the parent process has been forked + // but before the child process calls `exec`. + // + // This closure runs in the child process after a `fork`, which primarily means that any + // modifications made to memory on behalf of this closure will **not** be visible to the + // parent process. This environment is often very constrained. Normal operations--such + // as using `malloc`, accessing environment variables through [`std::env`] or acquiring + // a mutex--are not guaranteed to work, because other threads may still be running at + // the time of `fork`. + // + // This closure is considered safe for the following reasons: + // 1. The [`ResourcesLimits`] struct is fully constructed and moved into the closure. + // 2. No heap allocations occur in the `set` method. + // 3. `setrlimit` is an async-signal-safe system call, which means it is safe to invoke + // after `fork`. + command.pre_exec(move || self.set()) + } + #[cfg(not(unix))] + // Not implemented for Windows. + unimplemented!("Resource limits are not implemented for Windows.") + } +} diff --git a/crates/starknet_sierra_compile/src/resource_limits_test.rs b/crates/starknet_sierra_compile/src/resource_limits_test.rs new file mode 100644 index 0000000000..7d3b84bcef --- /dev/null +++ b/crates/starknet_sierra_compile/src/resource_limits_test.rs @@ -0,0 +1,43 @@ +use std::process::Command; +use std::time::Instant; + +use rstest::rstest; + +use crate::resource_limits::ResourcesLimits; + +#[rstest] +fn test_cpu_time_limit() { + let cpu_limit = 1; // 1 second + let cpu_time_rlimit = ResourcesLimits::new(Some(1), None, None); + + let start = Instant::now(); + let mut command = Command::new("bash"); + command.args(["-c", "while true; do :; done;"]); + cpu_time_rlimit.apply(&mut command); + command.spawn().expect("Failed to start CPU consuming process").wait().unwrap(); + assert!(start.elapsed().as_secs() <= 1); +} + +#[rstest] +fn test_memory_size_limit() { + let memory_limit = 100 * 1024; // 100 KB + let memory_size_rlimit = ResourcesLimits::new(None, None, Some(memory_limit)); + + let mut command = Command::new("bash"); + command.args(["-c", "a=(); while true; do a+=0; done;"]); + memory_size_rlimit.apply(&mut command); + command.spawn().expect("Failed to start memory consuming process").wait().unwrap(); +} + +#[rstest] +fn test_file_size_limit() { + let file_limit = 10; // 10 bytes + let file_size_rlimit = ResourcesLimits::new(None, Some(file_limit), None); + + let mut command = Command::new("bash"); + command.args(["-c", "echo 0 > /tmp/file.txt; while true; do echo 0 >> /tmp/file.txt; done;"]); + file_size_rlimit.apply(&mut command); + command.spawn().expect("Failed to start disk consuming process").wait().unwrap(); + assert!(std::fs::metadata("/tmp/file.txt").unwrap().len() <= file_limit); + std::fs::remove_file("/tmp/file.txt").unwrap(); +}