feat: allow limiting system resources in compilation processes

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"

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

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)]

crates/starknet_sierra_compile/src/resource_limits.rs

@@ -0,0 +1,93 @@
+#![allow(dead_code)]
+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;
+
+struct RLimit {
+    resource: Resource,
+    soft_limit: u64,
+    hard_limit: u64,
+    units: String,
+}
+
+impl RLimit {
+    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_limit, self.units);
+        setrlimit(self.resource, self.soft_limit, self.hard_limit)
+    }
+}
+
+struct ResourceLimits {
+    cpu_time: Option<RLimit>,
+    file_size: Option<RLimit>,
+    memory_size: Option<RLimit>,
+}
+
+impl ResourceLimits {
+    fn new(
+        cpu_time: Option<u64>,
+        file_size: Option<u64>,
+        memory_size: Option<u64>,
+    ) -> ResourceLimits {
+        ResourceLimits {
+            cpu_time: cpu_time.map(|t| RLimit {
+                resource: Resource::CPU,
+                soft_limit: t,
+                hard_limit: t,
+                units: "seconds".to_string(),
+            }),
+            file_size: file_size.map(|x| RLimit {
+                resource: Resource::FSIZE,
+                soft_limit: x,
+                hard_limit: x,
+                units: "bytes".to_string(),
+            }),
+            memory_size: memory_size.map(|y| RLimit {
+                resource: Resource::AS,
+                soft_limit: y,
+                hard_limit: y,
+                units: "bytes".to_string(),
+            }),
+        }
+    }
+
+    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())
+    }
+
+    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 [`ResourceLimits`] 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.")
+    }
+}

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::ResourceLimits;
+
+#[rstest]
+fn test_cpu_time_limit() {
+    let cpu_limit = 1; // 1 second
+    let cpu_time_rlimit = ResourceLimits::new(Some(cpu_limit), 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() <= cpu_limit);
+}
+
+#[rstest]
+fn test_memory_size_limit() {
+    let memory_limit = 100 * 1024; // 100 KB
+    let memory_size_rlimit = ResourceLimits::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 = ResourceLimits::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();
+}