diff --git a/Cargo.lock b/Cargo.lock index 050c36691c9..26a54bf2264 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10700,6 +10700,7 @@ dependencies = [ "infra_utils", "mempool_test_utils", "papyrus_config", + "rlimit", "rstest", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index bfac7a76f69..c034bacb931 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 bd559d226ef..e5eea128931 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 80a543d24b3..11773f51f5d 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 00000000000..6908618e750 --- /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 forking the parent process and + // before the exec of the child process. + // + // This closure will be run in the context of the child process after a `fork`. This + // primarily means that any modifications made to memory on behalf of this closure will + // **not** be visible to the parent process. This is often a very constrained + // environment where normal operations like `malloc`, accessing environment variables + // through [`std::env`] or acquiring a mutex are not guaranteed to work (due to other + // threads perhaps still running when the `fork` was run). + // + // The current closure is 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 00000000000..69182b213cb --- /dev/null +++ b/crates/starknet_sierra_compile/src/resource_limits_test.rs @@ -0,0 +1,40 @@ +use std::process::Command; +use std::time::Instant; + +use rstest::rstest; + +use crate::resource_limits::ResourcesLimits; + +#[rstest] +fn test_cpu_time_limit() { + let cpu_time_limits = 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_limits.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_size_limits = ResourcesLimits::new(None, None, Some(10000)); + + let mut command = Command::new("bash"); + command.args(["-c", "a=(); while true; do a+=0; done;"]); + memory_size_limits.apply(&mut command); + command.spawn().expect("Failed to start memory consuming process").wait().unwrap(); +} + +#[rstest] +fn test_file_size_limit() { + let file_size_limits = ResourcesLimits::new(None, Some(100), 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_limits.apply(&mut command); + command.spawn().expect("Failed to start disk consuming process").wait().unwrap(); + assert!(std::fs::metadata("/tmp/file.txt").unwrap().len() <= 100); + std::fs::remove_file("/tmp/file.txt").unwrap(); +}