Port Spam filter to Rust - part 6

This commit is contained in:
mdecimus
2024-12-15 18:01:17 +01:00
parent 08c43e2f15
commit 07449a3622
75 changed files with 1742 additions and 1101 deletions

View File

@@ -28,6 +28,7 @@ imap_proto = { path = "../crates/imap-proto" }
pop3 = { path = "../crates/pop3", features = ["test_mode"] }
smtp = { path = "../crates/smtp", features = ["test_mode"] }
common = { path = "../crates/common", features = ["test_mode", "enterprise"] }
spam-filter = { path = "../crates/spam-filter", features = ["test_mode", "enterprise"] }
trc = { path = "../crates/trc" }
managesieve = { path = "../crates/managesieve", features = ["test_mode", "enterprise"] }
smtp-proto = { version = "0.1" }

View File

@@ -21,7 +21,7 @@ use rustls::ServerConfig;
use rustls_pemfile::{certs, pkcs8_private_keys};
use rustls_pki_types::PrivateKeyDer;
use std::{borrow::Cow, io::BufReader, sync::Arc};
use store::{LookupStore, Store, Stores};
use store::{Store, Stores};
use tokio_rustls::TlsAcceptor;
use crate::{store::TempDir, AssertConfig};
@@ -307,7 +307,7 @@ fields.full-name = "name"
"#;
pub struct DirectoryStore {
pub store: LookupStore,
pub store: Store,
}
pub struct DirectoryTest {

View File

@@ -11,7 +11,7 @@ use directory::{
use mail_send::Credentials;
#[allow(unused_imports)]
use store::{LookupStore, Store};
use store::{InMemoryStore, Store};
use crate::directory::{
map_account_id, map_account_ids, DirectoryTest, IntoTestPrincipal, TestPrincipal,
@@ -37,7 +37,7 @@ async fn sql_directory() {
println!("Testing SQL directory {:?}", directory_id);
let handle = config.directories.directories.remove(directory_id).unwrap();
let store = DirectoryStore {
store: config.stores.lookup_stores.remove(directory_id).unwrap(),
store: config.stores.stores.remove(directory_id).unwrap(),
};
let base_store = config.stores.stores.get(directory_id).unwrap();
let core = config.server;
@@ -356,7 +356,7 @@ impl DirectoryStore {
// Create tables
for table in ["accounts", "group_members", "emails"] {
self.store
.query::<usize>(&format!("DROP TABLE IF EXISTS {table}"), vec![])
.sql_query::<usize>(&format!("DROP TABLE IF EXISTS {table}"), vec![])
.await
.unwrap();
}
@@ -383,7 +383,7 @@ impl DirectoryStore {
};
self.store
.query::<usize>(&query, vec![])
.sql_query::<usize>(&query, vec![])
.await
.unwrap_or_else(|_| panic!("failed for {query}"));
}
@@ -396,7 +396,7 @@ impl DirectoryStore {
"individual"
};
self.store
.query::<usize>(
.sql_query::<usize>(
if self.is_postgresql() {
concat!(
"INSERT INTO accounts (name, secret, description, ",
@@ -439,7 +439,7 @@ impl DirectoryStore {
pub async fn create_test_group(&self, login: &str, name: &str) {
self.store
.query::<usize>(
.sql_query::<usize>(
if self.is_postgresql() {
concat!(
"INSERT INTO accounts (name, description, ",
@@ -469,7 +469,7 @@ impl DirectoryStore {
pub async fn link_test_address(&self, login: &str, address: &str, typ: &str) {
self.store
.query::<usize>(
.sql_query::<usize>(
if self.is_postgresql() {
"INSERT INTO emails (name, address, type) VALUES ($1, $2, $3) ON CONFLICT (name, address) DO NOTHING"
} else if self.is_mysql() {
@@ -485,7 +485,7 @@ impl DirectoryStore {
pub async fn set_test_quota(&self, login: &str, quota: u32) {
self.store
.query::<usize>(
.sql_query::<usize>(
if self.is_postgresql() {
"UPDATE accounts SET quota = $1 where name = $2"
} else {
@@ -499,7 +499,7 @@ impl DirectoryStore {
pub async fn add_to_group(&self, login: &str, group: &str) {
self.store
.query::<usize>(
.sql_query::<usize>(
if self.is_postgresql() {
"INSERT INTO group_members (name, member_of) VALUES ($1, $2)"
} else {
@@ -513,7 +513,7 @@ impl DirectoryStore {
pub async fn remove_from_group(&self, login: &str, group: &str) {
self.store
.query::<usize>(
.sql_query::<usize>(
if self.is_postgresql() {
"DELETE FROM group_members WHERE name = $1 AND member_of = $2"
} else {
@@ -527,7 +527,7 @@ impl DirectoryStore {
pub async fn remove_test_alias(&self, login: &str, alias: &str) {
self.store
.query::<usize>(
.sql_query::<usize>(
if self.is_postgresql() {
"DELETE FROM emails WHERE name = $1 AND address = $2"
} else {
@@ -542,7 +542,7 @@ impl DirectoryStore {
fn is_mysql(&self) -> bool {
#[cfg(feature = "mysql")]
{
matches!(self.store, LookupStore::Store(Store::MySQL(_)))
matches!(self.store, Store::MySQL(_))
}
#[cfg(not(feature = "mysql"))]
{
@@ -553,7 +553,7 @@ impl DirectoryStore {
fn is_postgresql(&self) -> bool {
#[cfg(feature = "postgres")]
{
matches!(self.store, LookupStore::Store(Store::PostgreSQL(_)))
matches!(self.store, Store::PostgreSQL(_))
}
#[cfg(not(feature = "postgres"))]
{
@@ -565,7 +565,7 @@ impl DirectoryStore {
fn is_sqlite(&self) -> bool {
#[cfg(feature = "sqlite")]
{
matches!(self.store, LookupStore::Store(Store::SQLite(_)))
matches!(self.store, Store::SQLite(_))
}
#[cfg(not(feature = "sqlite"))]
{

View File

@@ -393,7 +393,7 @@ pub async fn test(params: &mut JMAPTest) {
.core
.storage
.lookup
.purge_lookup_store()
.purge_in_memory_store()
.await
.unwrap();
params.client.set_default_account_id(john_id);

View File

@@ -108,6 +108,7 @@ pub async fn test(params: &mut JMAPTest) {
metrics_alerts: parse_metric_alerts(&mut config),
logo_url: None,
ai_apis: Default::default(),
spam_filter_llm: None,
}
.into();
config.assert_no_errors();
@@ -173,6 +174,7 @@ impl EnterpriseCore for Core {
metrics_alerts: vec![],
logo_url: None,
ai_apis: Default::default(),
spam_filter_llm: None,
}
.into();
self

View File

@@ -219,7 +219,7 @@ async fn antispam() {
.enable_enterprise();
core.enterprise.as_mut().unwrap().ai_apis.insert(
"dummy".to_string(),
AiApiConfig::parse(&mut config, "dummy").unwrap(),
AiApiConfig::parse(&mut config, "dummy").unwrap().into(),
);
crate::AssertConfig::assert_no_errors(config);

View File

@@ -119,7 +119,7 @@ async fn lookup_sql() {
// Obtain directory handle
let handle = DirectoryStore {
store: core.storage.lookups.get("sql").unwrap().clone(),
store: core.storage.stores.get("sql").unwrap().clone(),
};
let test = TestSMTP::from_core(core);
@@ -162,7 +162,7 @@ async fn lookup_sql() {
] {
handle
.store
.query::<usize>(query, Vec::new())
.sql_query::<usize>(query, Vec::new())
.await
.unwrap();
}

View File

@@ -6,7 +6,7 @@
use std::time::Duration;
use store::{LookupStore, Stores};
use store::{dispatch::lookup::KeyValue, InMemoryStore, Stores};
use utils::config::{Config, Rate};
use crate::{
@@ -27,14 +27,14 @@ pub async fn lookup_tests() {
period: Duration::from_secs(1),
};
for (store_id, store) in stores.lookup_stores {
for (store_id, store) in stores.in_memory_stores {
println!("Testing lookup store {}...", store_id);
if let LookupStore::Store(store) = &store {
if let InMemoryStore::Store(store) = &store {
store.destroy().await;
} else {
// Reset redis counter
store
.key_set("abc".as_bytes().to_vec(), "0".as_bytes().to_vec(), None)
.key_set(KeyValue::new("abc", "0".as_bytes().to_vec()))
.await
.unwrap();
}
@@ -42,10 +42,10 @@ pub async fn lookup_tests() {
// Test key
let key = "xyz".as_bytes().to_vec();
store
.key_set(key.clone(), "world".to_string().into_bytes(), None)
.key_set(KeyValue::new(key.clone(), "world".to_string().into_bytes()))
.await
.unwrap();
store.purge_lookup_store().await.unwrap();
store.purge_in_memory_store().await.unwrap();
assert_eq!(
store.key_get::<String>(key.clone()).await.unwrap(),
Some("world".to_string())
@@ -53,7 +53,7 @@ pub async fn lookup_tests() {
// Test value expiry
store
.key_set(key.clone(), "hello".to_string().into_bytes(), 1.into())
.key_set(KeyValue::new(key.clone(), "hello".to_string().into_bytes()).expires(1))
.await
.unwrap();
assert_eq!(
@@ -63,25 +63,25 @@ pub async fn lookup_tests() {
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
assert_eq!(None, store.key_get::<String>(key.clone()).await.unwrap());
store.purge_lookup_store().await.unwrap();
if let LookupStore::Store(store) = &store {
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store.assert_is_empty(store.clone().into()).await;
}
// Test counter
let key = "abc".as_bytes().to_vec();
store
.counter_incr(key.clone(), 1, None, false)
.counter_incr(KeyValue::new(key.clone(), 1))
.await
.unwrap();
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
store
.counter_incr(key.clone(), 2, None, false)
.counter_incr(KeyValue::new(key.clone(), 2))
.await
.unwrap();
assert_eq!(3, store.counter_get(key.clone()).await.unwrap());
store
.counter_incr(key.clone(), -3, None, false)
.counter_incr(KeyValue::new(key.clone(), -3))
.await
.unwrap();
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
@@ -89,34 +89,34 @@ pub async fn lookup_tests() {
// Test counter expiry
let key = "fgh".as_bytes().to_vec();
store
.counter_incr(key.clone(), 1, 1.into(), false)
.counter_incr(KeyValue::new(key.clone(), 1).expires(1))
.await
.unwrap();
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
store.purge_lookup_store().await.unwrap();
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("rate".as_bytes(), &rate, false)
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
.await
.unwrap()
.is_none());
assert!(store
.is_rate_allowed("rate".as_bytes(), &rate, false)
.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("rate".as_bytes(), &rate, false)
.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_lookup_store().await.unwrap();
if let LookupStore::Store(store) = &store {
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store.assert_is_empty(store.clone().into()).await;
}
}