Reset quota endpoint

stalwartlabs · Jan 17, 2025 · 61e63d1 · 61e63d1
1 parent 2eb3886
commit 61e63d1
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 4 deletions.
diff --git a/crates/common/src/core.rs b/crates/common/src/core.rs
@@ -18,8 +18,8 @@ use store::{
         key::DeserializeBigEndian, log::ChangeLogBuilder, now, BatchBuilder, BitmapClass, BlobOp,
         DirectoryClass, QueueClass, TagValue, ValueClass,
     },
-    BitmapKey, BlobClass, BlobStore, Deserialize, FtsStore, InMemoryStore, IterateParams, LogKey,
-    Serialize, Store, ValueKey, U32_LEN,
+    BitmapKey, BlobClass, BlobStore, Deserialize, FtsStore, InMemoryStore, IndexKey, IterateParams,
+    LogKey, Serialize, Store, ValueKey, U32_LEN,
 };
 use trc::AddContext;
 use utils::BlobHash;
@@ -191,6 +191,54 @@ impl Server {
             .add_context(|err| err.caused_by(trc::location!()).account_id(account_id))
     }
 
+    pub async fn recalculate_quota(&self, account_id: u32) -> trc::Result<()> {
+        let mut quota = 0i64;
+
+        self.store()
+            .iterate(
+                IterateParams::new(
+                    IndexKey {
+                        account_id,
+                        collection: Collection::Email.into(),
+                        document_id: 0,
+                        field: Property::Size.into(),
+                        key: 0u32.serialize(),
+                    },
+                    IndexKey {
+                        account_id,
+                        collection: Collection::Email.into(),
+                        document_id: u32::MAX,
+                        field: Property::Size.into(),
+                        key: u32::MAX.serialize(),
+                    },
+                )
+                .no_values()
+                .ascending(),
+                |key, _| {
+                    let value = key
+                        .get(key.len() - (U32_LEN * 2)..key.len() - U32_LEN)
+                        .ok_or_else(|| trc::Error::corrupted_key(key, None, trc::location!()))
+                        .and_then(u32::deserialize)?;
+
+                    quota += value as i64;
+
+                    Ok(true)
+                },
+            )
+            .await
+            .caused_by(trc::location!())?;
+
+        let mut batch = BatchBuilder::new();
+        batch
+            .clear(DirectoryClass::UsedQuota(account_id))
+            .add(DirectoryClass::UsedQuota(account_id), quota);
+        self.store()
+            .write(batch)
+            .await
+            .caused_by(trc::location!())
+            .map(|_| ())
+    }
+
     pub async fn has_available_quota(
         &self,
         quotas: &ResourceToken,

diff --git a/crates/jmap/src/api/management/stores.rs b/crates/jmap/src/api/management/stores.rs
@@ -287,6 +287,26 @@ impl ManageStore for Server {
                 }))
                 .into_http_response())
             }
+            (Some("quota"), Some(account_id), None, method @ (&Method::GET | &Method::DELETE)) => {
+                let account_id = self
+                    .core
+                    .storage
+                    .data
+                    .get_principal_id(decode_path_element(account_id).as_ref())
+                    .await?
+                    .ok_or_else(|| trc::ManageEvent::NotFound.into_err())?;
+
+                if method == Method::DELETE {
+                    self.recalculate_quota(account_id).await?;
+                }
+
+                let result = self.get_used_quota(account_id).await?;
+
+                Ok(JsonResponse::new(json!({
+                    "data": result,
+                }))
+                .into_http_response())
+            }
             _ => Err(trc::ResourceEvent::NotFound.into_err()),
         }
     }

diff --git a/tests/src/jmap/mod.rs b/tests/src/jmap/mod.rs
@@ -392,10 +392,10 @@ pub async fn jmap_tests() {
     sieve_script::test(&mut params).await;
     vacation_response::test(&mut params).await;
     email_submission::test(&mut params).await;
-    websocket::test(&mut params).await;
+    websocket::test(&mut params).await;*/
     quota::test(&mut params).await;
     crypto::test(&mut params).await;
-    blob::test(&mut params).await;*/
+    blob::test(&mut params).await;
     permissions::test(&params).await;
     purge::test(&mut params).await;
     enterprise::test(&mut params).await;

diff --git a/tests/src/jmap/quota.rs b/tests/src/jmap/quota.rs
@@ -227,6 +227,23 @@ pub async fn test(params: &mut JMAPTest) {
     let create_id = create_item.create_id().unwrap();
     assert_over_quota(request.send_set_email().await.unwrap().created(&create_id));
 
+    // Recalculate quota
+    let prev_quota = server
+        .get_used_quota(account_id.document_id())
+        .await
+        .unwrap();
+    server
+        .recalculate_quota(account_id.document_id())
+        .await
+        .unwrap();
+    assert_eq!(
+        server
+            .get_used_quota(account_id.document_id())
+            .await
+            .unwrap(),
+        prev_quota
+    );
+
     // Delete messages and check available quota
     for message_id in message_ids {
         client.email_destroy(&message_id).await.unwrap();
@@ -333,6 +350,7 @@ pub async fn test(params: &mut JMAPTest) {
             .len(),
         1,
     );
+
     DISABLE_UPLOAD_QUOTA.store(true, std::sync::atomic::Ordering::Relaxed);
 
     // Remove test data