video thumbnails

This commit is contained in:
2025-11-25 22:34:05 +01:00
parent c4603a0b3a
commit 757abe9a0f
9 changed files with 389 additions and 15 deletions
+21
View File
@@ -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",
+1
View File
@@ -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"] }
+1
View File
@@ -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
+21
View File
@@ -27,6 +27,8 @@ pub trait ObjectStorage: Send + Sync + 'static {
async fn get_object(&self, key: &str) -> Result<Vec<u8>>;
async fn get_object_range(&self, key: &str, start: u64, end: Option<u64>) -> Result<Vec<u8>>;
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<u64>) -> Result<Vec<u8>> {
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<u64>,
) -> Result<Vec<u8>> {
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
+20
View File
@@ -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<u64>,
) -> Result<Vec<u8>> {
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);
+115 -4
View File
@@ -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<AppState>,
}
const MIME_SNIFF_BYTES: usize = 8192;
impl AnalyzePlanner {
fn new(force: bool, state: Arc<AppState>) -> Self {
Self { force, state }
@@ -114,6 +119,8 @@ impl TaskPlanner<DocumentVersionTaskContext> for AnalyzePlanner {
let document = ctx.document().await?.clone();
let mut tasks: Vec<BoxedTask<DocumentVersionTaskContext>> = 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<DocumentVersionTaskContext> for AnalyzePlanner {
}
}
struct EnsureMimeTask;
#[async_trait]
impl Task<DocumentVersionTaskContext> 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<Option<String>> {
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<String> {
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<String>) {
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"));
}
}
+29 -1
View File
@@ -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<String>) {
if let Some(document) = self.document.as_mut() {
document.mime_type = mime;
}
}
pub fn state(&self) -> &Arc<AppState> {
&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<Vec<u8>> {
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(());
+174 -4
View File
@@ -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<DocumentVersionTaskContext> 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<GeneratedAss
}
}
async fn generate_video_thumbnail(
ctx: &DocumentVersionTaskContext,
context: &ThumbnailContext,
) -> TaskResult<GeneratedAssets> {
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<GeneratedAsset, String> {
let reader = ImageReader::new(Cursor::new(bytes))
.with_guessed_format()
@@ -352,6 +400,107 @@ fn generate_pdf_assets(bytes: &[u8]) -> Result<PdfGeneratedAssets, String> {
})
}
#[derive(Deserialize)]
struct FfprobeOutput {
format: Option<FfprobeFormat>,
}
#[derive(Deserialize)]
struct FfprobeFormat {
duration: Option<String>,
}
struct VideoProbe {
duration: Option<f64>,
}
async fn probe_video_metadata(url: &str) -> Result<VideoProbe, String> {
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::<f64>().ok());
Ok(VideoProbe { duration })
}
fn pick_thumbnail_timestamp(duration: Option<f64>) -> 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<Vec<u8>, 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<GeneratedImage, String> {
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
+7 -6
View File
@@ -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<DocumentVersion>,
}
#[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);