460 lines
16 KiB
Rust
460 lines
16 KiB
Rust
//! One bounded, path-free document request per resource-limited child process.
|
|
//! The configured executable is trusted application code, not a search result.
|
|
use crate::{
|
|
extraction_protocol::{ExtractionResponse, RequestHeader},
|
|
extraction_transport,
|
|
};
|
|
use std::{
|
|
io::{Read, Write},
|
|
path::PathBuf,
|
|
process::{Child, Command, Stdio},
|
|
sync::{
|
|
atomic::{AtomicBool, Ordering},
|
|
mpsc,
|
|
},
|
|
thread,
|
|
time::{Duration, Instant},
|
|
};
|
|
|
|
/// Parser limits are fixed; callers may shorten, but never extend, the deadline.
|
|
#[derive(Debug, Clone)]
|
|
pub struct ProcessConfig {
|
|
pub executable: PathBuf,
|
|
pub timeout: Duration,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum ProcessError {
|
|
Cancelled,
|
|
TimedOut,
|
|
Failed(String),
|
|
}
|
|
impl std::fmt::Display for ProcessError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
Self::Cancelled => f.write_str("Document extraction cancelled"),
|
|
Self::TimedOut => f.write_str("Document extraction timed out"),
|
|
Self::Failed(s) => write!(f, "Document extraction failed: {s}"),
|
|
}
|
|
}
|
|
}
|
|
impl std::error::Error for ProcessError {}
|
|
fn failed(error: impl std::fmt::Display) -> ProcessError {
|
|
ProcessError::Failed(error.to_string())
|
|
}
|
|
enum Event {
|
|
Written,
|
|
Response(ExtractionResponse),
|
|
Stderr,
|
|
}
|
|
|
|
/// Includes pipe writes, response EOF and process exit in the same deadline.
|
|
/// No parser fallback occurs when limits, transport or worker execution fail.
|
|
pub fn run(
|
|
config: &ProcessConfig,
|
|
request: RequestHeader,
|
|
input: Vec<u8>,
|
|
cancelled: &AtomicBool,
|
|
) -> Result<ExtractionResponse, ProcessError> {
|
|
if config.timeout.is_zero() || config.timeout > Duration::from_secs(30) {
|
|
return Err(failed("Invalid extraction deadline"));
|
|
}
|
|
request.validate_payload(&input).map_err(failed)?;
|
|
if cancelled.load(Ordering::Relaxed) {
|
|
return Err(ProcessError::Cancelled);
|
|
}
|
|
let started = Instant::now();
|
|
let mut command = Command::new(&config.executable);
|
|
command
|
|
.stdin(Stdio::piped())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
#[cfg(unix)]
|
|
{
|
|
use std::os::unix::process::CommandExt;
|
|
command.process_group(0);
|
|
}
|
|
#[cfg(windows)]
|
|
{
|
|
use std::os::windows::process::CommandExt;
|
|
command.creation_flags(0x08000000);
|
|
}
|
|
let child = command.spawn().map_err(failed)?;
|
|
let mut owned = OwnedChild {
|
|
child,
|
|
threads: Vec::new(),
|
|
completed: false,
|
|
#[cfg(windows)]
|
|
job: None,
|
|
};
|
|
// The writer cannot transmit even the admission prefix until assignment succeeds.
|
|
#[cfg(windows)]
|
|
{
|
|
owned.job = Some(Job::assign(&owned.child).map_err(failed)?);
|
|
}
|
|
let mut stdin = owned
|
|
.child
|
|
.stdin
|
|
.take()
|
|
.ok_or_else(|| failed("Missing stdin"))?;
|
|
let mut stdout = owned
|
|
.child
|
|
.stdout
|
|
.take()
|
|
.ok_or_else(|| failed("Missing stdout"))?;
|
|
let mut stderr = owned
|
|
.child
|
|
.stderr
|
|
.take()
|
|
.ok_or_else(|| failed("Missing stderr"))?;
|
|
let (tx, rx) = mpsc::sync_channel(3);
|
|
let writer_tx = tx.clone();
|
|
let writer_request = request.clone();
|
|
owned.spawn("document-input", move || {
|
|
let result = (|| {
|
|
extraction_transport::write_request(&mut stdin, &writer_request, &input)
|
|
.map_err(failed)?;
|
|
stdin.flush().map_err(failed)?;
|
|
Ok(Event::Written)
|
|
})();
|
|
drop(stdin);
|
|
let _ = writer_tx.send(result);
|
|
})?;
|
|
let reader_tx = tx.clone();
|
|
owned.spawn("document-output", move || {
|
|
let result = (|| {
|
|
let response =
|
|
extraction_transport::read_response(&mut stdout, &request).map_err(failed)?;
|
|
let mut extra = [0];
|
|
if stdout.read(&mut extra).map_err(failed)? != 0 {
|
|
return Err(failed("Trailing extraction output"));
|
|
}
|
|
Ok(Event::Response(response))
|
|
})();
|
|
let _ = reader_tx.send(result);
|
|
})?;
|
|
owned.spawn("document-errors", move || {
|
|
// Do not log document-derived stderr. Bound total bytes, not just retention.
|
|
let result = (|| {
|
|
let mut total = 0usize;
|
|
let mut buffer = [0; 1024];
|
|
loop {
|
|
let count = stderr.read(&mut buffer).map_err(failed)?;
|
|
if count == 0 {
|
|
break;
|
|
}
|
|
total += count;
|
|
if total > 16 * 1024 {
|
|
return Err(failed("Worker stderr exceeds budget"));
|
|
}
|
|
}
|
|
Ok(Event::Stderr)
|
|
})();
|
|
let _ = tx.send(result);
|
|
})?;
|
|
let mut received = 0;
|
|
let mut response = None;
|
|
loop {
|
|
if cancelled.load(Ordering::Relaxed) {
|
|
return Err(ProcessError::Cancelled);
|
|
}
|
|
if started.elapsed() >= config.timeout {
|
|
return Err(ProcessError::TimedOut);
|
|
}
|
|
match rx.recv_timeout(
|
|
Duration::from_millis(5).min(config.timeout.saturating_sub(started.elapsed())),
|
|
) {
|
|
Ok(Ok(event)) => {
|
|
received += 1;
|
|
if let Event::Response(value) = event {
|
|
response = Some(value);
|
|
}
|
|
}
|
|
Ok(Err(error)) => return Err(error),
|
|
Err(mpsc::RecvTimeoutError::Timeout) => {}
|
|
Err(mpsc::RecvTimeoutError::Disconnected) if received != 3 => {
|
|
return Err(failed("Extraction I/O worker stopped"))
|
|
}
|
|
Err(mpsc::RecvTimeoutError::Disconnected) => thread::sleep(Duration::from_millis(1)),
|
|
}
|
|
// Keep the leader unreaped while pipes may still be inherited by descendants.
|
|
// Its PID then anchors the owned process-group identity during error cleanup.
|
|
if received == 3 {
|
|
if let Some(status) = owned.child.try_wait().map_err(failed)? {
|
|
owned.completed = true;
|
|
if !status.success() {
|
|
return Err(failed("Extraction worker exited unsuccessfully"));
|
|
}
|
|
let response = response.ok_or_else(|| failed("Missing extraction response"))?;
|
|
return Ok(response);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
struct OwnedChild {
|
|
child: Child,
|
|
threads: Vec<thread::JoinHandle<()>>,
|
|
completed: bool,
|
|
#[cfg(windows)]
|
|
job: Option<Job>,
|
|
}
|
|
impl OwnedChild {
|
|
fn spawn(
|
|
&mut self,
|
|
name: &str,
|
|
work: impl FnOnce() + Send + 'static,
|
|
) -> Result<(), ProcessError> {
|
|
self.threads.push(
|
|
thread::Builder::new()
|
|
.name(name.into())
|
|
.spawn(work)
|
|
.map_err(failed)?,
|
|
);
|
|
Ok(())
|
|
}
|
|
}
|
|
impl Drop for OwnedChild {
|
|
fn drop(&mut self) {
|
|
#[cfg(windows)]
|
|
if let Some(job) = self.job.take() {
|
|
drop(job);
|
|
}
|
|
if !self.completed {
|
|
#[cfg(unix)]
|
|
unsafe {
|
|
libc::kill(-(self.child.id() as i32), libc::SIGKILL);
|
|
}
|
|
let _ = self.child.kill();
|
|
let _ = self.child.wait();
|
|
}
|
|
for handle in self.threads.drain(..) {
|
|
let _ = handle.join();
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(windows)]
|
|
struct Job(windows_sys::Win32::Foundation::HANDLE);
|
|
#[cfg(windows)]
|
|
impl Job {
|
|
fn assign(child: &Child) -> std::io::Result<Self> {
|
|
use std::os::windows::io::AsRawHandle;
|
|
use windows_sys::Win32::System::JobObjects::*;
|
|
unsafe {
|
|
let handle = CreateJobObjectW(std::ptr::null(), std::ptr::null());
|
|
if handle.is_null() {
|
|
return Err(std::io::Error::last_os_error());
|
|
}
|
|
let job = Self(handle);
|
|
let mut limits: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed();
|
|
limits.BasicLimitInformation.LimitFlags =
|
|
JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE | JOB_OBJECT_LIMIT_PROCESS_MEMORY;
|
|
limits.ProcessMemoryLimit = 512 * 1024 * 1024;
|
|
if SetInformationJobObject(
|
|
handle,
|
|
JobObjectExtendedLimitInformation,
|
|
(&limits as *const JOBOBJECT_EXTENDED_LIMIT_INFORMATION).cast(),
|
|
std::mem::size_of_val(&limits) as u32,
|
|
) == 0
|
|
|| AssignProcessToJobObject(handle, child.as_raw_handle()) == 0
|
|
{
|
|
return Err(std::io::Error::last_os_error());
|
|
}
|
|
Ok(job)
|
|
}
|
|
}
|
|
}
|
|
#[cfg(windows)]
|
|
impl Drop for Job {
|
|
fn drop(&mut self) {
|
|
unsafe {
|
|
windows_sys::Win32::Foundation::CloseHandle(self.0);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::extraction_protocol::{DocumentFormat, ExtractionOutcome};
|
|
|
|
// These tests compile real executables. Avoid concurrent rustc/linker and
|
|
// executable-admission load contaminating the deliberately short deadlines.
|
|
static PROCESS_FIXTURE_SLOT: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
|
fn fixture_slot() -> std::sync::MutexGuard<'static, ()> {
|
|
PROCESS_FIXTURE_SLOT
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner())
|
|
}
|
|
|
|
fn fixtures() -> tempfile::TempDir {
|
|
let directory = tempfile::tempdir().unwrap();
|
|
let source = directory.path().join("fake.rs");
|
|
std::fs::write(&source, r#"
|
|
use std::{io::{Read, Write},time::Duration};
|
|
fn main() {
|
|
let mode=std::env::current_exe().unwrap().file_stem().unwrap().to_string_lossy().into_owned();
|
|
if mode=="blocked" { std::thread::sleep(Duration::from_secs(60)); return; }
|
|
if mode=="crash" { std::process::exit(7); }
|
|
let mut bytes=Vec::new(); std::io::stdin().read_to_end(&mut bytes).unwrap();
|
|
if mode=="sleep" { std::thread::sleep(Duration::from_secs(60)); return; }
|
|
if mode=="oversize" { std::io::stdout().write_all(&u32::MAX.to_le_bytes()).unwrap(); return; }
|
|
let size=u32::from_le_bytes(bytes[..4].try_into().unwrap()) as usize;
|
|
let header=std::str::from_utf8(&bytes[4..4+size]).unwrap();
|
|
let digest=header.split("\"input_sha256\":\"").nth(1).unwrap().split('"').next().unwrap();
|
|
let json=format!("{{\"version\":1,\"input_sha256\":\"{}\",\"outcome\":{{\"status\":\"rejected\",\"reason\":\"controlled fixture\"}}}}",digest);
|
|
let mut out=std::io::stdout(); out.write_all(&(json.len() as u32).to_le_bytes()).unwrap();out.write_all(json.as_bytes()).unwrap();
|
|
if mode=="stderr" { std::io::stderr().write_all(&vec![b'x';17000]).unwrap(); }
|
|
if mode=="extra" { out.write_all(b"extra").unwrap(); }
|
|
}
|
|
"#).unwrap();
|
|
let executable = directory
|
|
.path()
|
|
.join(format!("valid{}", std::env::consts::EXE_SUFFIX));
|
|
let output = Command::new("rustc")
|
|
.arg("--edition=2021")
|
|
.arg(&source)
|
|
.arg("-o")
|
|
.arg(&executable)
|
|
.output()
|
|
.unwrap();
|
|
assert!(
|
|
output.status.success(),
|
|
"{}",
|
|
String::from_utf8_lossy(&output.stderr)
|
|
);
|
|
for mode in ["blocked", "crash", "sleep", "oversize", "stderr", "extra"] {
|
|
std::fs::copy(
|
|
&executable,
|
|
directory
|
|
.path()
|
|
.join(format!("{mode}{}", std::env::consts::EXE_SUFFIX)),
|
|
)
|
|
.unwrap();
|
|
}
|
|
directory
|
|
}
|
|
fn config(directory: &tempfile::TempDir, mode: &str) -> ProcessConfig {
|
|
ProcessConfig {
|
|
executable: directory
|
|
.path()
|
|
.join(format!("{mode}{}", std::env::consts::EXE_SUFFIX)),
|
|
timeout: Duration::from_secs(2),
|
|
}
|
|
}
|
|
fn request(bytes: &[u8]) -> RequestHeader {
|
|
RequestHeader::from_bytes(DocumentFormat::Pdf, bytes).unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn valid_response_requires_complete_input_eof_frame_and_exit() {
|
|
let _slot = fixture_slot();
|
|
let directory = fixtures();
|
|
let result = run(
|
|
&config(&directory, "valid"),
|
|
request(b"fixture"),
|
|
b"fixture".to_vec(),
|
|
&AtomicBool::new(false),
|
|
)
|
|
.unwrap();
|
|
assert!(matches!(result.outcome, ExtractionOutcome::Rejected { .. }));
|
|
}
|
|
#[test]
|
|
fn crashes_oversized_trailing_output_and_stderr_fail_closed() {
|
|
let _slot = fixture_slot();
|
|
let directory = fixtures();
|
|
for mode in ["crash", "oversize", "extra", "stderr"] {
|
|
let result = run(
|
|
&config(&directory, mode),
|
|
request(b"fixture"),
|
|
b"fixture".to_vec(),
|
|
&AtomicBool::new(false),
|
|
);
|
|
assert!(
|
|
matches!(result, Err(ProcessError::Failed(_))),
|
|
"{mode}: {result:?}"
|
|
);
|
|
}
|
|
}
|
|
#[test]
|
|
fn deadline_covers_blocked_stdin_and_missing_response() {
|
|
let _slot = fixture_slot();
|
|
let directory = fixtures();
|
|
for mode in ["blocked", "sleep"] {
|
|
let bytes = vec![0; 1024 * 1024];
|
|
let started = Instant::now();
|
|
assert!(matches!(
|
|
run(
|
|
&config(&directory, mode),
|
|
request(&bytes),
|
|
bytes,
|
|
&AtomicBool::new(false)
|
|
),
|
|
Err(ProcessError::TimedOut)
|
|
));
|
|
assert!(started.elapsed() < Duration::from_secs(5));
|
|
}
|
|
}
|
|
#[test]
|
|
fn cancellation_reaps_a_blocked_writer_and_invalid_limits_do_not_spawn() {
|
|
let _slot = fixture_slot();
|
|
let directory = fixtures();
|
|
let cancelled = AtomicBool::new(false);
|
|
thread::scope(|scope| {
|
|
scope.spawn(|| {
|
|
thread::sleep(Duration::from_millis(50));
|
|
cancelled.store(true, Ordering::Relaxed);
|
|
});
|
|
let bytes = vec![0; 1024 * 1024];
|
|
assert!(matches!(
|
|
run(
|
|
&config(&directory, "blocked"),
|
|
request(&bytes),
|
|
bytes,
|
|
&cancelled
|
|
),
|
|
Err(ProcessError::Cancelled)
|
|
));
|
|
});
|
|
let mut invalid = config(&directory, "valid");
|
|
invalid.timeout = Duration::from_secs(31);
|
|
assert!(matches!(
|
|
run(&invalid, request(b""), Vec::new(), &AtomicBool::new(false)),
|
|
Err(ProcessError::Failed(_))
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires a separately built SEARCHI_EXTRACT_WORKER executable"]
|
|
fn real_worker_parses_exact_docx_bytes_in_parent_owned_limits() {
|
|
use std::io::Cursor;
|
|
let mut archive = zip::ZipWriter::new(Cursor::new(Vec::new()));
|
|
archive
|
|
.start_file(
|
|
"word/document.xml",
|
|
zip::write::SimpleFileOptions::default(),
|
|
)
|
|
.unwrap();
|
|
archive.write_all(b"<w:document xmlns:w=\"http://schemas.openxmlformats.org/wordprocessingml/2006/main\"><w:body><w:p><w:r><w:t>Supplier cancellation notice</w:t></w:r></w:p></w:body></w:document>").unwrap();
|
|
let input = archive.finish().unwrap().into_inner();
|
|
let header = RequestHeader::from_bytes(DocumentFormat::Docx, &input).unwrap();
|
|
let config = ProcessConfig {
|
|
executable: std::env::var_os("SEARCHI_EXTRACT_WORKER")
|
|
.expect("Set actual worker executable")
|
|
.into(),
|
|
timeout: Duration::from_secs(5),
|
|
};
|
|
let response = run(&config, header, input, &AtomicBool::new(false)).unwrap();
|
|
match response.outcome {
|
|
ExtractionOutcome::Extracted { passages } => {
|
|
assert_eq!(passages.len(), 1);
|
|
assert_eq!(passages[0].text, "Supplier cancellation notice");
|
|
assert_eq!(passages[0].source.kind, "paragraph");
|
|
assert_eq!(passages[0].source.number, 1);
|
|
}
|
|
_ => panic!("Actual DOCX unexpectedly rejected"),
|
|
}
|
|
}
|
|
}
|