Save sizes of media attachments and other files to database

CHANGELOG.md

@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
 
+### Added
+
+- Save sizes of media attachments and other files to database.
+
 ### Security
 
 - Validate emoji name before saving.

migrations/V0040__media_attachment__add_file_size.sql

@@ -0,0 +1 @@
+ALTER TABLE media_attachment ADD COLUMN file_size INTEGER;

migrations/schema.sql

@@ -106,8 +106,9 @@ CREATE TABLE post_reaction (
 CREATE TABLE media_attachment (
     id UUID PRIMARY KEY,
     owner_id UUID NOT NULL REFERENCES actor_profile (id) ON DELETE CASCADE,
-    media_type VARCHAR(50),
     file_name VARCHAR(200) NOT NULL,
+    file_size INTEGER,
+    media_type VARCHAR(50),
     ipfs_cid VARCHAR(200),
     post_id UUID REFERENCES post (id) ON DELETE CASCADE,
     created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now()

src/activitypub/actors/helpers.rs

@@ -9,9 +9,9 @@ use crate::activitypub::{
 };
 use crate::config::Instance;
 use crate::database::DatabaseClient;
-use crate::models::profiles::{
-    queries::{create_profile, update_profile},
-    types::{
+use crate::models::{
+    profiles::queries::{create_profile, update_profile},
+    profiles::types::{
         DbActorProfile,
         ProfileImage,
         ProfileCreateData,
@@ -36,9 +36,10 @@ async fn fetch_actor_images(
             ACTOR_IMAGE_MAX_SIZE,
             media_dir,
         ).await {
-            Ok((file_name, maybe_media_type)) => {
+            Ok((file_name, file_size, maybe_media_type)) => {
                 let image = ProfileImage::new(
                     file_name,
+                    file_size,
                     maybe_media_type,
                 );
                 Some(image)
@@ -59,9 +60,10 @@ async fn fetch_actor_images(
             ACTOR_IMAGE_MAX_SIZE,
             media_dir,
         ).await {
-            Ok((file_name, maybe_media_type)) => {
+            Ok((file_name, file_size, maybe_media_type)) => {
                 let image = ProfileImage::new(
                     file_name,
+                    file_size,
                     maybe_media_type,
                 );
                 Some(image)

src/activitypub/fetcher/fetchers.rs

@@ -110,7 +110,7 @@ pub async fn fetch_file(
     maybe_media_type: Option<&str>,
     file_max_size: usize,
     output_dir: &Path,
-) -> Result<(String, Option<String>), FetchError> {
+) -> Result<(String, usize, Option<String>), FetchError> {
     let client = build_client(instance)?;
     let request_builder =
         build_request(instance, client, Method::GET, url);
@@ -123,7 +123,8 @@ pub async fn fetch_file(
         };
     };
     let file_data = response.bytes().await?;
-    if file_data.len() > file_max_size {
+    let file_size = file_data.len();
+    if file_size > file_max_size {
         return Err(FetchError::OtherError("file is too large"));
     };
     let maybe_media_type = maybe_media_type
@@ -148,7 +149,7 @@ pub async fn fetch_file(
         output_dir,
         maybe_media_type.as_deref(),
     )?;
-    Ok((file_name, maybe_media_type))
+    Ok((file_name, file_size, maybe_media_type))
 }
 
 pub async fn perform_webfinger_query(

src/activitypub/handlers/create.rs

@@ -215,7 +215,7 @@ pub async fn handle_note(
             };
             let attachment_url = attachment.url
                 .ok_or(ValidationError("attachment URL is missing"))?;
-            let (file_name, maybe_media_type) = fetch_file(
+            let (file_name, file_size, maybe_media_type) = fetch_file(
                 instance,
                 &attachment_url,
                 attachment.media_type.as_deref(),
@@ -227,18 +227,19 @@ pub async fn handle_note(
                     ValidationError("failed to fetch attachment")
                 })?;
             log::info!("downloaded attachment {}", attachment_url);
-            downloaded.push((file_name, maybe_media_type));
+            downloaded.push((file_name, file_size, maybe_media_type));
             // Stop downloading if limit is reached
             if downloaded.len() >= ATTACHMENTS_MAX_NUM {
                 log::warn!("too many attachments");
                 break;
             };
         };
-        for (file_name, maybe_media_type) in downloaded {
+        for (file_name, file_size, maybe_media_type) in downloaded {
             let db_attachment = create_attachment(
                 db_client,
                 &author.id,
                 file_name,
+                file_size,
                 maybe_media_type,
             ).await?;
             attachments.push(db_attachment.id);
@@ -411,7 +412,7 @@ pub async fn handle_note(
                     Err(DatabaseError::NotFound("emoji")) => None,
                     Err(other_error) => return Err(other_error.into()),
                 };
-                let (file_name, maybe_media_type) = match fetch_file(
+                let (file_name, file_size, maybe_media_type) = match fetch_file(
                     instance,
                     &tag.icon.url,
                     tag.icon.media_type.as_deref(),
@@ -437,7 +438,7 @@ pub async fn handle_note(
                     },
                 };
                 log::info!("downloaded emoji {}", tag.icon.url);
-                let image = EmojiImage { file_name, media_type };
+                let image = EmojiImage { file_name, file_size, media_type };
                 let emoji = if let Some(emoji_id) = maybe_emoji_id {
                     update_emoji(
                         db_client,

src/mastodon_api/accounts/types.rs

@@ -237,7 +237,7 @@ fn process_b64_image_field_value(
                 None
             } else {
                 // Decode and save file
-                let (file_name, media_type) = save_b64_file(
+                let (file_name, file_size, media_type) = save_b64_file(
                     &b64_data,
                     form_media_type,
                     output_dir,
@@ -245,6 +245,7 @@ fn process_b64_image_field_value(
                 )?;
                 let image = ProfileImage::new(
                     file_name,
+                    file_size,
                     Some(media_type),
                 );
                 Some(image)
@@ -487,7 +488,7 @@ mod tests {
     #[test]
     fn test_create_account_from_profile() {
         let profile = DbActorProfile {
-            avatar: Some(ProfileImage::new("test".to_string(), None)),
+            avatar: Some(ProfileImage::new("test".to_string(), 1000, None)),
             ..Default::default()
         };
         let account = Account::from_profile(profile, INSTANCE_URL);

src/mastodon_api/media/views.rs

@@ -19,7 +19,7 @@ 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, media_type) = save_b64_file(
+    let (file_name, file_size, media_type) = save_b64_file(
         &attachment_data.file,
         attachment_data.media_type.clone(),
         &config.media_dir(),
@@ -29,6 +29,7 @@ async fn create_attachment_view(
         db_client,
         &current_user.id,
         file_name,
+        file_size,
         Some(media_type),
     ).await?;
     let attachment = Attachment::from_db(

src/mastodon_api/uploads.rs

@@ -40,13 +40,14 @@ pub fn save_b64_file(
     maybe_media_type: Option<String>,
     output_dir: &Path,
     maybe_expected_prefix: Option<&str>,
-) -> Result<(String, String), UploadError> {
-    let data = base64::decode(b64data)?;
-    if data.len() > UPLOAD_MAX_SIZE {
+) -> Result<(String, usize, String), UploadError> {
+    let file_data = base64::decode(b64data)?;
+    let file_size = file_data.len();
+    if file_size > UPLOAD_MAX_SIZE {
         return Err(UploadError::TooLarge);
     };
     // Sniff media type if not provided
-    let media_type = maybe_media_type.or(sniff_media_type(&data))
+    let media_type = maybe_media_type.or(sniff_media_type(&file_data))
         .ok_or(UploadError::InvalidMediaType)?;
     if !SUPPORTED_MEDIA_TYPES.contains(&media_type.as_str()) {
         return Err(UploadError::InvalidMediaType);
@@ -57,9 +58,9 @@ pub fn save_b64_file(
         };
     };
     let file_name = save_file(
-        data,
+        file_data,
         output_dir,
         Some(&media_type),
     )?;
-    Ok((file_name, media_type))
+    Ok((file_name, file_size, media_type))
 }

src/models/attachments/queries.rs

@@ -14,16 +14,31 @@ pub async fn create_attachment(
     db_client: &impl DatabaseClient,
     owner_id: &Uuid,
     file_name: String,
+    file_size: usize,
     media_type: Option<String>,
 ) -> Result<DbMediaAttachment, DatabaseError> {
     let attachment_id = new_uuid();
+    let file_size: i32 = file_size.try_into()
+        .expect("value should be within bounds");
     let inserted_row = db_client.query_one(
         "
-        INSERT INTO media_attachment (id, owner_id, media_type, file_name)
-        VALUES ($1, $2, $3, $4)
+        INSERT INTO media_attachment (
+            id,
+            owner_id,
+            file_name,
+            file_size,
+            media_type
+        )
+        VALUES ($1, $2, $3, $4, $5)
         RETURNING media_attachment
         ",
-        &[&attachment_id, &owner_id, &media_type, &file_name],
+        &[
+            &attachment_id,
+            &owner_id,
+            &file_name,
+            &file_size,
+            &media_type,
+        ],
     ).await?;
     let db_attachment: DbMediaAttachment = inserted_row.try_get("media_attachment")?;
     Ok(db_attachment)

src/models/attachments/types.rs

@@ -7,8 +7,9 @@ use uuid::Uuid;
 pub struct DbMediaAttachment {
     pub id: Uuid,
     pub owner_id: Uuid,
-    pub media_type: Option<String>,
     pub file_name: String,
+    pub file_size: Option<i32>,
+    pub media_type: Option<String>,
     pub ipfs_cid: Option<String>,
     pub post_id: Option<Uuid>,
     pub created_at: DateTime<Utc>,

src/models/emojis/queries.rs

@@ -129,6 +129,7 @@ mod tests {
         let hostname = "example.org";
         let image = EmojiImage {
             file_name: "test.png".to_string(),
+            file_size: 10000,
             media_type: "image/png".to_string(),
         };
         let object_id = "https://example.org/emojis/test";
@@ -156,6 +157,7 @@ mod tests {
         let db_client = &create_test_database().await;
         let image = EmojiImage {
             file_name: "test.png".to_string(),
+            file_size: 10000,
             media_type: "image/png".to_string(),
         };
         let emoji = create_emoji(

src/models/emojis/types.rs

@@ -4,10 +4,15 @@ use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
 use crate::database::json_macro::{json_from_sql, json_to_sql};
+use super::validators::EMOJI_MAX_SIZE;
+
+fn default_emoji_file_size() -> usize { EMOJI_MAX_SIZE }
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct EmojiImage {
     pub file_name: String,
+    #[serde(default = "default_emoji_file_size")]
+    pub file_size: usize,
     pub media_type: String,
 }
 

src/models/profiles/types.rs

@@ -34,16 +34,19 @@ use super::validators::{
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct ProfileImage {
     pub file_name: String,
+    pub file_size: Option<usize>,
     pub media_type: Option<String>,
 }
 
 impl ProfileImage {
     pub fn new(
         file_name: String,
+        file_size: usize,
         media_type: Option<String>,
     ) -> Self {
         Self {
             file_name,
+            file_size: Some(file_size),
             media_type,
         }
     }