feat(infra): add tools service with document scanner, image & PDF tools
Implement self-hosted document scanner and media processing tools as an alternative to CamScanner/ilovepdf without third-party uploads. Backend: Rust Axum gateway + worker pool with NATS JetStream queue Frontend: Next.js 16 + shadcn/ui + Tailwind v4 + Framer Motion Pipeline: Canny edge detection -> DLT homography warp -> Sauvola binarization -> Hough deskew -> Tesseract OCR -> searchable PDF Phase 1 (MVP) delivers: - Document scanner with perspective correction and OCR - Image compress/resize/convert tools - PDF merge/split/compress tools - Real-time WebSocket progress updates - Rate limiting, auto-cleanup, Prometheus metrics - Full CI/CD pipeline with Docker multi-stage build Co-Authored-By: Kilo <kilo@kilo.ai>
This commit is contained in:
asepharyana
@@ -0,0 +1,25 @@
|
||||
[package]
|
||||
name = "tools-gateway"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
tools-common = { path = "../common" }
|
||||
|
||||
axum = { version = "0.8", features = ["multipart", "ws"] }
|
||||
tower = "0.5"
|
||||
tower-http = { version = "0.6", features = ["cors", "trace", "limit"] }
|
||||
tokio.workspace = true
|
||||
tokio-util = { version = "0.7", features = ["io"] }
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
uuid.workspace = true
|
||||
chrono.workspace = true
|
||||
tracing.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
async-nats.workspace = true
|
||||
redis.workspace = true
|
||||
thiserror.workspace = true
|
||||
anyhow.workspace = true
|
||||
futures = "0.3"
|
||||
@@ -0,0 +1,76 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// Application configuration loaded from environment variables.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AppConfig {
|
||||
pub port: u16,
|
||||
pub nats_url: String,
|
||||
pub redis_url: String,
|
||||
pub storage_path: PathBuf,
|
||||
pub max_file_size_mb: u64,
|
||||
pub job_ttl_seconds: u64,
|
||||
pub rate_limit_per_minute: u32,
|
||||
pub rust_log: String,
|
||||
}
|
||||
|
||||
impl AppConfig {
|
||||
/// Load configuration from environment variables with sensible defaults.
|
||||
pub fn from_env() -> Self {
|
||||
Self {
|
||||
port: env_or_default("GATEWAY_PORT", "3001")
|
||||
.parse()
|
||||
.unwrap_or(3001),
|
||||
nats_url: env_or_default("NATS_URL", "nats://localhost:4222"),
|
||||
redis_url: env_or_default("REDIS_URL", "redis://localhost:6379"),
|
||||
storage_path: PathBuf::from(env_or_default("STORAGE_PATH", "/data/tools")),
|
||||
max_file_size_mb: env_or_default("MAX_FILE_SIZE_MB", "50")
|
||||
.parse()
|
||||
.unwrap_or(50),
|
||||
job_ttl_seconds: env_or_default("JOB_TTL_SECONDS", "3600")
|
||||
.parse()
|
||||
.unwrap_or(3600),
|
||||
rate_limit_per_minute: env_or_default("RATE_LIMIT_PER_MINUTE", "30")
|
||||
.parse()
|
||||
.unwrap_or(30),
|
||||
rust_log: env_or_default("RUST_LOG", "info"),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn max_file_size_bytes(&self) -> u64 {
|
||||
self.max_file_size_mb * 1024 * 1024
|
||||
}
|
||||
}
|
||||
|
||||
fn env_or_default(key: &str, default: &str) -> String {
|
||||
std::env::var(key).unwrap_or_else(|_| default.to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_default_config() {
|
||||
let config = AppConfig::from_env();
|
||||
assert_eq!(config.port, 3001);
|
||||
assert_eq!(config.nats_url, "nats://localhost:4222");
|
||||
assert_eq!(config.redis_url, "redis://localhost:6379");
|
||||
assert_eq!(config.max_file_size_mb, 50);
|
||||
assert_eq!(config.job_ttl_seconds, 3600);
|
||||
assert_eq!(config.rate_limit_per_minute, 30);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_file_size_bytes() {
|
||||
let config = AppConfig::from_env();
|
||||
assert_eq!(config.max_file_size_bytes(), 50 * 1024 * 1024);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_env_override() {
|
||||
std::env::set_var("GATEWAY_PORT", "9999");
|
||||
let config = AppConfig::from_env();
|
||||
assert_eq!(config.port, 9999);
|
||||
std::env::remove_var("GATEWAY_PORT");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
use tower_http::cors::{Any, CorsLayer};
|
||||
use tower_http::limit::RequestBodyLimitLayer;
|
||||
use tower_http::trace::TraceLayer;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
mod config;
|
||||
mod metrics;
|
||||
mod middleware;
|
||||
mod nats;
|
||||
mod redis;
|
||||
mod routes;
|
||||
|
||||
use config::AppConfig;
|
||||
use metrics::Metrics;
|
||||
use routes::health::AppState;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
// Load config
|
||||
let config = AppConfig::from_env();
|
||||
|
||||
// Init logging
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(EnvFilter::new(&config.rust_log))
|
||||
.init();
|
||||
|
||||
tracing::info!("Starting tools-gateway...");
|
||||
|
||||
// Init Redis client
|
||||
let redis_client = redis::create_client(&config.redis_url)
|
||||
.expect("Failed to create Redis client");
|
||||
tracing::info!("Redis client created for {}", config.redis_url);
|
||||
|
||||
// Init NATS connection
|
||||
let nats = nats::publisher::NatsPublisher::connect(&config.nats_url)
|
||||
.await
|
||||
.expect("Failed to connect to NATS");
|
||||
tracing::info!("Connected to NATS at {}", config.nats_url);
|
||||
|
||||
// Ensure NATS streams exist
|
||||
if let Err(e) = ensure_nats_streams(&nats).await {
|
||||
tracing::warn!("Failed to create NATS streams: {}", e);
|
||||
}
|
||||
|
||||
// Init metrics
|
||||
let metrics = Metrics::new();
|
||||
|
||||
// Shared state
|
||||
let state = Arc::new(AppState {
|
||||
redis: redis_client,
|
||||
nats,
|
||||
config: config.clone(),
|
||||
metrics,
|
||||
});
|
||||
|
||||
// Build router
|
||||
let app = Router::new()
|
||||
.route("/api/upload", post(routes::upload::upload_handler))
|
||||
.route("/api/job/{id}", get(routes::job::job_status_handler))
|
||||
.route(
|
||||
"/api/job/{id}/preview",
|
||||
get(routes::job::job_preview_handler),
|
||||
)
|
||||
.route("/api/job/{id}/ws", get(routes::ws::ws_handler))
|
||||
.route("/api/download/{id}", get(routes::download::download_handler))
|
||||
.route("/health", get(routes::health::health_handler))
|
||||
.route("/metrics", get(routes::health::metrics_handler))
|
||||
.layer(TraceLayer::new_for_http())
|
||||
.layer(CorsLayer::new().allow_origin(Any))
|
||||
.layer(RequestBodyLimitLayer::new(
|
||||
((config.max_file_size_mb + 1) * 1024 * 1024) as usize,
|
||||
))
|
||||
.with_state(state);
|
||||
|
||||
// Start server
|
||||
let addr = SocketAddr::from(([0, 0, 0, 0], config.port));
|
||||
tracing::info!("Gateway listening on {}", addr);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(shutdown_signal())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
/// Ensure required NATS JetStream streams exist.
|
||||
async fn ensure_nats_streams(
|
||||
nats: &async_nats::Client,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let js = async_nats::jetstream::new(nats.clone());
|
||||
|
||||
match js
|
||||
.get_or_create_stream(tools_common::nats::jobs_stream_config())
|
||||
.await
|
||||
{
|
||||
Ok(_) => tracing::info!("NATS stream 'tools-jobs' ready"),
|
||||
Err(e) => tracing::warn!("Failed to create tools-jobs stream: {}", e),
|
||||
}
|
||||
|
||||
match js
|
||||
.get_or_create_stream(tools_common::nats::progress_stream_config())
|
||||
.await
|
||||
{
|
||||
Ok(_) => tracing::info!("NATS stream 'tools-progress' ready"),
|
||||
Err(e) => tracing::warn!("Failed to create tools-progress stream: {}", e),
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Handle graceful shutdown on SIGINT/SIGTERM.
|
||||
async fn shutdown_signal() {
|
||||
let ctrl_c = async {
|
||||
tokio::signal::ctrl_c()
|
||||
.await
|
||||
.expect("Failed to install Ctrl+C handler");
|
||||
};
|
||||
|
||||
#[cfg(unix)]
|
||||
let terminate = async {
|
||||
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
|
||||
.expect("Failed to install SIGTERM handler")
|
||||
.recv()
|
||||
.await;
|
||||
};
|
||||
|
||||
#[cfg(not(unix))]
|
||||
let terminate = std::future::pending::<()>();
|
||||
|
||||
tokio::select! {
|
||||
_ = ctrl_c => {},
|
||||
_ = terminate => {},
|
||||
}
|
||||
|
||||
tracing::info!("Shutting down gateway...");
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::Mutex;
|
||||
|
||||
/// Simple Prometheus metrics collector.
|
||||
pub struct Metrics {
|
||||
/// Counter: tools_jobs_total{tool, status}
|
||||
jobs_total: Mutex<HashMap<(String, String), AtomicU64>>,
|
||||
/// Counter: tools_uploaded_files_total{tool, status}
|
||||
uploaded_files_total: Mutex<HashMap<(String, String), AtomicU64>>,
|
||||
/// Histogram buckets for processing duration (ms)
|
||||
duration_buckets: Vec<f64>,
|
||||
/// Histogram: tools_processing_duration_ms{tool}
|
||||
duration_histogram: Mutex<HashMap<String, Vec<AtomicU64>>>,
|
||||
/// Gauge: tools_queue_depth{tool}
|
||||
queue_depth: Mutex<HashMap<String, AtomicU64>>,
|
||||
/// Counter: tools_rate_limit_hits{tool}
|
||||
rate_limit_hits: Mutex<HashMap<String, AtomicU64>>,
|
||||
/// Counter: cleanup deleted files
|
||||
cleanup_deleted_files: AtomicU64,
|
||||
}
|
||||
|
||||
impl Metrics {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
jobs_total: Mutex::new(HashMap::new()),
|
||||
uploaded_files_total: Mutex::new(HashMap::new()),
|
||||
duration_buckets: vec![
|
||||
100.0, 250.0, 500.0, 1000.0, 2000.0, 4000.0, 8000.0, 16000.0, 32000.0,
|
||||
],
|
||||
duration_histogram: Mutex::new(HashMap::new()),
|
||||
queue_depth: Mutex::new(HashMap::new()),
|
||||
rate_limit_hits: Mutex::new(HashMap::new()),
|
||||
cleanup_deleted_files: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn increment_jobs_total(&self, tool: &str, status: &str) {
|
||||
if let Ok(mut map) = self.jobs_total.lock() {
|
||||
map.entry((tool.to_string(), status.to_string()))
|
||||
.or_insert_with(|| AtomicU64::new(0))
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn increment_uploaded_files(&self, tool: &str, status: &str) {
|
||||
if let Ok(mut map) = self.uploaded_files_total.lock() {
|
||||
map.entry((tool.to_string(), status.to_string()))
|
||||
.or_insert_with(|| AtomicU64::new(0))
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(unused)]
|
||||
pub fn record_duration(&self, tool: &str, duration_ms: f64) {
|
||||
if let Ok(mut map) = self.duration_histogram.lock() {
|
||||
let entry = map
|
||||
.entry(tool.to_string())
|
||||
.or_insert_with(|| {
|
||||
(0..self.duration_buckets.len())
|
||||
.map(|_| AtomicU64::new(0))
|
||||
.collect()
|
||||
});
|
||||
for (i, bucket) in self.duration_buckets.iter().enumerate() {
|
||||
if duration_ms <= *bucket {
|
||||
if let Some(b) = entry.get(i) {
|
||||
b.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_queue_depth(&self, tool: &str, depth: u64) {
|
||||
if let Ok(mut map) = self.queue_depth.lock() {
|
||||
map.entry(tool.to_string())
|
||||
.or_insert_with(|| AtomicU64::new(0))
|
||||
.store(depth, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(unused)]
|
||||
pub fn increment_rate_limit_hits(&self, tool: &str) {
|
||||
if let Ok(mut map) = self.rate_limit_hits.lock() {
|
||||
map.entry(tool.to_string())
|
||||
.or_insert_with(|| AtomicU64::new(0))
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(unused)]
|
||||
pub fn increment_cleanup_deleted(&self) {
|
||||
self.cleanup_deleted_files.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Format all metrics as Prometheus text format.
|
||||
pub fn format(&self) -> String {
|
||||
let mut output = String::new();
|
||||
|
||||
output.push_str("# HELP tools_jobs_total Total jobs processed\n");
|
||||
output.push_str("# TYPE tools_jobs_total counter\n");
|
||||
if let Ok(map) = self.jobs_total.lock() {
|
||||
for ((tool, status), count) in map.iter() {
|
||||
let val = count.load(Ordering::Relaxed);
|
||||
output.push_str(&format!(
|
||||
"tools_jobs_total{{tool=\"{}\",status=\"{}\"}} {}\n",
|
||||
tool, status, val
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
output.push_str("# HELP tools_uploaded_files_total Total uploaded files\n");
|
||||
output.push_str("# TYPE tools_uploaded_files_total counter\n");
|
||||
if let Ok(map) = self.uploaded_files_total.lock() {
|
||||
for ((tool, status), count) in map.iter() {
|
||||
let val = count.load(Ordering::Relaxed);
|
||||
output.push_str(&format!(
|
||||
"tools_uploaded_files_total{{tool=\"{}\",status=\"{}\"}} {}\n",
|
||||
tool, status, val
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
output.push_str("# HELP tools_processing_duration_ms Processing duration histogram\n");
|
||||
output.push_str("# TYPE tools_processing_duration_ms histogram\n");
|
||||
if let Ok(map) = self.duration_histogram.lock() {
|
||||
for (tool, buckets) in map.iter() {
|
||||
for (i, bucket) in self.duration_buckets.iter().enumerate() {
|
||||
if let Some(b) = buckets.get(i) {
|
||||
let val = b.load(Ordering::Relaxed);
|
||||
if val > 0 {
|
||||
output.push_str(&format!(
|
||||
"tools_processing_duration_ms_bucket{{tool=\"{}\",le=\"{}\"}} {}\n",
|
||||
tool, bucket, val
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
output.push_str("# HELP tools_queue_depth Current queue depth\n");
|
||||
output.push_str("# TYPE tools_queue_depth gauge\n");
|
||||
if let Ok(map) = self.queue_depth.lock() {
|
||||
for (tool, depth) in map.iter() {
|
||||
let val = depth.load(Ordering::Relaxed);
|
||||
output.push_str(&format!(
|
||||
"tools_queue_depth{{tool=\"{}\"}} {}\n",
|
||||
tool, val
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
output.push_str("# HELP tools_rate_limit_hits Total rate limit violations\n");
|
||||
output.push_str("# TYPE tools_rate_limit_hits counter\n");
|
||||
if let Ok(map) = self.rate_limit_hits.lock() {
|
||||
for (tool, count) in map.iter() {
|
||||
let val = count.load(Ordering::Relaxed);
|
||||
output.push_str(&format!(
|
||||
"tools_rate_limit_hits{{tool=\"{}\"}} {}\n",
|
||||
tool, val
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
output.push_str("# HELP tools_cleanup_deleted_files Total files deleted by cleanup\n");
|
||||
output.push_str("# TYPE tools_cleanup_deleted_files counter\n");
|
||||
output.push_str(&format!(
|
||||
"tools_cleanup_deleted_files {}\n",
|
||||
self.cleanup_deleted_files.load(Ordering::Relaxed)
|
||||
));
|
||||
|
||||
output
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Metrics {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
use axum::{
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
/// Unified JSON error response format.
|
||||
#[derive(Debug)]
|
||||
pub struct AppError {
|
||||
pub status_code: StatusCode,
|
||||
pub code: String,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
impl AppError {
|
||||
pub fn bad_request(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
status_code: StatusCode::BAD_REQUEST,
|
||||
code: "bad_request".to_string(),
|
||||
message: message.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn not_found(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
status_code: StatusCode::NOT_FOUND,
|
||||
code: "not_found".to_string(),
|
||||
message: message.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn too_large(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
status_code: StatusCode::PAYLOAD_TOO_LARGE,
|
||||
code: "file_too_large".to_string(),
|
||||
message: message.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn rate_limited(retry_after: u64) -> Self {
|
||||
Self {
|
||||
status_code: StatusCode::TOO_MANY_REQUESTS,
|
||||
code: "rate_limit_exceeded".to_string(),
|
||||
message: format!("Rate limit exceeded. Retry after {} seconds", retry_after),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn internal(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
status_code: StatusCode::INTERNAL_SERVER_ERROR,
|
||||
code: "internal_error".to_string(),
|
||||
message: message.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoResponse for AppError {
|
||||
fn into_response(self) -> Response {
|
||||
let body = json!({
|
||||
"error": self.message,
|
||||
"code": self.code,
|
||||
});
|
||||
(self.status_code, Json(body)).into_response()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod error_handler;
|
||||
pub mod request_id;
|
||||
@@ -0,0 +1,73 @@
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use axum::{extract::Request, response::Response};
|
||||
use tower::{Layer, Service};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Middleware that adds a unique X-Request-Id header to every request.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct RequestIdLayer;
|
||||
|
||||
impl<S> Layer<S> for RequestIdLayer {
|
||||
type Service = RequestIdMiddleware<S>;
|
||||
|
||||
fn layer(&self, inner: S) -> Self::Service {
|
||||
RequestIdMiddleware {
|
||||
inner,
|
||||
counter: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RequestIdMiddleware<S> {
|
||||
inner: S,
|
||||
counter: AtomicU64,
|
||||
}
|
||||
|
||||
impl<S: Clone> Clone for RequestIdMiddleware<S> {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
inner: self.inner.clone(),
|
||||
counter: AtomicU64::new(self.counter.load(Ordering::Relaxed)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for RequestIdMiddleware<S>
|
||||
where
|
||||
S: Service<Request<ReqBody>, Response = Response<ResBody>>,
|
||||
S::Future: Send + 'static,
|
||||
S::Error: 'static,
|
||||
ReqBody: Send + 'static,
|
||||
ResBody: Default + Send + 'static,
|
||||
{
|
||||
type Response = Response<ResBody>;
|
||||
type Error = S::Error;
|
||||
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
|
||||
|
||||
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
self.inner.poll_ready(cx)
|
||||
}
|
||||
|
||||
fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
|
||||
let request_id = Uuid::new_v4().to_string();
|
||||
let (mut parts, body) = req.into_parts();
|
||||
parts
|
||||
.headers
|
||||
.insert("x-request-id", request_id.parse().unwrap());
|
||||
|
||||
let req = Request::from_parts(parts, body);
|
||||
let fut = self.inner.call(req);
|
||||
|
||||
Box::pin(async move {
|
||||
let mut response: Response<ResBody> = fut.await?;
|
||||
response
|
||||
.headers_mut()
|
||||
.insert("x-request-id", request_id.parse().unwrap());
|
||||
Ok(response)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod publisher;
|
||||
@@ -0,0 +1,64 @@
|
||||
use async_nats::Client;
|
||||
use tools_common::error::NatsError;
|
||||
use tools_common::nats;
|
||||
use tools_common::types::{Job, JobProgress, Tool};
|
||||
|
||||
/// NATS publisher for job and progress messages.
|
||||
pub struct NatsPublisher;
|
||||
|
||||
impl NatsPublisher {
|
||||
/// Connect to NATS server.
|
||||
pub async fn connect(url: &str) -> Result<Client, NatsError> {
|
||||
async_nats::connect(url)
|
||||
.await
|
||||
.map_err(|e| NatsError::Connection(e.to_string()))
|
||||
}
|
||||
|
||||
/// Publish a job to the appropriate NATS subject.
|
||||
pub async fn publish_job(nats: &Client, tool: &Tool, job: &Job) -> Result<(), NatsError> {
|
||||
let prefix = tool.subject_prefix();
|
||||
let subject = nats::job_subject(prefix, &job.id.to_string());
|
||||
let payload = serde_json::to_vec(job)
|
||||
.map_err(|e| NatsError::Publish(e.to_string()))?;
|
||||
|
||||
nats.publish(subject, payload.into())
|
||||
.await
|
||||
.map_err(|e| NatsError::Publish(e.to_string()))?;
|
||||
|
||||
tracing::debug!(
|
||||
job_id = %job.id,
|
||||
tool = %tool.as_str(),
|
||||
"Published job to NATS"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Publish a progress update to the NATS progress subject.
|
||||
pub async fn publish_progress(
|
||||
nats: &Client,
|
||||
progress: &JobProgress,
|
||||
) -> Result<(), NatsError> {
|
||||
let tool_prefix = ""; // We need the tool from somewhere — stored in progress
|
||||
let subject = format!("tools.*.progress.{}", progress.job_id);
|
||||
let payload = serde_json::to_vec(progress)
|
||||
.map_err(|e| NatsError::Publish(e.to_string()))?;
|
||||
|
||||
nats.publish(subject, payload.into())
|
||||
.await
|
||||
.map_err(|e| NatsError::Publish(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Subscribe to NATS progress updates for a specific job.
|
||||
pub async fn subscribe_progress(
|
||||
nats: &Client,
|
||||
job_id: &str,
|
||||
) -> Result<async_nats::Subscriber, NatsError> {
|
||||
let subject = format!("tools.*.progress.{}", job_id);
|
||||
nats.subscribe(subject)
|
||||
.await
|
||||
.map_err(|e| NatsError::Subscribe(e.to_string()))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
use redis::{AsyncCommands, RedisError};
|
||||
use uuid::Uuid;
|
||||
|
||||
use tools_common::types::Job;
|
||||
|
||||
/// Repository for job CRUD operations on Redis.
|
||||
pub struct JobRepository;
|
||||
|
||||
impl JobRepository {
|
||||
/// Create a new job record in Redis with TTL.
|
||||
pub async fn create(
|
||||
conn: &mut impl AsyncCommands,
|
||||
job: &Job,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
let key = format!("job:{}", job.id);
|
||||
let json = serde_json::to_string(job)?;
|
||||
let _: () = conn
|
||||
.set_ex(key, json, job.ttl_seconds)
|
||||
.await
|
||||
.map_err(|e: RedisError| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get a job by ID from Redis.
|
||||
pub async fn get(
|
||||
conn: &mut impl AsyncCommands,
|
||||
job_id: Uuid,
|
||||
) -> Result<Job, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let key = format!("job:{}", job_id);
|
||||
let json: String = conn.get(&key).await.map_err(|_| {
|
||||
Box::new(std::io::Error::new(
|
||||
std::io::ErrorKind::NotFound,
|
||||
format!("Job {} not found", job_id),
|
||||
)) as Box<dyn std::error::Error + Send + Sync>
|
||||
})?;
|
||||
let job: Job = serde_json::from_str(&json)?;
|
||||
Ok(job)
|
||||
}
|
||||
|
||||
/// Update the status of a job in Redis and refresh TTL.
|
||||
pub async fn update_status(
|
||||
conn: &mut impl AsyncCommands,
|
||||
job_id: Uuid,
|
||||
status: &tools_common::types::JobStatus,
|
||||
result_path: Option<String>,
|
||||
ttl_seconds: u64,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
let key = format!("job:{}", job_id);
|
||||
let json: String = conn
|
||||
.get(&key)
|
||||
.await
|
||||
.map_err(|e: RedisError| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?;
|
||||
let mut job: Job = serde_json::from_str(&json)?;
|
||||
job.status = status.clone();
|
||||
if let Some(path) = result_path {
|
||||
job.result_path = Some(path);
|
||||
}
|
||||
let json = serde_json::to_string(&job)?;
|
||||
let _: () = conn
|
||||
.set_ex(key, json, ttl_seconds)
|
||||
.await
|
||||
.map_err(|e: RedisError| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Delete a job from Redis.
|
||||
pub async fn delete(
|
||||
conn: &mut impl AsyncCommands,
|
||||
job_id: Uuid,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
let key = format!("job:{}", job_id);
|
||||
let _: usize = conn
|
||||
.del(key)
|
||||
.await
|
||||
.map_err(|e: RedisError| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
pub mod job;
|
||||
pub mod ratelimit;
|
||||
|
||||
use redis::Client;
|
||||
|
||||
/// Create a Redis client.
|
||||
pub fn create_client(url: &str) -> Result<Client, redis::RedisError> {
|
||||
Client::open(url)
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
use redis::AsyncCommands;
|
||||
|
||||
/// Sliding window rate limiter using Redis sorted sets.
|
||||
pub struct RateLimiter;
|
||||
|
||||
impl RateLimiter {
|
||||
/// Check if a request is within the rate limit.
|
||||
pub async fn check(
|
||||
conn: &mut impl AsyncCommands,
|
||||
ip: &str,
|
||||
tool: &str,
|
||||
max_per_minute: u32,
|
||||
) -> Result<bool, Box<dyn std::error::Error>> {
|
||||
let key = format!("ratelimit:{}:{}", ip, tool);
|
||||
let now = chrono::Utc::now().timestamp_millis();
|
||||
let window_start = now - 60_000;
|
||||
|
||||
// Remove entries outside the window
|
||||
let _: usize = conn.zrembyscore(&key, 0, window_start).await?;
|
||||
|
||||
// Add current entry
|
||||
let _: usize = conn
|
||||
.zadd(&key, format!("{}:{}", ip, now), now as f64)
|
||||
.await?;
|
||||
|
||||
// Set TTL on the key (cleanup)
|
||||
let _: usize = conn.expire(&key, 120).await?;
|
||||
|
||||
// Count entries in window
|
||||
let count: u32 = conn.zcount(&key, window_start, now).await?;
|
||||
|
||||
Ok(count <= max_per_minute)
|
||||
}
|
||||
|
||||
/// Get remaining requests within the current window.
|
||||
pub async fn remaining(
|
||||
conn: &mut impl AsyncCommands,
|
||||
ip: &str,
|
||||
tool: &str,
|
||||
max_per_minute: u32,
|
||||
) -> Result<u32, Box<dyn std::error::Error>> {
|
||||
let key = format!("ratelimit:{}:{}", ip, tool);
|
||||
let now = chrono::Utc::now().timestamp_millis();
|
||||
let window_start = now - 60_000;
|
||||
|
||||
let count: u32 = conn.zcount(&key, window_start, now).await?;
|
||||
|
||||
Ok(max_per_minute.saturating_sub(count))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
http::{header, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use tokio_util::io::ReaderStream;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::routes::health::AppState;
|
||||
use tools_common::types::JobStatus;
|
||||
|
||||
/// Handle GET /api/download/{id}
|
||||
pub async fn download_handler(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(id): Path<Uuid>,
|
||||
) -> Result<Response, (StatusCode, JsonResponse)> {
|
||||
let mut conn = state.redis.get_multiplexed_async_connection().await.map_err(|_| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
JsonResponse(serde_json::json!({ "error": "Redis connection failed" })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let job = crate::redis::job::JobRepository::get(&mut conn, id)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
(
|
||||
StatusCode::NOT_FOUND,
|
||||
JsonResponse(serde_json::json!({ "error": "Job not found or expired" })),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Verify job is completed
|
||||
if job.status != JobStatus::Completed {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
JsonResponse(serde_json::json!({
|
||||
"error": "Job is not completed yet",
|
||||
"status": match job.status {
|
||||
JobStatus::Queued => "queued",
|
||||
JobStatus::Processing { .. } => "processing",
|
||||
JobStatus::NeedsManualCrop => "needs_manual_crop",
|
||||
JobStatus::Failed(_) => "failed",
|
||||
_ => "unknown",
|
||||
}
|
||||
})),
|
||||
));
|
||||
}
|
||||
|
||||
let result_path = job.result_path.ok_or_else(|| {
|
||||
(
|
||||
StatusCode::NOT_FOUND,
|
||||
JsonResponse(serde_json::json!({ "error": "Result file path not found" })),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Open file
|
||||
let file = tokio::fs::File::open(&result_path).await.map_err(|e| {
|
||||
(
|
||||
StatusCode::NOT_FOUND,
|
||||
JsonResponse(serde_json::json!({ "error": format!("File not found: {}", e) })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let metadata = file.metadata().await.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
JsonResponse(serde_json::json!({
|
||||
"error": format!("Failed to read metadata: {}", e)
|
||||
})),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Determine content type
|
||||
let ext = result_path.rsplit('.').next().unwrap_or("bin").to_string();
|
||||
let content_type = match ext.as_str() {
|
||||
"pdf" => "application/pdf",
|
||||
"jpg" | "jpeg" => "image/jpeg",
|
||||
"png" => "image/png",
|
||||
"webp" => "image/webp",
|
||||
"mp4" => "video/mp4",
|
||||
"mp3" => "audio/mpeg",
|
||||
"zip" => "application/zip",
|
||||
_ => "application/octet-stream",
|
||||
};
|
||||
|
||||
// Generate filename for download
|
||||
let file_name = format!(
|
||||
"{}_{}.{}",
|
||||
job.tool.as_str(),
|
||||
job.id.to_string().split('-').next().unwrap_or("result"),
|
||||
ext
|
||||
);
|
||||
|
||||
// Stream the file
|
||||
let stream = ReaderStream::new(file);
|
||||
let body = axum::body::Body::from_stream(stream);
|
||||
|
||||
let response = Response::builder()
|
||||
.header(header::CONTENT_TYPE, content_type)
|
||||
.header(
|
||||
header::CONTENT_DISPOSITION,
|
||||
format!("attachment; filename=\"{}\"", file_name),
|
||||
)
|
||||
.header(header::CONTENT_LENGTH, metadata.len().to_string())
|
||||
.body(body)
|
||||
.unwrap();
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// Wrapper for JSON error responses.
|
||||
pub struct JsonResponse(pub serde_json::Value);
|
||||
|
||||
impl IntoResponse for JsonResponse {
|
||||
fn into_response(self) -> Response {
|
||||
(StatusCode::OK, axum::Json(self.0)).into_response()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{extract::State, Json};
|
||||
use redis::AsyncCommands;
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::metrics::Metrics;
|
||||
|
||||
/// Shared application state accessible from all handlers.
|
||||
pub struct AppState {
|
||||
pub redis: redis::Client,
|
||||
pub nats: async_nats::Client,
|
||||
pub config: crate::config::AppConfig,
|
||||
pub metrics: Metrics,
|
||||
}
|
||||
|
||||
/// Health check response.
|
||||
#[derive(Serialize)]
|
||||
pub struct HealthResponse {
|
||||
pub status: String,
|
||||
pub version: String,
|
||||
pub redis: String,
|
||||
pub nats: String,
|
||||
pub uptime_seconds: u64,
|
||||
}
|
||||
|
||||
/// Handle GET /health
|
||||
pub async fn health_handler(
|
||||
State(state): State<Arc<AppState>>,
|
||||
) -> Json<HealthResponse> {
|
||||
let redis_status = {
|
||||
match state.redis.get_multiplexed_async_connection().await {
|
||||
Ok(mut conn) => match redis::cmd("PING").query_async::<String>(&mut conn).await {
|
||||
Ok(_) => "connected".to_string(),
|
||||
Err(_) => "error".to_string(),
|
||||
},
|
||||
Err(_) => "disconnected".to_string(),
|
||||
}
|
||||
};
|
||||
|
||||
let nats_status = if state
|
||||
.nats
|
||||
.publish("tools.health.check", b"ping".to_vec().into())
|
||||
.await
|
||||
.is_ok()
|
||||
{
|
||||
"connected".to_string()
|
||||
} else {
|
||||
"disconnected".to_string()
|
||||
};
|
||||
|
||||
Json(HealthResponse {
|
||||
status: "ok".to_string(),
|
||||
version: env!("CARGO_PKG_VERSION").to_string(),
|
||||
redis: redis_status,
|
||||
nats: nats_status,
|
||||
uptime_seconds: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0),
|
||||
})
|
||||
}
|
||||
|
||||
/// Handle GET /metrics
|
||||
pub async fn metrics_handler(
|
||||
State(state): State<Arc<AppState>>,
|
||||
) -> Result<String, (axum::http::StatusCode, axum::Json<serde_json::Value>)> {
|
||||
Ok(state.metrics.format())
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
Json,
|
||||
};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::routes::health::AppState;
|
||||
use tools_common::types::*;
|
||||
|
||||
/// Handle GET /api/job/{id}
|
||||
pub async fn job_status_handler(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(id): Path<Uuid>,
|
||||
) -> Result<Json<JobStatusResponse>, (StatusCode, Json<serde_json::Value>)> {
|
||||
let mut conn = state.redis.get_multiplexed_async_connection().await.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({ "error": format!("Redis error: {}", e) })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let job = crate::redis::job::JobRepository::get(&mut conn, id)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
(
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(serde_json::json!({ "error": "Job not found or expired" })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let status_str = match &job.status {
|
||||
JobStatus::Queued => "queued",
|
||||
JobStatus::Processing { .. } => "processing",
|
||||
JobStatus::Completed => "completed",
|
||||
JobStatus::NeedsManualCrop => "needs_manual_crop",
|
||||
JobStatus::Failed(_) => "failed",
|
||||
};
|
||||
|
||||
let (progress, stage, message) = match &job.status {
|
||||
JobStatus::Processing { stage, progress } => (*progress, stage.clone(), String::new()),
|
||||
JobStatus::Failed(msg) => (0, String::new(), msg.clone()),
|
||||
JobStatus::Completed => (100, "complete".to_string(), "Processing complete".to_string()),
|
||||
JobStatus::Queued => (0, "queued".to_string(), "Waiting in queue".to_string()),
|
||||
JobStatus::NeedsManualCrop => {
|
||||
(0, "manual_crop".to_string(), "Manual crop needed".to_string())
|
||||
}
|
||||
};
|
||||
|
||||
let result = if job.status == JobStatus::Completed {
|
||||
let file_name = job
|
||||
.result_path
|
||||
.as_ref()
|
||||
.and_then(|p| std::path::Path::new(p).file_name())
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or("result")
|
||||
.to_string();
|
||||
Some(ResultInfo {
|
||||
download_url: format!("/api/download/{}", job.id),
|
||||
file_size: job.file_size,
|
||||
file_name,
|
||||
preview_url: Some(format!("/api/job/{}/preview", job.id)),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let error = match &job.status {
|
||||
JobStatus::Failed(msg) => Some(msg.clone()),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
Ok(Json(JobStatusResponse {
|
||||
job_id: job.id,
|
||||
status: status_str.to_string(),
|
||||
tool: job.tool.as_str().to_string(),
|
||||
progress,
|
||||
stage,
|
||||
message,
|
||||
result,
|
||||
created_at: job.created_at,
|
||||
error,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Handle GET /api/job/{id}/preview
|
||||
pub async fn job_preview_handler(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(id): Path<Uuid>,
|
||||
) -> Result<(StatusCode, [(String, String); 2], Vec<u8>), (StatusCode, Json<serde_json::Value>)> {
|
||||
let mut conn = state.redis.get_multiplexed_async_connection().await.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({ "error": format!("Redis error: {}", e) })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let job = crate::redis::job::JobRepository::get(&mut conn, id)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
(
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(serde_json::json!({ "error": "Job not found or expired" })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let result_path = job.result_path.ok_or_else(|| {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({ "error": "No result available yet" })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let data = tokio::fs::read(&result_path).await.map_err(|e| {
|
||||
(
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(serde_json::json!({ "error": format!("File not found: {}", e) })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let ext = result_path.rsplit('.').next().unwrap_or("bin").to_string();
|
||||
let content_type = match ext.as_str() {
|
||||
"jpg" | "jpeg" => "image/jpeg",
|
||||
"png" => "image/png",
|
||||
"webp" => "image/webp",
|
||||
"pdf" => "application/pdf",
|
||||
_ => "application/octet-stream",
|
||||
};
|
||||
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
[
|
||||
("Content-Type".to_string(), content_type.to_string()),
|
||||
(
|
||||
"Cache-Control".to_string(),
|
||||
"private, max-age=300".to_string(),
|
||||
),
|
||||
],
|
||||
data,
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod download;
|
||||
pub mod health;
|
||||
pub mod job;
|
||||
pub mod upload;
|
||||
pub mod ws;
|
||||
@@ -0,0 +1,261 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
extract::{Multipart, State},
|
||||
http::StatusCode,
|
||||
Json,
|
||||
};
|
||||
use chrono::Utc;
|
||||
use tokio::fs;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::routes::health::AppState;
|
||||
use tools_common::error::UploadError;
|
||||
use tools_common::types::*;
|
||||
|
||||
/// Handle POST /api/upload
|
||||
pub async fn upload_handler(
|
||||
State(state): State<Arc<AppState>>,
|
||||
mut multipart: Multipart,
|
||||
) -> Result<Json<UploadResponse>, (StatusCode, Json<serde_json::Value>)> {
|
||||
let config = &state.config;
|
||||
let max_size = config.max_file_size_bytes();
|
||||
|
||||
// Extract fields from multipart
|
||||
let mut file_data: Option<(String, Vec<u8>)> = None;
|
||||
let mut tool_str: Option<String> = None;
|
||||
let mut options: serde_json::Value = serde_json::Value::Null;
|
||||
|
||||
while let Ok(Some(field)) = multipart.next_field().await {
|
||||
let name = field.name().unwrap_or("").to_string();
|
||||
match name.as_str() {
|
||||
"file" => {
|
||||
let filename = field.file_name().unwrap_or("unknown").to_string();
|
||||
let data = field.bytes().await.map_err(|e| {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({
|
||||
"error": "Failed to read file",
|
||||
"detail": e.to_string()
|
||||
})),
|
||||
)
|
||||
})?;
|
||||
file_data = Some((filename, data.to_vec()));
|
||||
}
|
||||
"tool" => {
|
||||
tool_str = Some(field.text().await.unwrap_or_default());
|
||||
}
|
||||
"options" => {
|
||||
let text = field.text().await.unwrap_or_default();
|
||||
if !text.is_empty() {
|
||||
options = serde_json::from_str(&text).unwrap_or(serde_json::Value::Null);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// Validate fields
|
||||
let (filename, data) = file_data.ok_or_else(|| {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({ "error": "No file provided" })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let tool_str = tool_str.ok_or_else(|| {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({ "error": "No tool specified" })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let tool = Tool::from_str(&tool_str).ok_or_else(|| {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({ "error": format!("Unknown tool: {}", tool_str) })),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Validate file size
|
||||
let file_size = data.len() as u64;
|
||||
if file_size > max_size {
|
||||
return Err((
|
||||
StatusCode::PAYLOAD_TOO_LARGE,
|
||||
Json(serde_json::json!({
|
||||
"error": format!("File too large: {} bytes (max {} bytes)", file_size, max_size)
|
||||
})),
|
||||
));
|
||||
}
|
||||
|
||||
// Validate MIME type based on tool
|
||||
let ext = filename.rsplit('.').next().unwrap_or("").to_lowercase();
|
||||
validate_mime(&tool, &ext).map_err(|e| {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({ "error": e.to_string() })),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Verify magic bytes
|
||||
if !verify_magic_bytes(&data, &ext) {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({ "error": "File content does not match extension" })),
|
||||
));
|
||||
}
|
||||
|
||||
// Create directories
|
||||
let upload_dir = config.storage_path.join("upload");
|
||||
fs::create_dir_all(&upload_dir).await.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({ "error": format!("Storage error: {}", e) })),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Generate job ID and save file
|
||||
let job_id = Uuid::new_v4();
|
||||
let storage_filename = format!("{}.{}", job_id, ext);
|
||||
let file_path = upload_dir.join(&storage_filename);
|
||||
fs::write(&file_path, &data).await.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({ "error": format!("Failed to save file: {}", e) })),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Create job record
|
||||
let job = Job {
|
||||
id: job_id,
|
||||
tool: tool.clone(),
|
||||
status: JobStatus::Queued,
|
||||
file_path: file_path.to_string_lossy().to_string(),
|
||||
result_path: None,
|
||||
file_size,
|
||||
options: options.clone(),
|
||||
created_at: Utc::now(),
|
||||
ttl_seconds: config.job_ttl_seconds,
|
||||
};
|
||||
|
||||
// Save to Redis
|
||||
{
|
||||
let mut conn = state.redis.get_multiplexed_async_connection().await.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({ "error": format!("Redis error: {}", e) })),
|
||||
)
|
||||
})?;
|
||||
crate::redis::job::JobRepository::create(&mut conn, &job)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({ "error": format!("Failed to create job: {}", e) })),
|
||||
)
|
||||
})?;
|
||||
}
|
||||
|
||||
// Publish to NATS
|
||||
crate::nats::publisher::NatsPublisher::publish_job(&state.nats, &tool, &job)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!("Failed to publish job to NATS: {}", e);
|
||||
});
|
||||
|
||||
// Update metrics
|
||||
state.metrics.increment_jobs_total(tool.as_str(), "queued");
|
||||
|
||||
// Return response
|
||||
Ok(Json(UploadResponse {
|
||||
job_id,
|
||||
status: "queued".to_string(),
|
||||
tool: tool_str,
|
||||
ws_url: format!("/api/job/{}/ws", job_id),
|
||||
created_at: job.created_at,
|
||||
estimated_seconds: match tool {
|
||||
Tool::Scan => 5,
|
||||
_ => 3,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
fn validate_mime(tool: &Tool, ext: &str) -> Result<(), UploadError> {
|
||||
let image_exts = ["jpg", "jpeg", "png", "webp", "heic", "bmp", "tiff", "tif"];
|
||||
let pdf_exts = ["pdf"];
|
||||
let video_exts = ["mp4", "webm", "avi", "mov", "mkv"];
|
||||
let audio_exts = ["mp3", "wav", "flac", "aac", "ogg", "m4a"];
|
||||
|
||||
match tool {
|
||||
Tool::Scan
|
||||
| Tool::ImageCompress
|
||||
| Tool::ImageResize
|
||||
| Tool::ImageConvert
|
||||
| Tool::RemoveBg => {
|
||||
if !image_exts.contains(&ext) {
|
||||
return Err(UploadError::InvalidMime(format!(
|
||||
"Expected image file, got .{}",
|
||||
ext
|
||||
)));
|
||||
}
|
||||
}
|
||||
Tool::PdfMerge | Tool::PdfSplit | Tool::PdfCompress | Tool::PdfToImages => {
|
||||
if !pdf_exts.contains(&ext) {
|
||||
return Err(UploadError::InvalidMime(format!(
|
||||
"Expected PDF file, got .{}",
|
||||
ext
|
||||
)));
|
||||
}
|
||||
}
|
||||
Tool::ImagesToPdf => {
|
||||
if !image_exts.contains(&ext) {
|
||||
return Err(UploadError::InvalidMime(format!(
|
||||
"Expected image file, got .{}",
|
||||
ext
|
||||
)));
|
||||
}
|
||||
}
|
||||
Tool::VideoCompress | Tool::VideoTrim | Tool::GifMaker => {
|
||||
if !video_exts.contains(&ext) {
|
||||
return Err(UploadError::InvalidMime(format!(
|
||||
"Expected video file, got .{}",
|
||||
ext
|
||||
)));
|
||||
}
|
||||
}
|
||||
Tool::AudioExtract => {
|
||||
if !video_exts.contains(&ext) {
|
||||
return Err(UploadError::InvalidMime(format!(
|
||||
"Expected video file, got .{}",
|
||||
ext
|
||||
)));
|
||||
}
|
||||
}
|
||||
Tool::AudioConvert => {
|
||||
if !audio_exts.contains(&ext) {
|
||||
return Err(UploadError::InvalidMime(format!(
|
||||
"Expected audio file, got .{}",
|
||||
ext
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn verify_magic_bytes(data: &[u8], ext: &str) -> bool {
|
||||
if data.is_empty() {
|
||||
return false;
|
||||
}
|
||||
match ext {
|
||||
"jpg" | "jpeg" => data.starts_with(&[0xFF, 0xD8, 0xFF]),
|
||||
"png" => data.starts_with(&[0x89, 0x50, 0x4E, 0x47]),
|
||||
"webp" => data.len() > 12 && &data[0..4] == b"RIFF" && &data[8..12] == b"WEBP",
|
||||
"gif" => data.starts_with(b"GIF8"),
|
||||
"bmp" => data.starts_with(b"BM"),
|
||||
"pdf" => data.starts_with(b"%PDF"),
|
||||
"mp4" => data.len() > 8 && (&data[4..8] == b"ftyp" || &data[4..8] == b"ftyp"),
|
||||
"heic" => data.len() > 12 && &data[4..12] == b"ftypheic",
|
||||
_ => true,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
extract::{
|
||||
ws::{Message, WebSocket},
|
||||
Path, State, WebSocketUpgrade,
|
||||
},
|
||||
response::IntoResponse,
|
||||
};
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use tokio::sync::mpsc;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::routes::health::AppState;
|
||||
use tools_common::types::{JobProgress, JobStatus};
|
||||
|
||||
/// Handle WebSocket upgrade at /api/job/{id}/ws.
|
||||
pub async fn ws_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(job_id): Path<Uuid>,
|
||||
) -> impl IntoResponse {
|
||||
ws.on_upgrade(move |socket| handle_ws(socket, state, job_id))
|
||||
}
|
||||
|
||||
async fn handle_ws(ws: WebSocket, state: Arc<AppState>, job_id: Uuid) {
|
||||
let (mut sender, mut receiver) = ws.split();
|
||||
|
||||
// Subscribe to NATS progress updates
|
||||
let nats = state.nats.clone();
|
||||
let subject = format!("tools.*.progress.{}", job_id);
|
||||
let mut subscriber = match nats.subscribe(subject).await {
|
||||
Ok(sub) => sub,
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to subscribe to NATS: {}", e);
|
||||
let _ = sender
|
||||
.send(Message::Text(
|
||||
serde_json::json!({
|
||||
"type": "error",
|
||||
"job_id": job_id,
|
||||
"status": "failed",
|
||||
"error": format!("Connection error: {}", e)
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Send initial status from Redis
|
||||
if let Ok(mut conn) = state.redis.get_multiplexed_async_connection().await {
|
||||
if let Ok(job) = crate::redis::job::JobRepository::get(&mut conn, job_id).await {
|
||||
let init_msg = serde_json::json!({
|
||||
"type": "status",
|
||||
"job_id": job_id,
|
||||
"status": match &job.status {
|
||||
JobStatus::Queued => "queued",
|
||||
JobStatus::Processing { .. } => "processing",
|
||||
JobStatus::Completed => "completed",
|
||||
JobStatus::NeedsManualCrop => "needs_manual_crop",
|
||||
JobStatus::Failed(_) => "failed",
|
||||
},
|
||||
"progress": match &job.status {
|
||||
JobStatus::Processing { progress, .. } => *progress,
|
||||
JobStatus::Completed => 100,
|
||||
_ => 0,
|
||||
},
|
||||
});
|
||||
let _ = sender
|
||||
.send(Message::Text(init_msg.to_string().into()))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
// Channel for NATS messages
|
||||
let (tx, mut rx) = mpsc::channel::<String>(32);
|
||||
|
||||
// Spawn NATS listener
|
||||
let tx_clone = tx.clone();
|
||||
let nats_listener = tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
msg = subscriber.next() => {
|
||||
match msg {
|
||||
Some(nats_msg) => {
|
||||
if let Ok(progress) = serde_json::from_slice::<JobProgress>(&nats_msg.payload) {
|
||||
let json = serde_json::json!({
|
||||
"type": "progress",
|
||||
"job_id": progress.job_id,
|
||||
"status": match &progress.status {
|
||||
JobStatus::Queued => "queued",
|
||||
JobStatus::Processing { .. } => "processing",
|
||||
JobStatus::Completed => "completed",
|
||||
JobStatus::NeedsManualCrop => "needs_manual_crop",
|
||||
JobStatus::Failed(_) => "failed",
|
||||
},
|
||||
"progress": progress.progress,
|
||||
"stage": progress.stage,
|
||||
"message": progress.message,
|
||||
});
|
||||
let _ = tx_clone.send(json.to_string()).await;
|
||||
}
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
_ = tokio::time::sleep(tokio::time::Duration::from_secs(30)) => {
|
||||
// Keepalive ping
|
||||
let _ = tx_clone.send(serde_json::json!({"type": "ping"}).to_string()).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Forward messages from channel to WebSocket
|
||||
let ws_sender = tokio::spawn(async move {
|
||||
while let Some(msg) = rx.recv().await {
|
||||
if sender.send(Message::Text(msg.into())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Listen for client close
|
||||
let ws_receiver = tokio::spawn(async move {
|
||||
while let Some(Ok(_)) = receiver.next().await {
|
||||
// Client messages ignored (we only forward server→client)
|
||||
}
|
||||
});
|
||||
|
||||
// Wait for either task to complete (connection closed)
|
||||
tokio::select! {
|
||||
_ = ws_sender => {},
|
||||
_ = ws_receiver => {},
|
||||
}
|
||||
|
||||
// Cancel NATS listener
|
||||
nats_listener.abort();
|
||||
}
|
||||
Reference in New Issue
Block a user