Improve file limit handling (#4242)

* Improve file limit handling

* Oops

* Update PostgreSQL migration

* Review comments

---------

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

src/db/models/attachment.rs

@@ -1,5 +1,6 @@
 use std::io::ErrorKind;
 
+use bigdecimal::{BigDecimal, ToPrimitive};
 use serde_json::Value;
 
 use crate::CONFIG;
@@ -13,14 +14,14 @@ db_object! {
         pub id: String,
         pub cipher_uuid: String,
         pub file_name: String, // encrypted
-        pub file_size: i32,
+        pub file_size: i64,
         pub akey: Option<String>,
     }
 }
 
 /// Local methods
 impl Attachment {
-    pub const fn new(id: String, cipher_uuid: String, file_name: String, file_size: i32, akey: Option<String>) -> Self {
+    pub const fn new(id: String, cipher_uuid: String, file_name: String, file_size: i64, akey: Option<String>) -> Self {
         Self {
             id,
             cipher_uuid,
@@ -145,13 +146,18 @@ impl Attachment {
 
     pub async fn size_by_user(user_uuid: &str, conn: &mut DbConn) -> i64 {
         db_run! { conn: {
-            let result: Option<i64> = attachments::table
+            let result: Option<BigDecimal> = attachments::table
                 .left_join(ciphers::table.on(ciphers::uuid.eq(attachments::cipher_uuid)))
                 .filter(ciphers::user_uuid.eq(user_uuid))
                 .select(diesel::dsl::sum(attachments::file_size))
                 .first(conn)
                 .expect("Error loading user attachment total size");
-            result.unwrap_or(0)
+
+            match result.map(|r| r.to_i64()) {
+                Some(Some(r)) => r,
+                Some(None) => i64::MAX,
+                None => 0
+            }
         }}
     }
 
@@ -168,13 +174,18 @@ impl Attachment {
 
     pub async fn size_by_org(org_uuid: &str, conn: &mut DbConn) -> i64 {
         db_run! { conn: {
-            let result: Option<i64> = attachments::table
+            let result: Option<BigDecimal> = attachments::table
                 .left_join(ciphers::table.on(ciphers::uuid.eq(attachments::cipher_uuid)))
                 .filter(ciphers::organization_uuid.eq(org_uuid))
                 .select(diesel::dsl::sum(attachments::file_size))
                 .first(conn)
                 .expect("Error loading user attachment total size");
-            result.unwrap_or(0)
+
+            match result.map(|r| r.to_i64()) {
+                Some(Some(r)) => r,
+                Some(None) => i64::MAX,
+                None => 0
+            }
         }}
     }
 

src/db/models/send.rs

@@ -172,6 +172,7 @@ use crate::db::DbConn;
 
 use crate::api::EmptyResult;
 use crate::error::MapResult;
+use crate::util::NumberOrString;
 
 impl Send {
     pub async fn save(&mut self, conn: &mut DbConn) -> EmptyResult {
@@ -286,6 +287,36 @@ impl Send {
         }}
     }
 
+    pub async fn size_by_user(user_uuid: &str, conn: &mut DbConn) -> Option<i64> {
+        let sends = Self::find_by_user(user_uuid, conn).await;
+
+        #[allow(non_snake_case)]
+        #[derive(serde::Deserialize, Default)]
+        struct FileData {
+            Size: Option<NumberOrString>,
+            size: Option<NumberOrString>,
+        }
+
+        let mut total: i64 = 0;
+        for send in sends {
+            if send.atype == SendType::File as i32 {
+                let data: FileData = serde_json::from_str(&send.data).unwrap_or_default();
+
+                let size = match (data.size, data.Size) {
+                    (Some(s), _) => s.into_i64(),
+                    (_, Some(s)) => s.into_i64(),
+                    (None, None) => continue,
+                };
+
+                if let Ok(size) = size {
+                    total = total.checked_add(size)?;
+                };
+            }
+        }
+
+        Some(total)
+    }
+
     pub async fn find_by_org(org_uuid: &str, conn: &mut DbConn) -> Vec<Self> {
         db_run! {conn: {
             sends::table

src/db/schemas/mysql/schema.rs

@@ -3,7 +3,7 @@ table! {
         id -> Text,
         cipher_uuid -> Text,
         file_name -> Text,
-        file_size -> Integer,
+        file_size -> BigInt,
         akey -> Nullable<Text>,
     }
 }

src/db/schemas/postgresql/schema.rs

@@ -3,7 +3,7 @@ table! {
         id -> Text,
         cipher_uuid -> Text,
         file_name -> Text,
-        file_size -> Integer,
+        file_size -> BigInt,
         akey -> Nullable<Text>,
     }
 }

src/db/schemas/sqlite/schema.rs

@@ -3,7 +3,7 @@ table! {
         id -> Text,
         cipher_uuid -> Text,
         file_name -> Text,
-        file_size -> Integer,
+        file_size -> BigInt,
         akey -> Nullable<Text>,
     }
 }