From e5c4681e04feb8c762fc224bdf4925548de5cbba Mon Sep 17 00:00:00 2001
From: Michael Debertol <michael.debertol@gmail.com>
Date: Wed, 2 Jun 2021 18:06:56 +0200
Subject: [PATCH] tests: add the ability to set resource limits

---
 .../workspace.wordlist.txt                    |  1 +
 Cargo.lock                                    | 11 +++++++
 Cargo.toml                                    |  1 +
 tests/common/util.rs                          | 31 +++++++++++++++++++
 4 files changed, 44 insertions(+)

diff --git a/.vscode/cspell.dictionaries/workspace.wordlist.txt b/.vscode/cspell.dictionaries/workspace.wordlist.txt
index b567a6c21..e2a864f9c 100644
--- a/.vscode/cspell.dictionaries/workspace.wordlist.txt
+++ b/.vscode/cspell.dictionaries/workspace.wordlist.txt
@@ -30,6 +30,7 @@ peekreader
 quickcheck
 rand_chacha
 ringbuffer
+rlimit
 smallvec
 tempdir
 tempfile
diff --git a/Cargo.lock b/Cargo.lock
index 17fa9e2b7..0455d7118 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -233,6 +233,7 @@ dependencies = [
  "pretty_assertions",
  "rand 0.7.3",
  "regex",
+ "rlimit",
  "sha1",
  "tempfile",
  "textwrap",
@@ -1410,6 +1411,16 @@ version = "0.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e9c17925a9027d298a4603d286befe3f9dc0e8ed02523141914eb628798d6e5b"
 
+[[package]]
+name = "rlimit"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49b02d62c38353a6fce45c25ca19783e25dd5f495ca681c674a4ee15aa4c1536"
+dependencies = [
+ "cfg-if 0.1.10",
+ "libc",
+]
+
 [[package]]
 name = "rust-ini"
 version = "0.13.0"
diff --git a/Cargo.toml b/Cargo.toml
index 9a6e8bd46..5f89a4077 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -354,6 +354,7 @@ walkdir = "2.2"
 atty = "0.2.14"
 
 [target.'cfg(unix)'.dev-dependencies]
+rlimit = "0.4.0"
 rust-users = { version="0.10", package="users" }
 unix_socket = "0.5.0"
 
diff --git a/tests/common/util.rs b/tests/common/util.rs
index d466ba1e9..2f7d7dcc4 100644
--- a/tests/common/util.rs
+++ b/tests/common/util.rs
@@ -1,5 +1,10 @@
+//spell-checker: ignore (linux) rlimit prlimit Rlim
+
 #![allow(dead_code)]
+
 use pretty_assertions::assert_eq;
+#[cfg(target_os = "linux")]
+use rlimit::{prlimit, rlim};
 use std::env;
 #[cfg(not(windows))]
 use std::ffi::CString;
@@ -724,6 +729,8 @@ pub struct UCommand {
     stdout: Option<Stdio>,
     stderr: Option<Stdio>,
     bytes_into_stdin: Option<Vec<u8>>,
+    #[cfg(target_os = "linux")]
+    limits: Vec<(rlimit::Resource, rlim, rlim)>,
 }
 
 impl UCommand {
@@ -758,6 +765,8 @@ impl UCommand {
             stdin: None,
             stdout: None,
             stderr: None,
+            #[cfg(target_os = "linux")]
+            limits: vec![],
         }
     }
 
@@ -855,6 +864,17 @@ impl UCommand {
         self
     }
 
+    #[cfg(target_os = "linux")]
+    pub fn with_limit(
+        &mut self,
+        resource: rlimit::Resource,
+        soft_limit: rlim,
+        hard_limit: rlim,
+    ) -> &mut Self {
+        self.limits.push((resource, soft_limit, hard_limit));
+        self
+    }
+
     /// Spawns the command, feeds the stdin if any, and returns the
     /// child process immediately.
     pub fn run_no_wait(&mut self) -> Child {
@@ -871,6 +891,17 @@ impl UCommand {
             .spawn()
             .unwrap();
 
+        #[cfg(target_os = "linux")]
+        for &(resource, soft_limit, hard_limit) in &self.limits {
+            prlimit(
+                child.id() as i32,
+                resource,
+                Some((soft_limit, hard_limit)),
+                None,
+            )
+            .unwrap();
+        }
+
         if let Some(ref input) = self.bytes_into_stdin {
             let write_result = child
                 .stdin