From 7539533b69b9c43ecdfcfd579fdcb8293f553450 Mon Sep 17 00:00:00 2001
From: silverpill <silverpill@firemail.cc>
Date: Fri, 6 Jan 2023 01:50:30 +0000
Subject: [PATCH] Remove ability to upload non-images using /api/v1/media
 endpoint

---
 CHANGELOG.md                    |  4 ++++
 src/mastodon_api/media/views.rs |  5 +++--
 src/mastodon_api/uploads.rs     | 28 +++++++++++++++++++---------
 src/utils/files.rs              |  7 +++++++
 4 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index aad6ead..cb1bb8b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Updated installation instructions, default mitra config and recommended nginx config.
 - Limited the number of requests made during the processing of a thread.
 
+### Removed
+
+- Removed ability to upload non-images using `/api/v1/media` endpoint.
+
 ### Fixed
 
 - Fixed post and profile page redirections.
diff --git a/src/mastodon_api/media/views.rs b/src/mastodon_api/media/views.rs
index dca716e..10db7bf 100644
--- a/src/mastodon_api/media/views.rs
+++ b/src/mastodon_api/media/views.rs
@@ -19,16 +19,17 @@ async fn create_attachment_view(
 ) -> Result<HttpResponse, HttpError> {
     let db_client = &**get_database_client(&db_pool).await?;
     let current_user = get_current_user(db_client, auth.token()).await?;
-    let (file_name, maybe_media_type) = save_b64_file(
+    let (file_name, media_type) = save_b64_file(
         &attachment_data.file,
         attachment_data.media_type.clone(),
         &config.media_dir(),
+        Some("image/"),
     )?;
     let db_attachment = create_attachment(
         db_client,
         &current_user.id,
         file_name,
-        maybe_media_type,
+        Some(media_type),
     ).await?;
     let attachment = Attachment::from_db(
         db_attachment,
diff --git a/src/mastodon_api/uploads.rs b/src/mastodon_api/uploads.rs
index c419253..47dd95e 100644
--- a/src/mastodon_api/uploads.rs
+++ b/src/mastodon_api/uploads.rs
@@ -1,7 +1,11 @@
 use std::path::Path;
 
 use crate::errors::HttpError;
-use crate::utils::files::{save_file, sniff_media_type};
+use crate::utils::files::{
+    save_file,
+    sniff_media_type,
+    SUPPORTED_MEDIA_TYPES,
+};
 
 pub const UPLOAD_MAX_SIZE: usize = 1024 * 1024 * 5;
 
@@ -33,25 +37,31 @@ impl From<UploadError> for HttpError {
 
 pub fn save_b64_file(
     b64data: &str,
-    mut maybe_media_type: Option<String>,
+    maybe_media_type: Option<String>,
     output_dir: &Path,
-) -> Result<(String, Option<String>), UploadError> {
+    maybe_expected_prefix: Option<&str>,
+) -> Result<(String, String), UploadError> {
     let data = base64::decode(b64data)?;
     if data.len() > UPLOAD_MAX_SIZE {
         return Err(UploadError::TooLarge);
     };
     // Sniff media type if not provided
-    maybe_media_type = maybe_media_type.or(sniff_media_type(&data));
-    if maybe_media_type.as_deref() == Some("image/svg+xml") {
-        // Don't treat SVG files as images
-        maybe_media_type = None;
+    let media_type = maybe_media_type.or(sniff_media_type(&data))
+        .ok_or(UploadError::InvalidMediaType)?;
+    if !SUPPORTED_MEDIA_TYPES.contains(&media_type.as_str()) {
+        return Err(UploadError::InvalidMediaType);
+    };
+    if let Some(expected_prefix) = maybe_expected_prefix {
+        if !media_type.starts_with(expected_prefix) {
+            return Err(UploadError::InvalidMediaType);
+        };
     };
     let file_name = save_file(
         data,
         output_dir,
-        maybe_media_type.as_deref(),
+        Some(&media_type),
     )?;
-    Ok((file_name, maybe_media_type))
+    Ok((file_name, media_type))
 }
 
 pub fn save_validated_b64_file(
diff --git a/src/utils/files.rs b/src/utils/files.rs
index f4d499f..31dccdf 100644
--- a/src/utils/files.rs
+++ b/src/utils/files.rs
@@ -13,6 +13,13 @@ use mime_guess::get_mime_extensions_str;
 use mime_sniffer::MimeTypeSniffer;
 use sha2::{Digest, Sha256};
 
+pub const SUPPORTED_MEDIA_TYPES: [&str; 4] = [
+    "image/jpeg",
+    "image/png",
+    "image/gif",
+    "video/mp4",
+];
+
 pub fn sniff_media_type(data: &[u8]) -> Option<String> {
     data.sniff_mime_type().map(|val| val.to_string())
 }