From 5449bb82f1d6575b688893c183b9cfc610b4a639 Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
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<Option<InputFormat>, 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?,
                 )),
             ))
         }