Add resource-checked extraction worker and bounded parent process owner

This commit is contained in:
karl committed 2026-10-13 15:06:40 +02:00
1 parent 688b5dc519
commit 844fdbf89f
11 files changed
+849 -3

No files matched your search

+20
View File
@@ -0,0 +1,20 @@
[package]
name = "searchi-extract-worker"
version.workspace = true
edition.workspace = true
license.workspace = true
[dependencies]
anyhow.workspace = true
searchi-core = { path = "../searchi-core" }
[target.'cfg(unix)'.dependencies]
libc = "0.2"
[target.'cfg(windows)'.dependencies]
windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_System_JobObjects", "Win32_System_Threading"] }
[dev-dependencies]
serde_json.workspace = true
zip = { version = "8", default-features = false, features = ["deflate"] }
+40
View File
@@ -0,0 +1,40 @@
//! Single-request, path-free document parser helper. Parent owns deadline,
//! cancellation and Windows Job lifetime; this helper never spawns children.
mod resources;
use anyhow::{ensure, Result};
use searchi_core::{extraction_protocol::extract_document, extraction_transport};
use std::io::{self, Read, Write};
fn run() -> Result<()> {
ensure!(std::env::args_os().len() == 1, "Arguments are unsupported");
let mut input = io::stdin().lock();
#[cfg(unix)]
resources::apply()?;
// Prefix arrival synchronizes with Windows parent job assignment. No
// attacker-declared allocation or parser runs before the policy check.
let mut prefix = [0u8; 4];
input.read_exact(&mut prefix)?;
#[cfg(not(unix))]
resources::apply()?;
let mut input = io::Cursor::new(prefix).chain(&mut input);
let (request, bytes) = extraction_transport::read_request(&mut input)?;
// Parent closes stdin after the single frame. Extra data never reaches a parser.
let mut trailing = [0u8; 1];
ensure!(
input.read(&mut trailing)? == 0,
"Trailing input is unsupported"
);
let response = extract_document(&request, &bytes)?;
let mut output = io::stdout().lock();
extraction_transport::write_response(&mut output, &response, &request)?;
output.flush()?;
Ok(())
}
fn main() {
if run().is_err() {
// Do not echo document contents, external paths or untrusted protocol values.
eprintln!("searchi extraction worker rejected request or resource policy");
std::process::exit(1);
}
}
@@ -0,0 +1,61 @@
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")
}
@@ -0,0 +1,63 @@
//! Exercise the real helper through the shared engine's process owner.
use searchi_core::{
extraction_process::{run, ProcessConfig, ProcessError},
extraction_protocol::{DocumentFormat, ExtractionOutcome, RequestHeader},
};
use std::{
io::{Cursor, Write},
path::PathBuf,
sync::atomic::AtomicBool,
time::Duration,
};
fn config() -> ProcessConfig {
ProcessConfig {
executable: PathBuf::from(env!("CARGO_BIN_EXE_searchi-extract-worker")),
timeout: Duration::from_secs(10),
}
}
#[test]
fn actual_parent_and_worker_return_original_docx_passage_and_digest() {
let mut document = zip::ZipWriter::new(Cursor::new(Vec::new()));
document
.start_file(
"word/document.xml",
zip::write::SimpleFileOptions::default(),
)
.unwrap();
document.write_all(br#"<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:p><w:r><w:t>Cancel with sixty days written notice.</w:t></w:r></w:p></w:body></w:document>"#).unwrap();
let input = document.finish().unwrap().into_inner();
let request = RequestHeader::from_bytes(DocumentFormat::Docx, &input).unwrap();
let expected_hash = request.input_sha256.clone();
let response = run(&config(), request, input, &AtomicBool::new(false)).unwrap();
assert_eq!(response.input_sha256, expected_hash);
let ExtractionOutcome::Extracted { passages } = response.outcome else {
panic!("Valid document rejected");
};
assert_eq!(passages.len(), 1);
assert_eq!(passages[0].text, "Cancel with sixty days written notice.");
assert_eq!(passages[0].source.kind, "paragraph");
assert_eq!(passages[0].source.number, 1);
}
#[test]
fn parser_rejection_and_cancellation_remain_distinct() {
let input = b"invalid PDF".to_vec();
let request = RequestHeader::from_bytes(DocumentFormat::Pdf, &input).unwrap();
let response = run(
&config(),
request.clone(),
input.clone(),
&AtomicBool::new(false),
)
.unwrap();
assert!(matches!(
response.outcome,
ExtractionOutcome::Rejected { .. }
));
assert!(matches!(
run(&config(), request, input, &AtomicBool::new(true)),
Err(ProcessError::Cancelled)
));
}
@@ -0,0 +1,153 @@
use searchi_core::{
extraction_protocol::{DocumentFormat, ExtractionOutcome, RequestHeader, MAX_HEADER_BYTES},
extraction_transport::{read_response, write_request},
};
use std::{
io::{Cursor, Write},
process::{Command, Stdio},
time::{Duration, Instant},
};
#[cfg(windows)]
struct Job(windows_sys::Win32::Foundation::HANDLE);
#[cfg(windows)]
impl Job {
fn attach(child: &std::process::Child) -> Self {
use std::os::windows::io::AsRawHandle;
use windows_sys::Win32::System::JobObjects::*;
let handle = unsafe { CreateJobObjectW(std::ptr::null(), std::ptr::null()) };
assert!(!handle.is_null());
let guard = Self(handle);
let mut info: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { std::mem::zeroed() };
info.BasicLimitInformation.LimitFlags =
JOB_OBJECT_LIMIT_PROCESS_MEMORY | JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
info.ProcessMemoryLimit = 512 * 1024 * 1024;
assert_ne!(
unsafe {
SetInformationJobObject(
handle,
JobObjectExtendedLimitInformation,
&info as *const _ as *const _,
std::mem::size_of_val(&info) as u32,
)
},
0
);
assert_ne!(
unsafe { AssignProcessToJobObject(handle, child.as_raw_handle() as _) },
0
);
guard
}
}
#[cfg(windows)]
impl Drop for Job {
fn drop(&mut self) {
unsafe {
windows_sys::Win32::Foundation::CloseHandle(self.0);
}
}
}
fn exchange(wire: &[u8], arguments: &[&str], attach: bool) -> std::process::Output {
let mut child = Command::new(env!("CARGO_BIN_EXE_searchi-extract-worker"))
.args(arguments)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.unwrap();
#[cfg(windows)]
let _job = attach.then(|| Job::attach(&child));
#[cfg(not(windows))]
let _ = attach;
let mut input = child.stdin.take().unwrap();
// Unknown arguments may terminate before the request write.
let _ = input.write_all(wire);
drop(input);
let deadline = Instant::now() + Duration::from_secs(30);
while child.try_wait().unwrap().is_none() {
if Instant::now() >= deadline {
child.kill().unwrap();
child.wait().unwrap();
panic!("Worker subprocess exceeded deadline");
}
std::thread::sleep(Duration::from_millis(10));
}
child.wait_with_output().unwrap()
}
fn wire(format: DocumentFormat, input: &[u8]) -> (RequestHeader, Vec<u8>) {
let request = RequestHeader::from_bytes(format, input).unwrap();
let mut bytes = Vec::new();
write_request(&mut bytes, &request, input).unwrap();
(request, bytes)
}
#[test]
fn valid_docx_and_corrupt_pdf_are_real_worker_responses() {
let mut archive = zip::ZipWriter::new(Cursor::new(Vec::new()));
archive
.start_file(
"word/document.xml",
zip::write::SimpleFileOptions::default(),
)
.unwrap();
archive.write_all(br#"<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:p><w:r><w:t>supplier cancellation</w:t></w:r></w:p></w:body></w:document>"#).unwrap();
let input = archive.finish().unwrap().into_inner();
let (request, bytes) = wire(DocumentFormat::Docx, &input);
let output = exchange(&bytes, &[], true);
assert!(
output.status.success(),
"{}",
String::from_utf8_lossy(&output.stderr)
);
let response = read_response(&mut Cursor::new(output.stdout), &request).unwrap();
let ExtractionOutcome::Extracted { passages } = response.outcome else {
panic!("Valid DOCX rejected");
};
assert_eq!(passages[0].text, "supplier cancellation");
assert_eq!(passages[0].source.number, 1);
let (request, bytes) = wire(DocumentFormat::Pdf, b"not a PDF");
let output = exchange(&bytes, &[], true);
assert!(output.status.success());
let response = read_response(&mut Cursor::new(output.stdout), &request).unwrap();
assert!(matches!(
response.outcome,
ExtractionOutcome::Rejected { .. }
));
}
#[test]
fn rejects_unknown_arguments_trailing_input_hash_version_and_header_budget() {
let (request, mut bytes) = wire(DocumentFormat::Pdf, b"invalid PDF");
bytes.push(0);
let output = exchange(&bytes, &[], true);
assert!(!output.status.success() && output.stdout.is_empty());
for change_version in [false, true] {
let mut changed = request.clone();
if change_version {
changed.version += 1;
} else {
changed.input_sha256 = "0".repeat(64);
}
let header = serde_json::to_vec(&changed).unwrap();
let mut bytes = (header.len() as u32).to_le_bytes().to_vec();
bytes.extend(header);
bytes.extend(b"invalid PDF");
let output = exchange(&bytes, &[], true);
assert!(!output.status.success() && output.stdout.is_empty());
}
let output = exchange(&((MAX_HEADER_BYTES + 1) as u32).to_le_bytes(), &[], true);
assert!(!output.status.success() && output.stdout.is_empty());
let output = exchange(&[], &["--arbitrary-path"], true);
assert!(!output.status.success() && output.stdout.is_empty());
}
#[cfg(windows)]
#[test]
fn rejects_worker_without_parent_job_before_payload_parsing() {
let (_, bytes) = wire(DocumentFormat::Pdf, b"invalid PDF");
let output = exchange(&bytes, &[], false);
assert!(!output.status.success() && output.stdout.is_empty());
}