From 5449bb82f1d6575b688893c183b9cfc610b4a639 Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 25 Sep 2022 19:34:51 -0500 Subject: [PATCH] Better identify video uploads with ffprobe --- src/ffmpeg.rs | 67 +++++++++++++++++++++++++++++++++++++++----- src/queue/cleanup.rs | 2 +- src/validate.rs | 8 ++++-- 3 files changed, 67 insertions(+), 10 deletions(-) diff --git a/src/ffmpeg.rs b/src/ffmpeg.rs index 2452e33..ca70c4e 100644 --- a/src/ffmpeg.rs +++ b/src/ffmpeg.rs @@ -2,47 +2,58 @@ use crate::{ error::{Error, UploadError}, process::Process, store::Store, + magick::ValidInputType, }; use actix_web::web::Bytes; -use tokio::io::AsyncRead; +use tokio::io::{AsyncRead, AsyncReadExt}; use tracing::instrument; -#[derive(Debug)] +#[derive(Copy, Clone, Debug)] pub(crate) enum InputFormat { Gif, Mp4, + Webm, } -#[derive(Debug)] +#[derive(Copy, Clone, Debug)] pub(crate) enum ThumbnailFormat { Jpeg, // Webp, } impl InputFormat { - fn to_ext(&self) -> &'static str { + fn to_ext(self) -> &'static str { match self { InputFormat::Gif => ".gif", InputFormat::Mp4 => ".mp4", + InputFormat::Webm => ".webm", + } + } + + pub(crate) fn to_valid_input_type(self) -> ValidInputType { + match self { + Self::Gif => ValidInputType::Gif, + Self::Mp4 => ValidInputType::Mp4, + Self::Webm => ValidInputType::Webm, } } } impl ThumbnailFormat { - fn as_codec(&self) -> &'static str { + fn as_codec(self) -> &'static str { match self { ThumbnailFormat::Jpeg => "mjpeg", // ThumbnailFormat::Webp => "webp", } } - fn to_ext(&self) -> &'static str { + fn to_ext(self) -> &'static str { match self { ThumbnailFormat::Jpeg => ".jpeg", } } - fn as_format(&self) -> &'static str { + fn as_format(self) -> &'static str { match self { ThumbnailFormat::Jpeg => "image2", // ThumbnailFormat::Webp => "webp", @@ -50,6 +61,48 @@ impl ThumbnailFormat { } } +const FORMAT_MAPPINGS: &[(&str, InputFormat)] = &[ + ("gif", InputFormat::Gif), + ("mp4", InputFormat::Mp4), + ("webm", InputFormat::Webm), +]; + +pub(crate) async fn input_type_bytes( + input: Bytes, +) -> Result, Error> { + let input_file = crate::tmp_file::tmp_file(None); + let input_file_str = input_file.to_str().ok_or(UploadError::Path)?; + crate::store::file_store::safe_create_parent(&input_file).await?; + + let mut tmp_one = crate::file::File::create(&input_file).await?; + tmp_one.write_from_bytes(input).await?; + tmp_one.close().await?; + + let process = Process::run("ffprobe", &[ + "-v", + "quiet", + "-show_entries", + "format=format_name", + "-of", + "default=noprint_wrappers=1:nokey=1", + input_file_str, + ])?; + + let mut output = Vec::new(); + process.read().read_to_end(&mut output).await?; + let formats = String::from_utf8_lossy(&output); + + tracing::info!("FORMATS: {}", formats); + + for (k, v) in FORMAT_MAPPINGS { + if formats.contains(k) { + return Ok(Some(*v)) + } + } + + Ok(None) +} + #[tracing::instrument(name = "Convert to Mp4", skip(input))] pub(crate) async fn to_mp4_bytes( input: Bytes, diff --git a/src/queue/cleanup.rs b/src/queue/cleanup.rs index 687f79b..b0a588d 100644 --- a/src/queue/cleanup.rs +++ b/src/queue/cleanup.rs @@ -68,7 +68,7 @@ where let span = tracing::error_span!("Error deleting files"); span.in_scope(|| { for error in errors { - tracing::error!("{}", format!("{}" error)); + tracing::error!("{}", format!("{}", error)); } }); } diff --git a/src/validate.rs b/src/validate.rs index 316a340..b51a9ca 100644 --- a/src/validate.rs +++ b/src/validate.rs @@ -44,7 +44,11 @@ pub(crate) async fn validate_image_bytes( enable_full_video: bool, validate: bool, ) -> Result<(ValidInputType, impl AsyncRead + Unpin), Error> { - let input_type = crate::magick::input_type_bytes(bytes.clone()).await?; + let input_type = if let Some(input_type) = crate::ffmpeg::input_type_bytes(bytes.clone()).await? { + input_type.to_valid_input_type() + } else { + crate::magick::input_type_bytes(bytes.clone()).await? + }; if !validate { return Ok((input_type, Either::left(UnvalidatedBytes::new(bytes)))); @@ -80,7 +84,7 @@ pub(crate) async fn validate_image_bytes( Ok(( ValidInputType::Mp4, Either::right(Either::left( - crate::ffmpeg::to_mp4_bytes(bytes, InputFormat::Mp4, enable_full_video).await?, + crate::ffmpeg::to_mp4_bytes(bytes, InputFormat::Webm, enable_full_video).await?, )), )) }