diff --git a/backend/Cargo.lock b/backend/Cargo.lock index cede449..06b92ce 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -436,6 +436,17 @@ dependencies = [ "nom", ] +[[package]] +name = "cfb" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d38f2da7a0a2c4ccf0065be06397cc26a81f4e528be095826eee9d4adbb8c60f" +dependencies = [ + "byteorder", + "fnv", + "uuid", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -1568,6 +1579,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "infer" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a588916bfdfd92e71cacef98a63d9b1f0d74d6599980d11894290e7ddefffcf7" +dependencies = [ + "cfb", +] + [[package]] name = "ipnet" version = "2.11.0" @@ -2134,6 +2154,7 @@ dependencies = [ "http-body-util", "hyper", "image", + "infer", "jsonwebtoken", "mime_guess", "once_cell", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index f01dc58..e02976b 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -47,6 +47,7 @@ futures-util = "0.3" url = "2.5" once_cell = "1.19" regex = "1.11" +infer = "0.19" utoipa = { version = "4.2", default-features = false, features = ["chrono", "uuid", "preserve_order"] } clap = { version = "4.5", features = ["derive"] } diff --git a/backend/Dockerfile b/backend/Dockerfile index b81335f..ffce2ca 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -123,6 +123,7 @@ RUN apt-get update \ tesseract-ocr \ ghostscript \ qpdf \ + ffmpeg \ && rm -rf /var/lib/apt/lists/* \ && mkdir -p /usr/local/lib \ && useradd --system --create-home --uid 10001 appuser diff --git a/backend/src/storage.rs b/backend/src/storage.rs index 1009cbe..c70291b 100644 --- a/backend/src/storage.rs +++ b/backend/src/storage.rs @@ -27,6 +27,8 @@ pub trait ObjectStorage: Send + Sync + 'static { async fn get_object(&self, key: &str) -> Result>; + async fn get_object_range(&self, key: &str, start: u64, end: Option) -> Result>; + async fn delete_object(&self, key: &str) -> Result<()>; } @@ -104,6 +106,15 @@ impl ObjectStorage for S3Storage { Ok(data.into_bytes().to_vec()) } + async fn get_object_range(&self, key: &str, start: u64, end: Option) -> Result> { + let data = self + .bucket + .get_object_range(key, start, end) + .await + .context("failed to download ranged object from S3")?; + Ok(data.into_bytes().to_vec()) + } + async fn delete_object(&self, key: &str) -> Result<()> { self.bucket .delete_object(key) @@ -167,6 +178,16 @@ impl TenantStorage { self.inner.get_object(&qualified).await } + pub async fn get_object_range( + &self, + key: &str, + start: u64, + end: Option, + ) -> Result> { + let qualified = self.qualify(key); + self.inner.get_object_range(&qualified, start, end).await + } + pub async fn delete_object(&self, key: &str) -> Result<()> { let qualified = self.qualify(key); self.inner.delete_object(&qualified).await diff --git a/backend/src/test_support/mod.rs b/backend/src/test_support/mod.rs index 3ac95e5..50b8f38 100644 --- a/backend/src/test_support/mod.rs +++ b/backend/src/test_support/mod.rs @@ -112,6 +112,26 @@ impl ObjectStorage for FakeStorage { .ok_or_else(|| anyhow!("object {key} missing")) } + async fn get_object_range( + &self, + key: &str, + start: u64, + end: Option, + ) -> Result> { + let guard = self.objects.lock().await; + let bytes = guard + .get(key) + .map(|obj| obj.bytes.clone()) + .ok_or_else(|| anyhow!("object {key} missing"))?; + + let start_idx = start as usize; + let end_idx = end.map(|idx| idx.saturating_add(1) as usize).unwrap_or(bytes.len()); + if start_idx >= bytes.len() { + return Ok(Vec::new()); + } + Ok(bytes[start_idx..end_idx.min(bytes.len())].to_vec()) + } + async fn delete_object(&self, key: &str) -> Result<()> { let mut guard = self.objects.lock().await; guard.remove(key); diff --git a/backend/src/workers/analyze.rs b/backend/src/workers/analyze.rs index c213dc9..7344fcf 100644 --- a/backend/src/workers/analyze.rs +++ b/backend/src/workers/analyze.rs @@ -1,8 +1,10 @@ -use std::collections::HashSet; -use std::sync::Arc; +use std::{collections::HashSet, sync::Arc, time::Duration}; use async_trait::async_trait; +use diesel::prelude::*; +use infer; use serde::Deserialize; +use tokio::task; use uuid::Uuid; use crate::{ @@ -16,7 +18,8 @@ use super::{ job_execution_from_task_error, ocr::{GenerateOcrTask, OCR_TEXT_ASSET_TYPE}, taskflow::{ - document::DocumentVersionTaskContext, BoxedTask, TaskExecutor, TaskPlanner, TaskResult, + document::DocumentVersionTaskContext, BoxedTask, Task, TaskError, TaskExecutor, + TaskPlanner, TaskResult, }, thumbnails::GenerateThumbnailsTask, JobExecution, JobHandler, @@ -99,6 +102,8 @@ struct AnalyzePlanner { state: Arc, } +const MIME_SNIFF_BYTES: usize = 8192; + impl AnalyzePlanner { fn new(force: bool, state: Arc) -> Self { Self { force, state } @@ -114,6 +119,8 @@ impl TaskPlanner for AnalyzePlanner { let document = ctx.document().await?.clone(); let mut tasks: Vec> = Vec::new(); + tasks.push(Box::new(EnsureMimeTask)); + let (thumbnail_supported, _) = determine_thumbnail_support(&document); if thumbnail_supported { tasks.push(Box::new(GenerateThumbnailsTask::new(self.force))); @@ -142,6 +149,85 @@ impl TaskPlanner for AnalyzePlanner { } } +struct EnsureMimeTask; + +#[async_trait] +impl Task for EnsureMimeTask { + fn name(&self) -> &'static str { + "ensure-mime-type" + } + + async fn execute(&self, ctx: &mut DocumentVersionTaskContext) -> TaskResult<()> { + let document = ctx.document().await?.clone(); + let current = document.mime_type.clone(); + + let guessed = guess_mime_type(ctx, &document).await?; + let desired = match guessed { + Some(mime) if current.as_deref() != Some(mime.as_str()) => Some(mime), + _ => None, + }; + + if let Some(new_mime) = desired { + update_document_mime(ctx, document.id, new_mime).await?; + } + + Ok(()) + } +} + +async fn guess_mime_type( + ctx: &mut DocumentVersionTaskContext, + document: &Document, +) -> TaskResult> { + let bytes = ctx.object_head(MIME_SNIFF_BYTES).await?; + Ok(sniff_mime(&bytes, &document.original_name)) +} + +fn sniff_mime(bytes: &[u8], original_name: &str) -> Option { + if let Some(kind) = infer::get(bytes) { + return Some(kind.mime_type().to_string()); + } + + mime_guess::from_path(original_name) + .first_raw() + .map(|value| value.to_string()) +} + +async fn update_document_mime( + ctx: &mut DocumentVersionTaskContext, + document_id: Uuid, + mime_type: String, +) -> TaskResult<()> { + let tenant_id = ctx.tenant_id(); + let state = ctx.state().clone(); + let mime_type_clone = mime_type.clone(); + + task::spawn_blocking(move || -> Result<(), String> { + let mut conn = state + .db_for_tenant(tenant_id) + .map_err(|err| format!("{err:?}"))?; + + diesel::update( + crate::schema::documents::table.filter(crate::schema::documents::id.eq(document_id)), + ) + .set(crate::schema::documents::mime_type.eq(Some(mime_type_clone))) + .execute(&mut conn) + .map_err(|err| format!("{err:?}")) + .map(|_| ()) + }) + .await + .map_err(|err| { + TaskError::retry( + Duration::from_secs(60), + format!("mime update task panicked: {err}"), + ) + })? + .map_err(|err| TaskError::retry(Duration::from_secs(30), err))?; + + ctx.set_document_mime(Some(mime_type)); + Ok(()) +} + pub(crate) fn determine_thumbnail_support(document: &Document) -> (bool, Option) { let supported_mimes: HashSet<&'static str> = [ "image/jpeg", @@ -151,6 +237,12 @@ pub(crate) fn determine_thumbnail_support(document: &Document) -> (bool, Option< "image/bmp", "image/webp", "application/pdf", + "video/mp4", + "video/quicktime", + "video/webm", + "video/x-msvideo", + "video/x-ms-wmv", + "video/x-matroska", ] .into_iter() .collect(); @@ -168,7 +260,8 @@ pub(crate) fn determine_thumbnail_support(document: &Document) -> (bool, Option< .map(|ext| ext.to_ascii_lowercase()) { let supported_exts = [ - "jpg", "jpeg", "png", "gif", "tif", "tiff", "bmp", "webp", "pdf", + "jpg", "jpeg", "png", "gif", "tif", "tiff", "bmp", "webp", "pdf", "mp4", "m4v", "mov", + "webm", "mkv", "avi", "wmv", ]; if supported_exts.contains(&ext.as_str()) { return (true, None); @@ -195,3 +288,21 @@ fn document_supports_ocr(document: &Document) -> bool { .unwrap_or(false) }) } + +#[cfg(test)] +mod tests { + use super::sniff_mime; + + #[test] + fn sniff_mime_prefers_magic_bytes() { + const PNG_HEADER: [u8; 8] = [0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]; + let mime = sniff_mime(&PNG_HEADER, "file.txt"); + assert_eq!(mime.as_deref(), Some("image/png")); + } + + #[test] + fn sniff_mime_falls_back_to_extension() { + let mime = sniff_mime(b"not enough to detect", "video.mp4"); + assert_eq!(mime.as_deref(), Some("video/mp4")); + } +} diff --git a/backend/src/workers/taskflow/document.rs b/backend/src/workers/taskflow/document.rs index 772b62f..54f76ca 100644 --- a/backend/src/workers/taskflow/document.rs +++ b/backend/src/workers/taskflow/document.rs @@ -9,7 +9,7 @@ use crate::models::{Document, DocumentVersion}; use crate::state::AppState; use crate::storage::TenantStorage; use crate::workers::common::{load_document_version, load_version_assets, LoadedAsset}; -use crate::workers::{fetch_version_object, FetchVersionError}; +use crate::workers::{check_worker_document_limit, fetch_version_object, FetchVersionError}; use super::{TaskContext, TaskError, TaskResult}; @@ -86,6 +86,12 @@ impl DocumentVersionTaskContext { self.assets = None; } + pub fn set_document_mime(&mut self, mime: Option) { + if let Some(document) = self.document.as_mut() { + document.mime_type = mime; + } + } + pub fn state(&self) -> &Arc { &self.state } @@ -158,6 +164,28 @@ impl DocumentVersionTaskContext { Ok(self.object_bytes.as_deref().expect("bytes hydrated")) } + pub async fn object_head(&mut self, max_bytes: usize) -> TaskResult> { + let version = self.version().await?.clone(); + check_worker_document_limit(version.size_bytes, self.max_document_bytes).map_err( + |(size, limit)| { + TaskError::fail(format!( + "document size {size} bytes exceeds worker limit of {limit} bytes" + )) + }, + )?; + + let end = max_bytes.saturating_sub(1) as u64; + self.storage + .get_object_range(&version.s3_key, 0, Some(end)) + .await + .map_err(|err| { + TaskError::retry( + DEFAULT_RETRY_DELAY, + format!("failed to fetch ranged object: {err}"), + ) + }) + } + async fn ensure_document_loaded(&mut self) -> TaskResult<()> { if self.document.is_some() && self.version.is_some() { return Ok(()); diff --git a/backend/src/workers/thumbnails.rs b/backend/src/workers/thumbnails.rs index 5c42149..1c053e2 100644 --- a/backend/src/workers/thumbnails.rs +++ b/backend/src/workers/thumbnails.rs @@ -1,12 +1,13 @@ -use std::{convert::TryInto, io::Cursor, panic, sync::Arc, time::Duration}; +use std::{convert::TryInto, io::Cursor, panic, process::Stdio, sync::Arc, time::Duration}; use async_trait::async_trait; use chrono::Utc; use diesel::{pg::upsert::excluded, prelude::*}; use image::{GenericImageView, ImageFormat, ImageReader}; use pdfium_render::prelude::*; +use serde::Deserialize; use serde_json::{Map, Value}; -use tokio::task; +use tokio::{process::Command, task, time::timeout}; use tracing::{info, warn}; use uuid::Uuid; @@ -17,6 +18,7 @@ use crate::{ schema::{document_assets, document_versions}, state::AppState, utils::storage_paths::document_asset_key, + workers::check_worker_document_limit, }; use super::{ @@ -28,6 +30,8 @@ pub const THUMBNAIL_WIDTH: u32 = 512; pub const THUMBNAIL_HEIGHT: u32 = 512; const RENDER_WIDTH: u32 = THUMBNAIL_WIDTH * 4; const RENDER_HEIGHT: u32 = THUMBNAIL_HEIGHT * 4; +const VIDEO_PRESIGN_TTL: Duration = Duration::from_secs(300); +const FFMPEG_TIMEOUT: Duration = Duration::from_secs(30); pub const THUMBNAIL_ASSET_TYPE: &str = "thumbnail"; pub struct GenerateThumbnailsTask { @@ -54,8 +58,12 @@ impl Task for GenerateThumbnailsTask { return Ok(()); } - let bytes = ctx.buffered_object().await?; - let generation = generate_thumbnails(&context.document, bytes).map_err(TaskError::fail)?; + let generation = if document_is_video(&context.document) { + generate_video_thumbnail(ctx, &context).await? + } else { + let bytes = ctx.buffered_object().await?; + generate_thumbnails(&context.document, bytes).map_err(TaskError::fail)? + }; if let Some(page_count) = generation.page_count { let state = ctx.state().clone(); @@ -275,6 +283,46 @@ fn generate_thumbnails(document: &Document, bytes: &[u8]) -> Result TaskResult { + if let Err((size, limit)) = + check_worker_document_limit(context.version.size_bytes, ctx.max_document_bytes()) + { + return Err(TaskError::fail(format!( + "document size {size} bytes exceeds worker limit of {limit} bytes" + ))); + } + + let presigned_url = ctx + .storage() + .presign_get_object(&context.version.s3_key, VIDEO_PRESIGN_TTL, None) + .await + .map_err(|err| { + TaskError::retry( + Duration::from_secs(30), + format!("failed to presign video for thumbnail: {err}"), + ) + })?; + + let probe = probe_video_metadata(&presigned_url) + .await + .map_err(TaskError::fail)?; + let timestamp = pick_thumbnail_timestamp(probe.duration); + + let frame_bytes = extract_video_frame(&presigned_url, timestamp) + .await + .map_err(TaskError::fail)?; + + let thumbnail = generate_image_assets(&frame_bytes).map_err(TaskError::fail)?; + + Ok(GeneratedAssets { + thumbnail, + page_count: None, + }) +} + fn generate_image_assets(bytes: &[u8]) -> Result { let reader = ImageReader::new(Cursor::new(bytes)) .with_guessed_format() @@ -352,6 +400,107 @@ fn generate_pdf_assets(bytes: &[u8]) -> Result { }) } +#[derive(Deserialize)] +struct FfprobeOutput { + format: Option, +} + +#[derive(Deserialize)] +struct FfprobeFormat { + duration: Option, +} + +struct VideoProbe { + duration: Option, +} + +async fn probe_video_metadata(url: &str) -> Result { + let mut cmd = Command::new("ffprobe"); + cmd.arg("-v") + .arg("error") + .arg("-show_entries") + .arg("format=duration") + .arg("-of") + .arg("json") + .arg(url) + .stdout(Stdio::piped()); + + let output = timeout(FFMPEG_TIMEOUT, cmd.output()) + .await + .map_err(|_| "ffprobe timed out".to_string())? + .map_err(|err| format!("ffprobe failed to start: {err}"))?; + + if !output.status.success() { + return Err(format!( + "ffprobe exited with status {}: {}", + output.status, + String::from_utf8_lossy(&output.stderr) + )); + } + + let parsed: FfprobeOutput = serde_json::from_slice(&output.stdout) + .map_err(|err| format!("failed to parse ffprobe output: {err}"))?; + + let duration = parsed + .format + .and_then(|format| format.duration) + .and_then(|dur| dur.parse::().ok()); + + Ok(VideoProbe { duration }) +} + +fn pick_thumbnail_timestamp(duration: Option) -> f64 { + if let Some(duration) = duration { + if duration.is_finite() && duration > 0.0 { + let target = duration * 0.2; + let end = (duration - 1.0).max(0.0); + return target.max(2.0).min(end).max(0.0); + } + } + 2.0 +} + +async fn extract_video_frame(url: &str, timestamp_secs: f64) -> Result, String> { + let timestamp_arg = format!("{timestamp_secs:.3}"); + + let mut cmd = Command::new("ffmpeg"); + cmd.arg("-hide_banner") + .arg("-loglevel") + .arg("error") + .arg("-nostdin") + .arg("-ss") + .arg(timestamp_arg) + .arg("-i") + .arg(url) + .arg("-frames:v") + .arg("1") + .arg("-f") + .arg("image2pipe") + .arg("-vcodec") + .arg("png") + .arg("-") + .stdout(Stdio::piped()); + + let output = timeout(FFMPEG_TIMEOUT, cmd.output()) + .await + .map_err(|_| "ffmpeg timed out".to_string())? + .map_err(|err| format!("ffmpeg failed to start: {err}"))?; + + if !output.status.success() { + return Err(format!( + "ffmpeg exited with status {}: {}", + output.status, + String::from_utf8_lossy(&output.stderr) + )); + } + + if output.stdout.is_empty() { + return Err("ffmpeg produced no frame data".to_string()); + } + + Ok(output.stdout) +} + fn encode_dynamic_image(image: image::DynamicImage) -> Result { let (width, height) = image.dimensions(); let mut cursor = Cursor::new(Vec::new()); @@ -461,6 +610,27 @@ fn persist_document_page_count( Ok(()) } +fn document_is_video(document: &Document) -> bool { + const VIDEO_MIME_TYPES: [&str; 6] = [ + "video/mp4", + "video/quicktime", + "video/webm", + "video/x-msvideo", + "video/x-ms-wmv", + "video/x-matroska", + ]; + if let Some(mime) = document.mime_type.as_deref() { + if VIDEO_MIME_TYPES + .iter() + .any(|candidate| mime.eq_ignore_ascii_case(candidate)) + { + return true; + } + } + + false +} + fn document_is_pdf(document: &Document) -> bool { document .mime_type diff --git a/backend/tests/download_flow.rs b/backend/tests/download_flow.rs index d7a1204..148fcf6 100644 --- a/backend/tests/download_flow.rs +++ b/backend/tests/download_flow.rs @@ -3,22 +3,22 @@ use axum::http::StatusCode; use papercrate::test_support::{acquire_db_lock, body_to_vec, TestApp, TestUserRole}; use serde::Deserialize; -#[derive(Deserialize)] +#[derive(Clone, Deserialize)] struct DocumentDetail { document: DocumentInfo, } -#[derive(Deserialize)] +#[derive(Clone, Deserialize)] struct DocumentInfo { current_version: Option, } -#[derive(Deserialize)] +#[derive(Clone, Deserialize)] struct DocumentVersion { download: DownloadLink, } -#[derive(Deserialize)] +#[derive(Clone, Deserialize)] struct DownloadLink { url: String, expires_at: i64, @@ -47,14 +47,15 @@ async fn document_download_redirects_when_proxy_disabled() -> Result<()> { assert_eq!(upload.status(), StatusCode::CREATED); let body = body_to_vec(upload.into_body()).await?; let detail: DocumentDetail = serde_json::from_slice(&body)?; - let download_path = detail + let download_link = detail .document .current_version .as_ref() .expect("missing version") .download - .url .clone(); + assert!(download_link.expires_at > 0); + let download_path = download_link.url.clone(); let redirect = app.get(&download_path, None).await?; assert_eq!(redirect.status(), StatusCode::TEMPORARY_REDIRECT);