Disk quotas support.

This commit is contained in:
Mauro D
2023-06-08 12:53:29 +00:00
parent ab895b2fae
commit 891d39940d
44 changed files with 1169 additions and 532 deletions

View File

@@ -53,6 +53,7 @@ pub mod email_submission;
pub mod event_source;
pub mod mailbox;
pub mod push_subscription;
pub mod quota;
pub mod sieve_script;
pub mod stress_test;
pub mod thread_get;
@@ -150,6 +151,11 @@ max-concurrent = 8
[jmap.protocol.upload]
max-size = 5000000
max-concurrent = 4
ttl = "1m"
[jmap.protocol.upload.quota]
files = 3
size = 50000
[jmap.rate-limit]
account.rate = "1000/1m"
@@ -222,7 +228,7 @@ pub async fn jmap_tests() {
let delete = true;
let mut params = init_jmap_tests(delete).await;
email_query::test(params.server.clone(), &mut params.client, delete).await;
/*email_query::test(params.server.clone(), &mut params.client, delete).await;
email_get::test(params.server.clone(), &mut params.client).await;
email_set::test(params.server.clone(), &mut params.client).await;
email_parse::test(params.server.clone(), &mut params.client).await;
@@ -243,6 +249,7 @@ pub async fn jmap_tests() {
vacation_response::test(params.server.clone(), &mut params.client).await;
email_submission::test(params.server.clone(), &mut params.client).await;
websocket::test(params.server.clone(), &mut params.client).await;
quota::test(params.server.clone(), &mut params.client).await;*/
stress_test::test(params.server.clone(), params.client).await;
if delete {

303
tests/src/jmap/quota.rs Normal file
View File

@@ -0,0 +1,303 @@
use std::sync::Arc;
use jmap::{blob::upload::DISABLE_UPLOAD_QUOTA, mailbox::INBOX_ID, JMAP};
use jmap_client::{
client::Client,
core::set::{SetErrorType, SetObject},
email::EmailBodyPart,
};
use jmap_proto::types::{collection::Collection, id::Id};
use crate::{
directory::sql::{add_to_group, create_test_user_with_email, set_test_quota},
jmap::{delivery::SmtpConnection, mailbox::destroy_all_mailboxes, test_account_login},
};
pub async fn test(server: Arc<JMAP>, admin_client: &mut Client) {
println!("Running quota tests...");
let directory = server.directory.as_ref();
let other_account_id =
create_test_user_with_email(directory, "jdoe@example.com", "12345", "John Doe").await;
let account_id =
create_test_user_with_email(directory, "robert@example.com", "aabbcc", "Robert Foobar")
.await;
set_test_quota(directory, "robert@example.com", 1024).await;
add_to_group(directory, "robert@example.com", "jdoe@example.com").await;
// Delete temporary blobs from previous tests
server
.store
.delete_account_blobs(account_id.document_id())
.await
.unwrap();
server
.store
.delete_account_blobs(other_account_id.document_id())
.await
.unwrap();
// Test temporary blob quota (3 files)
DISABLE_UPLOAD_QUOTA.store(false, std::sync::atomic::Ordering::Relaxed);
let client = test_account_login("robert@example.com", "aabbcc").await;
for _ in 0..3 {
assert_eq!(
client
.upload(None, vec![b'A'; 1024], None)
.await
.unwrap()
.size(),
1024
);
}
match client
.upload(None, vec![b'A'; 1024], None)
.await
.unwrap_err()
{
jmap_client::Error::Problem(err) if err.detail().unwrap().contains("quota") => (),
other => panic!("Unexpected error: {:?}", other),
}
server
.store
.delete_account_blobs(account_id.document_id())
.await
.unwrap();
// Test temporary blob quota (50000 bytes)
for _ in 0..2 {
assert_eq!(
client
.upload(None, vec![b'A'; 25000], None)
.await
.unwrap()
.size(),
25000
);
}
match client
.upload(None, vec![b'A'; 1024], None)
.await
.unwrap_err()
{
jmap_client::Error::Problem(err) if err.detail().unwrap().contains("quota") => (),
other => panic!("Unexpected error: {:?}", other),
}
server
.store
.delete_account_blobs(account_id.document_id())
.await
.unwrap();
// Test Email/import quota
let inbox_id = Id::new(INBOX_ID as u64).to_string();
let mut message_ids = Vec::new();
for i in 0..2 {
message_ids.push(
client
.email_import(
create_message_with_size(
"jdoe@example.com",
"robert@example.com",
&format!("Test {i}"),
512,
),
vec![&inbox_id],
None::<Vec<String>>,
None,
)
.await
.unwrap()
.take_id(),
);
}
assert_over_quota(
client
.email_import(
create_message_with_size("test@example.com", "jdoe@example.com", "Test 3", 100),
vec![&inbox_id],
None::<Vec<String>>,
None,
)
.await,
);
// Delete messages and check available quota
for message_id in message_ids {
client.email_destroy(&message_id).await.unwrap();
}
assert_eq!(
server
.get_used_quota(account_id.document_id())
.await
.unwrap(),
0
);
// Test Email/set quota
let mut message_ids = Vec::new();
for i in 0..2 {
let mut request = client.build();
let create_item = request.set_email().create();
create_item
.mailbox_ids([&inbox_id])
.subject(format!("Test {i}"))
.from(["jdoe@example.com"])
.to(["robert@example.com"])
.body_value("a".to_string(), String::from_utf8(vec![b'A'; 200]).unwrap())
.text_body(EmailBodyPart::new().part_id("a"));
let create_id = create_item.create_id().unwrap();
message_ids.push(
request
.send_set_email()
.await
.unwrap()
.created(&create_id)
.unwrap()
.take_id(),
);
}
let mut request = client.build();
let create_item = request.set_email().create();
create_item
.mailbox_ids([&inbox_id])
.subject("Test 3")
.from(["jdoe@example.com"])
.to(["robert@example.com"])
.body_value("a".to_string(), String::from_utf8(vec![b'A'; 400]).unwrap())
.text_body(EmailBodyPart::new().part_id("a"));
let create_id = create_item.create_id().unwrap();
assert_over_quota(request.send_set_email().await.unwrap().created(&create_id));
// Delete messages and check available quota
for message_id in message_ids {
client.email_destroy(&message_id).await.unwrap();
}
assert_eq!(
server
.get_used_quota(account_id.document_id())
.await
.unwrap(),
0
);
// Test Email/copy quota
let other_client = test_account_login("jdoe@example.com", "12345").await;
let mut other_message_ids = Vec::new();
let mut message_ids = Vec::new();
for i in 0..3 {
other_message_ids.push(
other_client
.email_import(
create_message_with_size(
"jane@example.com",
"jdoe@example.com",
&format!("Other Test {i}"),
512,
),
vec![&inbox_id],
None::<Vec<String>>,
None,
)
.await
.unwrap()
.take_id(),
);
}
for id in other_message_ids.iter().take(2) {
message_ids.push(
client
.email_copy(
other_account_id.to_string(),
id,
vec![&inbox_id],
None::<Vec<String>>,
None,
)
.await
.unwrap()
.take_id(),
);
}
assert_over_quota(
client
.email_copy(
other_account_id.to_string(),
&other_message_ids[2],
vec![&inbox_id],
None::<Vec<String>>,
None,
)
.await,
);
// Delete messages and check available quota
for message_id in message_ids {
client.email_destroy(&message_id).await.unwrap();
}
assert_eq!(
server
.get_used_quota(account_id.document_id())
.await
.unwrap(),
0
);
// Test delivery quota
let mut lmtp = SmtpConnection::connect().await;
for i in 0..2 {
lmtp.ingest(
"jane@example.com",
&["robert@example.com"],
&String::from_utf8(create_message_with_size(
"jane@example.com",
"robert@example.com",
&format!("Ingest test {i}"),
100,
))
.unwrap(),
)
.await;
}
let quota = server
.get_used_quota(account_id.document_id())
.await
.unwrap();
assert!(quota > 0 && quota <= 1024, "Quota is {}", quota);
assert_eq!(
server
.get_document_ids(account_id.document_id(), Collection::Email)
.await
.unwrap()
.unwrap()
.len(),
1,
);
DISABLE_UPLOAD_QUOTA.store(true, std::sync::atomic::Ordering::Relaxed);
// Remove test data
for account_id in [&account_id, &other_account_id] {
admin_client.set_default_account_id(account_id.to_string());
destroy_all_mailboxes(admin_client).await;
}
server.store.assert_is_empty().await;
}
fn assert_over_quota<T: std::fmt::Debug>(result: Result<T, jmap_client::Error>) {
match result {
Ok(result) => panic!("Expected error, got {:?}", result),
Err(jmap_client::Error::Set(err)) if err.error() == &SetErrorType::OverQuota => (),
Err(err) => panic!("Expected OverQuota SetError, got {:?}", err),
}
}
fn create_message_with_size(from: &str, to: &str, subject: &str, size: usize) -> Vec<u8> {
let mut message = format!(
"From: {}\r\nTo: {}\r\nSubject: {}\r\n\r\n",
from, to, subject
);
for _ in 0..size - message.len() {
message.push('A');
}
message.into_bytes()
}

View File

@@ -1,4 +1,4 @@
use store::{BlobKind, Store};
use store::{write::now, BlobKind, Store};
use utils::config::Config;
use crate::store::TempDir;
@@ -34,7 +34,7 @@ path = "{TMP}"
const DATA: &[u8] = b"Lorem ipsum dolor sit amet, consectetur adipiscing elit. Fusce erat nisl, dignissim a porttitor id, varius nec arcu. Sed mauris.";
#[tokio::test]
pub async fn blob_s3_test() {
pub async fn blob_tests() {
let temp_dir = TempDir::new("blob_tests", true);
test_blob(
Store::open(
@@ -60,6 +60,12 @@ pub async fn blob_s3_test() {
}
async fn test_blob(store: Store) {
// Obtain temp quota
let (quota_items, quota_bytes) = store.get_tmp_blob_usage(2, 100).await.unwrap();
assert_eq!(quota_items, 0);
assert_eq!(quota_bytes, 0);
store.purge_tmp_blobs(0).await.unwrap();
// Store and fetch
let kind = BlobKind::LinkedMaildir {
account_id: 0,
@@ -104,32 +110,25 @@ async fn test_blob(store: Store) {
}
// Copy partial
let tmp_kind = BlobKind::Temporary {
account_id: 1,
creation_year: 2020,
creation_month: 12,
creation_day: 31,
seq: 0,
};
let tmp_kind2 = BlobKind::Temporary {
account_id: 1,
creation_year: 2021,
creation_month: 1,
creation_day: 1,
seq: 0,
};
assert!(store
.copy_blob(&src_kind, &tmp_kind, (0..11).into())
.await
.unwrap());
assert!(store
.copy_blob(&src_kind, &tmp_kind2, (0..11).into())
.await
.unwrap());
let now = now();
let mut tmp_kinds = Vec::new();
for i in 1..=3 {
let tmp_kind = BlobKind::Temporary {
account_id: 2,
timestamp: now - (i * 5),
seq: 0,
};
assert!(store
.copy_blob(&src_kind, &tmp_kind, (0..11).into())
.await
.unwrap());
tmp_kinds.push(tmp_kind);
}
assert_eq!(
String::from_utf8(
store
.get_blob(&tmp_kind, 0..u32::MAX)
.get_blob(&tmp_kinds[0], 0..u32::MAX)
.await
.unwrap()
.unwrap()
@@ -138,15 +137,17 @@ async fn test_blob(store: Store) {
std::str::from_utf8(&DATA[0..11]).unwrap()
);
// Obtain temp quota
let (quota_items, quota_bytes) = store.get_tmp_blob_usage(2, 100).await.unwrap();
assert_eq!(quota_items, 3);
assert_eq!(quota_bytes, 33);
let (quota_items, quota_bytes) = store.get_tmp_blob_usage(2, 12).await.unwrap();
assert_eq!(quota_items, 2);
assert_eq!(quota_bytes, 22);
// Delete range
store
.bulk_delete_blob(&BlobKind::LinkedMaildir {
account_id: 1,
document_id: 0,
})
.await
.unwrap();
store.bulk_delete_blob(&tmp_kind).await.unwrap();
store.delete_account_blobs(1).await.unwrap();
store.purge_tmp_blobs(7).await.unwrap();
// Make sure the blobs are deleted
for id in 0..4 {
@@ -162,11 +163,13 @@ async fn test_blob(store: Store) {
.unwrap()
.is_none());
}
assert!(store
.get_blob(&tmp_kind, 0..u32::MAX)
.await
.unwrap()
.is_none());
for i in [1, 2] {
assert!(store
.get_blob(&tmp_kinds[i], 0..u32::MAX)
.await
.unwrap()
.is_none());
}
// Make sure other blobs were not deleted
assert!(store
@@ -175,23 +178,26 @@ async fn test_blob(store: Store) {
.unwrap()
.is_some());
assert!(store
.get_blob(&tmp_kind2, 0..u32::MAX)
.get_blob(&tmp_kinds[0], 0..u32::MAX)
.await
.unwrap()
.is_some());
// Copying a non-existing blob should fail
assert!(!store.copy_blob(&tmp_kind, &src_kind, None).await.unwrap());
assert!(!store
.copy_blob(&tmp_kinds[1], &src_kind, None)
.await
.unwrap());
// Copy blob between buckets
assert!(store
.copy_blob(&src_kind, &tmp_kind, (10..20).into())
.copy_blob(&src_kind, &tmp_kinds[0], (10..20).into())
.await
.unwrap());
assert_eq!(
String::from_utf8(
store
.get_blob(&tmp_kind, 0..u32::MAX)
.get_blob(&tmp_kinds[0], 0..u32::MAX)
.await
.unwrap()
.unwrap()
@@ -201,7 +207,7 @@ async fn test_blob(store: Store) {
);
// Delete blobs
for blob_kind in [src_kind, tmp_kind, tmp_kind2] {
for blob_kind in [src_kind, tmp_kinds[0]] {
assert!(store.delete_blob(&blob_kind).await.unwrap());
assert!(store
.get_blob(&blob_kind, 0..u32::MAX)