From d45f1299d3df9be0688800f956a639964ec12a62 Mon Sep 17 00:00:00 2001 From: mdecimus <11444311+mdecimus@users.noreply.github.com> Date: Thu, 12 Mar 2026 18:26:35 +0100 Subject: [PATCH] Registry testing - part 2 --- crates/common/src/config/mailstore/email.rs | 2 +- crates/common/src/config/storage.rs | 11 +- crates/common/src/config/telemetry.rs | 10 + crates/common/src/enterprise/config.rs | 20 +- crates/common/src/manager/backup.rs | 4 +- crates/common/src/manager/defaults.rs | 108 ++- crates/email/src/sieve/ingest.rs | 5 +- crates/jmap/src/registry/mapping/task.rs | 5 +- crates/main/Cargo.toml | 4 +- .../src/backend/composite/read_replica.rs | 17 + crates/store/src/backend/foundationdb/read.rs | 16 +- .../store/src/backend/foundationdb/write.rs | 14 +- crates/store/src/backend/mysql/main.rs | 2 +- crates/store/src/backend/mysql/read.rs | 16 + crates/store/src/backend/mysql/write.rs | 67 +- crates/store/src/backend/postgres/read.rs | 16 + crates/store/src/backend/postgres/write.rs | 59 +- crates/store/src/backend/rocksdb/main.rs | 2 +- crates/store/src/backend/rocksdb/read.rs | 16 + crates/store/src/backend/sqlite/read.rs | 20 + crates/store/src/backend/sqlite/write.rs | 31 +- crates/store/src/build/data.rs | 1 + crates/store/src/build/registry.rs | 58 +- crates/store/src/dispatch/lookup.rs | 12 +- crates/store/src/dispatch/store.rs | 23 + crates/store/src/lib.rs | 10 +- crates/store/src/registry/get.rs | 8 +- crates/store/src/registry/write.rs | 20 +- crates/store/src/write/blob.rs | 14 +- crates/store/src/write/serialize.rs | 6 - tests/Cargo.toml | 4 +- tests/src/lib.rs | 4 +- tests/src/store/blob.rs | 832 +++++++++--------- tests/src/store/import_export.rs | 88 +- tests/src/store/lookup.rs | 479 +++++----- tests/src/store/mod.rs | 185 +--- tests/src/store/ops.rs | 79 +- tests/src/store/query.rs | 5 +- tests/src/utils/cleanup.rs | 2 +- tests/src/utils/mod.rs | 2 + tests/src/utils/registry.rs | 163 ++++ tests/src/utils/server.rs | 273 ++++++ 42 files changed, 1649 insertions(+), 1064 deletions(-) create mode 100644 tests/src/utils/registry.rs create mode 100644 tests/src/utils/server.rs diff --git a/crates/common/src/config/mailstore/email.rs b/crates/common/src/config/mailstore/email.rs index bbdb80f1..47383a75 100644 --- a/crates/common/src/config/mailstore/email.rs +++ b/crates/common/src/config/mailstore/email.rs @@ -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 diff --git a/crates/common/src/config/storage.rs b/crates/common/src/config/storage.rs index 8fa99a9c..6bc87b4f 100644 --- a/crates/common/src/config/storage.rs +++ b/crates/common/src/config/storage.rs @@ -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(), diff --git a/crates/common/src/config/telemetry.rs b/crates/common/src/config/telemetry.rs index aa3360f3..2d5a35a0 100644 --- a/crates/common/src/config/telemetry.rs +++ b/crates/common/src/config/telemetry.rs @@ -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()), diff --git a/crates/common/src/enterprise/config.rs b/crates/common/src/enterprise/config.rs index e28d354a..30027b2f 100644 --- a/crates/common/src/enterprise/config.rs +++ b/crates/common/src/enterprise/config.rs @@ -37,8 +37,17 @@ impl Enterprise { .await .default_hostname; let mut update_license = None; - let mut enterprise = bp.setting_infallible::().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); diff --git a/crates/common/src/manager/backup.rs b/crates/common/src/manager/backup.rs index c2da32f5..45224f48 100644 --- a/crates/common/src/manager/backup.rs +++ b/crates/common/src/manager/backup.rs @@ -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())) diff --git a/crates/common/src/manager/defaults.rs b/crates/common/src/manager/defaults.rs index d031797d..8058056a 100644 --- a/crates/common/src/manager/defaults.rs +++ b/crates/common/src/manager/defaults.rs @@ -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(()) } diff --git a/crates/email/src/sieve/ingest.rs b/crates/email/src/sieve/ingest.rs index fcd285f9..7e772c14 100644 --- a/crates/email/src/sieve/ingest.rs +++ b/crates/email/src/sieve/ingest.rs @@ -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() diff --git a/crates/jmap/src/registry/mapping/task.rs b/crates/jmap/src/registry/mapping/task.rs index 64b1a8c4..70110a66 100644 --- a/crates/jmap/src/registry/mapping/task.rs +++ b/crates/jmap/src/registry/mapping/task.rs @@ -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, diff --git a/crates/main/Cargo.toml b/crates/main/Cargo.toml index 3e106f6f..c78a8c42 100644 --- a/crates/main/Cargo.toml +++ b/crates/main/Cargo.toml @@ -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"] diff --git a/crates/store/src/backend/composite/read_replica.rs b/crates/store/src/backend/composite/read_replica.rs index 12c72d93..e6229a8a 100644 --- a/crates/store/src/backend/composite/read_replica.rs +++ b/crates/store/src/backend/composite/read_replica.rs @@ -121,6 +121,23 @@ impl SQLReadReplica { .await } + pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result { + 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( &self, params: IterateParams, diff --git a/crates/store/src/backend/foundationdb/read.rs b/crates/store/src/backend/foundationdb/read.rs index 4009b4b4..28b18225 100644 --- a/crates/store/src/backend/foundationdb/read.rs +++ b/crates/store/src/backend/foundationdb/read.rs @@ -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 { + 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( &self, params: IterateParams, diff --git a/crates/store/src/backend/foundationdb/write.rs b/crates/store/src/backend/foundationdb/write.rs index 18683b2d..23a7e689 100644 --- a/crates/store/src/backend/foundationdb/write.rs +++ b/crates/store/src/backend/foundationdb/write.rs @@ -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) diff --git a/crates/store/src/backend/mysql/main.rs b/crates/store/src/backend/mysql/main.rs index 319bd4ef..2acc56a6 100644 --- a/crates/store/src/backend/mysql/main.rs +++ b/crates/store/src/backend/mysql/main.rs @@ -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, }; diff --git a/crates/store/src/backend/mysql/read.rs b/crates/store/src/backend/mysql/read.rs index 46b5b97b..847f6033 100644 --- a/crates/store/src/backend/mysql/read.rs +++ b/crates/store/src/backend/mysql/read.rs @@ -35,6 +35,22 @@ impl MysqlStore { }) } + pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result { + 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::(&s, (&key,)) + .await + .map_err(into_error) + .map(|r| r.is_some()) + } + pub(crate) async fn iterate( &self, params: IterateParams, diff --git a/crates/store/src/backend/mysql/write.rs b/crates/store/src/backend/mysql/write.rs index a03c318f..543c4ca6 100644 --- a/crates/store/src/backend/mysql/write.rs +++ b/crates/store/src/backend/mysql/write.rs @@ -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 diff --git a/crates/store/src/backend/postgres/read.rs b/crates/store/src/backend/postgres/read.rs index 75d55b55..027f7005 100644 --- a/crates/store/src/backend/postgres/read.rs +++ b/crates/store/src/backend/postgres/read.rs @@ -37,6 +37,22 @@ impl PostgresStore { }) } + pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result { + 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( &self, params: IterateParams, diff --git a/crates/store/src/backend/postgres/write.rs b/crates/store/src/backend/postgres/write.rs index 4067d826..2b0f21bf 100644 --- a/crates/store/src/backend/postgres/write.rs +++ b/crates/store/src/backend/postgres/write.rs @@ -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) => { diff --git a/crates/store/src/backend/rocksdb/main.rs b/crates/store/src/backend/rocksdb/main.rs index 813ab744..87ed522e 100644 --- a/crates/store/src/backend/rocksdb/main.rs +++ b/crates/store/src/backend/rocksdb/main.rs @@ -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( diff --git a/crates/store/src/backend/rocksdb/read.rs b/crates/store/src/backend/rocksdb/read.rs index 0f5c5df6..a28bf5d7 100644 --- a/crates/store/src/backend/rocksdb/read.rs +++ b/crates/store/src/backend/rocksdb/read.rs @@ -36,6 +36,22 @@ impl RocksDbStore { .await } + pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result { + 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( &self, params: IterateParams, diff --git a/crates/store/src/backend/sqlite/read.rs b/crates/store/src/backend/sqlite/read.rs index b9af98dd..c4fe10df 100644 --- a/crates/store/src/backend/sqlite/read.rs +++ b/crates/store/src/backend/sqlite/read.rs @@ -34,6 +34,26 @@ impl SqliteStore { .await } + pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result { + 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( &self, params: IterateParams, diff --git a/crates/store/src/backend/sqlite/write.rs b/crates/store/src/backend/sqlite/write.rs index c4ba38ca..eb9860ef 100644 --- a/crates/store/src/backend/sqlite/write.rs +++ b/crates/store/src/backend/sqlite/write.rs @@ -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)?; diff --git a/crates/store/src/build/data.rs b/crates/store/src/build/data.rs index b8a3b45f..591e8a13 100644 --- a/crates/store/src/build/data.rs +++ b/crates/store/src/build/data.rs @@ -10,6 +10,7 @@ use registry::schema::{ structs::{DataStore, MetricsStore, TracingStore}, }; +#[allow(unreachable_patterns)] impl Store { pub async fn build(config: DataStore) -> Result { #[allow(unreachable_patterns)] diff --git a/crates/store/src/build/registry.rs b/crates/store/src/build/registry.rs index cfbf069c..46455f1f 100644 --- a/crates/store/src/build/registry.rs +++ b/crates/store/src/build/registry.rs @@ -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 { + // 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, + ) -> 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() + } } diff --git a/crates/store/src/dispatch/lookup.rs b/crates/store/src/dispatch/lookup.rs index f367e00c..46eabb9c 100644 --- a/crates/store/src/dispatch/lookup.rs +++ b/crates/store/src/dispatch/lookup.rs @@ -271,11 +271,11 @@ impl InMemoryStore { pub async fn key_exists(&self, key: impl Into>) -> trc::Result { match self { InMemoryStore::Store(store) => store - .get_value::>(ValueKey::from(ValueClass::InMemory( + .get_value::>(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 KeyValue { } } +struct Empty; + enum LookupValue { Value(T), None, @@ -655,6 +657,12 @@ impl Deserialize for LookupValue { } } +impl Deserialize for Empty { + fn deserialize(_bytes: &[u8]) -> trc::Result { + Ok(Empty) + } +} + impl From> for Option { fn from(value: LookupValue) -> Self { match value { diff --git a/crates/store/src/dispatch/store.rs b/crates/store/src/dispatch/store.rs index 2c31f0c2..924e2e66 100644 --- a/crates/store/src/dispatch/store.rs +++ b/crates/store/src/dispatch/store.rs @@ -45,6 +45,29 @@ impl Store { .caused_by(trc::location!()) } + pub async fn key_exists(&self, key: impl Key) -> trc::Result { + 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 + // 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( &self, params: IterateParams, diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs index e06329ea..6428b672 100644 --- a/crates/store/src/lib.rs +++ b/crates/store/src/lib.rs @@ -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 { @@ -282,6 +282,12 @@ impl From for InMemoryStore { } } +impl From for BlobStore { + fn from(store: Store) -> Self { + Self::Store(store) + } +} + impl Default for BlobStore { fn default() -> Self { Self::Store(Store::None) diff --git a/crates/store/src/registry/get.rs b/crates/store/src/registry/get.rs index eb84f9c3..23eda7d4 100644 --- a/crates/store/src/registry/get.rs +++ b/crates/store/src/registry/get.rs @@ -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::()) .map(|r| r.0) .ok_or_else(|| { diff --git a/crates/store/src/registry/write.rs b/crates/store/src/registry/write.rs index 7ac73e3b..1adcec0f 100644 --- a/crates/store/src/registry/write.rs +++ b/crates/store/src/registry/write.rs @@ -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( diff --git a/crates/store/src/write/blob.rs b/crates/store/src/write/blob.rs index 01e1b603..fd64d190 100644 --- a/crates/store/src/write/blob.rs +++ b/crates/store/src/write/blob.rs @@ -29,7 +29,7 @@ pub struct BlobQuota { impl Store { pub async fn blob_exists(&self, hash: impl AsRef + Sync + Send) -> trc::Result { - 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(()) } diff --git a/crates/store/src/write/serialize.rs b/crates/store/src/write/serialize.rs index 9e142139..db20190c 100644 --- a/crates/store/src/write/serialize.rs +++ b/crates/store/src/write/serialize.rs @@ -563,12 +563,6 @@ impl Deserialize for u32 { } } -impl Deserialize for () { - fn deserialize(_bytes: &[u8]) -> trc::Result { - Ok(()) - } -} - impl From> for Archive { fn from(_: Value<'static>) -> Self { unimplemented!() diff --git a/tests/Cargo.toml b/tests/Cargo.toml index f3051d54..7e024fe1 100644 --- a/tests/Cargo.toml +++ b/tests/Cargo.toml @@ -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"] diff --git a/tests/src/lib.rs b/tests/src/lib.rs index 4edeb8c6..dc122981 100644 --- a/tests/src/lib.rs +++ b/tests/src/lib.rs @@ -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; diff --git a/tests/src/store/blob.rs b/tests/src/store/blob.rs index 8d812e85..56d4c83b 100644 --- a/tests/src/store/blob.rs +++ b/tests/src/store/blob.rs @@ -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 diff --git a/tests/src/store/import_export.rs b/tests/src/store/import_export.rs index 491d41df..f7dd108a 100644 --- a/tests/src/store/import_export.rs +++ b/tests/src/store/import_export.rs @@ -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), diff --git a/tests/src/store/lookup.rs b/tests/src/store/lookup.rs index 0a17e714..23a009ec 100644 --- a/tests/src/store/lookup.rs +++ b/tests/src/store/lookup.rs @@ -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::(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::(key.clone()).await.unwrap(), + Some("hello".to_string()) + ); + tokio::time::sleep(tokio::time::Duration::from_secs(2)).await; + assert_eq!(None, store.key_get::(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::(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::(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::(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::(KeyValue::<()>::build_key(1, [0])) - .await - .unwrap() - .as_deref() - ); - for v in [0, 1000, 1001, 2000, 2001] { - assert_eq!( - Some("world"), - store - .key_get::(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::(KeyValue::<()>::build_key(1, [0])) - .await - .unwrap() - .as_deref() - ); - for v in [0, 1000, 1001, 2000, 2001] { - assert_eq!( - Some("world"), - store - .key_get::(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::(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::(KeyValue::<()>::build_key(1, [0])) + .key_get::(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::(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::(KeyValue::<()>::build_key(1, [0])) + .await + .unwrap() + .as_deref() + ); + for v in [0, 1000, 1001, 2000, 2001] { + assert_eq!( + Some("world"), + store + .key_get::(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::(KeyValue::<()>::build_key(1, [0])) + .await + .unwrap() + .as_deref() + ); + for v in [0, 1000, 1001, 2000, 2001] { assert_eq!( None, store - .key_get::(KeyValue::<()>::build_key(1, [0])) + .key_get::(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::(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; } } diff --git a/tests/src/store/mod.rs b/tests/src/store/mod.rs index 9e632287..6e3da06c 100644 --- a/tests/src/store/mod.rs +++ b/tests/src/store/mod.rs @@ -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= 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= 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= 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}'" - -"#; diff --git a/tests/src/store/ops.rs b/tests/src/store/ops.rs index 78bce056..831e85d0 100644 --- a/tests/src/store/ops.rs +++ b/tests/src/store/ops.rs @@ -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) -> Vec { 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 { diff --git a/tests/src/store/query.rs b/tests/src/store/query.rs index a51553f0..a892aec7 100644 --- a/tests/src/store/query.rs +++ b/tests/src/store/query.rs @@ -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() diff --git a/tests/src/utils/cleanup.rs b/tests/src/utils/cleanup.rs index 9e8972b0..7c90de4b 100644 --- a/tests/src/utils/cleanup.rs +++ b/tests/src/utils/cleanup.rs @@ -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; diff --git a/tests/src/utils/mod.rs b/tests/src/utils/mod.rs index d0ccfc7b..1f1d580f 100644 --- a/tests/src/utils/mod.rs +++ b/tests/src/utils/mod.rs @@ -7,4 +7,6 @@ pub mod account; pub mod cleanup; pub mod jmap; +pub mod registry; +pub mod server; pub mod storage; diff --git a/tests/src/utils/registry.rs b/tests/src/utils/registry.rs new file mode 100644 index 00000000..9f57901e --- /dev/null +++ b/tests/src/utils/registry.rs @@ -0,0 +1,163 @@ +/* + * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC + * + * 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; +} + +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}"), + } + } +} diff --git a/tests/src/utils/server.rs b/tests/src/utils/server.rs new file mode 100644 index 00000000..74132e1e --- /dev/null +++ b/tests/src/utils/server.rs @@ -0,0 +1,273 @@ +/* + * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC + * + * 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, +} + +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= 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) -> Self { + self.insert_object(object).await; + self + } + + pub async fn insert_object(&self, object: impl Into) -> 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); + } +}