mirror of
https://github.com/ruvnet/RuView
synced 2026-08-01 19:01:42 +00:00
Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'
This commit is contained in:
+308
@@ -0,0 +1,308 @@
|
||||
use axum::{
|
||||
extract::{Path, Query, State},
|
||||
http::StatusCode,
|
||||
response::{sse::Event, IntoResponse, Sse},
|
||||
Json,
|
||||
};
|
||||
use futures::stream::{self, Stream};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::{convert::Infallible, time::Duration};
|
||||
use tracing::{error, info, warn};
|
||||
use validator::Validate;
|
||||
|
||||
use super::{
|
||||
jobs::{JobStatus, PdfJob},
|
||||
requests::{LatexRequest, PdfRequest, StrokesRequest, TextRequest},
|
||||
responses::{ErrorResponse, PdfResponse, TextResponse},
|
||||
state::AppState,
|
||||
};
|
||||
|
||||
/// Health check handler
|
||||
pub async fn get_health() -> impl IntoResponse {
|
||||
#[derive(Serialize)]
|
||||
struct Health {
|
||||
status: &'static str,
|
||||
version: &'static str,
|
||||
}
|
||||
|
||||
Json(Health {
|
||||
status: "ok",
|
||||
version: env!("CARGO_PKG_VERSION"),
|
||||
})
|
||||
}
|
||||
|
||||
/// Process text/image OCR request
|
||||
/// Supports multipart/form-data, base64, and URL inputs
|
||||
///
|
||||
/// # Important
|
||||
/// This endpoint requires OCR models to be configured. If models are not available,
|
||||
/// returns a 503 Service Unavailable error with instructions.
|
||||
pub async fn process_text(
|
||||
State(_state): State<AppState>,
|
||||
Json(request): Json<TextRequest>,
|
||||
) -> Result<Json<TextResponse>, ErrorResponse> {
|
||||
info!("Processing text OCR request");
|
||||
|
||||
// Validate request
|
||||
request.validate().map_err(|e| {
|
||||
warn!("Invalid request: {:?}", e);
|
||||
ErrorResponse::validation_error(format!("Validation failed: {}", e))
|
||||
})?;
|
||||
|
||||
// Download or decode image
|
||||
let image_data = match request.get_image_data().await {
|
||||
Ok(data) => data,
|
||||
Err(e) => {
|
||||
error!("Failed to get image data: {:?}", e);
|
||||
return Err(ErrorResponse::internal_error("Failed to process image"));
|
||||
}
|
||||
};
|
||||
|
||||
// Validate image data is not empty
|
||||
if image_data.is_empty() {
|
||||
return Err(ErrorResponse::validation_error("Image data is empty"));
|
||||
}
|
||||
|
||||
// OCR processing requires models to be configured
|
||||
// Return informative error explaining how to set up the service
|
||||
Err(ErrorResponse::service_unavailable(
|
||||
"OCR service not fully configured. ONNX models are required for OCR processing. \
|
||||
Please download compatible models (PaddleOCR, TrOCR) and configure the model directory. \
|
||||
See documentation at /docs/MODEL_SETUP.md for setup instructions.",
|
||||
))
|
||||
}
|
||||
|
||||
/// Process digital ink strokes
|
||||
///
|
||||
/// # Important
|
||||
/// This endpoint requires OCR models to be configured.
|
||||
pub async fn process_strokes(
|
||||
State(_state): State<AppState>,
|
||||
Json(request): Json<StrokesRequest>,
|
||||
) -> Result<Json<TextResponse>, ErrorResponse> {
|
||||
info!(
|
||||
"Processing strokes request with {} strokes",
|
||||
request.strokes.len()
|
||||
);
|
||||
|
||||
request
|
||||
.validate()
|
||||
.map_err(|e| ErrorResponse::validation_error(format!("Validation failed: {}", e)))?;
|
||||
|
||||
// Validate we have stroke data
|
||||
if request.strokes.is_empty() {
|
||||
return Err(ErrorResponse::validation_error("No strokes provided"));
|
||||
}
|
||||
|
||||
// Stroke recognition requires models to be configured
|
||||
Err(ErrorResponse::service_unavailable(
|
||||
"Stroke recognition service not configured. ONNX models required for ink recognition.",
|
||||
))
|
||||
}
|
||||
|
||||
/// Process legacy LaTeX equation request
|
||||
///
|
||||
/// # Important
|
||||
/// This endpoint requires OCR models to be configured.
|
||||
pub async fn process_latex(
|
||||
State(_state): State<AppState>,
|
||||
Json(request): Json<LatexRequest>,
|
||||
) -> Result<Json<TextResponse>, ErrorResponse> {
|
||||
info!("Processing legacy LaTeX request");
|
||||
|
||||
request
|
||||
.validate()
|
||||
.map_err(|e| ErrorResponse::validation_error(format!("Validation failed: {}", e)))?;
|
||||
|
||||
// LaTeX recognition requires models to be configured
|
||||
Err(ErrorResponse::service_unavailable(
|
||||
"LaTeX recognition service not configured. ONNX models required.",
|
||||
))
|
||||
}
|
||||
|
||||
/// Create async PDF processing job
|
||||
pub async fn process_pdf(
|
||||
State(state): State<AppState>,
|
||||
Json(request): Json<PdfRequest>,
|
||||
) -> Result<Json<PdfResponse>, ErrorResponse> {
|
||||
info!("Creating PDF processing job");
|
||||
|
||||
request
|
||||
.validate()
|
||||
.map_err(|e| ErrorResponse::validation_error(format!("Validation failed: {}", e)))?;
|
||||
|
||||
// Create job
|
||||
let job = PdfJob::new(request);
|
||||
let job_id = job.id.clone();
|
||||
|
||||
// Queue job
|
||||
state.job_queue.enqueue(job).await.map_err(|e| {
|
||||
error!("Failed to enqueue job: {:?}", e);
|
||||
ErrorResponse::internal_error("Failed to create PDF job")
|
||||
})?;
|
||||
|
||||
let response = PdfResponse {
|
||||
pdf_id: job_id,
|
||||
status: JobStatus::Processing,
|
||||
message: Some("PDF processing started".to_string()),
|
||||
result: None,
|
||||
error: None,
|
||||
};
|
||||
|
||||
Ok(Json(response))
|
||||
}
|
||||
|
||||
/// Get PDF job status
|
||||
pub async fn get_pdf_status(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<PdfResponse>, ErrorResponse> {
|
||||
info!("Getting PDF job status: {}", id);
|
||||
|
||||
let status = state
|
||||
.job_queue
|
||||
.get_status(&id)
|
||||
.await
|
||||
.ok_or_else(|| ErrorResponse::not_found("Job not found"))?;
|
||||
|
||||
let response = PdfResponse {
|
||||
pdf_id: id.clone(),
|
||||
status: status.clone(),
|
||||
message: Some(format!("Job status: {:?}", status)),
|
||||
result: state.job_queue.get_result(&id).await,
|
||||
error: state.job_queue.get_error(&id).await,
|
||||
};
|
||||
|
||||
Ok(Json(response))
|
||||
}
|
||||
|
||||
/// Delete PDF job
|
||||
pub async fn delete_pdf_job(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<StatusCode, ErrorResponse> {
|
||||
info!("Deleting PDF job: {}", id);
|
||||
|
||||
state
|
||||
.job_queue
|
||||
.cancel(&id)
|
||||
.await
|
||||
.map_err(|_| ErrorResponse::not_found("Job not found"))?;
|
||||
|
||||
Ok(StatusCode::NO_CONTENT)
|
||||
}
|
||||
|
||||
/// Stream PDF processing results via SSE
|
||||
pub async fn stream_pdf_results(
|
||||
State(_state): State<AppState>,
|
||||
Path(_id): Path<String>,
|
||||
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
|
||||
info!("Streaming PDF results for job: {}", _id);
|
||||
|
||||
let stream = stream::unfold(0, move |page| {
|
||||
async move {
|
||||
if page > 10 {
|
||||
// Example: stop after 10 pages
|
||||
return None;
|
||||
}
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
|
||||
let event = Event::default()
|
||||
.json_data(serde_json::json!({
|
||||
"page": page,
|
||||
"text": format!("Content from page {}", page),
|
||||
"progress": (page as f32 / 10.0) * 100.0
|
||||
}))
|
||||
.ok()?;
|
||||
|
||||
Some((Ok(event), page + 1))
|
||||
}
|
||||
});
|
||||
|
||||
Sse::new(stream)
|
||||
}
|
||||
|
||||
/// Convert document to different format (MMD/DOCX/etc)
|
||||
///
|
||||
/// # Note
|
||||
/// Document conversion requires additional backend services to be configured.
|
||||
pub async fn convert_document(
|
||||
State(_state): State<AppState>,
|
||||
Json(_request): Json<serde_json::Value>,
|
||||
) -> Result<Json<serde_json::Value>, ErrorResponse> {
|
||||
info!("Converting document");
|
||||
|
||||
// Document conversion is not yet implemented
|
||||
Err(ErrorResponse::not_implemented(
|
||||
"Document conversion is not yet implemented. This feature requires additional backend services."
|
||||
))
|
||||
}
|
||||
|
||||
/// Get OCR processing history
|
||||
#[derive(Deserialize)]
|
||||
pub struct HistoryQuery {
|
||||
#[serde(default)]
|
||||
page: u32,
|
||||
#[serde(default = "default_limit")]
|
||||
limit: u32,
|
||||
}
|
||||
|
||||
fn default_limit() -> u32 {
|
||||
50
|
||||
}
|
||||
|
||||
/// Get OCR processing history
|
||||
///
|
||||
/// # Note
|
||||
/// History storage requires a database backend to be configured.
|
||||
/// Returns empty results if no database is available.
|
||||
pub async fn get_ocr_results(
|
||||
State(_state): State<AppState>,
|
||||
Query(params): Query<HistoryQuery>,
|
||||
) -> Result<Json<serde_json::Value>, ErrorResponse> {
|
||||
info!(
|
||||
"Getting OCR results history: page={}, limit={}",
|
||||
params.page, params.limit
|
||||
);
|
||||
|
||||
// History storage not configured - return empty results with notice
|
||||
Ok(Json(serde_json::json!({
|
||||
"results": [],
|
||||
"total": 0,
|
||||
"page": params.page,
|
||||
"limit": params.limit,
|
||||
"notice": "History storage not configured. Results are not persisted."
|
||||
})))
|
||||
}
|
||||
|
||||
/// Get OCR usage statistics
|
||||
///
|
||||
/// # Note
|
||||
/// Usage tracking requires a database backend to be configured.
|
||||
/// Returns zeros if no database is available.
|
||||
pub async fn get_ocr_usage(
|
||||
State(_state): State<AppState>,
|
||||
) -> Result<Json<serde_json::Value>, ErrorResponse> {
|
||||
info!("Getting OCR usage statistics");
|
||||
|
||||
// Usage tracking not configured - return zeros with notice
|
||||
Ok(Json(serde_json::json!({
|
||||
"requests_today": 0,
|
||||
"requests_month": 0,
|
||||
"quota_limit": null,
|
||||
"quota_remaining": null,
|
||||
"notice": "Usage tracking not configured. Statistics are not recorded."
|
||||
})))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_health_check() {
|
||||
let response = get_health().await.into_response();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
}
|
||||
+281
@@ -0,0 +1,281 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{mpsc, RwLock};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::requests::PdfRequest;
|
||||
|
||||
/// Job status enumeration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum JobStatus {
|
||||
/// Job is queued but not started
|
||||
Queued,
|
||||
/// Job is currently processing
|
||||
Processing,
|
||||
/// Job completed successfully
|
||||
Completed,
|
||||
/// Job failed with error
|
||||
Failed,
|
||||
/// Job was cancelled
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
/// PDF processing job
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PdfJob {
|
||||
/// Unique job identifier
|
||||
pub id: String,
|
||||
|
||||
/// Original request
|
||||
pub request: PdfRequest,
|
||||
|
||||
/// Current status
|
||||
pub status: JobStatus,
|
||||
|
||||
/// Creation timestamp
|
||||
pub created_at: DateTime<Utc>,
|
||||
|
||||
/// Last update timestamp
|
||||
pub updated_at: DateTime<Utc>,
|
||||
|
||||
/// Processing result
|
||||
pub result: Option<String>,
|
||||
|
||||
/// Error message (if failed)
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
impl PdfJob {
|
||||
/// Create a new PDF job
|
||||
pub fn new(request: PdfRequest) -> Self {
|
||||
let now = Utc::now();
|
||||
Self {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
request,
|
||||
status: JobStatus::Queued,
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
result: None,
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Update job status
|
||||
pub fn update_status(&mut self, status: JobStatus) {
|
||||
self.status = status;
|
||||
self.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
/// Set job result
|
||||
pub fn set_result(&mut self, result: String) {
|
||||
self.result = Some(result);
|
||||
self.status = JobStatus::Completed;
|
||||
self.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
/// Set job error
|
||||
pub fn set_error(&mut self, error: String) {
|
||||
self.error = Some(error);
|
||||
self.status = JobStatus::Failed;
|
||||
self.updated_at = Utc::now();
|
||||
}
|
||||
}
|
||||
|
||||
/// Async job queue with webhook support
|
||||
pub struct JobQueue {
|
||||
/// Job storage
|
||||
jobs: Arc<RwLock<HashMap<String, PdfJob>>>,
|
||||
|
||||
/// Job submission channel
|
||||
tx: mpsc::Sender<PdfJob>,
|
||||
|
||||
/// Job processing handle
|
||||
_handle: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl JobQueue {
|
||||
/// Create a new job queue
|
||||
pub fn new() -> Self {
|
||||
Self::with_capacity(1000)
|
||||
}
|
||||
|
||||
/// Create a job queue with specific capacity
|
||||
pub fn with_capacity(capacity: usize) -> Self {
|
||||
let jobs = Arc::new(RwLock::new(HashMap::new()));
|
||||
let (tx, rx) = mpsc::channel(capacity);
|
||||
|
||||
let queue_jobs = jobs.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
Self::process_jobs(queue_jobs, rx).await;
|
||||
});
|
||||
|
||||
Self {
|
||||
jobs,
|
||||
tx,
|
||||
_handle: Some(handle),
|
||||
}
|
||||
}
|
||||
|
||||
/// Enqueue a new job
|
||||
pub async fn enqueue(&self, mut job: PdfJob) -> anyhow::Result<()> {
|
||||
job.update_status(JobStatus::Queued);
|
||||
|
||||
// Store job
|
||||
{
|
||||
let mut jobs = self.jobs.write().await;
|
||||
jobs.insert(job.id.clone(), job.clone());
|
||||
}
|
||||
|
||||
// Send to processing queue
|
||||
self.tx.send(job).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get job status
|
||||
pub async fn get_status(&self, id: &str) -> Option<JobStatus> {
|
||||
let jobs = self.jobs.read().await;
|
||||
jobs.get(id).map(|job| job.status.clone())
|
||||
}
|
||||
|
||||
/// Get job result
|
||||
pub async fn get_result(&self, id: &str) -> Option<String> {
|
||||
let jobs = self.jobs.read().await;
|
||||
jobs.get(id).and_then(|job| job.result.clone())
|
||||
}
|
||||
|
||||
/// Get job error
|
||||
pub async fn get_error(&self, id: &str) -> Option<String> {
|
||||
let jobs = self.jobs.read().await;
|
||||
jobs.get(id).and_then(|job| job.error.clone())
|
||||
}
|
||||
|
||||
/// Cancel a job
|
||||
pub async fn cancel(&self, id: &str) -> anyhow::Result<()> {
|
||||
let mut jobs = self.jobs.write().await;
|
||||
if let Some(job) = jobs.get_mut(id) {
|
||||
job.update_status(JobStatus::Cancelled);
|
||||
Ok(())
|
||||
} else {
|
||||
anyhow::bail!("Job not found")
|
||||
}
|
||||
}
|
||||
|
||||
/// Background job processor
|
||||
async fn process_jobs(
|
||||
jobs: Arc<RwLock<HashMap<String, PdfJob>>>,
|
||||
mut rx: mpsc::Receiver<PdfJob>,
|
||||
) {
|
||||
while let Some(job) = rx.recv().await {
|
||||
let job_id = job.id.clone();
|
||||
|
||||
// Update status to processing
|
||||
{
|
||||
let mut jobs_lock = jobs.write().await;
|
||||
if let Some(stored_job) = jobs_lock.get_mut(&job_id) {
|
||||
stored_job.update_status(JobStatus::Processing);
|
||||
}
|
||||
}
|
||||
|
||||
// Simulate PDF processing
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
|
||||
|
||||
// Update with result
|
||||
{
|
||||
let mut jobs_lock = jobs.write().await;
|
||||
if let Some(stored_job) = jobs_lock.get_mut(&job_id) {
|
||||
stored_job.set_result("Processed PDF content".to_string());
|
||||
|
||||
// Send webhook if specified
|
||||
if let Some(webhook_url) = &stored_job.request.webhook_url {
|
||||
Self::send_webhook(webhook_url, stored_job).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Send webhook notification
|
||||
async fn send_webhook(url: &str, job: &PdfJob) {
|
||||
let client = reqwest::Client::new();
|
||||
let payload = serde_json::json!({
|
||||
"job_id": job.id,
|
||||
"status": job.status,
|
||||
"result": job.result,
|
||||
"error": job.error,
|
||||
});
|
||||
|
||||
if let Err(e) = client.post(url).json(&payload).send().await {
|
||||
tracing::error!("Failed to send webhook to {}: {:?}", url, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for JobQueue {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::api::requests::{PdfOptions, RequestMetadata};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_job_creation() {
|
||||
let request = PdfRequest {
|
||||
url: "https://example.com/test.pdf".to_string(),
|
||||
options: PdfOptions::default(),
|
||||
webhook_url: None,
|
||||
metadata: RequestMetadata::default(),
|
||||
};
|
||||
|
||||
let job = PdfJob::new(request);
|
||||
assert_eq!(job.status, JobStatus::Queued);
|
||||
assert!(job.result.is_none());
|
||||
assert!(job.error.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_job_queue_enqueue() {
|
||||
let queue = JobQueue::new();
|
||||
let request = PdfRequest {
|
||||
url: "https://example.com/test.pdf".to_string(),
|
||||
options: PdfOptions::default(),
|
||||
webhook_url: None,
|
||||
metadata: RequestMetadata::default(),
|
||||
};
|
||||
|
||||
let job = PdfJob::new(request);
|
||||
let job_id = job.id.clone();
|
||||
|
||||
queue.enqueue(job).await.unwrap();
|
||||
|
||||
let status = queue.get_status(&job_id).await;
|
||||
assert!(status.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_job_cancellation() {
|
||||
let queue = JobQueue::new();
|
||||
let request = PdfRequest {
|
||||
url: "https://example.com/test.pdf".to_string(),
|
||||
options: PdfOptions::default(),
|
||||
webhook_url: None,
|
||||
metadata: RequestMetadata::default(),
|
||||
};
|
||||
|
||||
let job = PdfJob::new(request);
|
||||
let job_id = job.id.clone();
|
||||
|
||||
queue.enqueue(job).await.unwrap();
|
||||
queue.cancel(&job_id).await.unwrap();
|
||||
|
||||
let status = queue.get_status(&job_id).await;
|
||||
assert_eq!(status, Some(JobStatus::Cancelled));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
use axum::{
|
||||
extract::{Request, State},
|
||||
http::HeaderMap,
|
||||
middleware::Next,
|
||||
response::Response,
|
||||
};
|
||||
use governor::{
|
||||
clock::DefaultClock,
|
||||
state::{InMemoryState, NotKeyed},
|
||||
Quota, RateLimiter,
|
||||
};
|
||||
use nonzero_ext::nonzero;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::sync::Arc;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use super::{responses::ErrorResponse, state::AppState};
|
||||
|
||||
/// Authentication middleware
|
||||
/// Validates app_id and app_key from headers or query parameters
|
||||
pub async fn auth_middleware(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
request: Request,
|
||||
next: Next,
|
||||
) -> Result<Response, ErrorResponse> {
|
||||
// Check if authentication is enabled
|
||||
if !state.auth_enabled {
|
||||
debug!("Authentication disabled, allowing request");
|
||||
return Ok(next.run(request).await);
|
||||
}
|
||||
|
||||
// Extract credentials from headers
|
||||
let app_id = headers
|
||||
.get("app_id")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.or_else(|| {
|
||||
// Fallback to query parameters
|
||||
request
|
||||
.uri()
|
||||
.query()
|
||||
.and_then(|q| extract_query_param(q, "app_id"))
|
||||
});
|
||||
|
||||
let app_key = headers
|
||||
.get("app_key")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.or_else(|| {
|
||||
request
|
||||
.uri()
|
||||
.query()
|
||||
.and_then(|q| extract_query_param(q, "app_key"))
|
||||
});
|
||||
|
||||
// Validate credentials
|
||||
match (app_id, app_key) {
|
||||
(Some(id), Some(key)) => {
|
||||
if validate_credentials(&state, id, key).await {
|
||||
debug!("Authentication successful for app_id: {}", id);
|
||||
Ok(next.run(request).await)
|
||||
} else {
|
||||
warn!("Invalid credentials for app_id: {}", id);
|
||||
Err(ErrorResponse::unauthorized("Invalid credentials"))
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
warn!("Missing authentication credentials");
|
||||
Err(ErrorResponse::unauthorized("Missing app_id or app_key"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Rate limiting middleware using token bucket algorithm
|
||||
pub async fn rate_limit_middleware(
|
||||
State(state): State<AppState>,
|
||||
request: Request,
|
||||
next: Next,
|
||||
) -> Result<Response, ErrorResponse> {
|
||||
// Check rate limit
|
||||
match state.rate_limiter.check() {
|
||||
Ok(_) => {
|
||||
debug!("Rate limit check passed");
|
||||
Ok(next.run(request).await)
|
||||
}
|
||||
Err(_) => {
|
||||
warn!("Rate limit exceeded");
|
||||
Err(ErrorResponse::rate_limited(
|
||||
"Rate limit exceeded. Please try again later.",
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate app credentials using secure comparison
|
||||
///
|
||||
/// SECURITY: This implementation:
|
||||
/// 1. Requires credentials to be pre-configured in AppState
|
||||
/// 2. Uses constant-time comparison to prevent timing attacks
|
||||
/// 3. Hashes the key before comparison
|
||||
async fn validate_credentials(state: &AppState, app_id: &str, app_key: &str) -> bool {
|
||||
// Reject empty credentials
|
||||
if app_id.is_empty() || app_key.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Get configured credentials from state
|
||||
let Some(expected_key_hash) = state.api_keys.get(app_id) else {
|
||||
warn!("Unknown app_id attempted authentication: {}", app_id);
|
||||
return false;
|
||||
};
|
||||
|
||||
// Hash the provided key
|
||||
let provided_key_hash = hash_api_key(app_key);
|
||||
|
||||
// Constant-time comparison to prevent timing attacks
|
||||
constant_time_compare(&provided_key_hash, expected_key_hash.as_str())
|
||||
}
|
||||
|
||||
/// Hash an API key using SHA-256
|
||||
fn hash_api_key(key: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(key.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
/// Constant-time string comparison to prevent timing attacks
|
||||
fn constant_time_compare(a: &str, b: &str) -> bool {
|
||||
if a.len() != b.len() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let mut result = 0u8;
|
||||
for (x, y) in a.bytes().zip(b.bytes()) {
|
||||
result |= x ^ y;
|
||||
}
|
||||
result == 0
|
||||
}
|
||||
|
||||
/// Extract query parameter from query string
|
||||
fn extract_query_param<'a>(query: &'a str, param: &str) -> Option<&'a str> {
|
||||
query.split('&').find_map(|pair| {
|
||||
let mut parts = pair.split('=');
|
||||
match (parts.next(), parts.next()) {
|
||||
(Some(k), Some(v)) if k == param => Some(v),
|
||||
_ => None,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a rate limiter with token bucket algorithm
|
||||
pub fn create_rate_limiter() -> Arc<RateLimiter<NotKeyed, InMemoryState, DefaultClock>> {
|
||||
// Allow 100 requests per minute
|
||||
let quota = Quota::per_minute(nonzero!(100u32));
|
||||
Arc::new(RateLimiter::direct(quota))
|
||||
}
|
||||
|
||||
/// Type alias for rate limiter
|
||||
pub type AppRateLimiter = Arc<RateLimiter<NotKeyed, InMemoryState, DefaultClock>>;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_extract_query_param() {
|
||||
let query = "app_id=123&app_key=secret&foo=bar";
|
||||
assert_eq!(extract_query_param(query, "app_id"), Some("123"));
|
||||
assert_eq!(extract_query_param(query, "app_key"), Some("secret"));
|
||||
assert_eq!(extract_query_param(query, "foo"), Some("bar"));
|
||||
assert_eq!(extract_query_param(query, "missing"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hash_api_key() {
|
||||
let key = "test_key_123";
|
||||
let hash1 = hash_api_key(key);
|
||||
let hash2 = hash_api_key(key);
|
||||
assert_eq!(hash1, hash2);
|
||||
assert_ne!(hash_api_key("different"), hash1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_constant_time_compare() {
|
||||
assert!(constant_time_compare("abc", "abc"));
|
||||
assert!(!constant_time_compare("abc", "abd"));
|
||||
assert!(!constant_time_compare("abc", "ab"));
|
||||
assert!(!constant_time_compare("", "a"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_validate_credentials_rejects_empty() {
|
||||
let state = AppState::new();
|
||||
assert!(!validate_credentials(&state, "", "key").await);
|
||||
assert!(!validate_credentials(&state, "test", "").await);
|
||||
assert!(!validate_credentials(&state, "", "").await);
|
||||
}
|
||||
}
|
||||
+91
@@ -0,0 +1,91 @@
|
||||
pub mod handlers;
|
||||
pub mod jobs;
|
||||
pub mod middleware;
|
||||
pub mod requests;
|
||||
pub mod responses;
|
||||
pub mod routes;
|
||||
pub mod state;
|
||||
|
||||
use anyhow::Result;
|
||||
use axum::Router;
|
||||
use std::net::SocketAddr;
|
||||
use tokio::signal;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use self::state::AppState;
|
||||
|
||||
/// Main API server structure
|
||||
pub struct ApiServer {
|
||||
state: AppState,
|
||||
addr: SocketAddr,
|
||||
}
|
||||
|
||||
impl ApiServer {
|
||||
/// Create a new API server instance
|
||||
pub fn new(state: AppState, addr: SocketAddr) -> Self {
|
||||
Self { state, addr }
|
||||
}
|
||||
|
||||
/// Start the API server with graceful shutdown
|
||||
pub async fn start(self) -> Result<()> {
|
||||
let app = self.create_router();
|
||||
|
||||
info!("Starting Scipix API server on {}", self.addr);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(self.addr).await?;
|
||||
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(shutdown_signal())
|
||||
.await?;
|
||||
|
||||
info!("Server shutdown complete");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Create the application router with all routes and middleware
|
||||
fn create_router(&self) -> Router {
|
||||
routes::router(self.state.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Graceful shutdown signal handler
|
||||
async fn shutdown_signal() {
|
||||
let ctrl_c = async {
|
||||
signal::ctrl_c()
|
||||
.await
|
||||
.expect("failed to install Ctrl+C handler");
|
||||
};
|
||||
|
||||
#[cfg(unix)]
|
||||
let terminate = async {
|
||||
signal::unix::signal(signal::unix::SignalKind::terminate())
|
||||
.expect("failed to install signal handler")
|
||||
.recv()
|
||||
.await;
|
||||
};
|
||||
|
||||
#[cfg(not(unix))]
|
||||
let terminate = std::future::pending::<()>();
|
||||
|
||||
tokio::select! {
|
||||
_ = ctrl_c => {
|
||||
warn!("Received Ctrl+C, shutting down...");
|
||||
},
|
||||
_ = terminate => {
|
||||
warn!("Received terminate signal, shutting down...");
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_server_creation() {
|
||||
let state = AppState::new();
|
||||
let addr = "127.0.0.1:3000".parse().unwrap();
|
||||
let server = ApiServer::new(state, addr);
|
||||
assert_eq!(server.addr, addr);
|
||||
}
|
||||
}
|
||||
+227
@@ -0,0 +1,227 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use validator::Validate;
|
||||
|
||||
/// Text/Image OCR request
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
|
||||
pub struct TextRequest {
|
||||
/// Image source (base64, URL, or multipart)
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub src: Option<String>,
|
||||
|
||||
/// Base64 encoded image data
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub base64: Option<String>,
|
||||
|
||||
/// Image URL
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
#[validate(url)]
|
||||
pub url: Option<String>,
|
||||
|
||||
/// Request metadata
|
||||
#[serde(default)]
|
||||
pub metadata: RequestMetadata,
|
||||
}
|
||||
|
||||
impl TextRequest {
|
||||
/// Get image data from request
|
||||
pub async fn get_image_data(&self) -> anyhow::Result<Vec<u8>> {
|
||||
if let Some(base64_data) = &self.base64 {
|
||||
// Decode base64
|
||||
use base64::Engine;
|
||||
let decoded = base64::engine::general_purpose::STANDARD.decode(base64_data)?;
|
||||
Ok(decoded)
|
||||
} else if let Some(url) = &self.url {
|
||||
// Download from URL
|
||||
let response = reqwest::get(url).await?;
|
||||
let bytes = response.bytes().await?;
|
||||
Ok(bytes.to_vec())
|
||||
} else {
|
||||
anyhow::bail!("No image data provided")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Digital ink strokes request
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
|
||||
pub struct StrokesRequest {
|
||||
/// Array of stroke data
|
||||
#[validate(length(min = 1))]
|
||||
pub strokes: Vec<Stroke>,
|
||||
|
||||
/// Request metadata
|
||||
#[serde(default)]
|
||||
pub metadata: RequestMetadata,
|
||||
}
|
||||
|
||||
/// Single stroke in digital ink
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Stroke {
|
||||
/// X coordinates
|
||||
pub x: Vec<f64>,
|
||||
|
||||
/// Y coordinates
|
||||
pub y: Vec<f64>,
|
||||
|
||||
/// Optional timestamps
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub t: Option<Vec<f64>>,
|
||||
}
|
||||
|
||||
/// Legacy LaTeX equation request
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
|
||||
pub struct LatexRequest {
|
||||
/// Image source
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub src: Option<String>,
|
||||
|
||||
/// Base64 encoded image
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub base64: Option<String>,
|
||||
|
||||
/// Image URL
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
#[validate(url)]
|
||||
pub url: Option<String>,
|
||||
|
||||
/// Request metadata
|
||||
#[serde(default)]
|
||||
pub metadata: RequestMetadata,
|
||||
}
|
||||
|
||||
/// PDF processing request
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
|
||||
pub struct PdfRequest {
|
||||
/// PDF file URL
|
||||
#[validate(url)]
|
||||
pub url: String,
|
||||
|
||||
/// Conversion options
|
||||
#[serde(default)]
|
||||
pub options: PdfOptions,
|
||||
|
||||
/// Webhook URL for completion notification
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
#[validate(url)]
|
||||
pub webhook_url: Option<String>,
|
||||
|
||||
/// Request metadata
|
||||
#[serde(default)]
|
||||
pub metadata: RequestMetadata,
|
||||
}
|
||||
|
||||
/// PDF processing options
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct PdfOptions {
|
||||
/// Output format
|
||||
#[serde(default = "default_format")]
|
||||
pub format: String,
|
||||
|
||||
/// Enable OCR
|
||||
#[serde(default)]
|
||||
pub enable_ocr: bool,
|
||||
|
||||
/// Include images
|
||||
#[serde(default = "default_true")]
|
||||
pub include_images: bool,
|
||||
|
||||
/// Page range (e.g., "1-5")
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub page_range: Option<String>,
|
||||
}
|
||||
|
||||
fn default_format() -> String {
|
||||
"mmd".to_string()
|
||||
}
|
||||
|
||||
fn default_true() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Request metadata
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct RequestMetadata {
|
||||
/// Output formats
|
||||
#[serde(default = "default_formats")]
|
||||
pub formats: Vec<String>,
|
||||
|
||||
/// Include confidence scores
|
||||
#[serde(default)]
|
||||
pub include_confidence: bool,
|
||||
|
||||
/// Enable math mode
|
||||
#[serde(default = "default_true")]
|
||||
pub enable_math: bool,
|
||||
|
||||
/// Language hint
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub language: Option<String>,
|
||||
}
|
||||
|
||||
fn default_formats() -> Vec<String> {
|
||||
vec!["text".to_string()]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_text_request_validation() {
|
||||
let request = TextRequest {
|
||||
src: None,
|
||||
base64: Some("SGVsbG8gV29ybGQ=".to_string()),
|
||||
url: None,
|
||||
metadata: RequestMetadata::default(),
|
||||
};
|
||||
|
||||
assert!(request.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strokes_request_validation() {
|
||||
let request = StrokesRequest {
|
||||
strokes: vec![Stroke {
|
||||
x: vec![0.0, 1.0, 2.0],
|
||||
y: vec![0.0, 1.0, 0.0],
|
||||
t: None,
|
||||
}],
|
||||
metadata: RequestMetadata::default(),
|
||||
};
|
||||
|
||||
assert!(request.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_strokes_validation() {
|
||||
let request = StrokesRequest {
|
||||
strokes: vec![],
|
||||
metadata: RequestMetadata::default(),
|
||||
};
|
||||
|
||||
assert!(request.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pdf_request_validation() {
|
||||
let request = PdfRequest {
|
||||
url: "https://example.com/document.pdf".to_string(),
|
||||
options: PdfOptions::default(),
|
||||
webhook_url: None,
|
||||
metadata: RequestMetadata::default(),
|
||||
};
|
||||
|
||||
assert!(request.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_url() {
|
||||
let request = PdfRequest {
|
||||
url: "not-a-url".to_string(),
|
||||
options: PdfOptions::default(),
|
||||
webhook_url: None,
|
||||
metadata: RequestMetadata::default(),
|
||||
};
|
||||
|
||||
assert!(request.validate().is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
use axum::{
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::jobs::JobStatus;
|
||||
|
||||
/// Standard text/OCR response
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TextResponse {
|
||||
/// Unique request identifier
|
||||
pub request_id: String,
|
||||
|
||||
/// Recognized text
|
||||
pub text: String,
|
||||
|
||||
/// Confidence score (0.0 - 1.0)
|
||||
pub confidence: f64,
|
||||
|
||||
/// LaTeX output (if requested)
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub latex: Option<String>,
|
||||
|
||||
/// MathML output (if requested)
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub mathml: Option<String>,
|
||||
|
||||
/// HTML output (if requested)
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub html: Option<String>,
|
||||
}
|
||||
|
||||
/// PDF processing response
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PdfResponse {
|
||||
/// PDF job identifier
|
||||
pub pdf_id: String,
|
||||
|
||||
/// Current job status
|
||||
pub status: JobStatus,
|
||||
|
||||
/// Status message
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub message: Option<String>,
|
||||
|
||||
/// Processing result (when completed)
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub result: Option<String>,
|
||||
|
||||
/// Error details (if failed)
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// Error response
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ErrorResponse {
|
||||
/// Error code
|
||||
pub error_code: String,
|
||||
|
||||
/// Human-readable error message
|
||||
pub message: String,
|
||||
|
||||
/// HTTP status code
|
||||
#[serde(skip)]
|
||||
pub status: StatusCode,
|
||||
}
|
||||
|
||||
impl ErrorResponse {
|
||||
/// Create a validation error response
|
||||
pub fn validation_error(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
error_code: "VALIDATION_ERROR".to_string(),
|
||||
message: message.into(),
|
||||
status: StatusCode::BAD_REQUEST,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an unauthorized error response
|
||||
pub fn unauthorized(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
error_code: "UNAUTHORIZED".to_string(),
|
||||
message: message.into(),
|
||||
status: StatusCode::UNAUTHORIZED,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a not found error response
|
||||
pub fn not_found(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
error_code: "NOT_FOUND".to_string(),
|
||||
message: message.into(),
|
||||
status: StatusCode::NOT_FOUND,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a rate limit error response
|
||||
pub fn rate_limited(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
error_code: "RATE_LIMIT_EXCEEDED".to_string(),
|
||||
message: message.into(),
|
||||
status: StatusCode::TOO_MANY_REQUESTS,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an internal error response
|
||||
pub fn internal_error(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
error_code: "INTERNAL_ERROR".to_string(),
|
||||
message: message.into(),
|
||||
status: StatusCode::INTERNAL_SERVER_ERROR,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a service unavailable error response
|
||||
/// Used when the service is not fully configured (e.g., missing models)
|
||||
pub fn service_unavailable(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
error_code: "SERVICE_UNAVAILABLE".to_string(),
|
||||
message: message.into(),
|
||||
status: StatusCode::SERVICE_UNAVAILABLE,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a not implemented error response
|
||||
pub fn not_implemented(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
error_code: "NOT_IMPLEMENTED".to_string(),
|
||||
message: message.into(),
|
||||
status: StatusCode::NOT_IMPLEMENTED,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoResponse for ErrorResponse {
|
||||
fn into_response(self) -> Response {
|
||||
let status = self.status;
|
||||
(status, Json(self)).into_response()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_text_response_serialization() {
|
||||
let response = TextResponse {
|
||||
request_id: "test-123".to_string(),
|
||||
text: "Hello World".to_string(),
|
||||
confidence: 0.95,
|
||||
latex: Some("x^2".to_string()),
|
||||
mathml: None,
|
||||
html: None,
|
||||
};
|
||||
|
||||
let json = serde_json::to_string(&response).unwrap();
|
||||
assert!(json.contains("request_id"));
|
||||
assert!(json.contains("test-123"));
|
||||
assert!(!json.contains("mathml"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_error_response_creation() {
|
||||
let error = ErrorResponse::validation_error("Invalid input");
|
||||
assert_eq!(error.status, StatusCode::BAD_REQUEST);
|
||||
assert_eq!(error.error_code, "VALIDATION_ERROR");
|
||||
|
||||
let error = ErrorResponse::unauthorized("Invalid credentials");
|
||||
assert_eq!(error.status, StatusCode::UNAUTHORIZED);
|
||||
|
||||
let error = ErrorResponse::rate_limited("Too many requests");
|
||||
assert_eq!(error.status, StatusCode::TOO_MANY_REQUESTS);
|
||||
}
|
||||
}
|
||||
+103
@@ -0,0 +1,103 @@
|
||||
use axum::{
|
||||
routing::{delete, get, post},
|
||||
Router,
|
||||
};
|
||||
use tower::ServiceBuilder;
|
||||
use tower_http::{
|
||||
compression::CompressionLayer,
|
||||
cors::CorsLayer,
|
||||
trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer},
|
||||
};
|
||||
use tracing::Level;
|
||||
|
||||
use super::{
|
||||
handlers::{
|
||||
convert_document, delete_pdf_job, get_health, get_ocr_results, get_ocr_usage,
|
||||
get_pdf_status, process_latex, process_pdf, process_strokes, process_text,
|
||||
stream_pdf_results,
|
||||
},
|
||||
middleware::{auth_middleware, rate_limit_middleware},
|
||||
state::AppState,
|
||||
};
|
||||
|
||||
/// Create the main application router with all routes and middleware
|
||||
pub fn router(state: AppState) -> Router {
|
||||
// API v3 routes
|
||||
let api_routes = Router::new()
|
||||
// Image processing
|
||||
.route("/v3/text", post(process_text))
|
||||
// Digital ink processing
|
||||
.route("/v3/strokes", post(process_strokes))
|
||||
// Legacy equation processing
|
||||
.route("/v3/latex", post(process_latex))
|
||||
// Async PDF processing
|
||||
.route("/v3/pdf", post(process_pdf))
|
||||
.route("/v3/pdf/:id", get(get_pdf_status))
|
||||
.route("/v3/pdf/:id", delete(delete_pdf_job))
|
||||
.route("/v3/pdf/:id/stream", get(stream_pdf_results))
|
||||
// Document conversion
|
||||
.route("/v3/converter", post(convert_document))
|
||||
// History and usage
|
||||
.route("/v3/ocr-results", get(get_ocr_results))
|
||||
.route("/v3/ocr-usage", get(get_ocr_usage))
|
||||
// Apply auth and rate limiting to all API routes
|
||||
.layer(
|
||||
ServiceBuilder::new()
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
state.clone(),
|
||||
auth_middleware,
|
||||
))
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
state.clone(),
|
||||
rate_limit_middleware,
|
||||
)),
|
||||
);
|
||||
|
||||
// Health check (no auth required)
|
||||
let health_routes = Router::new().route("/health", get(get_health));
|
||||
|
||||
// Combine all routes
|
||||
Router::new()
|
||||
.merge(api_routes)
|
||||
.merge(health_routes)
|
||||
.layer(
|
||||
ServiceBuilder::new()
|
||||
// Tracing layer
|
||||
.layer(
|
||||
TraceLayer::new_for_http()
|
||||
.make_span_with(DefaultMakeSpan::new().level(Level::INFO))
|
||||
.on_response(DefaultOnResponse::new().level(Level::INFO)),
|
||||
)
|
||||
// CORS layer
|
||||
.layer(CorsLayer::permissive())
|
||||
// Compression layer
|
||||
.layer(CompressionLayer::new()),
|
||||
)
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use tower::ServiceExt;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_health_endpoint() {
|
||||
let state = AppState::new();
|
||||
let app = router(state);
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri("/health")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
}
|
||||
+148
@@ -0,0 +1,148 @@
|
||||
use moka::future::Cache;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{
|
||||
jobs::JobQueue,
|
||||
middleware::{create_rate_limiter, AppRateLimiter},
|
||||
};
|
||||
|
||||
/// Shared application state
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
/// Job queue for async PDF processing
|
||||
pub job_queue: Arc<JobQueue>,
|
||||
|
||||
/// Result cache
|
||||
pub cache: Cache<String, String>,
|
||||
|
||||
/// Rate limiter
|
||||
pub rate_limiter: AppRateLimiter,
|
||||
|
||||
/// Whether authentication is enabled
|
||||
pub auth_enabled: bool,
|
||||
|
||||
/// Map of app_id -> hashed API key
|
||||
/// Keys should be stored as SHA-256 hashes, never in plaintext
|
||||
pub api_keys: Arc<HashMap<String, String>>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
/// Create a new application state instance with authentication disabled
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
job_queue: Arc::new(JobQueue::new()),
|
||||
cache: create_cache(),
|
||||
rate_limiter: create_rate_limiter(),
|
||||
auth_enabled: false,
|
||||
api_keys: Arc::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create state with custom configuration
|
||||
pub fn with_config(max_jobs: usize, cache_size: u64) -> Self {
|
||||
Self {
|
||||
job_queue: Arc::new(JobQueue::with_capacity(max_jobs)),
|
||||
cache: Cache::builder()
|
||||
.max_capacity(cache_size)
|
||||
.time_to_live(Duration::from_secs(3600))
|
||||
.time_to_idle(Duration::from_secs(600))
|
||||
.build(),
|
||||
rate_limiter: create_rate_limiter(),
|
||||
auth_enabled: false,
|
||||
api_keys: Arc::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create state with authentication enabled
|
||||
pub fn with_auth(api_keys: HashMap<String, String>) -> Self {
|
||||
// Hash all provided API keys
|
||||
let hashed_keys: HashMap<String, String> = api_keys
|
||||
.into_iter()
|
||||
.map(|(app_id, key)| (app_id, hash_api_key(&key)))
|
||||
.collect();
|
||||
|
||||
Self {
|
||||
job_queue: Arc::new(JobQueue::new()),
|
||||
cache: create_cache(),
|
||||
rate_limiter: create_rate_limiter(),
|
||||
auth_enabled: true,
|
||||
api_keys: Arc::new(hashed_keys),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add an API key (hashes the key before storing)
|
||||
pub fn add_api_key(&mut self, app_id: String, api_key: &str) {
|
||||
let hashed = hash_api_key(api_key);
|
||||
Arc::make_mut(&mut self.api_keys).insert(app_id, hashed);
|
||||
self.auth_enabled = true;
|
||||
}
|
||||
|
||||
/// Enable or disable authentication
|
||||
pub fn set_auth_enabled(&mut self, enabled: bool) {
|
||||
self.auth_enabled = enabled;
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AppState {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Hash an API key using SHA-256
|
||||
fn hash_api_key(key: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(key.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
/// Create a cache with default configuration
|
||||
fn create_cache() -> Cache<String, String> {
|
||||
Cache::builder()
|
||||
// Max 10,000 entries
|
||||
.max_capacity(10_000)
|
||||
// Time to live: 1 hour
|
||||
.time_to_live(Duration::from_secs(3600))
|
||||
// Time to idle: 10 minutes
|
||||
.time_to_idle(Duration::from_secs(600))
|
||||
.build()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_state_creation() {
|
||||
let state = AppState::new();
|
||||
assert!(Arc::strong_count(&state.job_queue) >= 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_state_with_config() {
|
||||
let state = AppState::with_config(100, 5000);
|
||||
assert!(Arc::strong_count(&state.job_queue) >= 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_operations() {
|
||||
let state = AppState::new();
|
||||
|
||||
// Insert value
|
||||
state
|
||||
.cache
|
||||
.insert("key1".to_string(), "value1".to_string())
|
||||
.await;
|
||||
|
||||
// Retrieve value
|
||||
let value = state.cache.get(&"key1".to_string()).await;
|
||||
assert_eq!(value, Some("value1".to_string()));
|
||||
|
||||
// Non-existent key
|
||||
let missing = state.cache.get(&"missing".to_string()).await;
|
||||
assert_eq!(missing, None);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user