From 7c07cd79bc56724a438055b20cc42007abfa2cf2 Mon Sep 17 00:00:00 2001
From: silverpill <silverpill@firemail.cc>
Date: Wed, 4 Jan 2023 20:04:10 +0000
Subject: [PATCH] Use "mediaType" property value to determine file extension
 when saving downloaded media

---
 CHANGELOG.md                        |  1 +
 src/activitypub/actors/helpers.rs   | 28 ++++++++++++++++++++++------
 src/activitypub/actors/types.rs     |  3 +++
 src/activitypub/fetcher/fetchers.rs |  6 +++++-
 src/activitypub/handlers/create.rs  |  1 +
 5 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 46bb05e..8b3bd4b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 
 - Save downloaded media as "unknown" if its media type is not supported.
+- Use `mediaType` property value to determine file extension when saving downloaded media.
 
 ### Removed
 
diff --git a/src/activitypub/actors/helpers.rs b/src/activitypub/actors/helpers.rs
index f859e08..159b469 100644
--- a/src/activitypub/actors/helpers.rs
+++ b/src/activitypub/actors/helpers.rs
@@ -27,9 +27,17 @@ async fn fetch_actor_images(
     default_banner: Option<ProfileImage>,
 ) -> (Option<ProfileImage>, Option<ProfileImage>) {
     let maybe_avatar = if let Some(icon) = &actor.icon {
-        match fetch_file(instance, &icon.url, media_dir).await {
-            Ok((file_name, _)) => {
-                let image = ProfileImage { file_name, media_type: None };
+        match fetch_file(
+            instance,
+            &icon.url,
+            icon.media_type.as_deref(),
+            media_dir,
+        ).await {
+            Ok((file_name, maybe_media_type)) => {
+                let image = ProfileImage {
+                    file_name,
+                    media_type: maybe_media_type,
+                };
                 Some(image)
             },
             Err(error) => {
@@ -41,9 +49,17 @@ async fn fetch_actor_images(
         None
     };
     let maybe_banner = if let Some(image) = &actor.image {
-        match fetch_file(instance, &image.url, media_dir).await {
-            Ok((file_name, _)) => {
-                let image = ProfileImage { file_name, media_type: None };
+        match fetch_file(
+            instance,
+            &image.url,
+            image.media_type.as_deref(),
+            media_dir,
+        ).await {
+            Ok((file_name, maybe_media_type)) => {
+                let image = ProfileImage {
+                    file_name,
+                    media_type: maybe_media_type,
+                };
                 Some(image)
             },
             Err(error) => {
diff --git a/src/activitypub/actors/types.rs b/src/activitypub/actors/types.rs
index 52f45b7..637028b 100644
--- a/src/activitypub/actors/types.rs
+++ b/src/activitypub/actors/types.rs
@@ -49,6 +49,7 @@ pub struct ActorImage {
     #[serde(rename = "type")]
     object_type: String,
     pub url: String,
+    pub media_type: Option<String>,
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
@@ -240,6 +241,7 @@ pub fn get_local_actor(
             let actor_image = ActorImage {
                 object_type: IMAGE.to_string(),
                 url: get_file_url(instance_url, &image.file_name),
+                media_type: None,
             };
             Some(actor_image)
         },
@@ -250,6 +252,7 @@ pub fn get_local_actor(
             let actor_image = ActorImage {
                 object_type: IMAGE.to_string(),
                 url: get_file_url(instance_url, &image.file_name),
+                media_type: None,
             };
             Some(actor_image)
         },
diff --git a/src/activitypub/fetcher/fetchers.rs b/src/activitypub/fetcher/fetchers.rs
index 38e97fc..f5237dd 100644
--- a/src/activitypub/fetcher/fetchers.rs
+++ b/src/activitypub/fetcher/fetchers.rs
@@ -107,6 +107,7 @@ const FILE_MAX_SIZE: u64 = 1024 * 1024 * 20;
 pub async fn fetch_file(
     instance: &Instance,
     url: &str,
+    maybe_media_type: Option<&str>,
     output_dir: &Path,
 ) -> Result<(String, Option<String>), FetchError> {
     let client = build_client(instance)?;
@@ -122,7 +123,10 @@ pub async fn fetch_file(
     if file_data.len() > FILE_MAX_SIZE as usize {
         return Err(FetchError::OtherError("file is too large"));
     };
-    let maybe_media_type = sniff_media_type(&file_data)
+    let maybe_media_type = maybe_media_type
+        .map(|media_type| media_type.to_string())
+        // Sniff media type if not provided
+        .or(sniff_media_type(&file_data))
         // Remove media type if it is not supported to prevent XSS
         .filter(|media_type| {
             if SUPPORTED_MEDIA_TYPES.contains(&media_type.as_str()) {
diff --git a/src/activitypub/handlers/create.rs b/src/activitypub/handlers/create.rs
index 73a7d02..8ddca30 100644
--- a/src/activitypub/handlers/create.rs
+++ b/src/activitypub/handlers/create.rs
@@ -178,6 +178,7 @@ pub async fn handle_note(
             let (file_name, maybe_media_type) = fetch_file(
                 instance,
                 &attachment_url,
+                attachment.media_type.as_deref(),
                 media_dir,
             ).await
                 .map_err(|err| {