filip/Stalwart
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Registry testing - part 2

Browse Source
This commit is contained in:
mdecimus
2026-03-12 18:26:35 +01:00
parent 7c143c85f8
commit d45f1299d3
42 changed files with 1649 additions and 1064 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/common/src/config/mailstore/email.rs
View File

@@ -95,7 +95,7 @@ impl EmailConfig {
default_domain.name
} else {
bp.build_error(
ObjectType::Authentication.singleton(),
ObjectType::SystemSettings.singleton(),
format!(
"Default domain with ID {} not found",
system.default_domain_id

11
crates/common/src/config/storage.rs
View File

@@ -6,6 +6,7 @@
use coordinator::Coordinator;
use directory::{Directories, Directory};
use registry::schema::prelude::ObjectType;
use std::{collections::HashMap, sync::Arc};
use store::{
BlobStore, InMemoryStore, RegistryStore, SearchStore, Store, registry::bootstrap::Bootstrap,
@@ -31,12 +32,20 @@ impl Storage {
pub async fn parse(bp: &mut Bootstrap) -> Self {
let memory = InMemoryStore::build(bp).await.unwrap_or_default();
let directory = Directories::build(bp).await;
let search = SearchStore::build(bp).await.unwrap_or_default();
if let Err(err) = search.create_indexes().await {
bp.build_warning(
ObjectType::SearchStore.singleton(),
format!("Failed to create search indexes: {err}"),
);
}
Storage {
registry: bp.registry.clone(),
data: bp.data_store.clone(),
blob: BlobStore::build(bp).await.unwrap_or_default(),
search: SearchStore::build(bp).await.unwrap_or_default(),
search,
coordinator: Coordinator::build(bp, &memory).await.unwrap_or_default(),
memory,
tracing: Store::build_tracing(bp).await.unwrap_or_default(),

10
crates/common/src/config/telemetry.rs
View File

@@ -178,6 +178,7 @@ impl Tracers {
let lossy;
let events;
let events_policy;
let enable;
let typ = match tracer {
Tracer::Log(tracer) if tracer.enable => {
@@ -185,6 +186,7 @@ impl Tracers {
lossy = tracer.lossy;
events = tracer.events;
events_policy = tracer.events_policy;
enable = tracer.enable;
TelemetrySubscriberType::LogTracer(LogTracer {
path: tracer.path,
@@ -204,6 +206,7 @@ impl Tracers {
lossy = tracer.lossy;
events = tracer.events;
events_policy = tracer.events_policy;
enable = tracer.enable;
if !tracers
.iter()
@@ -226,6 +229,7 @@ impl Tracers {
lossy = tracer.lossy;
events = tracer.events;
events_policy = tracer.events_policy;
enable = tracer.enable;
if !tracers
.iter()
@@ -260,6 +264,7 @@ impl Tracers {
lossy = tracer.lossy;
events = tracer.events;
events_policy = tracer.events_policy;
enable = tracer.enable;
let headers = match tracer
.http_auth
@@ -328,6 +333,7 @@ impl Tracers {
lossy = tracer.lossy;
events = tracer.events;
events_policy = tracer.events_policy;
enable = tracer.enable;
let mut span_exporter = SpanExporter::builder()
.with_tonic()
@@ -374,6 +380,10 @@ impl Tracers {
_ => continue,
};
if !enable {
continue;
}
// Create tracer
let mut tracer = TelemetrySubscriber {
id: format!("t_{}", id.id()),

20
crates/common/src/enterprise/config.rs
View File

@@ -37,8 +37,17 @@ impl Enterprise {
.await
.default_hostname;
let mut update_license = None;
let mut enterprise = bp.setting_infallible::<structs::Enterprise>().await;
// WARNING: TAMPERING WITH THIS FUNCTION IS STRICTLY PROHIBITED
// Any attempt to modify, bypass, or disable this license validation mechanism
// constitutes a severe violation of the Stalwart Enterprise License Agreement.
// Such actions may result in immediate termination of your license, legal action,
// and substantial financial penalties. Stalwart Labs LLC actively monitors for
// unauthorized modifications and will pursue all available legal remedies against
// violators to the fullest extent of the law, including but not limited to claims
// for copyright infringement, breach of contract, and fraud.
let license_result = match (
enterprise.license_key.secret().await,
enterprise.api_key.secret().await,
@@ -75,7 +84,16 @@ impl Enterprise {
result.key
}),
(Ok(None), Ok(None)) => {
#[cfg(not(feature = "test_mode"))]
return None;
#[cfg(feature = "test_mode")]
Ok(LicenseKey {
valid_to: store::write::now() + (86400 * 365),
valid_from: store::write::now() - 3600,
domain: server_hostname.to_string(),
accounts: 100,
})
}
(Err(err), _) => {
bp.build_error(ObjectType::Enterprise.singleton(), err);

4
crates/common/src/manager/backup.rs
View File

@@ -179,7 +179,9 @@ impl Core {
key: vec![u8::MAX; 32],
},
)
.set_values(subspace != SUBSPACE_INDEXES),
.set_values(
![SUBSPACE_INDEXES, SUBSPACE_REGISTRY_IDX].contains(&subspace),
),
|key, value| {
writer
.send((key.to_vec(), value.to_vec()))

108
crates/common/src/manager/defaults.rs
View File

@@ -4,30 +4,22 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::{auth::permissions::DefaultPermissions, network::dkim::generate_dkim_private_key};
use crate::auth::permissions::DefaultPermissions;
use registry::{
schema::{
enums::{DkimSignatureType, MtaInboundThrottleKey, MtaIpStrategy},
prelude::ObjectType,
structs::{
Authentication, Dkim1Signature, DkimPrivateKey, DkimSignature, Domain,
MtaConnectionStrategy, MtaDeliveryExpiration, MtaDeliveryExpirationAttempts,
MtaDeliveryExpirationTtl, MtaDeliverySchedule, MtaDeliveryScheduleInterval,
MtaDeliveryScheduleIntervals, MtaDeliveryScheduleIntervalsOrDefault,
MtaInboundThrottle, MtaQueueQuota, MtaRoute, MtaRouteCommon, MtaRouteMx,
MtaTlsStrategy, MtaVirtualQueue, OidcProvider, Rate, Role, SecretKey, SecretKeyValue,
SecretText, SecretTextValue, SystemSettings,
},
enums::*,
prelude::{ObjectType, SocketAddr},
structs::*,
},
types::{duration::Duration, error::Error, list::List, map::Map},
};
use std::str::FromStr;
use store::{
rand::{Rng, distr::Alphanumeric, rng},
registry::{
bootstrap::Bootstrap,
write::{RegistryWrite, RegistryWriteResult},
},
write::now,
};
use types::id::Id;
@@ -399,21 +391,24 @@ async fn insert_safe_defaults(bp: &mut Bootstrap) -> trc::Result<()> {
}
}
#[cfg(not(feature = "test_mode"))]
if let Some(domain_id) = default_domain_id {
let now = now();
let now = store::write::now();
let signature_rsa = DkimSignature::Dkim1RsaSha256(Dkim1Signature {
domain_id,
enabled: true,
selector: format!("rsa-{now}"),
private_key: DkimPrivateKey::Value(SecretTextValue {
secret: generate_dkim_private_key(DkimSignatureType::Dkim1RsaSha256)
.await?
.map_err(|err| {
trc::EventType::Dkim(trc::DkimEvent::BuildError)
.into_err()
.reason(err)
.caused_by(trc::location!())
})?,
secret: crate::network::dkim::generate_dkim_private_key(
DkimSignatureType::Dkim1RsaSha256,
)
.await?
.map_err(|err| {
trc::EventType::Dkim(trc::DkimEvent::BuildError)
.into_err()
.reason(err)
.caused_by(trc::location!())
})?,
}),
..Default::default()
});
@@ -422,14 +417,16 @@ async fn insert_safe_defaults(bp: &mut Bootstrap) -> trc::Result<()> {
enabled: true,
selector: format!("ed-{now}"),
private_key: DkimPrivateKey::Value(SecretTextValue {
secret: generate_dkim_private_key(DkimSignatureType::Dkim1Ed25519Sha256)
.await?
.map_err(|err| {
trc::EventType::Dkim(trc::DkimEvent::BuildError)
.into_err()
.reason(err)
.caused_by(trc::location!())
})?,
secret: crate::network::dkim::generate_dkim_private_key(
DkimSignatureType::Dkim1Ed25519Sha256,
)
.await?
.map_err(|err| {
trc::EventType::Dkim(trc::DkimEvent::BuildError)
.into_err()
.reason(err)
.caused_by(trc::location!())
})?,
}),
..Default::default()
});
@@ -455,5 +452,56 @@ async fn insert_safe_defaults(bp: &mut Bootstrap) -> trc::Result<()> {
.await?;
}
if bp
.registry
.count_object(ObjectType::NetworkListener)
.await?
== 0
{
for (protocol, name, port, use_tls) in [
(NetworkListenerProtocol::Smtp, "smtp", 25, false),
(NetworkListenerProtocol::Smtp, "submission", 587, false),
(NetworkListenerProtocol::Smtp, "submissions", 465, true),
(NetworkListenerProtocol::Imap, "imap", 143, false),
(NetworkListenerProtocol::Imap, "imaps", 993, true),
(NetworkListenerProtocol::Pop3, "pop3", 110, false),
(NetworkListenerProtocol::Pop3, "pop3s", 995, true),
(NetworkListenerProtocol::ManageSieve, "sieve", 4190, false),
(NetworkListenerProtocol::Http, "https", 443, true),
(NetworkListenerProtocol::Http, "http", 8080, false),
] {
bp.registry
.write(RegistryWrite::insert(
&NetworkListener {
bind: Map::new(vec![
SocketAddr::from_str(&format!("[::]:{port}")).unwrap(),
]),
name: name.to_string(),
protocol,
use_tls,
..Default::default()
}
.into(),
))
.await?;
}
}
if bp.registry.count_object(ObjectType::Tracer).await? == 0 {
bp.registry
.write(RegistryWrite::insert(
&Tracer::Log(TracerLog {
enable: true,
ansi: false,
prefix: "stalwart.log".into(),
rotate: LogRotateFrequency::Daily,
path: "/var/log/stalwart".into(),
..Default::default()
})
.into(),
))
.await?;
}
Ok(())
}

5
crates/email/src/sieve/ingest.rs
View File

@@ -256,10 +256,9 @@ impl SieveScriptIngest for Server {
} else {
let exists = self
.in_memory_store()
.key_get::<()>(id_hash.key())
.key_exists(id_hash.key())
.await
.caused_by(trc::location!())?
.is_some();
.caused_by(trc::location!())?;
if !exists || last {
self.in_memory_store()

5
crates/jmap/src/registry/mapping/task.rs
View File

@@ -104,10 +104,10 @@ pub(crate) async fn task_set(
object_id: foreign_id,
..
} = key
&& set
&& !set
.server
.store()
.get_value::<()>(ValueKey::from(ValueClass::Registry(
.key_exists(ValueKey::from(ValueClass::Registry(
RegistryClass::IndexId {
object_id: foreign_id.object().to_id(),
item_id: foreign_id.id().id(),
@@ -115,7 +115,6 @@ pub(crate) async fn task_set(
)))
.await
.caused_by(trc::location!())?
.is_none()
{
set.response.not_created.append(
id,

4
crates/main/Cargo.toml
View File

@@ -41,8 +41,8 @@ jemallocator = "0.5.0"
[features]
#default = ["sqlite", "postgres", "mysql", "rocks", "s3", "redis", "azure", "nats", "enterprise", "zenoh", "kafka"]
#default = ["sqlite", "postgres", "mysql", "rocks", "s3", "redis", "azure", "nats", "enterprise"]
default = ["rocks", "enterprise"]
default = ["sqlite", "postgres", "mysql", "rocks", "s3", "redis", "azure", "nats", "enterprise"]
#default = ["rocks", "enterprise"]
sqlite = ["store/sqlite", "directory/sqlite"]
foundationdb = ["store/foundation", "common/foundation"]
postgres = ["store/postgres", "directory/postgres"]

17
crates/store/src/backend/composite/read_replica.rs
View File

@@ -121,6 +121,23 @@ impl SQLReadReplica {
.await
}
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
self.run_op(move |store| {
let key = key.clone();
async move {
match store {
#[cfg(feature = "postgres")]
Store::PostgreSQL(store) => store.key_exists(key).await,
#[cfg(feature = "mysql")]
Store::MySQL(store) => store.key_exists(key).await,
_ => panic!("Invalid store type"),
}
}
})
.await
}
pub async fn iterate<T: Key>(
&self,
params: IterateParams<T>,

16
crates/store/src/backend/foundationdb/read.rs
View File

@@ -38,14 +38,26 @@ impl FdbStore {
let trx = self.read_trx().await?;
match read_chunked_value(&key, &trx, true).await? {
ChunkedValue::Single(bytes) => U::deserialize_with_key(&key, &bytes).map(Some),
ChunkedValue::Single(bytes) => {
U::deserialize_with_key(key.get(1..).unwrap_or_default(), &bytes).map(Some)
}
ChunkedValue::Chunked { bytes, .. } => {
U::deserialize_owned_with_key(&key, bytes).map(Some)
U::deserialize_owned_with_key(key.get(1..).unwrap_or_default(), bytes).map(Some)
}
ChunkedValue::None => Ok(None),
}
}
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
let key = key.serialize(WITH_SUBSPACE);
let trx = self.read_trx().await?;
match read_chunked_value(&key, &trx, true).await? {
ChunkedValue::Single(_) | ChunkedValue::Chunked { .. } => Ok(true),
ChunkedValue::None => Ok(false),
}
}
pub(crate) async fn iterate<T: Key>(
&self,
params: IterateParams<T>,

14
crates/store/src/backend/foundationdb/write.rs
View File

@@ -11,8 +11,8 @@ use super::{
use crate::{
backend::deserialize_i64_le,
write::{
AssignedIds, Batch, DirectoryClass, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult,
Operation, TaskQueueClass, TelemetryClass, ValueClass, ValueOp, key::KeySerializer,
AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation,
RegistryClass, TaskQueueClass, TelemetryClass, ValueClass, ValueOp, key::KeySerializer,
},
*,
};
@@ -150,7 +150,6 @@ impl FdbStore {
MergeResult::Skip => (),
}
}
ValueOp::AtomicAdd(by) => {
trx.atomic_op(&key, &by.to_le_bytes()[..], MutationType::Add);
}
@@ -181,16 +180,13 @@ impl FdbStore {
class,
ValueClass::Property(_)
| ValueClass::Queue(_)
| ValueClass::Report(_)
| ValueClass::Directory(DirectoryClass::Principal(_))
| ValueClass::Registry(RegistryClass::Item { .. })
| ValueClass::ShareNotification { .. }
| ValueClass::Telemetry(TelemetryClass::Metric { .. })
| ValueClass::TaskQueue(TaskQueueClass::SendImip {
is_payload: true,
..
})
| ValueClass::TaskQueue(TaskQueueClass::Task { .. })
| ValueClass::InMemory(_)
) {
// Clear range for potentially chunked values to avoid leaving orphaned chunks
trx.clear_range(
&key,
&KeySerializer::new(key.len() + 1)

2
crates/store/src/backend/mysql/main.rs
View File

@@ -13,7 +13,7 @@ use crate::{
},
*,
};
use ::registry::{schema::structs, utils::OrderedMap};
use ::registry::schema::structs;
use mysql_async::{
Conn, OptsBuilder, Pool, PoolConstraints, PoolOpts, SslOpts, prelude::Queryable,
};

16
crates/store/src/backend/mysql/read.rs
View File

@@ -35,6 +35,22 @@ impl MysqlStore {
})
}
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
let mut conn = self.conn_pool.get_conn().await.map_err(into_error)?;
let s = conn
.prep(format!(
"SELECT 1 FROM {} WHERE k = ?",
char::from(key.subspace())
))
.await
.map_err(into_error)?;
let key = key.serialize(0);
conn.exec_first::<u8, _, _>(&s, (&key,))
.await
.map_err(into_error)
.map(|r| r.is_some())
}
pub(crate) async fn iterate<T: Key>(
&self,
params: IterateParams<T>,

67
crates/store/src/backend/mysql/write.rs
View File

@@ -7,6 +7,7 @@
use super::{MysqlStore, into_error};
use crate::{
IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA,
SUBSPACE_REGISTRY_IDX,
write::{
AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation,
ValueClass, ValueOp,
@@ -126,47 +127,56 @@ impl MysqlStore {
}
Operation::Value { class, op } => {
let key = class.serialize(account_id, collection, document_id, 0);
let table = char::from(class.subspace(collection));
let subspace = class.subspace(collection);
let table = char::from(subspace);
match op {
ValueOp::Set(value) => {
let exists = asserted_values.get(&key);
let s = if let Some(exists) = exists {
if *exists {
trx.prep(format!("UPDATE {} SET v = :v WHERE k = :k", table))
if subspace != SUBSPACE_REGISTRY_IDX {
let exists = asserted_values.get(&key);
let s = if let Some(exists) = exists {
if *exists {
trx.prep(format!(
"UPDATE {} SET v = :v WHERE k = :k",
table
))
.await?
} else {
trx.prep(format!(
"INSERT INTO {} (k, v) VALUES (:k, :v)",
table
))
.await?
}
} else {
trx.prep(format!(
"INSERT INTO {} (k, v) VALUES (:k, :v)",
table
))
.await?
}
} else {
trx
trx
.prep(
format!("INSERT INTO {} (k, v) VALUES (:k, :v) ON DUPLICATE KEY UPDATE v = VALUES(v)", table),
)
.await?
};
};
match trx
.exec_drop(&s, params! {"k" => key, "v" => &*value})
.await
{
Ok(_) => {
if trx.affected_rows() == 0 {
match trx
.exec_drop(&s, params! {"k" => key, "v" => &*value})
.await
{
Ok(_) => {
if trx.affected_rows() == 0 {
trx.rollback().await?;
return Err(trc::StoreEvent::AssertValueFailed
.into_err()
.caused_by(trc::location!())
.into());
}
}
Err(err) => {
trx.rollback().await?;
return Err(trc::StoreEvent::AssertValueFailed
.into_err()
.caused_by(trc::location!())
.into());
return Err(err.into());
}
}
Err(err) => {
trx.rollback().await?;
return Err(err.into());
}
} else {
let s = trx.prep("INSERT IGNORE INTO b (k) VALUES (?)").await?;
trx.exec_drop(&s, (key,)).await?;
}
}
ValueOp::SetFnc(set_op) => {
@@ -256,7 +266,6 @@ impl MysqlStore {
_ => (),
}
}
ValueOp::AtomicAdd(by) => {
if *by >= 0 {
let s = trx

16
crates/store/src/backend/postgres/read.rs
View File

@@ -37,6 +37,22 @@ impl PostgresStore {
})
}
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let s = conn
.prepare_cached(&format!(
"SELECT 1 FROM {} WHERE k = $1",
char::from(key.subspace())
))
.await
.map_err(into_error)?;
let key = key.serialize(0);
conn.query_opt(&s, &[&key])
.await
.map_err(into_error)
.map(|r| r.is_some())
}
pub(crate) async fn iterate<T: Key>(
&self,
params: IterateParams<T>,

59
crates/store/src/backend/postgres/write.rs
View File

@@ -7,6 +7,7 @@
use super::{PostgresStore, into_error};
use crate::{
IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA,
SUBSPACE_REGISTRY_IDX,
backend::postgres::into_pool_error,
write::{
AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation,
@@ -130,40 +131,50 @@ impl PostgresStore {
}
Operation::Value { class, op } => {
let key = class.serialize(account_id, collection, document_id, 0);
let table = char::from(class.subspace(collection));
let subspace = class.subspace(collection);
let table = char::from(subspace);
match op {
ValueOp::Set(value) => {
let s = if let Some(exists) = asserted_values.get(&key) {
if *exists {
trx.prepare_cached(&format!(
"UPDATE {} SET v = $2 WHERE k = $1",
table
))
.await?
if subspace != SUBSPACE_REGISTRY_IDX {
let s = if let Some(exists) = asserted_values.get(&key) {
if *exists {
trx.prepare_cached(&format!(
"UPDATE {} SET v = $2 WHERE k = $1",
table
))
.await?
} else {
trx.prepare_cached(&format!(
"INSERT INTO {} (k, v) VALUES ($1, $2)",
table
))
.await?
}
} else {
trx.prepare_cached(&format!(
"INSERT INTO {} (k, v) VALUES ($1, $2)",
concat!(
"INSERT INTO {} (k, v) VALUES ($1, $2) ",
"ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v"
),
table
))
.await?
};
if trx.execute(&s, &[&key, &(*value)]).await? == 0 {
return Err(trc::StoreEvent::AssertValueFailed
.into_err()
.caused_by(trc::location!())
.into());
}
} else {
trx.prepare_cached(&format!(
concat!(
"INSERT INTO {} (k, v) VALUES ($1, $2) ",
"ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v"
),
table
))
.await?
};
if trx.execute(&s, &[&key, &(*value)]).await? == 0 {
return Err(trc::StoreEvent::AssertValueFailed
.into_err()
.caused_by(trc::location!())
.into());
let s = trx
.prepare_cached(
"INSERT INTO b (k) VALUES ($1) ON CONFLICT (k) DO NOTHING",
)
.await?;
trx.execute(&s, &[&key]).await?;
}
}
ValueOp::SetFnc(set_op) => {

2
crates/store/src/backend/rocksdb/main.rs
View File

@@ -65,7 +65,7 @@ impl RocksDbStore {
SUBSPACE_REGISTRY_PK,
SUBSPACE_DIRECTORY,
LEGACY_SUBSPACE_BITMAP_TEXT,
LEGACY_SUBSPACE_FTS_INDEX,
LEGACY_SUBSPACE_BITMAP_TAG,
] {
let cf_opts = Options::default();
cfs.push(ColumnFamilyDescriptor::new(

16
crates/store/src/backend/rocksdb/read.rs
View File

@@ -36,6 +36,22 @@ impl RocksDbStore {
.await
}
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
let db = self.db.clone();
self.spawn_worker(move || {
let subspace = &[key.subspace()];
let key = key.serialize(0);
db.get_pinned_cf(
&db.cf_handle(unsafe { std::str::from_utf8_unchecked(subspace.as_slice()) })
.unwrap(),
&key,
)
.map_err(into_error)
.map(|value| value.is_some())
})
.await
}
pub(crate) async fn iterate<T: Key>(
&self,
params: IterateParams<T>,

20
crates/store/src/backend/sqlite/read.rs
View File

@@ -34,6 +34,26 @@ impl SqliteStore {
.await
}
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
let manager = self.conn_pool.clone();
self.spawn_worker(move || {
let conn = manager.get().map_err(into_error)?;
let mut result = conn
.prepare_cached(&format!(
"SELECT 1 FROM {} WHERE k = ?",
char::from(key.subspace())
))
.map_err(into_error)?;
let key = key.serialize(0);
result
.query_row([&key], |_| Ok(()))
.optional()
.map(|opt| opt.is_some())
.map_err(into_error)
})
.await
}
pub(crate) async fn iterate<T: Key>(
&self,
params: IterateParams<T>,

31
crates/store/src/backend/sqlite/write.rs
View File

@@ -7,6 +7,7 @@
use super::{SqliteStore, into_error};
use crate::{
IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA,
SUBSPACE_REGISTRY_IDX,
write::{AssignedIds, Batch, MergeResult, Operation, ValueClass, ValueOp},
};
use rusqlite::{OptionalExtension, TransactionBehavior, params};
@@ -69,19 +70,29 @@ impl SqliteStore {
}
Operation::Value { class, op } => {
let key = class.serialize(account_id, collection, document_id, 0);
let table = char::from(class.subspace(collection));
let subspace = class.subspace(collection);
let table = char::from(subspace);
match op {
ValueOp::Set(value) => {
trx.prepare_cached(&format!(
"INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)",
table
))
.map_err(into_error)
.caused_by(trc::location!())?
.execute([&key, value])
.map_err(into_error)
.caused_by(trc::location!())?;
if subspace != SUBSPACE_REGISTRY_IDX {
trx.prepare_cached(&format!(
"INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)",
table
))
.map_err(into_error)
.caused_by(trc::location!())?
.execute([&key, value])
.map_err(into_error)
.caused_by(trc::location!())?;
} else {
trx.prepare_cached("INSERT OR IGNORE INTO b (k) VALUES (?)")
.map_err(into_error)
.caused_by(trc::location!())?
.execute([&key])
.map_err(into_error)
.caused_by(trc::location!())?;
}
}
ValueOp::SetFnc(set_op) => {
let value = (set_op.fnc)(&set_op.params, &result)?;

1
crates/store/src/build/data.rs
View File

@@ -10,6 +10,7 @@ use registry::schema::{
structs::{DataStore, MetricsStore, TracingStore},
};
#[allow(unreachable_patterns)]
impl Store {
pub async fn build(config: DataStore) -> Result<Self, String> {
#[allow(unreachable_patterns)]

58
crates/store/src/build/registry.rs
View File

@@ -15,7 +15,6 @@ use crate::{
};
use std::{path::PathBuf, time::Duration};
use trc::AddContext;
use utils::snowflake::SnowflakeIdGenerator;
const STALE_NODE_TIMEOUT: u64 = 60 * 60; // 1 hour
const DEAD_NODE_TIMEOUT: u64 = 60 * 60 * 24; // 24 hours
@@ -26,7 +25,18 @@ impl RegistryStore {
let mut inner = RegistryStoreInner::new(local);
// Build store
let store = Store::build(inner.read_data_store().await?).await?;
inner.store = Store::build(inner.read_data_store().await?).await?;
Self::from_inner(inner).await
}
pub async fn from_inner(mut inner: RegistryStoreInner) -> Result<Self, String> {
// Create tables (SQL only)
inner
.store
.create_tables()
.await
.map_err(|err| format!("Failed to create tables: {err}"))?;
// Obtain node id
let mut retry_count = 0;
@@ -35,7 +45,8 @@ impl RegistryStore {
let mut batch = BatchBuilder::new();
let now = now();
let mut node_ids = Vec::new();
store
inner
.store
.iterate(
IterateParams::new(
ValueKey::from(ValueClass::NodeId(0)),
@@ -113,7 +124,7 @@ impl RegistryStore {
);
}
match store.write(batch.build_all()).await {
match inner.store.write(batch.build_all()).await {
Ok(_) => break,
Err(err) => {
if err.is_assertion_failure() && retry_count < 5 {
@@ -126,7 +137,6 @@ impl RegistryStore {
}
}
inner.id_generator = SnowflakeIdGenerator::new();
Ok(Self(inner.into()))
}
@@ -157,23 +167,61 @@ impl RegistryStore {
.map(|_| ())
}
#[inline(always)]
pub fn recovery_admin(&self) -> Option<&(String, String)> {
self.0.env_recovery_admin.as_ref()
}
#[inline(always)]
pub fn cluster_role(&self) -> Option<&str> {
self.0.env_cluster_role.as_deref()
}
#[inline(always)]
pub fn cluster_push_shard(&self) -> u32 {
self.0.env_push_shard_id
}
#[inline(always)]
pub fn local_hostname(&self) -> &str {
&self.0.env_hostname
}
#[inline(always)]
pub fn is_recovery_mode(&self) -> bool {
self.0.env_recovery_mode
}
#[inline(always)]
pub fn path(&self) -> &PathBuf {
&self.0.local_path
}
#[inline(always)]
pub fn store(&self) -> &Store {
&self.0.store
}
#[cfg(feature = "test_mode")]
pub async fn new(
path: &str,
store: Store,
hostname: String,
push_shard_id: u32,
cluster_role: Option<String>,
) -> Self {
Self::from_inner(RegistryStoreInner {
local_path: PathBuf::from(path),
store,
node_id: 0,
env_recovery_mode: false,
env_recovery_admin: Some(("admin".to_string(), "popolna_zapora".to_string())),
env_cluster_role: cluster_role,
env_push_shard_id: push_shard_id,
env_hostname: hostname,
id_generator: utils::snowflake::SnowflakeIdGenerator::new(),
})
.await
.unwrap()
}
}

12
crates/store/src/dispatch/lookup.rs
View File

@@ -271,11 +271,11 @@ impl InMemoryStore {
pub async fn key_exists(&self, key: impl Into<LookupKey<'_>>) -> trc::Result<bool> {
match self {
InMemoryStore::Store(store) => store
.get_value::<LookupValue<()>>(ValueKey::from(ValueClass::InMemory(
.get_value::<LookupValue<Empty>>(ValueKey::from(ValueClass::InMemory(
InMemoryClass::Key(key.into().into_bytes()),
)))
.await
.map(|value| matches!(value, Some(LookupValue::Value(())))),
.map(|value| matches!(value, Some(LookupValue::Value(Empty)))),
#[cfg(feature = "redis")]
InMemoryStore::Redis(store) => store.key_exists(key.into().as_bytes()).await,
// SPDX-SnippetBegin
@@ -635,6 +635,8 @@ impl<T> KeyValue<T> {
}
}
struct Empty;
enum LookupValue<T> {
Value(T),
None,
@@ -655,6 +657,12 @@ impl<T: Deserialize> Deserialize for LookupValue<T> {
}
}
impl Deserialize for Empty {
fn deserialize(_bytes: &[u8]) -> trc::Result<Self> {
Ok(Empty)
}
}
impl<T> From<LookupValue<T>> for Option<T> {
fn from(value: LookupValue<T>) -> Self {
match value {

23
crates/store/src/dispatch/store.rs
View File

@@ -45,6 +45,29 @@ impl Store {
.caused_by(trc::location!())
}
pub async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
match self {
#[cfg(feature = "sqlite")]
Self::SQLite(store) => store.key_exists(key).await,
#[cfg(feature = "foundation")]
Self::FoundationDb(store) => store.key_exists(key).await,
#[cfg(feature = "postgres")]
Self::PostgreSQL(store) => store.key_exists(key).await,
#[cfg(feature = "mysql")]
Self::MySQL(store) => store.key_exists(key).await,
#[cfg(feature = "rocks")]
Self::RocksDb(store) => store.key_exists(key).await,
// SPDX-SnippetBegin
// SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
// SPDX-License-Identifier: LicenseRef-SEL
#[cfg(all(feature = "enterprise", any(feature = "postgres", feature = "mysql")))]
Self::SQLReadReplica(store) => store.key_exists(key).await,
// SPDX-SnippetEnd
Self::None => Err(trc::StoreEvent::NotConfigured.into()),
}
.caused_by(trc::location!())
}
pub async fn iterate<T: Key>(
&self,
params: IterateParams<T>,

10
crates/store/src/lib.rs
View File

@@ -110,7 +110,7 @@ pub const SUBSPACE_IN_MEMORY_COUNTER: u8 = b'y';
pub const SUBSPACE_PROPERTY: u8 = b'p';
pub const SUBSPACE_REGISTRY: u8 = b's';
pub const SUBSPACE_REGISTRY_IDX: u8 = b'b';
pub const SUBSPACE_REGISTRY_PK: u8 = b'c';
pub const SUBSPACE_REGISTRY_PK: u8 = b'g';
pub const SUBSPACE_DIRECTORY: u8 = b'd';
pub const SUBSPACE_QUEUE_MESSAGE: u8 = b'e';
pub const SUBSPACE_QUEUE_EVENT: u8 = b'q';
@@ -125,7 +125,7 @@ pub const SUBSPACE_SPAM_SAMPLES: u8 = b'w';
// TODO: Remove in v1.0
pub const LEGACY_SUBSPACE_BITMAP_TEXT: u8 = b'v';
pub const LEGACY_SUBSPACE_FTS_INDEX: u8 = b'g';
pub const LEGACY_SUBSPACE_BITMAP_TAG: u8 = b'c';
#[derive(Clone)]
pub struct IterateParams<T: Key> {
@@ -282,6 +282,12 @@ impl From<Store> for InMemoryStore {
}
}
impl From<Store> for BlobStore {
fn from(store: Store) -> Self {
Self::Store(store)
}
}
impl Default for BlobStore {
fn default() -> Self {
Self::Store(Store::None)

8
crates/store/src/registry/get.rs
View File

@@ -63,15 +63,13 @@ impl RegistryStore {
IterateParams::new(
ValueKey::from(ValueClass::Any(AnyClass {
subspace: SUBSPACE_REGISTRY,
key: KeySerializer::new(U16_LEN + 1)
.write(0u8)
key: KeySerializer::new(U16_LEN)
.write(object_type.to_id())
.finalize(),
})),
ValueKey::from(ValueClass::Any(AnyClass {
subspace: SUBSPACE_REGISTRY,
key: KeySerializer::new(U16_LEN + U64_LEN + 1)
.write(0u8)
key: KeySerializer::new(U16_LEN + U64_LEN)
.write(object_type.to_id())
.write(u64::MAX)
.finalize(),
@@ -79,7 +77,7 @@ impl RegistryStore {
),
|key, value| {
let id = key
.get(U16_LEN + 1..)
.get(U16_LEN..)
.and_then(|key| key.read_leb128::<u64>())
.map(|r| r.0)
.ok_or_else(|| {

20
crates/store/src/registry/write.rs
View File

@@ -196,23 +196,22 @@ impl RegistryStore {
} else {
RegistryClass::IndexId { object_id, item_id }
};
if self
if !self
.0
.store
.get_value::<()>(ValueKey::from(ValueClass::Registry(key)))
.key_exists(ValueKey::from(ValueClass::Registry(key)))
.await
.caused_by(trc::location!())?
.is_none()
{
return Ok(RegistryWriteResult::InvalidForeignKey {
object_id: *foreign_id,
});
} else if let Some(tenant_id) = tenant_id
&& (object_flags & OBJ_FILTER_TENANT) != 0
&& self
&& !self
.0
.store
.get_value::<()>(ValueKey::from(ValueClass::Registry(
.key_exists(ValueKey::from(ValueClass::Registry(
RegistryClass::Index {
index_id: Property::MemberTenantId.to_id(),
object_id,
@@ -222,17 +221,16 @@ impl RegistryStore {
)))
.await
.caused_by(trc::location!())?
.is_none()
{
return Ok(RegistryWriteResult::InvalidForeignKey {
object_id: *foreign_id,
});
} else if (object_flags & OBJ_FILTER_ACCOUNT) != 0
&& let Some(account_id) = account_id
&& self
&& !self
.0
.store
.get_value::<()>(ValueKey::from(ValueClass::Registry(
.key_exists(ValueKey::from(ValueClass::Registry(
RegistryClass::Index {
index_id: Property::AccountId.to_id(),
object_id,
@@ -242,7 +240,6 @@ impl RegistryStore {
)))
.await
.caused_by(trc::location!())?
.is_none()
{
return Ok(RegistryWriteResult::InvalidForeignKey {
object_id: *foreign_id,
@@ -312,7 +309,10 @@ impl RegistryStore {
out,
);
Ok(RegistryWriteResult::Success(Id::new(item_id)))
self.store()
.write(batch.build_all())
.await
.map(|_| RegistryWriteResult::Success(Id::new(item_id)))
}
async fn delete(

14
crates/store/src/write/blob.rs
View File

@@ -29,7 +29,7 @@ pub struct BlobQuota {
impl Store {
pub async fn blob_exists(&self, hash: impl AsRef<BlobHash> + Sync + Send) -> trc::Result<bool> {
self.get_value::<()>(ValueKey {
self.key_exists(ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
@@ -38,7 +38,6 @@ impl Store {
}),
})
.await
.map(|v| v.is_some())
.caused_by(trc::location!())
}
@@ -76,7 +75,14 @@ impl Store {
_ => return Ok(false),
};
self.get_value::<()>(key).await.map(|v| v.is_some())
self.key_exists(key).await
}
pub async fn purge_blobs_all_shards(&self, blob_store: BlobStore) -> trc::Result<()> {
for shard_index in 0u8..=255 {
self.purge_blobs(blob_store.clone(), shard_index).await?;
}
Ok(())
}
pub async fn purge_blobs(&self, blob_store: BlobStore, shard_index: u8) -> trc::Result<()> {
@@ -270,6 +276,8 @@ impl BlobPurgeState {
self.delete_registry
.push((account_id, ObjectId::deserialize(value)?));
}
} else {
self.last_hash_is_linked = true;
}
Ok(())
}

6
crates/store/src/write/serialize.rs
View File

@@ -563,12 +563,6 @@ impl Deserialize for u32 {
}
}
impl Deserialize for () {
fn deserialize(_bytes: &[u8]) -> trc::Result<Self> {
Ok(())
}
}
impl<T> From<Value<'static>> for Archive<T> {
fn from(_: Value<'static>) -> Self {
unimplemented!()

4
tests/Cargo.toml
View File

@@ -5,8 +5,8 @@ edition = "2024"
[features]
#default = ["sqlite", "postgres", "mysql", "rocks", "s3", "redis", "nats", "azure", "foundationdb"]
#default = ["sqlite", "postgres", "mysql", "rocks", "s3", "redis", "foundationdb"]
default = ["postgres"]
default = ["sqlite", "postgres", "mysql", "rocks", "s3", "redis", "foundationdb"]
#default = ["postgres"]
sqlite = ["store/sqlite", "directory/sqlite"]
foundationdb = ["store/foundation", "common/foundation"]
postgres = ["store/postgres", "directory/postgres"]

4
tests/src/lib.rs
View File

@@ -29,11 +29,11 @@ pub mod jmap;
#[cfg(test)]
pub mod smtp;
#[cfg(test)]
pub mod store;
#[cfg(test)]
pub mod webdav;
*/
#[cfg(test)]
pub mod store;
#[cfg(test)]
pub mod system;
#[cfg(test)]
pub mod utils;

832
tests/src/store/blob.rs
View File

@@ -4,457 +4,447 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::store::{CONFIG, TempDir, cleanup::store_destroy};
use crate::utils::{cleanup::store_destroy, server::TestServerBuilder};
use ahash::AHashMap;
use common::{Core, Inner, Server, config::storage::Storage};
use email::message::metadata::MessageMetadata;
use std::sync::Arc;
use registry::{
schema::{enums::CompressionAlgo, structs::Jmap},
types::duration::Duration,
};
use services::task_manager::destroy_account::destroy_account_blobs;
use store::{
BlobStore, Serialize, SerializeInfallible,
write::{Archiver, BatchBuilder, BlobLink, BlobOp, ValueClass, blob::BlobQuota, now},
write::{Archiver, BatchBuilder, BlobLink, BlobOp, ValueClass, now},
};
use types::{blob::BlobClass, blob_hash::BlobHash, collection::Collection, field::EmailField};
#[tokio::test]
pub async fn blob_tests() {
let temp_dir = TempDir::new("blob_tests", true);
let mut config =
Config::new(CONFIG.replace("{TMP}", temp_dir.path.as_path().to_str().unwrap())).unwrap();
let stores = Stores::parse_all(&mut config, false).await;
let test = TestServerBuilder::new("blob_tests", true)
.await
.with_object(Jmap {
upload_quota: 1024,
upload_ttl: Duration::from_millis(1000),
..Default::default()
})
.await
.build()
.await;
for (store_id, blob_store) in &stores.blob_stores {
println!("Testing blob store {}...", store_id);
test_store(blob_store.clone()).await;
let store = test.server.core.storage.data.clone();
let blob_store = test.server.core.storage.blob.clone();
println!(
"Testing blob store {} with data store {}...",
std::env::var("BLOB_STORE").unwrap_or_else(|_| "default".to_string()),
std::env::var("STORE").unwrap()
);
// Test blob quota
assert!(test.server.blob_has_quota(0, 1024).await.unwrap());
assert!(!test.server.blob_has_quota(0, 1024).await.unwrap());
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
assert!(test.server.blob_has_quota(0, 1024).await.unwrap());
// Test and reset store
test_store(blob_store.clone()).await;
store_destroy(&store).await;
// Blob hash exists
let hash = BlobHash::generate(b"abc".as_slice());
assert!(!store.blob_exists(&hash).await.unwrap());
// Reserve blob
let until = now() + 1;
store
.write(
BatchBuilder::new()
.with_account_id(0)
.set(
BlobOp::Link {
to: BlobLink::Temporary { until },
hash: hash.clone(),
},
1024u32.serialize(),
)
.build_all(),
)
.await
.unwrap();
// Uncommitted blob, should not exist
assert!(!store.blob_exists(&hash).await.unwrap());
// Write blob to store
blob_store
.put_blob(hash.as_ref(), b"abc", CompressionAlgo::Lz4)
.await
.unwrap();
// Commit blob
store
.write(
BatchBuilder::new()
.set(BlobOp::Commit { hash: hash.clone() }, Vec::new())
.build_all(),
)
.await
.unwrap();
// Blob hash should now exist
assert!(store.blob_exists(&hash).await.unwrap());
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_some()
);
// AccountId 0 should be able to read blob
assert!(
store
.blob_has_access(
&hash,
BlobClass::Reserved {
account_id: 0,
expires: until
}
)
.await
.unwrap()
);
// AccountId 1 should not be able to read blob
assert!(
!store
.blob_has_access(
&hash,
BlobClass::Reserved {
account_id: 1,
expires: until
}
)
.await
.unwrap()
);
// Purge expired blobs
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
store
.purge_blobs_all_shards(blob_store.clone())
.await
.unwrap();
// Blob hash should no longer exist
assert!(!store.blob_exists(&hash).await.unwrap());
// AccountId 0 should not be able to read blob
assert!(
!store
.blob_has_access(
&hash,
BlobClass::Reserved {
account_id: 0,
expires: until
}
)
.await
.unwrap()
);
// Blob should no longer be in store
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_none()
);
// Upload one linked blob to accountId 1, two linked blobs to accountId 0, and three unlinked (reserved) blobs to accountId 2
let expiry_times = AHashMap::from_iter([
(b"abc", now() - 10),
(b"efg", now() + 10),
(b"hij", now() + 10),
]);
for (document_id, (blob, _)) in [
(b"123", vec![]),
(b"456", vec![]),
(b"789", vec![]),
(b"abc", 5000u32.serialize()),
(b"efg", 1000u32.serialize()),
(b"hij", 2000u32.serialize()),
]
.into_iter()
.enumerate()
{
let hash = BlobHash::generate(blob.as_slice());
let mut batch = BatchBuilder::new();
batch
.with_account_id(if document_id > 0 { 0 } else { 1 })
.with_collection(Collection::Email)
.with_document(document_id as u32);
if let Some(until) = expiry_times.get(blob) {
batch.set(
BlobOp::Link {
hash: hash.clone(),
to: BlobLink::Temporary { until: *until },
},
vec![],
);
} else {
batch
.set(
BlobOp::Link {
hash: hash.clone(),
to: BlobLink::Document,
},
vec![],
)
.set(
ValueClass::Property(EmailField::Metadata.into()),
Archiver::new(MessageMetadata {
contents: Default::default(),
rcvd_attach: Default::default(),
blob_hash: hash.clone(),
blob_body_offset: Default::default(),
preview: Default::default(),
raw_headers: Default::default(),
})
.serialize()
.unwrap(),
);
};
batch.set(BlobOp::Commit { hash: hash.clone() }, vec![]);
store.write(batch.build_all()).await.unwrap();
blob_store
.put_blob(hash.as_ref(), blob.as_slice(), CompressionAlgo::Lz4)
.await
.unwrap();
}
for (store_id, store) in stores.stores {
println!("Testing blob management on store {}...", store_id);
// Init store
store_destroy(&store).await;
// Test internal blob store
let blob_store: BlobStore = store.clone().into();
let server = Server {
inner: Arc::new(Inner::default()),
core: Arc::new(Core {
storage: Storage {
data: store.clone(),
blob: blob_store.clone(),
..Default::default()
},
..Default::default()
}),
};
// Blob hash exists
let hash = BlobHash::generate(b"abc".as_slice());
assert!(!store.blob_exists(&hash).await.unwrap());
// Reserve blob
let until = now() + 1;
store
.write(
BatchBuilder::new()
.with_account_id(0)
.set(
BlobOp::Link {
to: BlobLink::Temporary { until },
hash: hash.clone(),
},
1024u32.serialize(),
)
.build_all(),
)
.await
.unwrap();
// Uncommitted blob, should not exist
assert!(!store.blob_exists(&hash).await.unwrap());
// Write blob to store
blob_store.put_blob(hash.as_ref(), b"abc").await.unwrap();
// Commit blob
store
.write(
BatchBuilder::new()
.set(BlobOp::Commit { hash: hash.clone() }, Vec::new())
.build_all(),
)
.await
.unwrap();
// Blob hash should now exist
assert!(store.blob_exists(&hash).await.unwrap());
// Purge expired blobs and make sure nothing else is deleted
store
.purge_blobs_all_shards(blob_store.clone())
.await
.unwrap();
for (pos, (blob, blob_class)) in [
(
b"abc",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"abc"],
},
),
(
b"123",
BlobClass::Linked {
account_id: 1,
collection: 0,
document_id: 0,
},
),
(
b"456",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 1,
},
),
(
b"789",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 2,
},
),
(
b"efg",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"efg"],
},
),
(
b"hij",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"hij"],
},
),
]
.into_iter()
.enumerate()
{
let hash = BlobHash::generate(blob.as_slice());
let ct = pos == 0;
assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_some()
^ ct
);
}
// AccountId 0 should be able to read blob
assert!(
store
.blob_has_access(
&hash,
BlobClass::Reserved {
account_id: 0,
expires: until
}
)
.await
.unwrap()
);
// AccountId 0 should not have access to accountId 1's blobs
assert!(
!store
.blob_has_access(
BlobHash::generate(b"123".as_slice()),
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 0,
}
)
.await
.unwrap()
);
// AccountId 1 should not be able to read blob
assert!(
!store
.blob_has_access(
&hash,
BlobClass::Reserved {
account_id: 1,
expires: until
}
)
.await
.unwrap()
);
// Unlink blob
store
.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(2)
.clear(BlobOp::Link {
hash: BlobHash::generate(b"789".as_slice()),
to: BlobLink::Document,
})
.build_all(),
)
.await
.unwrap();
// Blob already expired, quota should be 0
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
assert_eq!(
store.blob_quota(0).await.unwrap(),
BlobQuota { bytes: 0, count: 0 }
);
// Purge expired blobs
store.purge_blobs(blob_store.clone()).await.unwrap();
// Blob hash should no longer exist
assert!(!store.blob_exists(&hash).await.unwrap());
// AccountId 0 should not be able to read blob
assert!(
!store
.blob_has_access(
&hash,
BlobClass::Reserved {
account_id: 0,
expires: until
}
)
.await
.unwrap()
);
// Blob should no longer be in store
// Purge and make sure blob is deleted
store
.purge_blobs_all_shards(blob_store.clone())
.await
.unwrap();
for (pos, (blob, blob_class)) in [
(
b"789",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 2,
},
),
(
b"123",
BlobClass::Linked {
account_id: 1,
collection: 0,
document_id: 0,
},
),
(
b"456",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 1,
},
),
(
b"efg",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"efg"],
},
),
(
b"hij",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"hij"],
},
),
]
.into_iter()
.enumerate()
{
let ct = pos == 0;
let hash = BlobHash::generate(blob.as_slice());
assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_none()
.is_some()
^ ct
);
}
// Upload one linked blob to accountId 1, two linked blobs to accountId 0, and three unlinked (reserved) blobs to accountId 2
let expiry_times = AHashMap::from_iter([
(b"abc", now() - 10),
(b"efg", now() + 10),
(b"hij", now() + 10),
]);
for (document_id, (blob, blob_value)) in [
(b"123", vec![]),
(b"456", vec![]),
(b"789", vec![]),
(b"abc", 5000u32.serialize()),
(b"efg", 1000u32.serialize()),
(b"hij", 2000u32.serialize()),
]
.into_iter()
.enumerate()
{
let hash = BlobHash::generate(blob.as_slice());
let mut batch = BatchBuilder::new();
batch
.with_account_id(if document_id > 0 { 0 } else { 1 })
.with_collection(Collection::Email)
.with_document(document_id as u32);
if let Some(until) = expiry_times.get(blob) {
if !blob_value.is_empty() {
batch.set(
BlobOp::Quota {
hash: hash.clone(),
until: *until,
},
blob_value,
);
}
batch.set(
BlobOp::Link {
hash: hash.clone(),
to: BlobLink::Temporary { until: *until },
},
vec![],
);
} else {
batch
.set(
BlobOp::Link {
hash: hash.clone(),
to: BlobLink::Document,
},
vec![],
)
.set(
ValueClass::Property(EmailField::Metadata.into()),
Archiver::new(MessageMetadata {
contents: Default::default(),
rcvd_attach: Default::default(),
blob_hash: hash.clone(),
blob_body_offset: Default::default(),
preview: Default::default(),
raw_headers: Default::default(),
})
.serialize()
.unwrap(),
);
};
batch.set(BlobOp::Commit { hash: hash.clone() }, vec![]);
// Unlink all blobs from accountId 1 and purge
destroy_account_blobs(&test.server, 1).await.unwrap();
store
.purge_blobs_all_shards(blob_store.clone())
.await
.unwrap();
store.write(batch.build_all()).await.unwrap();
blob_store
.put_blob(hash.as_ref(), blob.as_slice())
.await
.unwrap();
}
// One of the reserved blobs expired and should not count towards quota
assert_eq!(
store.blob_quota(0).await.unwrap(),
BlobQuota {
bytes: 3000,
count: 2
}
);
assert_eq!(
store.blob_quota(1).await.unwrap(),
BlobQuota { bytes: 0, count: 0 }
);
// Purge expired blobs and make sure nothing else is deleted
store.purge_blobs(blob_store.clone()).await.unwrap();
for (pos, (blob, blob_class)) in [
(
b"abc",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"abc"],
},
),
(
b"123",
BlobClass::Linked {
account_id: 1,
collection: 0,
document_id: 0,
},
),
(
b"456",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 1,
},
),
(
b"789",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 2,
},
),
(
b"efg",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"efg"],
},
),
(
b"hij",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"hij"],
},
),
]
.into_iter()
.enumerate()
{
let ct = pos == 0;
let hash = BlobHash::generate(blob.as_slice());
assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_some()
^ ct
);
}
// AccountId 0 should not have access to accountId 1's blobs
// Make sure only accountId 0's blobs are left
for (pos, (blob, blob_class)) in [
(
b"123",
BlobClass::Linked {
account_id: 1,
collection: 0,
document_id: 0,
},
),
(
b"456",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 1,
},
),
(
b"efg",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"efg"],
},
),
(
b"hij",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"hij"],
},
),
]
.into_iter()
.enumerate()
{
let ct = pos == 0;
let hash = BlobHash::generate(blob.as_slice());
assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
assert!(
!store
.blob_has_access(
BlobHash::generate(b"123".as_slice()),
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 0,
}
)
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_some()
^ ct
);
// Unlink blob
store
.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(2)
.clear(BlobOp::Link {
hash: BlobHash::generate(b"789".as_slice()),
to: BlobLink::Document,
})
.build_all(),
)
.await
.unwrap();
// Purge and make sure blob is deleted
store.purge_blobs(blob_store.clone()).await.unwrap();
for (pos, (blob, blob_class)) in [
(
b"789",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 2,
},
),
(
b"123",
BlobClass::Linked {
account_id: 1,
collection: 0,
document_id: 0,
},
),
(
b"456",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 1,
},
),
(
b"efg",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"efg"],
},
),
(
b"hij",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"hij"],
},
),
]
.into_iter()
.enumerate()
{
let ct = pos == 0;
let hash = BlobHash::generate(blob.as_slice());
assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_some()
^ ct
);
}
// Unlink all blobs from accountId 1 and purge
destroy_account_blobs(&server, 1).await.unwrap();
store.purge_blobs(blob_store.clone()).await.unwrap();
// Make sure only accountId 0's blobs are left
for (pos, (blob, blob_class)) in [
(
b"123",
BlobClass::Linked {
account_id: 1,
collection: 0,
document_id: 0,
},
),
(
b"456",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 1,
},
),
(
b"efg",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"efg"],
},
),
(
b"hij",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"hij"],
},
),
]
.into_iter()
.enumerate()
{
let ct = pos == 0;
let hash = BlobHash::generate(blob.as_slice());
assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_some()
^ ct
);
}
}
temp_dir.delete();
test.temp_dir.delete();
}
async fn test_store(store: BlobStore) {
@@ -462,7 +452,10 @@ async fn test_store(store: BlobStore) {
const DATA: &[u8] = b"Lorem ipsum dolor sit amet, consectetur adipiscing elit. Fusce erat nisl, dignissim a porttitor id, varius nec arcu. Sed mauris.";
let hash = BlobHash::generate(DATA);
store.put_blob(hash.as_slice(), DATA).await.unwrap();
store
.put_blob(hash.as_slice(), DATA, CompressionAlgo::Lz4)
.await
.unwrap();
assert_eq!(
String::from_utf8(
store
@@ -502,7 +495,10 @@ async fn test_store(store: BlobStore) {
data.extend_from_slice(marker.as_bytes());
}
let hash = BlobHash::generate(&data);
store.put_blob(hash.as_slice(), &data).await.unwrap();
store
.put_blob(hash.as_slice(), &data, CompressionAlgo::Lz4)
.await
.unwrap();
assert_eq!(
String::from_utf8(
store

88
tests/src/store/import_export.rs
View File

@@ -4,17 +4,21 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::store::{
TempDir,
cleanup::{store_assert_is_empty, store_destroy},
use crate::{
store::TempDir,
utils::{
cleanup::{store_assert_is_empty, store_destroy},
server::TestServer,
},
};
use ::registry::schema::enums::CompressionAlgo;
use ahash::AHashSet;
use common::{Core, DATABASE_SCHEMA_VERSION, manager::backup::BackupParams};
use common::{DATABASE_SCHEMA_VERSION, manager::backup::BackupParams};
use store::{
rand,
write::{
AnyClass, AnyKey, BatchBuilder, BlobLink, BlobOp, Operation, QueueClass, QueueEvent,
ValueClass,
RegistryClass, ValueClass,
},
*,
};
@@ -24,15 +28,10 @@ use types::{
field::{Field, MailboxField},
};
pub async fn test(db: Store) {
let mut core = Core::default();
core.storage.data = db.clone();
core.storage.blob = db.clone().into();
core.storage.fts = db.clone().into();
core.storage.lookup = db.clone().into();
pub async fn test(test: &TestServer) {
// Make sure the store is empty
store_assert_is_empty(&db, db.clone().into(), true).await;
store_assert_is_empty(test.server.store(), test.server.blob_store().clone(), true).await;
let db = test.server.store().clone();
// Create blobs
println!("Creating blobs...");
@@ -49,9 +48,9 @@ pub async fn test(db: Store) {
let data = random_bytes(blob_size);
let hash = BlobHash::generate(data.as_slice());
blob_hashes.push(hash.clone());
core.storage
.blob
.put_blob(hash.as_ref(), &data)
test.server
.blob_store()
.put_blob(hash.as_ref(), &data, CompressionAlgo::Lz4)
.await
.unwrap();
batch.set(ValueClass::Blob(BlobOp::Commit { hash }), vec![]);
@@ -142,16 +141,11 @@ pub async fn test(db: Store) {
})),
random_bytes(idx),
);
/*batch.set(
ValueClass::InMemory(InMemoryClass::Key(random_bytes(idx))),
random_bytes(idx),
);
batch.add(
ValueClass::InMemory(InMemoryClass::Counter(random_bytes(idx))),
rand::random(),
);*/
batch.set(
ValueClass::Config(random_bytes(idx + 10)),
ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: 1,
}),
random_bytes(idx + 10),
);
}
@@ -167,40 +161,7 @@ pub async fn test(db: Store) {
for account_id in [1, 2, 3, 4, 5] {
batch
.with_document(account_id)
.add(
ValueClass::Directory(DirectoryClass::UsedQuota(account_id)),
rand::random(),
)
.set(
ValueClass::Directory(DirectoryClass::NameToId(random_bytes(
2 + account_id as usize,
))),
random_bytes(4),
)
.set(
ValueClass::Directory(DirectoryClass::EmailToId(random_bytes(
4 + account_id as usize,
))),
random_bytes(4),
)
.set(
ValueClass::Directory(DirectoryClass::Principal(account_id)),
random_bytes(30),
)
.set(
ValueClass::Directory(DirectoryClass::MemberOf {
principal_id: account_id,
member_of: rand::random(),
}),
random_bytes(15),
)
.set(
ValueClass::Directory(DirectoryClass::Members {
principal_id: account_id,
has_member: rand::random(),
}),
random_bytes(15),
);
.add(ValueClass::Quota, account_id as i64 * 1000);
}
db.write(batch.build_all()).await.unwrap();
@@ -212,7 +173,10 @@ pub async fn test(db: Store) {
// Export store
println!("Exporting store...");
let temp_dir = TempDir::new("art_vandelay_tests", true);
core.backup(BackupParams::new(temp_dir.path.clone())).await;
test.server
.core
.backup(BackupParams::new(temp_dir.path.clone()))
.await;
// Destroy store
println!("Destroying store...");
@@ -221,7 +185,7 @@ pub async fn test(db: Store) {
// Import store
println!("Importing store...");
core.restore(temp_dir.path.clone()).await;
test.server.core.restore(temp_dir.path.clone()).await;
// Verify hash
print!("Verifying store hash...");
@@ -266,6 +230,8 @@ impl Snapshot {
(SUBSPACE_IN_MEMORY_VALUE, true),
(SUBSPACE_PROPERTY, true),
(SUBSPACE_REGISTRY, true),
(SUBSPACE_REGISTRY_IDX, !is_sql),
(SUBSPACE_REGISTRY_PK, true),
(SUBSPACE_QUEUE_MESSAGE, true),
(SUBSPACE_QUEUE_EVENT, true),
(SUBSPACE_QUOTA, !is_sql),

479
tests/src/store/lookup.rs
View File

@@ -4,285 +4,284 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::{
AssertConfig,
store::{
CONFIG, TempDir,
cleanup::{store_assert_is_empty, store_destroy},
},
use crate::utils::{
cleanup::{store_assert_is_empty, store_destroy},
server::TestServerBuilder,
};
use std::time::Duration;
use registry::schema::structs::Rate;
use registry::types::duration::Duration;
use store::{InMemoryStore, dispatch::lookup::KeyValue};
#[tokio::test]
pub async fn lookup_tests() {
let temp_dir = TempDir::new("lookup_tests", true);
let mut config =
Config::new(CONFIG.replace("{TMP}", temp_dir.path.as_path().to_str().unwrap()))
.unwrap()
.assert_no_errors();
let stores = Stores::parse_all(&mut config, false).await;
let insert = std::env::var("NO_INSERT").is_err();
let test = TestServerBuilder::new("lookup_tests", insert)
.await
.build()
.await;
let store = test.server.in_memory_store().clone();
let rate = Rate {
requests: 1,
period: Duration::from_secs(1),
count: 1,
period: Duration::from_millis(1000),
};
for (store_id, store) in stores.in_memory_stores {
println!("Testing in-memory store {}...", store_id);
if let InMemoryStore::Store(store) = &store {
store_destroy(store).await;
} else {
// Reset redis counter
store
.key_set(KeyValue::new("abc", "0".as_bytes().to_vec()))
.await
.unwrap();
println!(
"Testing in-memory store {}...",
std::env::var("MEMORY_STORE").unwrap_or_else(|_| "default".to_string())
);
if let InMemoryStore::Store(store) = &store {
store_destroy(store).await;
} else {
// Reset redis counter
store
.key_set(KeyValue::new("abc", "0".as_bytes().to_vec()))
.await
.unwrap();
}
// Test key
let key = "xyz".as_bytes().to_vec();
store
.key_set(KeyValue::new(key.clone(), "world".to_string().into_bytes()))
.await
.unwrap();
store.purge_in_memory_store().await.unwrap();
assert_eq!(
store.key_get::<String>(key.clone()).await.unwrap(),
Some("world".to_string())
);
// Test value expiry
store
.key_set(KeyValue::new(key.clone(), "hello".to_string().into_bytes()).expires(1))
.await
.unwrap();
assert_eq!(
store.key_get::<String>(key.clone()).await.unwrap(),
Some("hello".to_string())
);
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
assert_eq!(None, store.key_get::<String>(key.clone()).await.unwrap());
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
// Test counter
let key = "abc".as_bytes().to_vec();
store
.counter_incr(KeyValue::new(key.clone(), 1), true)
.await
.unwrap();
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
store
.counter_incr(KeyValue::new(key.clone(), 2), true)
.await
.unwrap();
assert_eq!(3, store.counter_get(key.clone()).await.unwrap());
store
.counter_incr(KeyValue::new(key.clone(), -3), false)
.await
.unwrap();
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
// Test counter expiry
let key = "fgh".as_bytes().to_vec();
store
.counter_incr(KeyValue::new(key.clone(), 1).expires(1), false)
.await
.unwrap();
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
store.purge_in_memory_store().await.unwrap();
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
// Test rate limiter
assert!(
store
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
.await
.unwrap()
.is_none()
);
assert!(
store
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
.await
.unwrap()
.is_some()
);
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
assert!(
store
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
.await
.unwrap()
.is_none()
);
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
// Test locking
for iteration in [1, 2] {
let mut tasks = Vec::new();
for _ in 0..100 {
let store = store.clone();
tasks.push(tokio::spawn(async move {
store.try_lock(0, "lock".as_bytes(), 1).await.unwrap()
}));
}
// Only one should return true
let mut count = 0;
for task in tasks {
if task.await.unwrap() {
count += 1;
}
}
assert_eq!(1, count, "Iteration {}", iteration);
// Test key
let key = "xyz".as_bytes().to_vec();
store
.key_set(KeyValue::new(key.clone(), "world".to_string().into_bytes()))
.await
.unwrap();
store.purge_in_memory_store().await.unwrap();
assert_eq!(
store.key_get::<String>(key.clone()).await.unwrap(),
Some("world".to_string())
);
// Test value expiry
store
.key_set(KeyValue::new(key.clone(), "hello".to_string().into_bytes()).expires(1))
.await
.unwrap();
assert_eq!(
store.key_get::<String>(key.clone()).await.unwrap(),
Some("hello".to_string())
);
// Wait 2 seconds for the lock to expire
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
assert_eq!(None, store.key_get::<String>(key.clone()).await.unwrap());
}
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
// Test counter
let key = "abc".as_bytes().to_vec();
store
.counter_incr(KeyValue::new(key.clone(), 1), true)
.await
.unwrap();
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
store
.counter_incr(KeyValue::new(key.clone(), 2), true)
.await
.unwrap();
assert_eq!(3, store.counter_get(key.clone()).await.unwrap());
store
.counter_incr(KeyValue::new(key.clone(), -3), false)
.await
.unwrap();
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
// Test counter expiry
let key = "fgh".as_bytes().to_vec();
store
.counter_incr(KeyValue::new(key.clone(), 1).expires(1), false)
.await
.unwrap();
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
store.purge_in_memory_store().await.unwrap();
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
// Test rate limiter
assert!(
store
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
.await
.unwrap()
.is_none()
);
assert!(
store
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
.await
.unwrap()
.is_some()
);
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
assert!(
store
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
.await
.unwrap()
.is_none()
);
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
// Test locking
for iteration in [1, 2] {
let mut tasks = Vec::new();
for _ in 0..100 {
let store = store.clone();
tasks.push(tokio::spawn(async move {
store.try_lock(0, "lock".as_bytes(), 1).await.unwrap()
}));
}
// Only one should return true
let mut count = 0;
for task in tasks {
if task.await.unwrap() {
count += 1;
}
}
assert_eq!(1, count, "Iteration {}", iteration);
// Wait 2 seconds for the lock to expire
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
}
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
// Test prefix delete
// Test prefix delete
store
.key_set(KeyValue::with_prefix(
1,
[0],
"hello".to_string().into_bytes(),
))
.await
.unwrap();
for v in 0u32..2020u32 {
store
.key_set(KeyValue::with_prefix(
1,
[0],
"hello".to_string().into_bytes(),
0,
pack_u32(0, v),
"world".to_string().into_bytes(),
))
.await
.unwrap();
for v in 0u32..2020u32 {
store
.key_set(KeyValue::with_prefix(
0,
pack_u32(0, v),
"world".to_string().into_bytes(),
))
.await
.unwrap();
store
.counter_incr(
KeyValue::with_prefix(0, pack_u32(1, v), 123).expires(3600),
false,
)
.await
.unwrap();
}
// Make sure the keys are there
assert_eq!(
Some("hello"),
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.await
.unwrap()
.as_deref()
);
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
Some("world"),
store
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
.await
.unwrap()
.as_deref()
);
}
for v in [0, 1000, 1001, 2000, 2001] {
assert_ne!(
0,
store
.counter_get(KeyValue::<()>::build_key(0, pack_u32(1, v)))
.await
.unwrap()
);
}
// Delete [0, 0, 0, 0, 1] prefix and make sure only the keys with that prefix are gone
store
.key_delete_prefix(&KeyValue::<()>::build_key(0, 1u32.to_be_bytes()))
.counter_incr(
KeyValue::with_prefix(0, pack_u32(1, v), 123).expires(3600),
false,
)
.await
.unwrap();
}
assert_eq!(
Some("hello"),
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.await
.unwrap()
.as_deref()
);
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
Some("world"),
store
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
.await
.unwrap()
.as_deref()
);
}
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
0,
store
.counter_get(KeyValue::<()>::build_key(0, pack_u32(1, v)))
.await
.unwrap()
);
}
// Delete [0, 0, 0, 0, 0] prefix and make sure only the keys with that prefix are gone
// Make sure the keys are there
assert_eq!(
Some("hello"),
store
.key_delete_prefix(&KeyValue::<()>::build_key(0, 0u32.to_be_bytes()))
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.await
.unwrap();
.unwrap()
.as_deref()
);
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
Some("hello"),
Some("world"),
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
.await
.unwrap()
.as_deref()
);
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
None,
store
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
.await
.unwrap()
.as_deref()
);
}
}
for v in [0, 1000, 1001, 2000, 2001] {
assert_ne!(
0,
store
.counter_get(KeyValue::<()>::build_key(0, pack_u32(1, v)))
.await
.unwrap()
);
}
// Delete [1, ...] prefix and make sure it's all gone
store.key_delete_prefix(&[1u8]).await.unwrap();
// Delete [0, 0, 0, 0, 1] prefix and make sure only the keys with that prefix are gone
store
.key_delete_prefix(&KeyValue::<()>::build_key(0, 1u32.to_be_bytes()))
.await
.unwrap();
assert_eq!(
Some("hello"),
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.await
.unwrap()
.as_deref()
);
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
Some("world"),
store
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
.await
.unwrap()
.as_deref()
);
}
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
0,
store
.counter_get(KeyValue::<()>::build_key(0, pack_u32(1, v)))
.await
.unwrap()
);
}
// Delete [0, 0, 0, 0, 0] prefix and make sure only the keys with that prefix are gone
store
.key_delete_prefix(&KeyValue::<()>::build_key(0, 0u32.to_be_bytes()))
.await
.unwrap();
assert_eq!(
Some("hello"),
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.await
.unwrap()
.as_deref()
);
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
None,
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
.await
.unwrap()
.as_deref()
);
}
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
// Delete [1, ...] prefix and make sure it's all gone
store.key_delete_prefix(&[1u8]).await.unwrap();
assert_eq!(
None,
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.await
.unwrap()
.as_deref()
);
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
}

185
tests/src/store/mod.rs
View File

@@ -5,16 +5,12 @@
*/
pub mod blob;
pub mod cleanup;
pub mod import_export;
pub mod lookup;
pub mod ops;
pub mod query;
use crate::{
AssertConfig,
store::cleanup::{search_store_destroy, store_destroy},
};
use crate::utils::server::TestServerBuilder;
use std::io::Read;
pub struct TempDir {
@@ -23,60 +19,41 @@ pub struct TempDir {
#[tokio::test(flavor = "multi_thread")]
pub async fn store_tests() {
let insert = true;
let temp_dir = TempDir::new("store_tests", insert);
let mut config = Config::new(build_store_config(&temp_dir.path.to_string_lossy()))
.unwrap()
.assert_no_errors();
let stores = Stores::parse_all(&mut config, false).await;
let insert = std::env::var("NO_INSERT").is_err();
let test = TestServerBuilder::new("store_tests", insert)
.await
.build()
.await;
let store_id = std::env::var("STORE")
.expect("Missing store type. Try running `STORE=<store_type> cargo test`");
let store = stores
.stores
.get(&store_id)
.expect("Store not found")
.clone();
println!("Testing store {}...", std::env::var("STORE").unwrap());
println!("Testing store {}...", store_id);
if insert {
store_destroy(&store).await;
}
test.destroy_store().await;
import_export::test(store.clone()).await;
ops::test(store.clone()).await;
import_export::test(&test).await;
ops::test(&test).await;
if insert {
temp_dir.delete();
test.temp_dir.delete();
}
}
#[tokio::test(flavor = "multi_thread")]
pub async fn search_tests() {
let insert = std::env::var("NO_INSERT").is_err();
let temp_dir = TempDir::new("search_store_tests", insert);
let mut config = Config::new(build_store_config(&temp_dir.path.to_string_lossy()))
.unwrap()
.assert_no_errors();
let stores = Stores::parse_all(&mut config, false).await;
let test = TestServerBuilder::new("search_store_tests", insert)
.await
.build()
.await;
let store_id = std::env::var("SEARCH_STORE")
.expect("Missing store type. Try running `SEARCH_STORE=<store_type> cargo test`");
let store = stores
.search_stores
.get(&store_id)
.expect("Store not found")
.clone();
println!(
"Testing search store {}...",
std::env::var("SEARCH_STORE").unwrap_or("default".to_string())
);
println!("Testing store {}...", store_id);
if insert {
search_store_destroy(&store).await;
}
query::test(store, insert).await;
query::test(&test, insert).await;
if insert {
temp_dir.delete();
test.temp_dir.delete();
}
}
@@ -108,123 +85,3 @@ impl TempDir {
std::fs::remove_dir_all(&self.path).unwrap();
}
}
pub fn build_store_config(temp_dir: &str) -> String {
let store = std::env::var("STORE")
.expect("Missing store type. Try running `STORE=<store_type> cargo test`");
let fts_store = std::env::var("SEARCH_STORE").unwrap_or_else(|_| store.clone());
let blob_store = std::env::var("BLOB_STORE").unwrap_or_else(|_| store.clone());
let lookup_store = std::env::var("LOOKUP_STORE").unwrap_or_else(|_| store.clone());
CONFIG
.replace("{STORE}", &store)
.replace("{SEARCH_STORE}", &fts_store)
.replace("{BLOB_STORE}", &blob_store)
.replace("{LOOKUP_STORE}", &lookup_store)
.replace("{TMP}", temp_dir)
.replace(
"{ELASTIC_ENABLED}",
if fts_store != "elastic" {
"true"
} else {
"false"
},
)
.replace(
"{MEILI_ENABLED}",
if fts_store != "meili" {
"true"
} else {
"false"
},
)
}
const CONFIG: &str = r#"
[store."sqlite"]
type = "sqlite"
path = "{TMP}/sqlite.db"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/rocks.db"
[store."foundationdb"]
type = "foundationdb"
[store."postgresql"]
type = "postgresql"
host = "localhost"
port = 5432
database = "stalwart"
user = "postgres"
password = "mysecretpassword"
[store."mysql"]
type = "mysql"
host = "localhost"
port = 3307
database = "stalwart"
user = "root"
password = "password"
[store."elastic"]
type = "elasticsearch"
url = "https://localhost:9200"
tls.allow-invalid-certs = true
disable = {ELASTIC_ENABLED}
[store."elastic".auth]
username = "elastic"
secret = "changeme"
[store."meili"]
type = "meilisearch"
url = "http://localhost:7700"
tls.allow-invalid-certs = true
disable = {MEILI_ENABLED}
[store."meili".task]
poll-interval = "100ms"
#[store."meili".auth]
#username = "meili"
#secret = "changeme"
#[store."s3"]
#type = "s3"
#access-key = "minioadmin"
#secret-key = "minioadmin"
#region = "eu-central-1"
#endpoint = "http://localhost:9000"
#bucket = "tmp"
[store."fs"]
type = "fs"
path = "{TMP}"
[store."redis"]
type = "redis"
urls = "redis://127.0.0.1"
redis-type = "single"
#[store."psql-replica"]
#type = "sql-read-replica"
#primary = "postgresql"
#replicas = "postgresql"
[storage]
data = "{STORE}"
fts = "{SEARCH_STORE}"
blob = "{BLOB_STORE}"
lookup = "{LOOKUP_STORE}"
directory = "{STORE}"
[directory."{STORE}"]
type = "internal"
store = "{STORE}"
[session.rcpt]
directory = "'{STORE}'"
[session.auth]
directory = "'{STORE}'"
"#;

79
tests/src/store/ops.rs
View File

@@ -4,15 +4,18 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::store::cleanup::store_assert_is_empty;
use crate::utils::{cleanup::store_assert_is_empty, server::TestServer};
use ahash::AHashSet;
use std::collections::HashSet;
use store::Store;
use store::write::RegistryClass;
use store::{
Store, ValueKey,
ValueKey,
rand::{self, Rng},
write::{AlignedBytes, Archive, Archiver, BatchBuilder, MergeResult, Params, ValueClass},
};
use types::collection::{Collection, SyncCollection};
use types::collection::Collection;
use types::collection::SyncCollection;
// FDB max value
const MAX_VALUE_SIZE: usize = 100000;
@@ -26,17 +29,17 @@ fn value_gen(chunks: impl IntoIterator<Item = (u8, usize)>) -> Vec<u8> {
value
}
pub async fn test(db: Store) {
pub async fn test(test: &TestServer) {
let db = test.server.store().clone();
#[cfg(feature = "foundationdb")]
if matches!(db, Store::FoundationDb(_)) {
use types::collection::Collection;
println!("Running FoundationDB chunked iterator test...");
let kvs = [
("a", value_gen([(b'a', 1)])),
("b", value_gen([(b'b', MAX_VALUE_SIZE), (b'0', 1)])),
(1, value_gen([(b'a', 1)])),
(2, value_gen([(b'b', MAX_VALUE_SIZE), (b'0', 1)])),
(
"c",
3,
value_gen([
(b'c', MAX_VALUE_SIZE),
(b'1', MAX_VALUE_SIZE),
@@ -44,10 +47,10 @@ pub async fn test(db: Store) {
]),
),
(
"d",
4,
value_gen([(b'd', MAX_VALUE_SIZE), (b'3', MAX_VALUE_SIZE)]),
),
("e", value_gen([(b'e', 1)])),
(5, value_gen([(b'e', 1)])),
];
let mut batch = BatchBuilder::new();
batch
@@ -56,7 +59,13 @@ pub async fn test(db: Store) {
.with_document(0);
for (key, value) in &kvs {
batch.set(ValueClass::Config(key.as_bytes().to_vec()), value.clone());
batch.set(
ValueClass::Registry(RegistryClass::Item {
object_id: *key,
item_id: 0,
}),
value.clone(),
);
}
db.write(batch.build_all()).await.unwrap();
@@ -68,13 +77,19 @@ pub async fn test(db: Store) {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Config(b"".to_vec()),
class: ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: 0,
}),
},
ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Config(b"\xFF".to_vec()),
class: ValueClass::Registry(RegistryClass::Item {
object_id: u16::MAX,
item_id: u64::MAX,
}),
},
),
|key, value| {
@@ -92,13 +107,19 @@ pub async fn test(db: Store) {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Config(b"".to_vec()),
class: ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: 0,
}),
},
ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Config(b"\xFF".to_vec()),
class: ValueClass::Registry(RegistryClass::Item {
object_id: u16::MAX,
item_id: u64::MAX,
}),
},
)
.await
@@ -114,7 +135,10 @@ pub async fn test(db: Store) {
.with_document(0);
for n in 0..900000 {
batch.set(
ValueClass::Config(format!("key{n:10}").into_bytes()),
ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: n,
}),
format!("value{n:10}").into_bytes(),
);
@@ -139,13 +163,19 @@ pub async fn test(db: Store) {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Config(b"".to_vec()),
class: ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: 0,
}),
},
ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Config(b"\xFF".to_vec()),
class: ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: u64::MAX,
}),
},
),
|key, value| {
@@ -169,7 +199,10 @@ pub async fn test(db: Store) {
.with_collection(Collection::Email)
.with_document(0);
for n in 0..900000 {
batch.clear(ValueClass::Config(format!("key{n:10}").into_bytes()));
batch.clear(ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: n,
}));
if n % 10000 == 0 {
db.write(batch.build_all()).await.unwrap();
@@ -260,7 +293,7 @@ pub async fn test(db: Store) {
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0)
.add_and_get(ValueClass::Directory(DirectoryClass::UsedQuota(0)), 1);
.add_and_get(ValueClass::Quota, 1);
db.write(builder.build_all())
.await
.unwrap()
@@ -283,7 +316,7 @@ pub async fn test(db: Store) {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Directory(DirectoryClass::UsedQuota(0)),
class: ValueClass::Quota,
})
.await
.unwrap(),
@@ -471,7 +504,7 @@ pub async fn test(db: Store) {
.clear(ValueClass::Property(0))
.clear(ValueClass::Property(2))
.clear(ValueClass::Property(3))
.clear(ValueClass::Directory(DirectoryClass::UsedQuota(0)))
.clear(ValueClass::Quota)
.clear(ValueClass::ChangeId);
for document_id in 0..1000 {

5
tests/src/store/query.rs
View File

@@ -4,7 +4,7 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::store::deflate_test_resource;
use crate::{store::deflate_test_resource, utils::server::TestServer};
use ahash::AHashSet;
use nlp::language::Language;
use std::{
@@ -105,7 +105,8 @@ const ALL_IDS: &[&str] = &[
];
#[allow(clippy::mutex_atomic)]
pub async fn test(store: SearchStore, do_insert: bool) {
pub async fn test(test: &TestServer, do_insert: bool) {
let store = test.server.search_store().clone();
println!("Running Store query tests...");
let pool = rayon::ThreadPoolBuilder::new()

2
tests/src/utils/cleanup.rs
View File

@@ -259,7 +259,7 @@ pub async fn store_assert_is_empty(store: &Store, blob_store: BlobStore, include
(SUBSPACE_REGISTRY_PK, true),
(SUBSPACE_DIRECTORY, true),
] {
if (subspace == SUBSPACE_SEARCH_INDEX && store.is_pg_or_mysql())
if subspace == SUBSPACE_SEARCH_INDEX && store.is_pg_or_mysql()
//|| (subspace == directory && !include_directory)
{
continue;

2
tests/src/utils/mod.rs
View File

@@ -7,4 +7,6 @@
pub mod account;
pub mod cleanup;
pub mod jmap;
pub mod registry;
pub mod server;
pub mod storage;

163
tests/src/utils/registry.rs Normal file
View File

@@ -0,0 +1,163 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use registry::{
schema::{
enums::{BlobStoreType, DataStoreType, InMemoryStoreType, SearchStoreType},
prelude::Object,
structs::{
BlobStore, DataStore, ElasticSearchStore, FileSystemStore, FoundationDbStore, HttpAuth,
HttpAuthBasic, InMemoryStore, MeilisearchStore, MySqlStore, PostgreSqlStore,
RedisStore, RocksDbStore, S3Store, S3StoreCustomRegion, S3StoreRegion, SearchStore,
SecretKey, SecretKeyOptional, SecretKeyValue, SqliteStore,
},
},
types::{EnumImpl, duration::Duration},
};
use store::{
RegistryStore,
registry::write::{RegistryWrite, RegistryWriteResult},
};
use types::id::Id;
pub trait RegistryEnvStores {
fn insert_stores_from_env(&self) -> impl Future<Output = ()>;
}
impl RegistryEnvStores for RegistryStore {
async fn insert_stores_from_env(&self) {
let path = self.path().as_os_str().to_str().unwrap();
let search_store = std::env::var("SEARCH_STORE")
.map(|store| SearchStoreType::parse(&store).expect("Invalid store type"))
.map(|store| build_search_store(store, path))
.map(Object::from)
.ok();
let blob_store = std::env::var("BLOB_STORE")
.map(|store| BlobStoreType::parse(&store).expect("Invalid store type"))
.map(|store| build_blob_store(store, path))
.map(Object::from)
.ok();
let in_memory = std::env::var("MEMORY_STORE")
.map(|store| InMemoryStoreType::parse(&store).expect("Invalid store type"))
.map(|store| build_in_memory_store(store, path))
.map(Object::from)
.ok();
for store in [search_store, blob_store, in_memory].into_iter().flatten() {
self.write(RegistryWrite::insert(&store))
.await
.expect("Failed to insert store into registry")
.unwrap_id(trc::location!());
}
}
}
pub fn build_data_store(typ: DataStoreType, path: &str) -> DataStore {
match typ {
DataStoreType::RocksDb => DataStore::RocksDb(RocksDbStore {
path: format!("{path}/rocks.db"),
..Default::default()
}),
DataStoreType::Sqlite => DataStore::Sqlite(SqliteStore {
path: format!("{path}/sqlite.db"),
..Default::default()
}),
DataStoreType::FoundationDb => DataStore::FoundationDb(FoundationDbStore::default()),
DataStoreType::PostgreSql => DataStore::PostgreSql(PostgreSqlStore {
host: "localhost".into(),
port: 5432,
auth_username: "postgres".to_string().into(),
auth_secret: SecretKeyOptional::Value(SecretKeyValue {
secret: "mysecretpassword".into(),
}),
database: "stalwart".into(),
use_tls: false,
allow_invalid_certs: true,
..Default::default()
}),
DataStoreType::MySql => DataStore::MySql(MySqlStore {
host: "localhost".into(),
port: 3307,
auth_username: "root".to_string().into(),
auth_secret: SecretKeyOptional::Value(SecretKeyValue {
secret: "password".into(),
}),
database: "stalwart".into(),
use_tls: false,
allow_invalid_certs: true,
..Default::default()
}),
}
}
fn build_blob_store(typ: BlobStoreType, path: &str) -> BlobStore {
match typ {
BlobStoreType::S3 => BlobStore::S3(S3Store {
access_key: "minioadmin".to_string().into(),
bucket: "tmp".into(),
region: S3StoreRegion::Custom(S3StoreCustomRegion {
custom_endpoint: "http://localhost:9000".into(),
custom_region: "eu-central-1".into(),
}),
secret_key: SecretKeyOptional::Value(SecretKeyValue {
secret: "minioadmin".into(),
}),
allow_invalid_certs: true,
..Default::default()
}),
BlobStoreType::FileSystem => BlobStore::FileSystem(FileSystemStore {
path: path.to_string(),
..Default::default()
}),
_ => unreachable!(),
}
}
fn build_in_memory_store(typ: InMemoryStoreType, _path: &str) -> InMemoryStore {
match typ {
InMemoryStoreType::Redis => InMemoryStore::Redis(RedisStore {
url: "redis://127.0.0.1".into(),
..Default::default()
}),
_ => unreachable!(),
}
}
fn build_search_store(typ: SearchStoreType, _path: &str) -> SearchStore {
match typ {
SearchStoreType::ElasticSearch => SearchStore::ElasticSearch(ElasticSearchStore {
url: "https://localhost:9200".into(),
allow_invalid_certs: true,
http_auth: HttpAuth::Basic(HttpAuthBasic {
username: "elastic".into(),
secret: SecretKey::Value(SecretKeyValue {
secret: "changeme".into(),
}),
}),
..Default::default()
}),
SearchStoreType::Meilisearch => SearchStore::Meilisearch(MeilisearchStore {
url: "http://localhost:7700".into(),
allow_invalid_certs: true,
poll_interval: Duration::from_millis(100),
..Default::default()
}),
_ => unreachable!(),
}
}
pub trait UnwrapRegistryId {
fn unwrap_id(self, location: &str) -> Id;
}
impl UnwrapRegistryId for RegistryWriteResult {
fn unwrap_id(self, location: &str) -> Id {
match self {
RegistryWriteResult::Success(id) => id,
err => panic!("Expected success at {location} but got {err}"),
}
}
}

273
tests/src/utils/server.rs Normal file
View File

@@ -0,0 +1,273 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::{
AssertConfig,
store::TempDir,
utils::{
account::Account,
cleanup::{search_store_destroy, store_destroy},
registry::{RegistryEnvStores, UnwrapRegistryId, build_data_store},
storage::assert_is_empty,
},
};
use ahash::AHashMap;
use common::{
BuildServer, Caches, Core, Data, Inner, Server,
config::{
server::{Listeners, ServerProtocol},
storage::Storage,
telemetry::Telemetry,
},
manager::{boot::build_ipc, defaults::BootstrapDefaults},
};
use http::HttpSessionManager;
use imap::core::ImapSessionManager;
use managesieve::core::ManageSieveSessionManager;
use pop3::Pop3SessionManager;
use registry::{
schema::{
enums::{DataStoreType, EventPolicy, NetworkListenerProtocol, TracingLevel},
prelude::{Object, SocketAddr},
structs::{NetworkListener, Tracer, TracerStdout},
},
types::{EnumImpl, map::Map},
};
use services::{SpawnServices, broadcast::subscriber::spawn_broadcast_subscriber};
use smtp::{SpawnQueueManager, core::SmtpSessionManager};
use std::{str::FromStr, sync::Arc};
use store::{
RegistryStore, Store,
registry::{bootstrap::Bootstrap, write::RegistryWrite},
};
use tokio::sync::watch;
use trc::EventType;
use types::id::Id;
pub struct TestServer {
pub server: Server,
accounts: AHashMap<&'static str, Account>,
pub temp_dir: TempDir,
shutdown_tx: watch::Sender<bool>,
}
pub struct TestServerBuilder {
bootstrap: Bootstrap,
temp_dir: TempDir,
reset: bool,
}
impl TestServerBuilder {
pub async fn new(test_name: &str, reset: bool) -> Self {
let temp_dir = TempDir::new(test_name, reset);
let path = temp_dir.path.to_string_lossy().to_string();
let data_store = build_data_store(
std::env::var("STORE")
.map(|store| DataStoreType::parse(&store).expect("Invalid store type"))
.expect(concat!(
"Missing or invalid store type. Try ",
"running `STORE=<store_type> cargo test`"
)),
&path,
);
let store = Store::build(data_store).await.unwrap();
store.create_tables().await.unwrap();
// Delete old store if requested
if reset {
store_destroy(&store).await;
}
Self {
bootstrap: Bootstrap::new(
RegistryStore::new(&path, store, "mail.example.org".to_string(), 1, None).await,
)
.await,
temp_dir,
reset,
}
}
pub async fn with_listener(
self,
protocol: NetworkListenerProtocol,
name: &str,
port: u16,
use_tls: bool,
) -> Self {
self.insert_object(NetworkListener {
bind: Map::new(vec![SocketAddr::from_str(&format!("[::]:{port}")).unwrap()]),
name: name.to_string(),
protocol,
use_tls,
..Default::default()
})
.await;
self
}
pub async fn with_object(self, object: impl Into<Object>) -> Self {
self.insert_object(object).await;
self
}
pub async fn insert_object(&self, object: impl Into<Object>) -> Id {
self.bootstrap
.registry
.write(RegistryWrite::insert(&object.into()))
.await
.unwrap()
.unwrap_id(trc::location!())
}
pub async fn build(mut self) -> TestServer {
// Register stores from environment
self.bootstrap.registry.insert_stores_from_env().await;
// Enable logging if requested
let level = std::env::var("LOG")
.map(|log| TracingLevel::parse(&log).expect("Invalid log level"))
.ok();
self.bootstrap
.registry
.write(RegistryWrite::insert(
&Tracer::Stdout(TracerStdout {
enable: level.is_some(),
level: level.unwrap_or(TracingLevel::Info),
ansi: true,
multiline: false,
events: Map::new(
EventType::variants()
.iter()
.filter(|ev| {
let ev = ev.as_str();
ev.starts_with("network.")
|| ev == "telemetry.webhook-error"
|| ev == "http.request-body"
})
.copied()
.collect(),
),
events_policy: EventPolicy::Exclude,
..Default::default()
})
.into(),
))
.await
.unwrap()
.unwrap_id(trc::location!());
// Start listeners
let mut servers = Listeners::parse(&mut self.bootstrap).await;
servers.bind_and_drop_priv(&mut self.bootstrap);
// Parse storage
let storage = Storage::parse(&mut self.bootstrap).await;
// Reset search store
if self.reset {
search_store_destroy(&storage.search).await;
}
// Parse telemetry
let telemetry = Telemetry::parse(&mut self.bootstrap, &storage).await;
// Add safe defaults if missing
self.bootstrap.insert_safe_defaults().await;
// Parse components
let core = Box::pin(Core::parse(&mut self.bootstrap, storage)).await;
let data = Data::parse(&mut self.bootstrap).await;
let cache = Caches::parse(&mut self.bootstrap).await;
// Enable telemetry
telemetry.enable(true);
// Build inner
let (ipc, mut ipc_rxs) = build_ipc(!core.storage.coordinator.is_none());
let inner = Arc::new(Inner {
shared_core: core.into_shared(),
data,
ipc,
cache,
});
// Parse TCP acceptors
servers
.parse_tcp_acceptors(&mut self.bootstrap, inner.clone())
.await;
// Start services
self.bootstrap.assert_no_errors();
ipc_rxs.spawn_queue_manager(inner.clone());
ipc_rxs.spawn_services(inner.clone());
// Spawn servers
let (shutdown_tx, shutdown_rx) = servers.spawn(|server, acceptor, shutdown_rx| {
match &server.protocol {
ServerProtocol::Smtp | ServerProtocol::Lmtp => server.spawn(
SmtpSessionManager::new(inner.clone()),
inner.clone(),
acceptor,
shutdown_rx,
),
ServerProtocol::Http => server.spawn(
HttpSessionManager::new(inner.clone()),
inner.clone(),
acceptor,
shutdown_rx,
),
ServerProtocol::Imap => server.spawn(
ImapSessionManager::new(inner.clone()),
inner.clone(),
acceptor,
shutdown_rx,
),
ServerProtocol::Pop3 => server.spawn(
Pop3SessionManager::new(inner.clone()),
inner.clone(),
acceptor,
shutdown_rx,
),
ServerProtocol::ManageSieve => server.spawn(
ManageSieveSessionManager::new(inner.clone()),
inner.clone(),
acceptor,
shutdown_rx,
),
};
});
// Start broadcast subscriber
spawn_broadcast_subscriber(inner.clone(), shutdown_rx);
TestServer {
server: inner.build_server(),
temp_dir: self.temp_dir,
accounts: Default::default(),
shutdown_tx,
}
}
}
impl TestServer {
pub fn account(&self, name: &str) -> &Account {
self.accounts.get(name).unwrap()
}
pub async fn assert_is_empty(&self) {
assert_is_empty(&self.server).await;
}
pub async fn destroy_store(&self) {
store_destroy(self.server.store()).await;
}
pub fn shutdown(&self) {
let _ = self.shutdown_tx.send(true);
}
}