61 lines
2.1 KiB
Rust
61 lines
2.1 KiB
Rust
|
use anyhow::{ensure, Result};
|
||
|
|
|
||
|
|
const MEMORY_BYTES: usize = 512 * 1024 * 1024;
|
||
|
|
|
||
|
|
#[cfg(windows)]
|
||
|
|
pub fn apply() -> Result<()> {
|
||
|
|
use windows_sys::Win32::System::JobObjects::{
|
||
|
|
JobObjectExtendedLimitInformation, QueryInformationJobObject,
|
||
|
|
JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
|
||
|
|
JOB_OBJECT_LIMIT_PROCESS_MEMORY,
|
||
|
|
};
|
||
|
|
let mut information: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { std::mem::zeroed() };
|
||
|
|
// NULL addresses the job associated with the current process. The parent
|
||
|
|
// must assign this process before writing its first input bytes.
|
||
|
|
let success = unsafe {
|
||
|
|
QueryInformationJobObject(
|
||
|
|
std::ptr::null_mut(),
|
||
|
|
JobObjectExtendedLimitInformation,
|
||
|
|
&mut information as *mut _ as *mut _,
|
||
|
|
std::mem::size_of_val(&information) as u32,
|
||
|
|
std::ptr::null_mut(),
|
||
|
|
)
|
||
|
|
};
|
||
|
|
ensure!(
|
||
|
|
success != 0,
|
||
|
|
"Extraction worker requires parent resource job"
|
||
|
|
);
|
||
|
|
let required = JOB_OBJECT_LIMIT_PROCESS_MEMORY | JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
|
||
|
|
ensure!(
|
||
|
|
information.BasicLimitInformation.LimitFlags & required == required
|
||
|
|
&& information.ProcessMemoryLimit > 0
|
||
|
|
&& information.ProcessMemoryLimit <= MEMORY_BYTES,
|
||
|
|
"Extraction worker job resource policy is missing"
|
||
|
|
);
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(unix)]
|
||
|
|
pub fn apply() -> Result<()> {
|
||
|
|
for (resource, maximum) in [
|
||
|
|
(libc::RLIMIT_AS, MEMORY_BYTES as libc::rlim_t),
|
||
|
|
(libc::RLIMIT_CPU, 30 as libc::rlim_t),
|
||
|
|
] {
|
||
|
|
let value = libc::rlimit {
|
||
|
|
rlim_cur: maximum,
|
||
|
|
rlim_max: maximum,
|
||
|
|
};
|
||
|
|
// Applied before allocating input or invoking any parser. Platform
|
||
|
|
// enforcement semantics, especially macOS, require native evidence.
|
||
|
|
ensure!(
|
||
|
|
unsafe { libc::setrlimit(resource, &value) } == 0,
|
||
|
|
"Cannot install extraction resource limit"
|
||
|
|
);
|
||
|
|
}
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(not(any(windows, unix)))]
|
||
|
|
pub fn apply() -> Result<()> {
|
||
|
|
anyhow::bail!("Extraction resource policy unsupported on this platform")
|
||
|
|
}
|