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,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<ResourceLimits>,
+    file_size: Option<ResourceLimits>,
+    memory_size: Option<ResourceLimits>,
+}
+
+impl ResourcesLimits {
+    #[allow(dead_code)]
+    fn new(
+        cpu_time: Option<u64>,
+        file_size: Option<u64>,
+        memory_size: Option<u64>,
+    ) -> 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.")
+    }
+}

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();
+}