diff --git a/crates/common/src/core.rs b/crates/common/src/core.rs index 5ef3dab4..b44faea7 100644 --- 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 index 715cda17..1c8f716e 100644 --- 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 index dcacb63f..38dab09e 100644 --- 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(¶ms).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 index e1ad9839..7a2f039c 100644 --- 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