Improve file limit handling (#4242)

* Improve file limit handling

* Oops

* Update PostgreSQL migration

* Review comments

---------

Co-authored-by: BlackDex <black.dex@gmail.com>

src/api/admin.rs

@@ -15,7 +15,7 @@ use rocket::{
 use crate::{
     api::{
         core::{log_event, two_factor},
-        unregister_push_device, ApiResult, EmptyResult, JsonResult, Notify, NumberOrString,
+        unregister_push_device, ApiResult, EmptyResult, JsonResult, Notify,
     },
     auth::{decode_admin, encode_jwt, generate_admin_claims, ClientIp},
     config::ConfigBuilder,
@@ -24,6 +24,7 @@ use crate::{
     mail,
     util::{
         docker_base_image, format_naive_datetime_local, get_display_size, get_reqwest_client, is_running_in_docker,
+        NumberOrString,
     },
     CONFIG, VERSION,
 };
@@ -345,7 +346,7 @@ async fn users_overview(_token: AdminToken, mut conn: DbConn) -> ApiResult<Html<
         let mut usr = u.to_json(&mut conn).await;
         usr["cipher_count"] = json!(Cipher::count_owned_by_user(&u.uuid, &mut conn).await);
         usr["attachment_count"] = json!(Attachment::count_by_user(&u.uuid, &mut conn).await);
-        usr["attachment_size"] = json!(get_display_size(Attachment::size_by_user(&u.uuid, &mut conn).await as i32));
+        usr["attachment_size"] = json!(get_display_size(Attachment::size_by_user(&u.uuid, &mut conn).await));
         usr["user_enabled"] = json!(u.enabled);
         usr["created_at"] = json!(format_naive_datetime_local(&u.created_at, DT_FMT));
         usr["last_active"] = match u.last_active(&mut conn).await {
@@ -549,7 +550,7 @@ async fn organizations_overview(_token: AdminToken, mut conn: DbConn) -> ApiResu
         org["group_count"] = json!(Group::count_by_org(&o.uuid, &mut conn).await);
         org["event_count"] = json!(Event::count_by_org(&o.uuid, &mut conn).await);
         org["attachment_count"] = json!(Attachment::count_by_org(&o.uuid, &mut conn).await);
-        org["attachment_size"] = json!(get_display_size(Attachment::size_by_org(&o.uuid, &mut conn).await as i32));
+        org["attachment_size"] = json!(get_display_size(Attachment::size_by_org(&o.uuid, &mut conn).await));
         organizations_json.push(org);
     }
 

src/api/core/accounts.rs

@@ -6,12 +6,14 @@ use serde_json::Value;
 use crate::{
     api::{
         core::log_user_event, register_push_device, unregister_push_device, AnonymousNotify, EmptyResult, JsonResult,
-        JsonUpcase, Notify, NumberOrString, PasswordOrOtpData, UpdateType,
+        JsonUpcase, Notify, PasswordOrOtpData, UpdateType,
     },
     auth::{decode_delete, decode_invite, decode_verify_email, ClientHeaders, Headers},
     crypto,
     db::{models::*, DbConn},
-    mail, CONFIG,
+    mail,
+    util::NumberOrString,
+    CONFIG,
 };
 
 use rocket::{

src/api/core/ciphers.rs

@@ -1,6 +1,7 @@
 use std::collections::{HashMap, HashSet};
 
 use chrono::{NaiveDateTime, Utc};
+use num_traits::ToPrimitive;
 use rocket::fs::TempFile;
 use rocket::serde::json::Json;
 use rocket::{
@@ -956,7 +957,7 @@ async fn get_attachment(uuid: &str, attachment_id: &str, headers: Headers, mut c
 struct AttachmentRequestData {
     Key: String,
     FileName: String,
-    FileSize: i32,
+    FileSize: i64,
     AdminRequest: Option<bool>, // true when attaching from an org vault view
 }
 
@@ -985,8 +986,11 @@ async fn post_attachment_v2(
         err!("Cipher is not write accessible")
     }
 
-    let attachment_id = crypto::generate_attachment_id();
     let data: AttachmentRequestData = data.into_inner().data;
+    if data.FileSize < 0 {
+        err!("Attachment size can't be negative")
+    }
+    let attachment_id = crypto::generate_attachment_id();
     let attachment =
         Attachment::new(attachment_id.clone(), cipher.uuid.clone(), data.FileName, data.FileSize, Some(data.Key));
     attachment.save(&mut conn).await.expect("Error saving attachment");
@@ -1028,6 +1032,15 @@ async fn save_attachment(
     mut conn: DbConn,
     nt: Notify<'_>,
 ) -> Result<(Cipher, DbConn), crate::error::Error> {
+    let mut data = data.into_inner();
+
+    let Some(size) = data.data.len().to_i64() else {
+        err!("Attachment data size overflow");
+    };
+    if size < 0 {
+        err!("Attachment size can't be negative")
+    }
+
     let cipher = match Cipher::find_by_uuid(cipher_uuid, &mut conn).await {
         Some(cipher) => cipher,
         None => err!("Cipher doesn't exist"),
@@ -1040,19 +1053,29 @@ async fn save_attachment(
     // In the v2 API, the attachment record has already been created,
     // so the size limit needs to be adjusted to account for that.
     let size_adjust = match &attachment {
-        None => 0,                         // Legacy API
-        Some(a) => i64::from(a.file_size), // v2 API
+        None => 0,              // Legacy API
+        Some(a) => a.file_size, // v2 API
     };
 
     let size_limit = if let Some(ref user_uuid) = cipher.user_uuid {
         match CONFIG.user_attachment_limit() {
             Some(0) => err!("Attachments are disabled"),
             Some(limit_kb) => {
-                let left = (limit_kb * 1024) - Attachment::size_by_user(user_uuid, &mut conn).await + size_adjust;
+                let already_used = Attachment::size_by_user(user_uuid, &mut conn).await;
+                let left = limit_kb
+                    .checked_mul(1024)
+                    .and_then(|l| l.checked_sub(already_used))
+                    .and_then(|l| l.checked_add(size_adjust));
+
+                let Some(left) = left else {
+                    err!("Attachment size overflow");
+                };
+
                 if left <= 0 {
                     err!("Attachment storage limit reached! Delete some attachments to free up space")
                 }
-                Some(left as u64)
+
+                Some(left)
             }
             None => None,
         }
@@ -1060,11 +1083,21 @@ async fn save_attachment(
         match CONFIG.org_attachment_limit() {
             Some(0) => err!("Attachments are disabled"),
             Some(limit_kb) => {
-                let left = (limit_kb * 1024) - Attachment::size_by_org(org_uuid, &mut conn).await + size_adjust;
+                let already_used = Attachment::size_by_org(org_uuid, &mut conn).await;
+                let left = limit_kb
+                    .checked_mul(1024)
+                    .and_then(|l| l.checked_sub(already_used))
+                    .and_then(|l| l.checked_add(size_adjust));
+
+                let Some(left) = left else {
+                    err!("Attachment size overflow");
+                };
+
                 if left <= 0 {
                     err!("Attachment storage limit reached! Delete some attachments to free up space")
                 }
-                Some(left as u64)
+
+                Some(left)
             }
             None => None,
         }
@@ -1072,10 +1105,8 @@ async fn save_attachment(
         err!("Cipher is neither owned by a user nor an organization");
     };
 
-    let mut data = data.into_inner();
-
     if let Some(size_limit) = size_limit {
-        if data.data.len() > size_limit {
+        if size > size_limit {
             err!("Attachment storage limit exceeded with this file");
         }
     }
@@ -1085,20 +1116,19 @@ async fn save_attachment(
         None => crypto::generate_attachment_id(),  // Legacy API
     };
 
-    let folder_path = tokio::fs::canonicalize(&CONFIG.attachments_folder()).await?.join(cipher_uuid);
-    let file_path = folder_path.join(&file_id);
-    tokio::fs::create_dir_all(&folder_path).await?;
-
-    let size = data.data.len() as i32;
     if let Some(attachment) = &mut attachment {
         // v2 API
 
         // Check the actual size against the size initially provided by
         // the client. Upstream allows +/- 1 MiB deviation from this
         // size, but it's not clear when or why this is needed.
-        const LEEWAY: i32 = 1024 * 1024; // 1 MiB
-        let min_size = attachment.file_size - LEEWAY;
-        let max_size = attachment.file_size + LEEWAY;
+        const LEEWAY: i64 = 1024 * 1024; // 1 MiB
+        let Some(min_size) = attachment.file_size.checked_add(LEEWAY) else {
+            err!("Invalid attachment size min")
+        };
+        let Some(max_size) = attachment.file_size.checked_sub(LEEWAY) else {
+            err!("Invalid attachment size max")
+        };
 
         if min_size <= size && size <= max_size {
             if size != attachment.file_size {
@@ -1113,6 +1143,10 @@ async fn save_attachment(
         }
     } else {
         // Legacy API
+
+        // SAFETY: This value is only stored in the database and is not used to access the file system.
+        // As a result, the conditions specified by Rocket [0] are met and this is safe to use.
+        // [0]: https://docs.rs/rocket/latest/rocket/fs/struct.FileName.html#-danger-
         let encrypted_filename = data.data.raw_name().map(|s| s.dangerous_unsafe_unsanitized_raw().to_string());
 
         if encrypted_filename.is_none() {
@@ -1122,10 +1156,14 @@ async fn save_attachment(
             err!("No attachment key provided")
         }
         let attachment =
-            Attachment::new(file_id, String::from(cipher_uuid), encrypted_filename.unwrap(), size, data.key);
+            Attachment::new(file_id.clone(), String::from(cipher_uuid), encrypted_filename.unwrap(), size, data.key);
         attachment.save(&mut conn).await.expect("Error saving attachment");
     }
 
+    let folder_path = tokio::fs::canonicalize(&CONFIG.attachments_folder()).await?.join(cipher_uuid);
+    let file_path = folder_path.join(&file_id);
+    tokio::fs::create_dir_all(&folder_path).await?;
+
     if let Err(_err) = data.data.persist_to(&file_path).await {
         data.data.move_copy_to(file_path).await?
     }

src/api/core/emergency_access.rs

@@ -5,11 +5,13 @@ use serde_json::Value;
 use crate::{
     api::{
         core::{CipherSyncData, CipherSyncType},
-        EmptyResult, JsonResult, JsonUpcase, NumberOrString,
+        EmptyResult, JsonResult, JsonUpcase,
     },
     auth::{decode_emergency_access_invite, Headers},
     db::{models::*, DbConn, DbPool},
-    mail, CONFIG,
+    mail,
+    util::NumberOrString,
+    CONFIG,
 };
 
 pub fn routes() -> Vec<Route> {

src/api/core/organizations.rs

@@ -6,14 +6,13 @@ use serde_json::Value;
 use crate::{
     api::{
         core::{log_event, two_factor, CipherSyncData, CipherSyncType},
-        EmptyResult, JsonResult, JsonUpcase, JsonUpcaseVec, JsonVec, Notify, NumberOrString, PasswordOrOtpData,
-        UpdateType,
+        EmptyResult, JsonResult, JsonUpcase, JsonUpcaseVec, JsonVec, Notify, PasswordOrOtpData, UpdateType,
     },
     auth::{decode_invite, AdminHeaders, Headers, ManagerHeaders, ManagerHeadersLoose, OwnerHeaders},
     db::{models::*, DbConn},
     error::Error,
     mail,
-    util::convert_json_key_lcase_first,
+    util::{convert_json_key_lcase_first, NumberOrString},
     CONFIG,
 };
 

src/api/core/sends.rs

@@ -1,6 +1,7 @@
 use std::path::Path;
 
 use chrono::{DateTime, Duration, Utc};
+use num_traits::ToPrimitive;
 use rocket::form::Form;
 use rocket::fs::NamedFile;
 use rocket::fs::TempFile;
@@ -8,17 +9,17 @@ use rocket::serde::json::Json;
 use serde_json::Value;
 
 use crate::{
-    api::{ApiResult, EmptyResult, JsonResult, JsonUpcase, Notify, NumberOrString, UpdateType},
+    api::{ApiResult, EmptyResult, JsonResult, JsonUpcase, Notify, UpdateType},
     auth::{ClientIp, Headers, Host},
     db::{models::*, DbConn, DbPool},
-    util::SafeString,
+    util::{NumberOrString, SafeString},
     CONFIG,
 };
 
 const SEND_INACCESSIBLE_MSG: &str = "Send does not exist or is no longer available";
 
 // The max file size allowed by Bitwarden clients and add an extra 5% to avoid issues
-const SIZE_525_MB: u64 = 550_502_400;
+const SIZE_525_MB: i64 = 550_502_400;
 
 pub fn routes() -> Vec<rocket::Route> {
     routes![
@@ -216,30 +217,41 @@ async fn post_send_file(data: Form<UploadData<'_>>, headers: Headers, mut conn:
     } = data.into_inner();
     let model = model.into_inner().data;
 
+    let Some(size) = data.len().to_i64() else {
+        err!("Invalid send size");
+    };
+    if size < 0 {
+        err!("Send size can't be negative")
+    }
+
     enforce_disable_hide_email_policy(&model, &headers, &mut conn).await?;
 
-    let size_limit = match CONFIG.user_attachment_limit() {
+    let size_limit = match CONFIG.user_send_limit() {
         Some(0) => err!("File uploads are disabled"),
         Some(limit_kb) => {
-            let left = (limit_kb * 1024) - Attachment::size_by_user(&headers.user.uuid, &mut conn).await;
+            let Some(already_used) = Send::size_by_user(&headers.user.uuid, &mut conn).await else {
+                err!("Existing sends overflow")
+            };
+            let Some(left) = limit_kb.checked_mul(1024).and_then(|l| l.checked_sub(already_used)) else {
+                err!("Send size overflow");
+            };
             if left <= 0 {
-                err!("Attachment storage limit reached! Delete some attachments to free up space")
+                err!("Send storage limit reached! Delete some sends to free up space")
             }
-            std::cmp::Ord::max(left as u64, SIZE_525_MB)
+            i64::clamp(left, 0, SIZE_525_MB)
         }
         None => SIZE_525_MB,
     };
 
+    if size > size_limit {
+        err!("Send storage limit exceeded with this file");
+    }
+
     let mut send = create_send(model, headers.user.uuid)?;
     if send.atype != SendType::File as i32 {
         err!("Send content is not a file");
     }
 
-    let size = data.len();
-    if size > size_limit {
-        err!("Attachment storage limit exceeded with this file");
-    }
-
     let file_id = crate::crypto::generate_send_id();
     let folder_path = tokio::fs::canonicalize(&CONFIG.sends_folder()).await?.join(&send.uuid);
     let file_path = folder_path.join(&file_id);
@@ -253,7 +265,7 @@ async fn post_send_file(data: Form<UploadData<'_>>, headers: Headers, mut conn:
     if let Some(o) = data_value.as_object_mut() {
         o.insert(String::from("Id"), Value::String(file_id));
         o.insert(String::from("Size"), Value::Number(size.into()));
-        o.insert(String::from("SizeName"), Value::String(crate::util::get_display_size(size as i32)));
+        o.insert(String::from("SizeName"), Value::String(crate::util::get_display_size(size)));
     }
     send.data = serde_json::to_string(&data_value)?;
 
@@ -285,24 +297,32 @@ async fn post_send_file_v2(data: JsonUpcase<SendData>, headers: Headers, mut con
     enforce_disable_hide_email_policy(&data, &headers, &mut conn).await?;
 
     let file_length = match &data.FileLength {
-        Some(m) => Some(m.into_i32()?),
-        _ => None,
+        Some(m) => m.into_i64()?,
+        _ => err!("Invalid send length"),
     };
+    if file_length < 0 {
+        err!("Send size can't be negative")
+    }
 
-    let size_limit = match CONFIG.user_attachment_limit() {
+    let size_limit = match CONFIG.user_send_limit() {
         Some(0) => err!("File uploads are disabled"),
         Some(limit_kb) => {
-            let left = (limit_kb * 1024) - Attachment::size_by_user(&headers.user.uuid, &mut conn).await;
+            let Some(already_used) = Send::size_by_user(&headers.user.uuid, &mut conn).await else {
+                err!("Existing sends overflow")
+            };
+            let Some(left) = limit_kb.checked_mul(1024).and_then(|l| l.checked_sub(already_used)) else {
+                err!("Send size overflow");
+            };
             if left <= 0 {
-                err!("Attachment storage limit reached! Delete some attachments to free up space")
+                err!("Send storage limit reached! Delete some sends to free up space")
             }
-            std::cmp::Ord::max(left as u64, SIZE_525_MB)
+            i64::clamp(left, 0, SIZE_525_MB)
         }
         None => SIZE_525_MB,
     };
 
-    if file_length.is_some() && file_length.unwrap() as u64 > size_limit {
-        err!("Attachment storage limit exceeded with this file");
+    if file_length > size_limit {
+        err!("Send storage limit exceeded with this file");
     }
 
     let mut send = create_send(data, headers.user.uuid)?;
@@ -312,8 +332,8 @@ async fn post_send_file_v2(data: JsonUpcase<SendData>, headers: Headers, mut con
     let mut data_value: Value = serde_json::from_str(&send.data)?;
     if let Some(o) = data_value.as_object_mut() {
         o.insert(String::from("Id"), Value::String(file_id.clone()));
-        o.insert(String::from("Size"), Value::Number(file_length.unwrap().into()));
-        o.insert(String::from("SizeName"), Value::String(crate::util::get_display_size(file_length.unwrap())));
+        o.insert(String::from("Size"), Value::Number(file_length.into()));
+        o.insert(String::from("SizeName"), Value::String(crate::util::get_display_size(file_length)));
     }
     send.data = serde_json::to_string(&data_value)?;
     send.save(&mut conn).await?;

src/api/core/two_factor/authenticator.rs

@@ -5,7 +5,7 @@ use rocket::Route;
 use crate::{
     api::{
         core::log_user_event, core::two_factor::_generate_recover_code, EmptyResult, JsonResult, JsonUpcase,
-        NumberOrString, PasswordOrOtpData,
+        PasswordOrOtpData,
     },
     auth::{ClientIp, Headers},
     crypto,
@@ -13,6 +13,7 @@ use crate::{
         models::{EventType, TwoFactor, TwoFactorType},
         DbConn,
     },
+    util::NumberOrString,
 };
 
 pub use crate::config::CONFIG;

src/api/core/two_factor/mod.rs

@@ -7,12 +7,14 @@ use serde_json::Value;
 use crate::{
     api::{
         core::{log_event, log_user_event},
-        EmptyResult, JsonResult, JsonUpcase, NumberOrString, PasswordOrOtpData,
+        EmptyResult, JsonResult, JsonUpcase, PasswordOrOtpData,
     },
     auth::{ClientHeaders, Headers},
     crypto,
     db::{models::*, DbConn, DbPool},
-    mail, CONFIG,
+    mail,
+    util::NumberOrString,
+    CONFIG,
 };
 
 pub mod authenticator;

src/api/core/two_factor/webauthn.rs

@@ -7,7 +7,7 @@ use webauthn_rs::{base64_data::Base64UrlSafeData, proto::*, AuthenticationState,
 use crate::{
     api::{
         core::{log_user_event, two_factor::_generate_recover_code},
-        EmptyResult, JsonResult, JsonUpcase, NumberOrString, PasswordOrOtpData,
+        EmptyResult, JsonResult, JsonUpcase, PasswordOrOtpData,
     },
     auth::Headers,
     db::{
@@ -15,6 +15,7 @@ use crate::{
         DbConn,
     },
     error::Error,
+    util::NumberOrString,
     CONFIG,
 };
 

src/api/mod.rs

@@ -73,30 +73,3 @@ impl PasswordOrOtpData {
         Ok(())
     }
 }
-
-#[derive(Deserialize, Debug, Clone)]
-#[serde(untagged)]
-enum NumberOrString {
-    Number(i32),
-    String(String),
-}
-
-impl NumberOrString {
-    fn into_string(self) -> String {
-        match self {
-            NumberOrString::Number(n) => n.to_string(),
-            NumberOrString::String(s) => s,
-        }
-    }
-
-    #[allow(clippy::wrong_self_convention)]
-    fn into_i32(&self) -> ApiResult<i32> {
-        use std::num::ParseIntError as PIE;
-        match self {
-            NumberOrString::Number(n) => Ok(*n),
-            NumberOrString::String(s) => {
-                s.parse().map_err(|e: PIE| crate::Error::new("Can't convert to number", e.to_string()))
-            }
-        }
-    }
-}