From 01144261722548eb88a49a58ac78c913748f6ec3 Mon Sep 17 00:00:00 2001 From: Arnaud_Cayrol Date: Mon, 2 Feb 2026 21:12:05 +0100 Subject: [PATCH] Reduce code duplication, add utility functions --- src/config.rs | 10 +- src/job/mod.rs | 156 +----------------------------- src/job/processing.rs | 161 +++++++++++++++++++++++++++++++ src/models/dlib_landmarks.rs | 6 +- src/pipeline/mod.rs | 4 +- src/pipeline/steps/alignment.rs | 14 ++- src/pipeline/steps/brightness.rs | 14 +-- src/pipeline/steps/crop.rs | 8 +- src/pipeline/steps/head_pose.rs | 8 +- src/pipeline/steps/landmarks.rs | 16 ++- src/pipeline/steps/resize.rs | 8 +- src/pipeline/traits.rs | 31 ++++++ src/web/handlers/config.rs | 16 +-- src/web/handlers/output.rs | 122 ++++++++--------------- src/web/state.rs | 51 ++++++++++ 15 files changed, 330 insertions(+), 295 deletions(-) create mode 100644 src/job/processing.rs diff --git a/src/config.rs b/src/config.rs index ae7ae95..f3fac89 100644 --- a/src/config.rs +++ b/src/config.rs @@ -285,12 +285,10 @@ impl Default for EyeFilterConfig { impl EyeFilterConfig { /// Validate the configuration values. pub fn validate(&self) -> Result<()> { - if self.enabled { - if self.min_ear < 0.0 || self.min_ear > 0.5 { - return Err(Error::Config( - "Eye filter min_ear must be between 0.0 and 0.5".to_string(), - )); - } + if self.enabled && !(0.0..=0.5).contains(&self.min_ear) { + return Err(Error::Config( + "Eye filter min_ear must be between 0.0 and 0.5".to_string(), + )); } Ok(()) } diff --git a/src/job/mod.rs b/src/job/mod.rs index 682bdb4..e74961c 100644 --- a/src/job/mod.rs +++ b/src/job/mod.rs @@ -6,17 +6,17 @@ //! 3. Processing images through the extensible pipeline //! 4. Compiling processed images into a timelapse video -use crate::config::Config; +mod processing; + use crate::error::{Error, Result}; use crate::immich_api::{Asset, FaceData, ImmichClient}; -use crate::pipeline::{Pipeline, PipelineContext, PipelineResult}; +use crate::pipeline::Pipeline; use crate::utils::sanitize_folder_name; use crate::video::compile_timelapse; use crate::web::{AppState, AtomicSkipStats, JobStatus, Progress, SkipStats}; -use bytes::Bytes; -use image::ImageFormat; -use std::io::Cursor; +use processing::{process_single_asset, AssetProcessResult, DebugDirs, OutputDirs}; + use std::path::PathBuf; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::Arc; @@ -32,27 +32,6 @@ pub struct JobParams { pub date_to: Option, } -/// Output directories for a processing job. -#[derive(Debug, Clone)] -struct OutputDirs { - /// Directory for final processed images (used for video) - images: PathBuf, - /// Path for output video - video: PathBuf, - /// Optional debug directories for visualizing processing stages. - debug: Option, -} - -/// Debug output base directory for visualizing processing stages. -/// Subdirectories are created on-demand by the pipeline for each step: -/// - `{step_id}/passed/` - Images that passed the step -/// - `{step_id}/failed/` - Images that failed/skipped at the step -#[derive(Debug, Clone)] -struct DebugDirs { - /// Base directory for debug output (e.g., `output/PersonName/debug`) - base: PathBuf, -} - /// Run the complete processing pipeline. /// /// This is the main entry point for background job processing. @@ -413,131 +392,6 @@ fn find_face_for_person(asset: &Asset, person_id: &str) -> Option { None } -/// Result of processing a single asset. -/// -/// Fields are kept for debugging/logging purposes even if not currently read. -#[derive(Debug)] -#[allow(dead_code)] -enum AssetProcessResult { - /// Successfully processed. - Success { asset_id: String }, - /// Skipped for some reason. - Skipped { asset_id: String, reason: String }, - /// Error during processing. - Error { asset_id: String, error: String }, - /// Cancelled by user. - Cancelled { asset_id: String }, -} - -/// Process a single asset using the pipeline. -async fn process_single_asset( - client: &ImmichClient, - config: &Config, - asset: &Asset, - face_data: &FaceData, - output_dirs: &OutputDirs, - cancel_token: &CancellationToken, - skip_stats: &Arc, - pipeline: &Pipeline, -) -> AssetProcessResult { - let asset_id = &asset.id; - - // Generate timestamp-based filename for sorting - let timestamp = asset - .file_created_at - .as_ref() - .or(asset.local_date_time.as_ref()) - .cloned() - .unwrap_or_else(|| asset_id.clone()); - - // Create pipeline context - let mut ctx = PipelineContext::new( - asset_id.clone(), - timestamp.clone(), - face_data.clone(), - ); - - // Check before download (potentially slow) - if cancel_token.is_cancelled() { - return AssetProcessResult::Cancelled { - asset_id: asset_id.clone(), - }; - } - - // Download image - let image_bytes: Bytes = match client.download_asset(asset_id).await { - Ok(bytes) => bytes, - Err(e) => { - skip_stats.increment("download_failed"); - return AssetProcessResult::Error { - asset_id: asset_id.clone(), - error: format!("Download failed: {}", e), - }; - } - }; - - // Set raw bytes on context - ctx = ctx.with_bytes(image_bytes); - - // Determine debug directory - let debug_dir = output_dirs.debug.as_ref().map(|d| d.base.clone()); - - // Execute the pipeline - let result = pipeline.execute( - ctx, - config, - cancel_token, - skip_stats, - debug_dir.as_ref(), - ).await; - - match result { - PipelineResult::Success { image, asset_id, timestamp, .. } => { - // Sanitize timestamp for filename - let safe_timestamp: String = timestamp - .chars() - .map(|c| { - if c.is_alphanumeric() || c == '-' || c == '_' { - c - } else { - '_' - } - }) - .collect(); - - let filename = format!("{}_{}.jpg", safe_timestamp, asset_id); - let output_path = output_dirs.images.join(&filename); - - // Encode and save final image - let mut buffer = Cursor::new(Vec::new()); - if let Err(e) = image.write_to(&mut buffer, ImageFormat::Jpeg) { - return AssetProcessResult::Error { - asset_id, - error: format!("Failed to encode image: {}", e), - }; - } - - if let Err(e) = tokio::fs::write(&output_path, buffer.into_inner()).await { - return AssetProcessResult::Error { - asset_id, - error: format!("Failed to save image: {}", e), - }; - } - - AssetProcessResult::Success { asset_id } - } - PipelineResult::Skipped { asset_id, reason, .. } => { - AssetProcessResult::Skipped { asset_id, reason } - } - PipelineResult::Error { asset_id, error } => { - AssetProcessResult::Error { asset_id, error } - } - PipelineResult::Cancelled { asset_id } => { - AssetProcessResult::Cancelled { asset_id } - } - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/job/processing.rs b/src/job/processing.rs new file mode 100644 index 0000000..70d2779 --- /dev/null +++ b/src/job/processing.rs @@ -0,0 +1,161 @@ +//! Single asset processing logic. +//! +//! Contains the logic for processing individual images through the pipeline, +//! including downloading, running the pipeline, and saving results. + +use crate::config::Config; +use crate::immich_api::{Asset, FaceData, ImmichClient}; +use crate::pipeline::{Pipeline, PipelineContext, PipelineResult}; +use crate::web::AtomicSkipStats; + +use bytes::Bytes; +use image::ImageFormat; +use std::io::Cursor; +use std::path::PathBuf; +use std::sync::Arc; +use tokio_util::sync::CancellationToken; + +/// Output directories for a processing job. +#[derive(Debug, Clone)] +pub struct OutputDirs { + /// Directory for final processed images (used for video) + pub images: PathBuf, + /// Path for output video + pub video: PathBuf, + /// Optional debug directories for visualizing processing stages. + pub debug: Option, +} + +/// Debug output base directory for visualizing processing stages. +/// Subdirectories are created on-demand by the pipeline for each step: +/// - `{step_id}/passed/` - Images that passed the step +/// - `{step_id}/failed/` - Images that failed/skipped at the step +#[derive(Debug, Clone)] +pub struct DebugDirs { + /// Base directory for debug output (e.g., `output/PersonName/debug`) + pub base: PathBuf, +} + +/// Result of processing a single asset. +/// +/// Fields are kept for debugging/logging purposes even if not currently read. +#[derive(Debug)] +#[allow(dead_code)] +pub enum AssetProcessResult { + /// Successfully processed. + Success { asset_id: String }, + /// Skipped for some reason. + Skipped { asset_id: String, reason: String }, + /// Error during processing. + Error { asset_id: String, error: String }, + /// Cancelled by user. + Cancelled { asset_id: String }, +} + +/// Process a single asset using the pipeline. +/// +/// This function handles: +/// 1. Downloading the image from Immich +/// 2. Running it through the processing pipeline +/// 3. Saving the result to disk +#[allow(clippy::too_many_arguments)] +pub async fn process_single_asset( + client: &ImmichClient, + config: &Config, + asset: &Asset, + face_data: &FaceData, + output_dirs: &OutputDirs, + cancel_token: &CancellationToken, + skip_stats: &Arc, + pipeline: &Pipeline, +) -> AssetProcessResult { + let asset_id = &asset.id; + + // Generate timestamp-based filename for sorting + let timestamp = asset + .file_created_at + .as_ref() + .or(asset.local_date_time.as_ref()) + .cloned() + .unwrap_or_else(|| asset_id.clone()); + + // Create pipeline context + let mut ctx = PipelineContext::new(asset_id.clone(), timestamp.clone(), face_data.clone()); + + // Check before download (potentially slow) + if cancel_token.is_cancelled() { + return AssetProcessResult::Cancelled { + asset_id: asset_id.clone(), + }; + } + + // Download image + let image_bytes: Bytes = match client.download_asset(asset_id).await { + Ok(bytes) => bytes, + Err(e) => { + skip_stats.increment("download_failed"); + return AssetProcessResult::Error { + asset_id: asset_id.clone(), + error: format!("Download failed: {}", e), + }; + } + }; + + // Set raw bytes on context + ctx = ctx.with_bytes(image_bytes); + + // Determine debug directory + let debug_dir = output_dirs.debug.as_ref().map(|d| d.base.clone()); + + // Execute the pipeline + let result = pipeline + .execute(ctx, config, cancel_token, skip_stats, debug_dir.as_ref()) + .await; + + match result { + PipelineResult::Success { + image, + asset_id, + timestamp, + .. + } => { + // Sanitize timestamp for filename + let safe_timestamp: String = timestamp + .chars() + .map(|c| { + if c.is_alphanumeric() || c == '-' || c == '_' { + c + } else { + '_' + } + }) + .collect(); + + let filename = format!("{}_{}.jpg", safe_timestamp, asset_id); + let output_path = output_dirs.images.join(&filename); + + // Encode and save final image + let mut buffer = Cursor::new(Vec::new()); + if let Err(e) = image.write_to(&mut buffer, ImageFormat::Jpeg) { + return AssetProcessResult::Error { + asset_id, + error: format!("Failed to encode image: {}", e), + }; + } + + if let Err(e) = tokio::fs::write(&output_path, buffer.into_inner()).await { + return AssetProcessResult::Error { + asset_id, + error: format!("Failed to save image: {}", e), + }; + } + + AssetProcessResult::Success { asset_id } + } + PipelineResult::Skipped { + asset_id, reason, .. + } => AssetProcessResult::Skipped { asset_id, reason }, + PipelineResult::Error { asset_id, error } => AssetProcessResult::Error { asset_id, error }, + PipelineResult::Cancelled { asset_id } => AssetProcessResult::Cancelled { asset_id }, + } +} diff --git a/src/models/dlib_landmarks.rs b/src/models/dlib_landmarks.rs index fb11fb1..305840e 100644 --- a/src/models/dlib_landmarks.rs +++ b/src/models/dlib_landmarks.rs @@ -68,8 +68,8 @@ impl DlibLandmarks { /// * `height` - Image height /// * `pixels` - Raw RGB pixel data (width * height * 3 bytes) /// * `face_rect` - Optional face bounding box (x1, y1, x2, y2) in image coordinates. - /// If provided, uses this rectangle for landmark detection. - /// If None, attempts to detect the face or uses the whole image. + /// If provided, uses this rectangle for landmark detection. + /// If None, attempts to detect the face or uses the whole image. /// /// # Returns /// Landmarks struct containing the 68 facial landmark points, or an error @@ -107,7 +107,7 @@ impl DlibLandmarks { // No face rect provided - try to detect or use whole image let faces = detector.face_locations(&matrix); if !faces.is_empty() { - faces[0].clone() + faces[0] } else { // Fallback: use whole image with small margin let margin = 5; diff --git a/src/pipeline/mod.rs b/src/pipeline/mod.rs index cfed085..141aa40 100644 --- a/src/pipeline/mod.rs +++ b/src/pipeline/mod.rs @@ -22,7 +22,7 @@ use crate::config::Config; use crate::web::AtomicSkipStats; use image::{DynamicImage, ImageFormat}; use std::io::Cursor; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::Arc; use tokio_util::sync::CancellationToken; @@ -236,7 +236,7 @@ impl Pipeline { /// debug/{step_id}/failed/{filename}.jpg /// ``` async fn save_debug_images( - debug_base: &PathBuf, + debug_base: &Path, debug_images: &[(String, DebugImage)], timestamp: &str, asset_id: &str, diff --git a/src/pipeline/steps/alignment.rs b/src/pipeline/steps/alignment.rs index 9c08986..67de80b 100644 --- a/src/pipeline/steps/alignment.rs +++ b/src/pipeline/steps/alignment.rs @@ -48,11 +48,9 @@ impl ProcessingStep for AlignmentStep { } }; - let image = match ctx.image.take() { - Some(img) => img, - None => { - return StepOutcome::Error("No image available for alignment".to_string()); - } + let image = match ctx.take_image("alignment") { + Ok(img) => img, + Err(e) => return StepOutcome::Error(e), }; let (width, height) = image.dimensions(); @@ -207,12 +205,12 @@ impl ProcessingStep for AlignmentStep { /// Draw a marker (small filled square) at the given position. fn draw_marker(img: &mut RgbImage, x: u32, y: u32, color: Rgb) { let (width, height) = (img.width(), img.height()); - let size = 3; + let size: i32 = 3; for dy in 0..=size * 2 { for dx in 0..=size * 2 { - let px = (x as i32 + dx as i32 - size as i32) as u32; - let py = (y as i32 + dy as i32 - size as i32) as u32; + let px = (x as i32 + dx - size) as u32; + let py = (y as i32 + dy - size) as u32; if px < width && py < height { img.put_pixel(px, py, color); } diff --git a/src/pipeline/steps/brightness.rs b/src/pipeline/steps/brightness.rs index da6ac9f..fc3b972 100644 --- a/src/pipeline/steps/brightness.rs +++ b/src/pipeline/steps/brightness.rs @@ -55,13 +55,9 @@ impl ProcessingStep for BrightnessStep { async fn execute(&self, mut ctx: PipelineContext, config: &Config) -> StepOutcome { let step_config = &config.processing.brightness; - let image = match &ctx.image { - Some(img) => img, - None => { - return StepOutcome::Error( - "No image available for brightness check".to_string(), - ); - } + let image = match ctx.require_image("brightness check") { + Ok(img) => img, + Err(e) => return StepOutcome::Error(e), }; let brightness = Self::calculate_brightness(image); @@ -160,9 +156,9 @@ impl ProcessingStep for BrightnessStep { /// Convert brightness value to a color (red for dark/bright, green for good) fn brightness_to_color(brightness: f32) -> Rgb { // Very dark or very bright = red, middle range = green - if brightness < 0.15 || brightness > 0.85 { + if !(0.15..=0.85).contains(&brightness) { Rgb([255, 80, 80]) // Red - } else if brightness < 0.25 || brightness > 0.75 { + } else if !(0.25..=0.75).contains(&brightness) { Rgb([255, 200, 80]) // Yellow/orange } else { Rgb([80, 255, 80]) // Green diff --git a/src/pipeline/steps/crop.rs b/src/pipeline/steps/crop.rs index 18dc431..b62ebe1 100644 --- a/src/pipeline/steps/crop.rs +++ b/src/pipeline/steps/crop.rs @@ -26,11 +26,9 @@ impl ProcessingStep for CropFaceStep { } async fn execute(&self, mut ctx: PipelineContext, config: &Config) -> StepOutcome { - let image = match &ctx.image { - Some(img) => img, - None => { - return StepOutcome::Error("No image available for cropping".to_string()); - } + let image = match ctx.require_image("cropping") { + Ok(img) => img, + Err(e) => return StepOutcome::Error(e), }; // Crop returns CropResult with cropped images and face rectangle in crop coordinates. diff --git a/src/pipeline/steps/head_pose.rs b/src/pipeline/steps/head_pose.rs index 865d1d9..a2d38c6 100644 --- a/src/pipeline/steps/head_pose.rs +++ b/src/pipeline/steps/head_pose.rs @@ -33,11 +33,9 @@ impl ProcessingStep for HeadPoseStep { return StepOutcome::Continue(ctx); } - let image = match &ctx.image { - Some(img) => img, - None => { - return StepOutcome::Error("No image available for head pose estimation".to_string()); - } + let image = match ctx.require_image("head pose estimation") { + Ok(img) => img, + Err(e) => return StepOutcome::Error(e), }; // Load the DMHead model diff --git a/src/pipeline/steps/landmarks.rs b/src/pipeline/steps/landmarks.rs index 984859f..ab1dffd 100644 --- a/src/pipeline/steps/landmarks.rs +++ b/src/pipeline/steps/landmarks.rs @@ -38,13 +38,9 @@ impl ProcessingStep for LandmarksStep { return StepOutcome::Continue(ctx); } - let image = match &ctx.image { - Some(img) => img, - None => { - return StepOutcome::Error( - "No image available for landmark detection".to_string(), - ); - } + let image = match ctx.require_image("landmark detection") { + Ok(img) => img, + Err(e) => return StepOutcome::Error(e), }; // Get the global landmark predictor (loaded once, reused for all images) @@ -178,7 +174,7 @@ impl ProcessingStep for LandmarksStep { /// Draw a small cross at the given position. fn draw_cross(img: &mut RgbImage, x: u32, y: u32, color: Rgb) { let (width, height) = (img.width(), img.height()); - let size = 2; + let size: i32 = 2; // Check base coordinates are in bounds if x >= width || y >= height { @@ -186,13 +182,13 @@ fn draw_cross(img: &mut RgbImage, x: u32, y: u32, color: Rgb) { } for dx in 0..=size * 2 { - let px = (x as i32 + dx as i32 - size as i32) as u32; + let px = (x as i32 + dx - size) as u32; if px < width && y < height { img.put_pixel(px, y, color); } } for dy in 0..=size * 2 { - let py = (y as i32 + dy as i32 - size as i32) as u32; + let py = (y as i32 + dy - size) as u32; if x < width && py < height { img.put_pixel(x, py, color); } diff --git a/src/pipeline/steps/resize.rs b/src/pipeline/steps/resize.rs index 1670bdc..4248737 100644 --- a/src/pipeline/steps/resize.rs +++ b/src/pipeline/steps/resize.rs @@ -24,11 +24,9 @@ impl ProcessingStep for ResizeStep { } async fn execute(&self, mut ctx: PipelineContext, config: &Config) -> StepOutcome { - let image = match ctx.image.take() { - Some(img) => img, - None => { - return StepOutcome::Error("No image available for resizing".to_string()); - } + let image = match ctx.take_image("resizing") { + Ok(img) => img, + Err(e) => return StepOutcome::Error(e), }; let output_size = config.processing.output.size; diff --git a/src/pipeline/traits.rs b/src/pipeline/traits.rs index 96f4846..95131df 100644 --- a/src/pipeline/traits.rs +++ b/src/pipeline/traits.rs @@ -175,6 +175,37 @@ impl PipelineContext { self.computed.get(key) } + /// Get a reference to the image, returning an error message if not available. + /// + /// Use this in pipeline steps that need to read the image without modifying it. + /// + /// # Example + /// ```ignore + /// let image = ctx.require_image("brightness check")?; + /// ``` + pub fn require_image(&self, step_name: &str) -> Result<&DynamicImage, String> { + self.image + .as_ref() + .ok_or_else(|| format!("No image available for {}", step_name)) + } + + /// Take ownership of the image, returning an error message if not available. + /// + /// Use this in pipeline steps that need to transform the image (alignment, resize). + /// The step should set `ctx.image` to the transformed result before returning. + /// + /// # Example + /// ```ignore + /// let image = ctx.take_image("alignment")?; + /// // ... transform image ... + /// ctx.image = Some(transformed); + /// ``` + pub fn take_image(&mut self, step_name: &str) -> Result { + self.image + .take() + .ok_or_else(|| format!("No image available for {}", step_name)) + } + /// Add a debug image for a step. /// /// # Arguments diff --git a/src/web/handlers/config.rs b/src/web/handlers/config.rs index b581019..44b3dfe 100644 --- a/src/web/handlers/config.rs +++ b/src/web/handlers/config.rs @@ -4,7 +4,7 @@ use crate::config::{ AlignmentConfig, BrightnessConfig, FaceResolutionConfig, OutputConfig, ProcessingConfig, VideoConfig, }; -use crate::web::state::{AppState, JobStatus}; +use crate::web::state::AppState; use axum::{extract::State, http::StatusCode, response::Json}; use serde::{Deserialize, Serialize}; @@ -174,16 +174,10 @@ pub async fn update_config( State(state): State, Json(update): Json, ) -> Result, (StatusCode, String)> { - // Check if a job is running - { - let progress = state.progress.read().await; - if progress.status == JobStatus::Running || progress.status == JobStatus::CompilingVideo { - return Err(( - StatusCode::CONFLICT, - "Cannot update config while a job is running".to_string(), - )); - } - } + state + .ensure_no_job_running() + .await + .map_err(|e| (StatusCode::CONFLICT, e.to_string()))?; // Validate input before updating if let Some(ref proc) = update.processing { diff --git a/src/web/handlers/output.rs b/src/web/handlers/output.rs index 66977ba..514065a 100644 --- a/src/web/handlers/output.rs +++ b/src/web/handlers/output.rs @@ -1,6 +1,6 @@ //! Output folder and image management endpoints. -use crate::web::state::{AppState, JobStatus, Progress, SkipStats}; +use crate::web::state::{validate_path_component, AppState, JobStatus, Progress, SkipStats}; use axum::{ extract::{Path, State}, http::StatusCode, @@ -120,20 +120,14 @@ pub async fn list_output_folders( pub async fn cleanup_all_output( State(state): State, ) -> Result, (StatusCode, String)> { + state + .ensure_no_job_running() + .await + .map_err(|e| (StatusCode::CONFLICT, e.to_string()))?; + let config = state.config.read().await; let output_dir = &config.output_dir; - // Check if a job is running - { - let progress = state.progress.read().await; - if progress.status == JobStatus::Running || progress.status == JobStatus::CompilingVideo { - return Err(( - StatusCode::CONFLICT, - "Cannot cleanup while a job is running".to_string(), - )); - } - } - // Remove all contents of output directory if output_dir.exists() { let mut entries = tokio::fs::read_dir(output_dir).await.map_err(|e| { @@ -181,24 +175,16 @@ pub async fn cleanup_output_folder( State(state): State, Path(folder_name): Path, ) -> Result, (StatusCode, String)> { + state + .ensure_no_job_running() + .await + .map_err(|e| (StatusCode::CONFLICT, e.to_string()))?; + + validate_path_component(&folder_name, "folder name") + .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; + let config = state.config.read().await; - // Check if a job is running - { - let progress = state.progress.read().await; - if progress.status == JobStatus::Running || progress.status == JobStatus::CompilingVideo { - return Err(( - StatusCode::CONFLICT, - "Cannot cleanup while a job is running".to_string(), - )); - } - } - - // Sanitize folder name to prevent path traversal - if folder_name.contains("..") || folder_name.contains('/') || folder_name.contains('\\') { - return Err((StatusCode::BAD_REQUEST, "Invalid folder name".to_string())); - } - let folder_path = config.output_dir.join(&folder_name); if !folder_path.exists() { @@ -235,12 +221,10 @@ pub async fn list_folder_images( State(state): State, Path(folder_name): Path, ) -> Result, (StatusCode, String)> { - let config = state.config.read().await; + validate_path_component(&folder_name, "folder name") + .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; - // Sanitize folder name to prevent path traversal - if folder_name.contains("..") || folder_name.contains('/') || folder_name.contains('\\') { - return Err((StatusCode::BAD_REQUEST, "Invalid folder name".to_string())); - } + let config = state.config.read().await; let images_dir = config.output_dir.join(&folder_name).join("images"); @@ -311,26 +295,17 @@ pub async fn delete_single_image( State(state): State, Path((folder_name, filename)): Path<(String, String)>, ) -> Result, (StatusCode, String)> { - // Check if a job is running - { - let progress = state.progress.read().await; - if progress.status == JobStatus::Running || progress.status == JobStatus::CompilingVideo { - return Err(( - StatusCode::CONFLICT, - "Cannot delete images while a job is running".to_string(), - )); - } - } + state + .ensure_no_job_running() + .await + .map_err(|e| (StatusCode::CONFLICT, e.to_string()))?; + + validate_path_component(&folder_name, "folder name") + .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; + validate_path_component(&filename, "filename") + .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; let config = state.config.read().await; - - // Sanitize folder name and filename to prevent path traversal - if folder_name.contains("..") || folder_name.contains('/') || folder_name.contains('\\') { - return Err((StatusCode::BAD_REQUEST, "Invalid folder name".to_string())); - } - if filename.contains("..") || filename.contains('/') || filename.contains('\\') { - return Err((StatusCode::BAD_REQUEST, "Invalid filename".to_string())); - } if !filename.ends_with(".jpg") { return Err(( StatusCode::BAD_REQUEST, @@ -372,24 +347,16 @@ pub async fn delete_images_bulk( Path(folder_name): Path, Json(request): Json, ) -> Result, (StatusCode, String)> { - // Check if a job is running - { - let progress = state.progress.read().await; - if progress.status == JobStatus::Running || progress.status == JobStatus::CompilingVideo { - return Err(( - StatusCode::CONFLICT, - "Cannot delete images while a job is running".to_string(), - )); - } - } + state + .ensure_no_job_running() + .await + .map_err(|e| (StatusCode::CONFLICT, e.to_string()))?; + + validate_path_component(&folder_name, "folder name") + .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; let config = state.config.read().await; - // Sanitize folder name - if folder_name.contains("..") || folder_name.contains('/') || folder_name.contains('\\') { - return Err((StatusCode::BAD_REQUEST, "Invalid folder name".to_string())); - } - let images_dir = config.output_dir.join(&folder_name).join("images"); if !images_dir.exists() { @@ -403,8 +370,8 @@ pub async fn delete_images_bulk( let mut failed_count = 0u32; for filename in &request.filenames { - // Sanitize each filename - if filename.contains("..") || filename.contains('/') || filename.contains('\\') { + // Validate each filename + if validate_path_component(filename, "filename").is_err() { failed_count += 1; continue; } @@ -457,21 +424,16 @@ pub async fn compile_folder_video( State(state): State, Path(folder_name): Path, ) -> Result, (StatusCode, String)> { - // Check if a job is running - { - let progress = state.progress.read().await; - if progress.status == JobStatus::Running || progress.status == JobStatus::CompilingVideo { - return Err((StatusCode::CONFLICT, "A job is already running".to_string())); - } - } + state + .ensure_no_job_running() + .await + .map_err(|e| (StatusCode::CONFLICT, e.to_string()))?; + + validate_path_component(&folder_name, "folder name") + .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; let config = state.config.read().await; - // Sanitize folder name - if folder_name.contains("..") || folder_name.contains('/') || folder_name.contains('\\') { - return Err((StatusCode::BAD_REQUEST, "Invalid folder name".to_string())); - } - let folder_path = config.output_dir.join(&folder_name); let images_dir = folder_path.join("images"); diff --git a/src/web/state.rs b/src/web/state.rs index 4699a23..0eda17b 100644 --- a/src/web/state.rs +++ b/src/web/state.rs @@ -244,4 +244,55 @@ impl AppState { pub async fn clear_cancel_token(&self) { *self.cancel_token.write().await = None; } + + /// Check if a job is currently running and return an error if so. + /// + /// Use this at the start of handlers that cannot run while a job is in progress. + pub async fn ensure_no_job_running(&self) -> Result<(), JobRunningError> { + let progress = self.progress.read().await; + if progress.status == JobStatus::Running || progress.status == JobStatus::CompilingVideo { + Err(JobRunningError) + } else { + Ok(()) + } + } } + +/// Error returned when an operation cannot proceed because a job is running. +#[derive(Debug, Clone, Copy)] +pub struct JobRunningError; + +impl std::fmt::Display for JobRunningError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "Cannot perform this operation while a job is running") + } +} + +impl std::error::Error for JobRunningError {} + +/// Validate that a path component (folder name or filename) is safe. +/// +/// Returns an error if the path contains traversal sequences or separators. +pub fn validate_path_component(name: &str, component_type: &str) -> Result<(), PathValidationError> { + if name.contains("..") || name.contains('/') || name.contains('\\') { + Err(PathValidationError { + component_type: component_type.to_string(), + }) + } else { + Ok(()) + } +} + +/// Error returned when a path component fails validation. +#[derive(Debug, Clone)] +pub struct PathValidationError { + pub component_type: String, +} + +impl std::fmt::Display for PathValidationError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "Invalid {}", self.component_type) + } +} + +impl std::error::Error for PathValidationError {}