use std::io;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};
use serde_json::{Value, json};
use tiny_http::{Header, Request, Response, Server, StatusCode};
use rapid_structure_rs::{ProviderPreference, WordOutputMode};
use super::admit::{
self, AdmissionError, BodySource, ChunkOutcome, HttpMethod, QueueAdmission, RequestDescriptor,
RouteDecision,
};
use super::cli::ProviderChoice;
use super::error::ServeError;
use super::evaluate;
use super::export::ExportFormat;
use super::flowlog::RequestFlow;
use super::jobs::JobQueue;
use super::queue::QueueClass;
use super::run::ServeStartError;
use super::security::{self, SECURITY_HEADERS};
use super::server::{
Body, OcrJobParams, READ_CHUNK_BYTES, ServeContext, ServeRuntime, ServeShared,
};
const ACCEPT_TIMEOUT: Duration = Duration::from_millis(200);
const BODY_READ_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_SIDE_CEILING: u32 = 32_768;
pub(super) const EXPORT_CSP: &str =
"default-src 'none'; style-src 'unsafe-inline'; img-src data:; script-src 'none'; sandbox";
pub(super) fn page_csp(nonce: &str) -> String {
format!(
"default-src 'none'; script-src 'nonce-{nonce}'; style-src 'nonce-{nonce}'; font-src \
'self'; img-src 'self' blob: data:; connect-src 'self'; base-uri 'none'; form-action \
'none'; frame-ancestors 'none'"
)
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum Route {
Page,
Static(String),
Status,
Models,
ModelsDownload,
ModelsReverify,
EngineReload,
Ocr,
Evaluate,
Job(String),
JobResult(String),
JobCancel(String),
JobAnnotated(String),
JobExport(String),
}
impl Route {
fn method(&self) -> HttpMethod {
match self {
Self::Page
| Self::Status
| Self::Models
| Self::Job(_)
| Self::JobResult(_)
| Self::JobAnnotated(_)
| Self::JobExport(_)
| Self::Static(_) => HttpMethod::Get,
Self::ModelsDownload
| Self::ModelsReverify
| Self::EngineReload
| Self::Ocr
| Self::Evaluate
| Self::JobCancel(_) => HttpMethod::Post,
}
}
fn allow(&self) -> &'static str {
match self.method() {
HttpMethod::Get => "GET",
_ => "POST",
}
}
fn requires_token(&self) -> bool {
!matches!(self, Self::Page | Self::Static(_))
}
}
pub(super) struct BoundServer {
server: Arc<Server>,
local_addr: SocketAddr,
}
impl BoundServer {
pub fn bind(port: u16) -> Result<Self, ServeStartError> {
let requested = security::bind_address(port);
let server = Server::http(requested).map_err(|error| ServeStartError::Bind {
address: requested,
detail: error.to_string(),
})?;
let bound = server.server_addr().to_ip().ok_or(ServeStartError::Bind {
address: requested,
detail: "the listener does not expose an IPv4 address".to_string(),
})?;
let local_addr = security::assert_loopback(bound)?;
Ok(Self {
server: Arc::new(server),
local_addr,
})
}
pub fn local_addr(&self) -> SocketAddr {
self.local_addr
}
}
pub(super) struct ServeHandle {
server: Arc<Server>,
runtime: ServeRuntime,
accept: Option<JoinHandle<()>>,
local_addr: SocketAddr,
}
impl ServeHandle {
pub fn start(bound: BoundServer, context: ServeContext) -> Result<Self, ServeStartError> {
let BoundServer { server, local_addr } = bound;
let runtime = ServeRuntime::new(context).map_err(ServeStartError::Runtime)?;
let accept = {
let server = Arc::clone(&server);
let shared = Arc::clone(runtime.shared());
thread::Builder::new()
.name("serve-accept".to_string())
.spawn(move || accept_loop(&server, &shared))
.map_err(ServeStartError::Runtime)?
};
Ok(Self {
server,
runtime,
accept: Some(accept),
local_addr,
})
}
pub fn local_addr(&self) -> SocketAddr {
self.local_addr
}
pub fn shared(&self) -> &Arc<ServeShared> {
self.runtime.shared()
}
pub fn accept_loop(&mut self) {
if let Some(accept) = self.accept.take() {
let _ = accept.join();
}
}
}
impl Drop for ServeHandle {
fn drop(&mut self) {
self.shared().shutdown();
self.server.unblock();
if let Some(accept) = self.accept.take() {
let _ = accept.join();
}
self.runtime.stop();
}
}
fn accept_loop(server: &Arc<Server>, shared: &Arc<ServeShared>) {
while !shared.is_shutting_down() {
match server.recv_timeout(ACCEPT_TIMEOUT) {
Ok(Some(request)) => handle(shared, request),
Ok(None) => {}
Err(error) => {
if !shared.is_shutting_down() {
eprintln!("serve: accept failed: {error}");
}
}
}
}
}
fn handle(shared: &Arc<ServeShared>, mut request: Request) {
let method = HttpMethod::parse(request.method().as_str());
let url = request.url().to_string();
let (path, query) = split_url(&url);
let (decision, route) = route_of(method, path);
let flow = RequestFlow::begin(shared.flow(), method.name(), path);
if decision != RouteDecision::Matched {
let headers = match (&route, decision) {
(Some(known), RouteDecision::MethodNotAllowed) => {
vec![("Allow".to_string(), known.allow().to_string())]
}
_ => Vec::new(),
};
let _ = respond(
request,
route_error(decision, method, path),
&headers,
&flow,
);
return;
}
let route = route.expect("a matched route decision always carries its route");
let queue_raw = query_param(query, "queue");
let class = match route {
Route::Ocr => shared.routing().class_for(queue_raw).ok(),
_ => None,
};
let table_raw = query_param(query, "table");
let table = match route {
Route::Ocr => shared.routing().table_requested(table_raw).ok(),
_ => None,
};
let slot_source = class.map(|class| QueueSlotSource { shared, class });
let token_header = header(&request, "X-RapidStructure-Token").map(str::to_string);
let host_header = header(&request, "Host").map(str::to_string);
let origin_header = header(&request, "Origin").map(str::to_string);
let content_type_header = header(&request, "Content-Type").map(str::to_string);
let content_length = request.body_length().map(|length| length as u64);
let is_chunked = header(&request, "Transfer-Encoding").is_some();
let descriptor = RequestDescriptor {
method,
path,
route: decision,
has_token: !route.requires_token()
|| token_header
.as_deref()
.is_some_and(|candidate| shared.token().matches(candidate)),
host: host_header.as_deref(),
origin: origin_header.as_deref(),
queue: match &slot_source {
Some(source) => QueueAdmission::Reserve(source),
None => QueueAdmission::NotQueued,
},
content_length,
body_so_far: 0,
content_type: content_type_header.as_deref(),
is_chunked,
};
let admitted = match admit::admit(&descriptor, shared.local(), shared.limits().max_body_bytes) {
Ok(admitted) => admitted,
Err(error) => {
match &error {
AdmissionError::Route {
path: route_path,
method: route_method,
decision: route_decision,
} => {
let body = route_error(*route_decision, *route_method, route_path);
let _ = respond(request, body, &[], &flow);
}
AdmissionError::Rejected(_) => {
let serve_error = error
.rejected()
.expect("a Rejected admission error always carries a ServeError");
let body = shared.error_body(serve_error);
if matches!(route, Route::Ocr) {
flow.job_rejected(
queue_label(class, queue_raw),
serve_error.status_code(),
serve_error.code(),
&serve_error.detail_text(),
);
}
let _ = respond(request, body, &[], &flow);
}
}
return;
}
};
if matches!(route, Route::Ocr)
&& class == Some(QueueClass::Formula)
&& !shared.formula_models_ready()
{
let body = shared.formula_blocked_body();
let (code, detail) = shared.formula_blocked_reason();
flow.job_rejected(queue_label(class, queue_raw), body.status, code, &detail);
let _ = respond(request, body, &[], &flow);
return;
}
if matches!(route, Route::Ocr) && table == Some(true) && !shared.table_models_ready() {
let body = shared.table_blocked_body();
let (code, detail) = shared.table_blocked_reason();
flow.job_rejected(queue_label(class, queue_raw), body.status, code, &detail);
let _ = respond(request, body, &[], &flow);
return;
}
let outcome = dispatch(
shared,
&mut request,
&route,
query,
&descriptor,
admitted,
&flow,
);
match outcome {
Ok(Dispatch::Respond(body, headers)) => {
let _ = respond(request, body, &headers, &flow);
}
Ok(Dispatch::EngineWork(work)) => {
if let Err(failure) = spawn_engine_work(shared, work, request, flow) {
let (error, request, flow) = *failure;
let _ = respond(request, Body::error(&error), &[], &flow);
}
}
Ok(Dispatch::Evaluate(manifest)) => {
if let Err(failure) = spawn_evaluation(shared, manifest, request, flow) {
let (error, request, flow) = *failure;
let _ = respond(request, Body::error(&error), &[], &flow);
}
}
Err(error) => {
let body = shared.error_body(&error);
let _ = respond(request, body, &[], &flow);
}
}
}
fn queue_label(class: Option<QueueClass>, raw: Option<&str>) -> &str {
match class {
Some(QueueClass::Text) => QueueClass::Text.name(),
Some(QueueClass::Formula) => QueueClass::Formula.name(),
None => raw.filter(|raw| !raw.trim().is_empty()).unwrap_or("text"),
}
}
struct QueueSlotSource<'a> {
shared: &'a Arc<ServeShared>,
class: QueueClass,
}
impl admit::QueueSlots for QueueSlotSource<'_> {
fn reserve(&self) -> Option<admit::QueueReservation> {
self.shared.reserve_queue_slot(self.class)
}
}
enum Dispatch {
Respond(Body, Vec<(String, String)>),
EngineWork(EngineWork),
Evaluate(PathBuf),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum EngineWork {
Reload(Option<ProviderPreference>),
ReverifyModels,
}
fn dispatch(
shared: &Arc<ServeShared>,
request: &mut Request,
route: &Route,
query: &str,
descriptor: &RequestDescriptor<'_>,
admitted: admit::Admit,
flow: &RequestFlow,
) -> Result<Dispatch, ServeError> {
match route {
Route::Page => Ok(Dispatch::Respond(
Body::html(shared.page().as_bytes().to_vec()),
vec![(
"Content-Security-Policy".to_string(),
page_csp(shared.nonce()),
)],
)),
Route::Static(path) => match super::assets::lookup(path) {
Some((content_type, bytes)) => Ok(Dispatch::Respond(
Body::typed(200, content_type, bytes.to_vec()),
Vec::new(),
)),
None => Ok(Dispatch::Respond(
route_error(RouteDecision::NotFound, HttpMethod::Get, path),
Vec::new(),
)),
},
Route::Status => Ok(Dispatch::Respond(
json_body(200, shared.status_json())?,
Vec::new(),
)),
Route::Models => Ok(Dispatch::Respond(
json_body(200, shared.models_json())?,
Vec::new(),
)),
Route::Ocr => {
descriptor.check_content_type()?;
let bytes = read_body(request, &admitted)?;
let max_side = parse_max_side(query_param(query, "max_side"), shared.min_side_len())?;
let class = shared.routing().class_for(query_param(query, "queue"))?;
let table = shared
.routing()
.table_requested(query_param(query, "table"))?;
let words = parse_words(query_param(query, "words"))?;
let params = OcrJobParams {
max_side,
table,
words,
};
let value = shared.submit_ocr(bytes, class, params, admitted.reservation, flow)?;
Ok(Dispatch::Respond(json_body(202, value)?, Vec::new()))
}
Route::ModelsDownload => {
check_json_content_type(descriptor.content_type)?;
let bytes = read_body(request, &admitted)?;
let set_id = parse_set_id(&bytes)?;
let value = match shared.submit_download(&set_id, flow) {
Ok(value) => value,
Err(error) => {
flow.job_rejected(
JobQueue::Download.name(),
error.status_code(),
error.code(),
&error.detail_text(),
);
return Err(error);
}
};
Ok(Dispatch::Respond(json_body(202, value)?, Vec::new()))
}
Route::ModelsReverify => {
let body = read_body(request, &admitted)?;
if !body.is_empty() {
return Err(ServeError::BadRequest);
}
Ok(Dispatch::EngineWork(EngineWork::ReverifyModels))
}
Route::Evaluate => {
check_json_content_type(descriptor.content_type)?;
let bytes = read_body(request, &admitted)?;
let manifest = evaluate::parse_manifest_body(&bytes)?;
Ok(Dispatch::Evaluate(manifest))
}
Route::EngineReload => {
let body = read_body(request, &admitted)?;
let provider = parse_provider_body(body)?;
Ok(Dispatch::EngineWork(EngineWork::Reload(provider)))
}
Route::Job(id) => Ok(Dispatch::Respond(
json_body(200, shared.job_view(id)?)?,
Vec::new(),
)),
Route::JobResult(id) => Ok(Dispatch::Respond(shared.job_result(id)?, Vec::new())),
Route::JobCancel(id) => Ok(Dispatch::Respond(
json_body(200, shared.cancel_job(id)?)?,
Vec::new(),
)),
Route::JobAnnotated(id) => Ok(Dispatch::Respond(shared.annotated_png(id)?, Vec::new())),
Route::JobExport(id) => {
let Some(format) = export_format(query_param(query, "format")) else {
return Err(ServeError::BadRequest);
};
let body = shared.export(id, format)?;
let mut headers = vec![(
"Content-Disposition".to_string(),
content_disposition(id, format.extension()),
)];
if format == ExportFormat::Html {
headers.push((
"Content-Security-Policy".to_string(),
EXPORT_CSP.to_string(),
));
}
Ok(Dispatch::Respond(body, headers))
}
}
}
fn spawn_engine_work(
shared: &Arc<ServeShared>,
work: EngineWork,
request: Request,
flow: RequestFlow,
) -> Result<(), Box<(ServeError, Request, RequestFlow)>> {
let Some(guard) = shared.begin_provider_switch() else {
return Err(Box::new((ServeError::Busy, request, flow)));
};
let shared = Arc::clone(shared);
let (tx, rx) = std::sync::mpsc::sync_channel::<Request>(1);
let thread_flow = flow.clone();
let spawned = thread::Builder::new()
.name("serve-engine-work".to_string())
.spawn(move || {
let _guard = guard;
let Ok(request) = rx.recv() else {
return;
};
let result = match work {
EngineWork::Reload(provider) => shared.reload_engine(provider),
EngineWork::ReverifyModels => shared.reverify_models(),
};
let body = match result {
Ok(value) => match serde_json::to_vec(&value) {
Ok(bytes) => Body::json(200, bytes),
Err(_) => Body::error(&ServeError::Internal),
},
Err(error) => Body::error(&error),
};
let _ = respond(request, body, &[], &thread_flow);
});
match spawned {
Ok(_) => {
let _ = tx.send(request);
Ok(())
}
Err(error) => {
eprintln!("serve: cannot start the engine-work thread: {error}");
Err(Box::new((ServeError::Internal, request, flow)))
}
}
}
fn spawn_evaluation(
shared: &Arc<ServeShared>,
manifest: PathBuf,
request: Request,
flow: RequestFlow,
) -> Result<(), Box<(ServeError, Request, RequestFlow)>> {
let Some(guard) = shared.begin_evaluation() else {
return Err(Box::new((ServeError::Busy, request, flow)));
};
let shared = Arc::clone(shared);
let (tx, rx) = std::sync::mpsc::sync_channel::<Request>(1);
let thread_flow = flow.clone();
let spawned = thread::Builder::new()
.name("serve-evaluate".to_string())
.spawn(move || {
let _guard = guard;
let Ok(request) = rx.recv() else {
return;
};
let body = match evaluate::run(&shared, &manifest) {
Ok(value) => match serde_json::to_vec(&value) {
Ok(bytes) => Body::json(200, bytes),
Err(_) => Body::error(&ServeError::Internal),
},
Err(error) => shared.error_body(&error),
};
let _ = respond(request, body, &[], &thread_flow);
});
match spawned {
Ok(_) => {
let _ = tx.send(request);
Ok(())
}
Err(error) => {
eprintln!("serve: cannot start the evaluation thread: {error}");
Err(Box::new((ServeError::Internal, request, flow)))
}
}
}
fn route_error(decision: RouteDecision, method: HttpMethod, path: &str) -> Body {
let (status, code, message) = match decision {
RouteDecision::NotFound => (404, "not_found", "no route matches this path".to_string()),
RouteDecision::MethodNotAllowed => (
405,
"method_not_allowed",
format!("{} is not allowed on this path", method.name()),
),
RouteDecision::Matched => (
500,
"internal",
"the router returned no route for a matched path".to_string(),
),
};
let body = json!({
"code": code,
"message": message,
"detail": { "path": path, "method": method.name(), "route": decision.name() },
});
Body::json(
status,
serde_json::to_vec(&body).expect("route error serialization cannot fail"),
)
}
fn json_body(status: u16, value: Value) -> Result<Body, ServeError> {
let bytes = serde_json::to_vec(&value).map_err(|_| ServeError::Internal)?;
Ok(Body::json(status, bytes))
}
fn read_body(request: &mut Request, admitted: &admit::Admit) -> Result<Vec<u8>, ServeError> {
let mut source = HttpBody {
request,
deadline: Instant::now() + BODY_READ_TIMEOUT,
};
admit::read_body(&mut source, admitted.expected_body, admitted.max_body)
}
struct HttpBody<'a> {
request: &'a mut Request,
deadline: Instant,
}
impl BodySource for HttpBody<'_> {
fn read_timed_out(&self) -> bool {
Instant::now() >= self.deadline
}
fn next_chunk(&mut self) -> ChunkOutcome {
let mut buffer = vec![0_u8; READ_CHUNK_BYTES];
match self.request.as_reader().read(&mut buffer) {
Ok(0) => ChunkOutcome::End,
Ok(read) => {
buffer.truncate(read);
ChunkOutcome::Data(buffer)
}
Err(error)
if matches!(
error.kind(),
io::ErrorKind::TimedOut | io::ErrorKind::WouldBlock
) =>
{
ChunkOutcome::TimedOut
}
Err(error) => ChunkOutcome::Failed(error.to_string()),
}
}
}
fn parse_max_side(raw: Option<&str>, min_side: u32) -> Result<Option<u32>, ServeError> {
let Some(raw) = raw else {
return Ok(None);
};
let value: u32 = raw.trim().parse().map_err(|_| ServeError::BadRequest)?;
if value < min_side || value > MAX_SIDE_CEILING {
return Err(ServeError::BadRequest);
}
Ok(Some(value))
}
fn parse_words(raw: Option<&str>) -> Result<WordOutputMode, ServeError> {
match raw {
None => Ok(WordOutputMode::Off),
Some(raw) => WordOutputMode::from_name(raw.trim()).ok_or(ServeError::BadRequest),
}
}
fn parse_set_id(bytes: &[u8]) -> Result<String, ServeError> {
let value: Value = serde_json::from_slice(bytes).map_err(|_| ServeError::BadRequest)?;
let object = value.as_object().ok_or(ServeError::BadRequest)?;
if object.len() != 1 {
return Err(ServeError::BadRequest);
}
let set_id = object
.get("set_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|id| !id.is_empty())
.ok_or(ServeError::BadRequest)?;
Ok(set_id.to_string())
}
fn check_json_content_type(content_type: Option<&str>) -> Result<(), ServeError> {
let Some(raw) = content_type else {
return Ok(());
};
let essence = raw
.split(';')
.next()
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
if essence == "application/json" {
return Ok(());
}
Err(ServeError::BadRequest)
}
fn parse_provider_body(bytes: Vec<u8>) -> Result<Option<ProviderPreference>, ServeError> {
if bytes.is_empty() {
return Ok(None);
}
let value: Value = serde_json::from_slice(&bytes).map_err(|_| ServeError::BadRequest)?;
let object = value.as_object().ok_or(ServeError::BadRequest)?;
if object.len() != 1 {
return Err(ServeError::BadRequest);
}
let raw = object
.get("provider")
.and_then(Value::as_str)
.map(str::trim)
.ok_or(ServeError::BadRequest)?;
let provider = match raw.to_ascii_lowercase().as_str() {
"cpu" => ProviderChoice::Cpu,
"directml" => ProviderChoice::Directml,
"cuda" => ProviderChoice::Cuda,
_ => return Err(ServeError::BadRequest),
};
Ok(Some(provider.preference()))
}
fn export_format(raw: Option<&str>) -> Option<ExportFormat> {
match raw? {
"json" => Some(ExportFormat::Json),
"md" => Some(ExportFormat::Markdown),
"html" => Some(ExportFormat::Html),
_ => None,
}
}
fn content_disposition(job_id: &str, extension: &str) -> String {
format!("attachment; filename=\"ocr-{job_id}.{extension}\"")
}
fn route_of(method: HttpMethod, path: &str) -> (RouteDecision, Option<Route>) {
let candidate = match path {
"/" => Some(Route::Page),
"/api/status" => Some(Route::Status),
"/api/models" => Some(Route::Models),
"/api/models/download" => Some(Route::ModelsDownload),
"/api/models/reverify" => Some(Route::ModelsReverify),
"/api/engine/reload" => Some(Route::EngineReload),
"/api/ocr" => Some(Route::Ocr),
"/api/evaluate" => Some(Route::Evaluate),
_ if path.starts_with("/static/") => Some(Route::Static(path.to_string())),
_ => job_route(path),
};
match candidate {
None => (RouteDecision::NotFound, None),
Some(route) => {
let expected = route.method();
let matched =
method == expected || (method == HttpMethod::Head && expected == HttpMethod::Get);
if matched {
(RouteDecision::Matched, Some(route))
} else {
(RouteDecision::MethodNotAllowed, Some(route))
}
}
}
}
fn job_route(path: &str) -> Option<Route> {
let rest = path.strip_prefix("/api/jobs/")?;
let mut parts = rest.split('/');
let id = parts.next()?;
if id.is_empty() {
return None;
}
match (parts.next(), parts.next()) {
(None, None) => Some(Route::Job(id.to_string())),
(Some("result"), None) => Some(Route::JobResult(id.to_string())),
(Some("cancel"), None) => Some(Route::JobCancel(id.to_string())),
(Some("annotated.png"), None) => Some(Route::JobAnnotated(id.to_string())),
(Some("export"), None) => Some(Route::JobExport(id.to_string())),
_ => None,
}
}
fn split_url(url: &str) -> (&str, &str) {
match url.split_once('?') {
Some((path, query)) => (path, query),
None => (url, ""),
}
}
fn query_param<'a>(query: &'a str, name: &str) -> Option<&'a str> {
query.split('&').find_map(|pair| {
let (key, value) = pair.split_once('=')?;
(key == name).then_some(value)
})
}
fn header<'a>(request: &'a Request, name: &'static str) -> Option<&'a str> {
request
.headers()
.iter()
.find(|header| header.field.equiv(name))
.map(|header| header.value.as_str())
}
fn respond(
request: Request,
body: Body,
extra: &[(String, String)],
flow: &RequestFlow,
) -> io::Result<()> {
flow.response(body.status, body.bytes.len());
let mut response = Response::from_data(body.bytes).with_status_code(StatusCode(body.status));
response.add_header(header_of("Content-Type", body.content_type));
for (name, value) in SECURITY_HEADERS {
response.add_header(header_of(name, value));
}
for (name, value) in extra {
response.add_header(header_of(name, value));
}
request.respond(response)
}
fn header_of(name: &str, value: &str) -> Header {
Header::from_bytes(name.as_bytes(), value.as_bytes())
.unwrap_or_else(|()| panic!("serve: header `{name}` is not ASCII"))
}
#[cfg(test)]
mod tests {
use super::{
EXPORT_CSP, HttpMethod, Route, RouteDecision, check_json_content_type, content_disposition,
export_format, page_csp, parse_max_side, parse_provider_body, parse_set_id, parse_words,
query_param, route_of, split_url,
};
use crate::serve::error::ServeError;
use crate::serve::export::ExportFormat;
use rapid_structure_rs::{ProviderPreference, WordOutputMode};
#[test]
fn the_router_matches_the_documented_endpoint_set_and_nothing_else() {
let cases = [
(HttpMethod::Get, "/", Route::Page),
(HttpMethod::Get, "/api/status", Route::Status),
(HttpMethod::Get, "/api/models", Route::Models),
(
HttpMethod::Post,
"/api/models/download",
Route::ModelsDownload,
),
(
HttpMethod::Post,
"/api/models/reverify",
Route::ModelsReverify,
),
(HttpMethod::Post, "/api/engine/reload", Route::EngineReload),
(HttpMethod::Post, "/api/ocr", Route::Ocr),
(HttpMethod::Post, "/api/evaluate", Route::Evaluate),
(
HttpMethod::Get,
"/api/jobs/job-1",
Route::Job("job-1".into()),
),
(
HttpMethod::Get,
"/api/jobs/job-1/result",
Route::JobResult("job-1".into()),
),
(
HttpMethod::Post,
"/api/jobs/job-1/cancel",
Route::JobCancel("job-1".into()),
),
(
HttpMethod::Get,
"/api/jobs/job-1/annotated.png",
Route::JobAnnotated("job-1".into()),
),
(
HttpMethod::Get,
"/api/jobs/job-1/export",
Route::JobExport("job-1".into()),
),
(
HttpMethod::Get,
"/static/katex/katex.min.js",
Route::Static("/static/katex/katex.min.js".into()),
),
];
for (method, path, expected) in cases {
let (decision, route) = route_of(method, path);
assert_eq!(decision, RouteDecision::Matched, "{path}");
assert_eq!(route.as_ref(), Some(&expected), "{path}");
}
for path in [
"/api/models/download/cancel",
"/api/jobs/job-1/annotated",
"/api/jobs/job-1/annotated.png/extra",
"/api/jobs/bogus/export/data",
"/favicon.ico",
"/api",
"/api/jobs/",
] {
let (decision, route) = route_of(HttpMethod::Get, path);
assert_eq!(decision, RouteDecision::NotFound, "{path}");
assert!(route.is_none(), "{path}");
}
let (decision, route) = route_of(HttpMethod::Get, "/api/ocr");
assert_eq!(decision, RouteDecision::MethodNotAllowed);
assert_eq!(route.expect("known path").allow(), "POST");
let (decision, route) = route_of(HttpMethod::Get, "/api/evaluate");
assert_eq!(decision, RouteDecision::MethodNotAllowed);
assert_eq!(route.expect("known path").allow(), "POST");
let (decision, route) = route_of(HttpMethod::Post, "/api/status");
assert_eq!(decision, RouteDecision::MethodNotAllowed);
assert_eq!(route.expect("known path").allow(), "GET");
for path in ["/api/jobs/job-1/annotated.png", "/api/jobs/job-1/export"] {
let (decision, route) = route_of(HttpMethod::Post, path);
assert_eq!(decision, RouteDecision::MethodNotAllowed, "{path}");
assert_eq!(route.expect("known path").allow(), "GET");
}
assert_eq!(route_of(HttpMethod::Head, "/").0, RouteDecision::Matched);
assert_eq!(
route_of(HttpMethod::Head, "/api/ocr").0,
RouteDecision::MethodNotAllowed
);
}
#[test]
fn static_routes_are_get_only_and_outside_the_token_boundary() {
let (decision, route) = route_of(HttpMethod::Get, "/static/mathjax/tex-svg.js");
assert_eq!(decision, RouteDecision::Matched);
assert!(
!route.expect("static route").requires_token(),
"/static/* must not require the token"
);
let (decision, route) = route_of(HttpMethod::Post, "/static/katex/katex.min.css");
assert_eq!(decision, RouteDecision::MethodNotAllowed);
assert_eq!(route.expect("known path").allow(), "GET");
assert!(!Route::Page.requires_token());
assert!(Route::Status.requires_token());
assert!(Route::JobExport("job-1".into()).requires_token());
}
#[test]
fn only_the_three_documented_export_formats_are_accepted() {
assert_eq!(export_format(Some("json")), Some(ExportFormat::Json));
assert_eq!(export_format(Some("md")), Some(ExportFormat::Markdown));
assert_eq!(export_format(Some("html")), Some(ExportFormat::Html));
for bad in [None, Some(""), Some("JSON"), Some("pdf"), Some("mdx")] {
assert_eq!(export_format(bad), None, "{bad:?}");
}
assert_eq!(ExportFormat::Json.extension(), "json");
assert_eq!(ExportFormat::Markdown.extension(), "md");
assert_eq!(ExportFormat::Html.extension(), "html");
assert_eq!(
ExportFormat::Html.content_type(),
"text/html; charset=utf-8"
);
assert_eq!(
content_disposition("job-1", ExportFormat::Html.extension()),
"attachment; filename=\"ocr-job-1.html\""
);
}
#[test]
fn the_reload_body_carries_at_most_a_provider() {
assert_eq!(parse_provider_body(Vec::new()).expect("no body"), None);
assert_eq!(
parse_provider_body(br#"{"provider":"cpu"}"#.to_vec()).expect("cpu"),
Some(ProviderPreference::Cpu)
);
assert_eq!(
parse_provider_body(br#"{"provider":" directml "}"#.to_vec()).expect("directml"),
Some(ProviderPreference::DirectMl { device_id: 0 })
);
assert_eq!(
parse_provider_body(br#"{"provider":"CUDA"}"#.to_vec()).expect("cuda"),
Some(ProviderPreference::Cuda { device_id: 0 })
);
for bad in [
&br#"{"provider":"tpu"}"#[..],
&br#"{"provider":null}"#[..],
&br#"{"provider":""}"#[..],
&br#"{"provider":"cpu","allow_provider_fallback":true}"#[..],
&br#"{"other":"cpu"}"#[..],
&br#"{}"#[..],
&br#"[]"#[..],
&br#"not json"#[..],
] {
assert!(
matches!(
parse_provider_body(bad.to_vec()),
Err(ServeError::BadRequest)
),
"{}",
String::from_utf8_lossy(bad)
);
}
}
#[test]
fn the_export_csp_is_the_documented_one_and_differs_from_the_page_csp() {
assert_eq!(
EXPORT_CSP,
"default-src 'none'; style-src 'unsafe-inline'; img-src data:; script-src 'none'; \
sandbox"
);
assert!(EXPORT_CSP.contains("script-src 'none'"), "{EXPORT_CSP}");
assert!(EXPORT_CSP.contains("img-src data:"), "{EXPORT_CSP}");
assert!(!EXPORT_CSP.contains("nonce"), "{EXPORT_CSP}");
let page = page_csp("abc");
assert!(!page.contains("unsafe-inline"), "{page}");
assert_ne!(page, EXPORT_CSP);
}
#[test]
fn urls_and_query_parameters_are_split_without_decoding() {
assert_eq!(
split_url("/api/ocr?max_side=2000"),
("/api/ocr", "max_side=2000")
);
assert_eq!(split_url("/"), ("/", ""));
assert_eq!(
query_param("max_side=2000&queue=text", "max_side"),
Some("2000")
);
assert_eq!(
query_param("max_side=2000&queue=text", "queue"),
Some("text")
);
assert_eq!(query_param("a=1", "b"), None);
assert_eq!(query_param("noseparator", "noseparator"), None);
}
#[test]
fn the_documented_max_side_parameter_is_validated() {
assert_eq!(parse_max_side(None, 30).expect("ok"), None);
assert_eq!(parse_max_side(Some("2000"), 30).expect("ok"), Some(2000));
assert_eq!(parse_max_side(Some(" 736 "), 30).expect("ok"), Some(736));
for bad in ["0", "29", "abc", "-1", "999999999999999999999"] {
assert!(parse_max_side(Some(bad), 30).is_err(), "{bad}");
}
}
#[test]
fn the_documented_words_parameter_is_validated() {
assert_eq!(parse_words(None).expect("ok"), WordOutputMode::Off);
for (raw, expected) in [
("off", WordOutputMode::Off),
("words", WordOutputMode::Words),
("chars", WordOutputMode::Chars),
(" words ", WordOutputMode::Words),
] {
assert_eq!(parse_words(Some(raw)).expect("ok"), expected, "{raw}");
}
for bad in [
"Off", "CHARS", "char", "word", "", "1", "0", "true", "table",
] {
assert!(parse_words(Some(bad)).is_err(), "{bad}");
}
}
#[test]
fn the_download_body_carries_only_a_set_id() {
assert_eq!(
parse_set_id(br#"{"set_id":"PP-OCRv6-small-ch"}"#).expect("ok"),
"PP-OCRv6-small-ch"
);
for bad in [
&br#"{"url":"https://evil.example/x.onnx"}"#[..],
&br#"{"set_id":"a","url":"https://evil.example/x.onnx"}"#[..],
&br#"{"set_id":""}"#[..],
&br#"{"set_id":null}"#[..],
&br#"not json"#[..],
&br#"{}"#[..],
&br#"[]"#[..],
] {
assert!(
matches!(parse_set_id(bad), Err(ServeError::BadRequest)),
"{}",
String::from_utf8_lossy(bad)
);
}
assert!(check_json_content_type(None).is_ok());
assert!(check_json_content_type(Some("application/json; charset=utf-8")).is_ok());
assert!(check_json_content_type(Some("application/octet-stream")).is_err());
assert!(check_json_content_type(Some("multipart/form-data")).is_err());
}
#[test]
fn the_main_page_csp_is_nonce_based_and_never_allows_inline() {
let csp = page_csp("abc123");
assert!(csp.contains("script-src 'nonce-abc123'"), "{csp}");
assert!(csp.contains("style-src 'nonce-abc123'"), "{csp}");
assert!(!csp.contains("unsafe-inline"), "{csp}");
assert!(csp.contains("default-src 'none'"), "{csp}");
assert!(csp.contains("frame-ancestors 'none'"), "{csp}");
assert!(csp.contains("font-src 'self'"), "{csp}");
assert!(!csp.contains("font-src *"), "{csp}");
assert!(!csp.contains("font-src 'none'"), "{csp}");
}
}