use tokio::process::Command;
#[cfg(not(unix))]
use tracing::warn;
#[derive(Debug, Clone, Default)]
pub struct RlimitConfig {
pub max_open_files: Option<u64>,
pub max_file_size_bytes: Option<u64>,
pub disable_core_dumps: bool,
}
impl RlimitConfig {
#[inline]
pub fn is_empty(&self) -> bool {
self.max_open_files.is_none()
&& self.max_file_size_bytes.is_none()
&& !self.disable_core_dumps
}
}
pub fn attach_rlimits(cmd: &mut Command, config: &RlimitConfig) {
if config.is_empty() {
return;
}
#[cfg(unix)]
{
unix_impl::attach_rlimits(cmd, config);
}
#[cfg(not(unix))]
{
warn!(
?config,
"rlimit-based process limits requested on a non-Unix OS; limits will be ignored"
);
}
}
#[cfg(unix)]
mod unix_impl {
use super::RlimitConfig;
use crate::utils::log::{pre_exec_log, pre_exec_log_errno};
use std::io;
use tokio::process::Command;
pub fn attach_rlimits(cmd: &mut Command, config: &RlimitConfig) {
let max_file_size_bytes = config.max_file_size_bytes;
let disable_core_dumps = config.disable_core_dumps;
let max_open_files = config.max_open_files;
unsafe {
cmd.pre_exec(move || {
if let Some(nofile) = max_open_files
&& let Err(e) = apply_rlimit(NOFILE, nofile)
{
pre_exec_log(b"solti-exec: failed to set RLIMIT_NOFILE: ");
if let Some(code) = e.raw_os_error() {
pre_exec_log_errno(code);
}
return Err(e);
}
if let Some(fsize) = max_file_size_bytes
&& let Err(e) = apply_rlimit(FSIZE, fsize)
{
pre_exec_log(b"solti-exec: failed to set RLIMIT_FSIZE: ");
if let Some(code) = e.raw_os_error() {
pre_exec_log_errno(code);
}
return Err(e);
}
if disable_core_dumps && let Err(e) = apply_rlimit(CORE, 0) {
pre_exec_log(b"solti-exec: failed to set RLIMIT_CORE: ");
if let Some(code) = e.raw_os_error() {
pre_exec_log_errno(code);
}
return Err(e);
}
Ok(())
});
}
}
#[cfg(any(target_os = "linux", target_os = "android"))]
type RlimitResource = libc::__rlimit_resource_t;
#[cfg(not(any(target_os = "linux", target_os = "android")))]
type RlimitResource = libc::c_int;
const NOFILE: RlimitResource = libc::RLIMIT_NOFILE as RlimitResource;
const FSIZE: RlimitResource = libc::RLIMIT_FSIZE as RlimitResource;
const CORE: RlimitResource = libc::RLIMIT_CORE as RlimitResource;
fn apply_rlimit(resource: RlimitResource, value: u64) -> io::Result<()> {
let mut current = libc::rlimit {
rlim_cur: 0,
rlim_max: 0,
};
if unsafe { libc::getrlimit(resource, &mut current) } != 0 {
return Err(io::Error::last_os_error());
}
let requested = value as libc::rlim_t;
let new_soft = if current.rlim_max == libc::RLIM_INFINITY {
requested
} else if requested > current.rlim_max {
current.rlim_max
} else {
requested
};
let rlim = libc::rlimit {
rlim_cur: new_soft,
rlim_max: current.rlim_max,
};
if unsafe { libc::setrlimit(resource, &rlim) } != 0 {
Err(io::Error::last_os_error())
} else {
Ok(())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn empty_config_is_noop() {
let config = RlimitConfig::default();
assert!(config.is_empty());
let mut cmd = Command::new("sh");
attach_rlimits(&mut cmd, &config);
}
#[cfg(unix)]
#[test]
fn non_empty_config_attaches_pre_exec_hook() {
let config = RlimitConfig {
max_open_files: Some(1024),
max_file_size_bytes: Some(10 * 1024 * 1024),
disable_core_dumps: true,
};
let mut cmd = Command::new("sh");
attach_rlimits(&mut cmd, &config);
}
#[cfg(not(unix))]
#[test]
fn non_empty_config_is_ignored_on_non_unix() {
let config = RlimitConfig {
max_open_files: Some(512),
max_file_size_bytes: None,
disable_core_dumps: true,
};
let mut cmd = Command::new("sh");
attach_rlimits(&mut cmd, &config);
}
#[cfg(unix)]
#[tokio::test]
async fn rlimits_can_be_applied() {
let config = RlimitConfig {
max_open_files: Some(512),
max_file_size_bytes: Some(1024 * 1024),
disable_core_dumps: true,
};
let mut cmd = Command::new("sh");
cmd.arg("-c").arg("ulimit -a");
attach_rlimits(&mut cmd, &config);
let result = cmd.status().await;
assert!(result.is_ok(), "rlimits should be applied successfully");
assert!(result.unwrap().success());
}
}