searchi/crates/searchi-extract-worker/tests/protocol.rs

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());
}