136 lines
4.8 KiB
Rust
136 lines
4.8 KiB
Rust
//! Bounded single-request framing. Process deadlines, EOF/status validation and
|
|
//! cancellation belong to the caller; these blocking I/O helpers do not isolate
|
|
//! parsing or make an unresponsive pipe safe on their own.
|
|
use crate::extraction_protocol::{
|
|
ExtractionResponse, RequestHeader, MAX_HEADER_BYTES, MAX_RESPONSE_BYTES,
|
|
};
|
|
use anyhow::{ensure, Context, Result};
|
|
use std::io::{Read, Write};
|
|
|
|
fn read_frame(reader: &mut impl Read, limit: usize) -> Result<Vec<u8>> {
|
|
let mut prefix = [0; 4];
|
|
reader
|
|
.read_exact(&mut prefix)
|
|
.context("Missing extraction frame length")?;
|
|
let length = u32::from_le_bytes(prefix) as usize;
|
|
ensure!(
|
|
length > 0 && length <= limit,
|
|
"Extraction frame exceeds budget"
|
|
);
|
|
let mut bytes = Vec::new();
|
|
bytes
|
|
.try_reserve_exact(length)
|
|
.context("Cannot allocate bounded extraction frame")?;
|
|
bytes.resize(length, 0);
|
|
reader
|
|
.read_exact(&mut bytes)
|
|
.context("Truncated extraction frame")?;
|
|
Ok(bytes)
|
|
}
|
|
|
|
fn write_frame(writer: &mut impl Write, bytes: &[u8], limit: usize) -> Result<()> {
|
|
ensure!(
|
|
!bytes.is_empty() && bytes.len() <= limit,
|
|
"Extraction frame exceeds budget"
|
|
);
|
|
writer.write_all(&u32::try_from(bytes.len())?.to_le_bytes())?;
|
|
writer.write_all(bytes)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Header length + JSON header + exact raw payload. No filename is transmitted.
|
|
pub fn write_request(writer: &mut impl Write, header: &RequestHeader, input: &[u8]) -> Result<()> {
|
|
header.validate_payload(input)?;
|
|
write_frame(writer, &serde_json::to_vec(header)?, MAX_HEADER_BYTES)?;
|
|
writer.write_all(input)?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn read_request(reader: &mut impl Read) -> Result<(RequestHeader, Vec<u8>)> {
|
|
let header = RequestHeader::decode(&read_frame(reader, MAX_HEADER_BYTES)?)?;
|
|
// Validate the declaration before allocating or reading any raw payload.
|
|
let mut input = Vec::new();
|
|
input
|
|
.try_reserve_exact(header.input_bytes)
|
|
.context("Cannot allocate bounded document input")?;
|
|
input.resize(header.input_bytes, 0);
|
|
reader
|
|
.read_exact(&mut input)
|
|
.context("Truncated document input")?;
|
|
header.validate_payload(&input)?;
|
|
Ok((header, input))
|
|
}
|
|
|
|
pub fn write_response(
|
|
writer: &mut impl Write,
|
|
response: &ExtractionResponse,
|
|
request: &RequestHeader,
|
|
) -> Result<()> {
|
|
response.validate(request)?;
|
|
write_frame(writer, &serde_json::to_vec(response)?, MAX_RESPONSE_BYTES)
|
|
}
|
|
|
|
pub fn read_response(
|
|
reader: &mut impl Read,
|
|
request: &RequestHeader,
|
|
) -> Result<ExtractionResponse> {
|
|
ExtractionResponse::decode(&read_frame(reader, MAX_RESPONSE_BYTES)?, request)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::extraction_protocol::{DocumentFormat, ExtractionOutcome, VERSION};
|
|
use std::io::Cursor;
|
|
|
|
#[test]
|
|
fn binary_input_preserves_every_byte_without_base64_or_paths() -> Result<()> {
|
|
let input = b"\0\xff\r\nPDF binary input";
|
|
let header = RequestHeader::from_bytes(DocumentFormat::Pdf, input)?;
|
|
let mut wire = Vec::new();
|
|
write_request(&mut wire, &header, input)?;
|
|
assert!(wire.ends_with(input));
|
|
let (decoded, bytes) = read_request(&mut Cursor::new(wire))?;
|
|
assert_eq!(bytes, input);
|
|
assert_eq!(decoded.input_sha256, header.input_sha256);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn oversized_and_truncated_frames_fail_without_payload_allocation() {
|
|
let oversized = u32::MAX.to_le_bytes();
|
|
assert!(read_request(&mut Cursor::new(oversized)).is_err());
|
|
assert!(read_request(&mut Cursor::new([1, 0])).is_err());
|
|
let header = RequestHeader::from_bytes(DocumentFormat::Pdf, b"input").unwrap();
|
|
let mut wire = Vec::new();
|
|
write_request(&mut wire, &header, b"input").unwrap();
|
|
wire.pop();
|
|
assert!(read_request(&mut Cursor::new(wire)).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn wrong_input_digest_is_rejected_and_writer_emits_nothing() {
|
|
let header = RequestHeader::from_bytes(DocumentFormat::Pdf, b"first").unwrap();
|
|
let mut wire = Vec::new();
|
|
assert!(write_request(&mut wire, &header, b"other").is_err());
|
|
assert!(wire.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn response_identity_and_size_are_checked() -> Result<()> {
|
|
let header = RequestHeader::from_bytes(DocumentFormat::Pdf, b"input")?;
|
|
let response = ExtractionResponse {
|
|
version: VERSION,
|
|
input_sha256: header.input_sha256.clone(),
|
|
outcome: ExtractionOutcome::Rejected {
|
|
reason: "Invalid PDF".into(),
|
|
},
|
|
};
|
|
let mut wire = Vec::new();
|
|
write_response(&mut wire, &response, &header)?;
|
|
read_response(&mut Cursor::new(wire), &header)?;
|
|
assert!(read_response(&mut Cursor::new(u32::MAX.to_le_bytes()), &header).is_err());
|
|
Ok(())
|
|
}
|
|
}
|