From 07449a36227ee867cc1bece5120471214321d1d9 Mon Sep 17 00:00:00 2001 From: mdecimus Date: Sun, 15 Dec 2024 18:01:17 +0100 Subject: [PATCH] Port Spam filter to Rust - part 6 --- Cargo.lock | 15 +- crates/common/src/config/inner.rs | 19 +- crates/common/src/config/mod.rs | 10 +- crates/common/src/config/scripts.rs | 2 +- crates/common/src/config/spamfilter.rs | 27 + crates/common/src/config/storage.rs | 6 +- crates/common/src/core.rs | 20 +- crates/common/src/enterprise/config.rs | 5 +- crates/common/src/enterprise/mod.rs | 20 +- crates/common/src/expr/functions/asynch.rs | 25 +- crates/common/src/ipc.rs | 4 +- crates/common/src/lib.rs | 48 +- crates/common/src/listener/acme/order.rs | 25 +- crates/common/src/listener/acme/resolver.rs | 2 +- crates/common/src/listener/blocked.rs | 29 +- crates/common/src/manager/reload.rs | 6 +- crates/common/src/scripts/plugins/lookup.rs | 18 +- crates/common/src/scripts/plugins/query.rs | 12 +- crates/directory/src/backend/sql/config.rs | 17 +- crates/directory/src/backend/sql/lookup.rs | 38 +- crates/directory/src/backend/sql/mod.rs | 4 +- crates/directory/src/core/mod.rs | 2 +- crates/directory/src/lib.rs | 2 +- crates/imap/src/core/client.rs | 9 +- crates/jmap/src/api/form.rs | 6 +- crates/jmap/src/api/management/stores.rs | 2 +- crates/jmap/src/auth/oauth/auth.rs | 34 +- crates/jmap/src/auth/rate_limit.rs | 13 +- crates/jmap/src/email/delete.rs | 19 +- crates/jmap/src/services/housekeeper.rs | 9 +- crates/main/Cargo.toml | 8 +- crates/managesieve/src/core/client.rs | 8 +- crates/nlp/Cargo.toml | 5 +- crates/nlp/src/bayes/cache.rs | 112 ----- crates/nlp/src/bayes/mod.rs | 114 ++++- crates/nlp/src/bayes/tokenize.rs | 165 ++++--- crates/nlp/src/bayes/train.rs | 16 +- crates/nlp/src/lib.rs | 61 ++- crates/nlp/src/tokenizers/japanese.rs | 251 +++++++++- crates/nlp/src/tokenizers/osb.rs | 32 +- crates/pop3/src/client.rs | 8 +- crates/smtp/src/core/throttle.rs | 58 ++- crates/smtp/src/inbound/rcpt.rs | 38 +- crates/smtp/src/queue/throttle.rs | 36 +- crates/spam-filter/src/analysis/bayes.rs | 52 ++ crates/spam-filter/src/analysis/init.rs | 24 - crates/spam-filter/src/analysis/llm.rs | 90 ++++ crates/spam-filter/src/analysis/mod.rs | 9 + crates/spam-filter/src/analysis/reputation.rs | 69 ++- crates/spam-filter/src/analysis/score.rs | 97 ++++ .../spam-filter/src/analysis/trusted_reply.rs | 83 ++++ crates/spam-filter/src/lib.rs | 2 + crates/spam-filter/src/modules/bayes.rs | 463 +++++++++++------- crates/spam-filter/src/modules/mod.rs | 15 +- crates/store/src/backend/memory/mod.rs | 12 +- crates/store/src/config.rs | 63 +-- crates/store/src/dispatch/lookup.rs | 218 ++++----- crates/store/src/dispatch/store.rs | 32 +- crates/store/src/lib.rs | 21 +- crates/store/src/write/purge.rs | 6 +- .../spamfilter/scripts/bayes_classify.sieve | 17 - .../config/spamfilter/scripts/epilogue.sieve | 28 -- resources/config/spamfilter/scripts/llm.sieve | 41 -- .../spamfilter/scripts/replies_in.sieve | 12 - .../spamfilter/scripts/replies_out.sieve | 12 - .../config/spamfilter/scripts/scores.sieve | 27 - .../config/spamfilter/scripts/spamtrap.sieve | 9 - tests/Cargo.toml | 1 + tests/src/directory/mod.rs | 4 +- tests/src/directory/sql.rs | 28 +- tests/src/jmap/auth_oauth.rs | 2 +- tests/src/jmap/enterprise.rs | 2 + tests/src/smtp/inbound/antispam.rs | 2 +- tests/src/smtp/lookup/sql.rs | 4 +- tests/src/store/lookup.rs | 38 +- 75 files changed, 1742 insertions(+), 1101 deletions(-) delete mode 100644 crates/nlp/src/bayes/cache.rs create mode 100644 crates/spam-filter/src/analysis/bayes.rs create mode 100644 crates/spam-filter/src/analysis/llm.rs create mode 100644 crates/spam-filter/src/analysis/score.rs create mode 100644 crates/spam-filter/src/analysis/trusted_reply.rs delete mode 100644 resources/config/spamfilter/scripts/bayes_classify.sieve delete mode 100644 resources/config/spamfilter/scripts/epilogue.sieve delete mode 100644 resources/config/spamfilter/scripts/llm.sieve delete mode 100644 resources/config/spamfilter/scripts/replies_in.sieve delete mode 100644 resources/config/spamfilter/scripts/replies_out.sieve delete mode 100644 resources/config/spamfilter/scripts/scores.sieve delete mode 100644 resources/config/spamfilter/scripts/spamtrap.sieve diff --git a/Cargo.lock b/Cargo.lock index 433053a5..71224739 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3880,6 +3880,7 @@ dependencies = [ "managesieve", "pop3", "smtp", + "spam-filter", "store", "tokio", "trc", @@ -4173,14 +4174,15 @@ dependencies = [ "farmhash", "jieba-rs", "lru-cache", + "maplit", "nohash", "parking_lot", "phf", "psl", + "radix_trie", "rust-stemmers", "serde", "siphasher 1.0.1", - "tinysegmenter", "tokio", "whatlang", "xxhash-rust", @@ -6773,6 +6775,7 @@ dependencies = [ "sieve-rs", "smtp", "smtp-proto", + "spam-filter", "store", "tokio", "tokio-rustls 0.26.0", @@ -6861,16 +6864,6 @@ dependencies = [ "crunchy", ] -[[package]] -name = "tinysegmenter" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1755695d17d470baf2d937a59ab4e86de3034b056fc8700e21411b0efca36497" -dependencies = [ - "lazy_static", - "maplit", -] - [[package]] name = "tinystr" version = "0.7.6" diff --git a/crates/common/src/config/inner.rs b/crates/common/src/config/inner.rs index 54b216cd..ac6ad199 100644 --- a/crates/common/src/config/inner.rs +++ b/crates/common/src/config/inner.rs @@ -4,13 +4,12 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use std::{sync::Arc, time::Duration}; +use std::sync::Arc; use ahash::{AHashMap, AHashSet, RandomState}; use arc_swap::ArcSwap; use dashmap::DashMap; use mail_send::smtp::tls::build_tls_connector; -use nlp::bayes::cache::BayesTokenCache; use parking_lot::RwLock; use utils::{ config::Config, @@ -105,17 +104,6 @@ impl Data { shard_amount, ), smtp_connectors: TlsConnectors::default(), - bayes_cache: BayesTokenCache::new( - config - .property_or_default("cache.bayes.capacity", "8192") - .unwrap_or(8192), - config - .property_or_default("cache.bayes.ttl.positive", "1h") - .unwrap_or_else(|| Duration::from_secs(3600)), - config - .property_or_default("cache.bayes.ttl.negative", "1h") - .unwrap_or_else(|| Duration::from_secs(3600)), - ), remote_lists: Default::default(), asn_geo_data: Default::default(), } @@ -148,11 +136,6 @@ impl Default for Data { smtp_session_throttle: Default::default(), smtp_queue_throttle: Default::default(), smtp_connectors: Default::default(), - bayes_cache: BayesTokenCache::new( - 8192, - Duration::from_secs(3600), - Duration::from_secs(3600), - ), asn_geo_data: Default::default(), } } diff --git a/crates/common/src/config/mod.rs b/crates/common/src/config/mod.rs index eb115b89..f22fed07 100644 --- a/crates/common/src/config/mod.rs +++ b/crates/common/src/config/mod.rs @@ -15,7 +15,7 @@ use hyper::{ }; use ring::signature::{EcdsaKeyPair, RsaKeyPair}; use spamfilter::SpamFilterConfig; -use store::{BlobBackend, BlobStore, FtsStore, LookupStore, Store, Stores}; +use store::{BlobBackend, BlobStore, FtsStore, InMemoryStore, Store, Stores}; use telemetry::Metrics; use utils::config::{utils::AsKey, Config}; @@ -111,7 +111,7 @@ impl Core { .value_require("storage.lookup") .map(|id| id.to_string()) .and_then(|id| { - if let Some(store) = stores.lookup_stores.get(&id) { + if let Some(store) = stores.in_memory_stores.get(&id) { store.clone().into() } else { config.new_parse_error( @@ -161,12 +161,12 @@ impl Core { // If any of the stores are missing, disable all stores to avoid data loss if matches!(data, Store::None) || matches!(&blob.backend, BlobBackend::Store(Store::None)) - || matches!(lookup, LookupStore::Store(Store::None)) + || matches!(lookup, InMemoryStore::Store(Store::None)) || matches!(fts, FtsStore::Store(Store::None)) { data = Store::default(); blob = BlobStore::default(); - lookup = LookupStore::default(); + lookup = InMemoryStore::default(); fts = FtsStore::default(); config.new_build_error( "storage.*", @@ -196,7 +196,7 @@ impl Core { purge_schedules: stores.purge_schedules, config: config_manager, stores: stores.stores, - lookups: stores.lookup_stores, + lookups: stores.in_memory_stores, blobs: stores.blob_stores, ftss: stores.fts_stores, }, diff --git a/crates/common/src/config/scripts.rs b/crates/common/src/config/scripts.rs index eb2b59ab..254cd5b9 100644 --- a/crates/common/src/config/scripts.rs +++ b/crates/common/src/config/scripts.rs @@ -239,7 +239,7 @@ impl Scripting { ) .with_max_header_size(10240) .with_valid_notification_uri("mailto") - .with_valid_ext_lists(stores.lookup_stores.keys().map(|k| k.to_string())) + .with_valid_ext_lists(stores.in_memory_stores.keys().map(|k| k.to_string())) .with_functions(&mut fnc_map_trusted) .with_max_redirects( config diff --git a/crates/common/src/config/spamfilter.rs b/crates/common/src/config/spamfilter.rs index fb313b91..1032aa44 100644 --- a/crates/common/src/config/spamfilter.rs +++ b/crates/common/src/config/spamfilter.rs @@ -8,6 +8,7 @@ use std::{net::SocketAddr, time::Duration}; use ahash::AHashSet; use mail_parser::HeaderName; +use nlp::bayes::BayesClassifier; use utils::{ config::Config, glob::{GlobMap, GlobSet}, @@ -21,10 +22,16 @@ pub struct SpamFilterConfig { pub max_rbl_domain_checks: usize, pub max_rbl_email_checks: usize, pub max_rbl_url_checks: usize, + pub trusted_reply: Option, pub greylist_duration: Option, pub pyzor: Option, pub reputation: Option, + pub bayes: Option, + + pub score_reject_threshold: f64, + pub score_discard_threshold: f64, + pub score_spam_threshold: f64, pub list_dmarc_allow: GlobSet, pub list_spf_dkim_allow: GlobSet, @@ -33,11 +40,31 @@ pub struct SpamFilterConfig { pub list_trusted_domains: GlobSet, pub list_url_redirectors: GlobSet, pub list_file_extensions: GlobMap, + pub list_scores: GlobMap>, + pub list_spamtraps: GlobSet, pub remote_lists: Vec, pub dnsbls: Vec, } +#[derive(Debug, Clone)] +pub enum SpamFilterAction { + Allow(T), + Discard, + Reject, +} + +#[derive(Debug, Clone, Default)] +pub struct BayesConfig { + pub classifier: BayesClassifier, + pub auto_learn: bool, + pub auto_learn_reply_ham: bool, + pub auto_learn_spam_threshold: f64, + pub auto_learn_ham_threshold: f64, + pub score_spam: f64, + pub score_ham: f64, +} + #[derive(Debug, Clone, Default)] pub struct ReputationConfig { pub expiry: u64, diff --git a/crates/common/src/config/storage.rs b/crates/common/src/config/storage.rs index 3ab8ec8b..ef42256f 100644 --- a/crates/common/src/config/storage.rs +++ b/crates/common/src/config/storage.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use ahash::AHashMap; use directory::Directory; -use store::{write::purge::PurgeSchedule, BlobStore, FtsStore, LookupStore, Store}; +use store::{write::purge::PurgeSchedule, BlobStore, FtsStore, InMemoryStore, Store}; use crate::manager::config::ConfigManager; @@ -17,7 +17,7 @@ pub struct Storage { pub data: Store, pub blob: BlobStore, pub fts: FtsStore, - pub lookup: LookupStore, + pub lookup: InMemoryStore, pub directory: Arc, pub directories: AHashMap>, pub purge_schedules: Vec, @@ -25,6 +25,6 @@ pub struct Storage { pub stores: AHashMap, pub blobs: AHashMap, - pub lookups: AHashMap, + pub lookups: AHashMap, pub ftss: AHashMap, } diff --git a/crates/common/src/core.rs b/crates/common/src/core.rs index 31ed0bf7..34c13d01 100644 --- a/crates/common/src/core.rs +++ b/crates/common/src/core.rs @@ -10,7 +10,7 @@ use directory::{backend::internal::manage::ManageDirectory, Directory, Type}; use sieve::Sieve; use store::{ write::{QueueClass, ValueClass}, - BlobStore, FtsStore, IterateParams, LookupStore, Store, ValueKey, + BlobStore, FtsStore, InMemoryStore, IterateParams, Store, ValueKey, }; use trc::AddContext; @@ -39,7 +39,7 @@ impl Server { } #[inline(always)] - pub fn lookup_store(&self) -> &LookupStore { + pub fn in_memory_store(&self) -> &InMemoryStore { &self.core.storage.lookup } @@ -66,7 +66,7 @@ impl Server { }) } - pub fn get_lookup_store(&self, name: &str, session_id: u64) -> &LookupStore { + pub fn get_in_memory_store(&self, name: &str, session_id: u64) -> &InMemoryStore { self.core.storage.lookups.get(name).unwrap_or_else(|| { if !name.is_empty() { trc::event!( @@ -80,6 +80,20 @@ impl Server { }) } + pub fn get_data_store(&self, name: &str, session_id: u64) -> &Store { + self.core.storage.stores.get(name).unwrap_or_else(|| { + if !name.is_empty() { + trc::event!( + Eval(trc::EvalEvent::StoreNotFound), + Id = name.to_string(), + SpanId = session_id, + ); + } + + &self.core.storage.data + }) + } + pub fn get_arc_sealer(&self, name: &str, session_id: u64) -> Option<&ArcSealer> { self.core .smtp diff --git a/crates/common/src/enterprise/config.rs b/crates/common/src/enterprise/config.rs index 60403013..f8924931 100644 --- a/crates/common/src/enterprise/config.rs +++ b/crates/common/src/enterprise/config.rs @@ -183,10 +183,12 @@ impl Enterprise { .collect::>() { if let Some(api) = AiApiConfig::parse(config, &id) { - ai_apis.insert(id, api); + ai_apis.insert(id, api.into()); } } + todo!("implement spam filter llm config"); + Some(Enterprise { license, undelete: config @@ -198,6 +200,7 @@ impl Enterprise { metrics_store, metrics_alerts: parse_metric_alerts(config), ai_apis, + spam_filter_llm: None, }) } } diff --git a/crates/common/src/enterprise/mod.rs b/crates/common/src/enterprise/mod.rs index e9325c18..a5b986f1 100644 --- a/crates/common/src/enterprise/mod.rs +++ b/crates/common/src/enterprise/mod.rs @@ -14,9 +14,9 @@ pub mod license; pub mod llm; pub mod undelete; -use std::time::Duration; +use std::{sync::Arc, time::Duration}; -use ahash::AHashMap; +use ahash::{AHashMap, AHashSet}; use directory::{ backend::internal::{lookup::DirectoryStore, PrincipalField}, QueryBy, Type, @@ -38,7 +38,21 @@ pub struct Enterprise { pub trace_store: Option, pub metrics_store: Option, pub metrics_alerts: Vec, - pub ai_apis: AHashMap, + pub ai_apis: AHashMap>, + pub spam_filter_llm: Option, +} + +#[derive(Debug, Clone)] +pub struct SpamFilterLlmConfig { + pub model: Arc, + pub temperature: f64, + pub prompt: String, + pub separator: char, + pub index_category: usize, + pub index_confidence: Option, + pub index_explanation: Option, + pub categories: AHashSet, + pub confidence: AHashSet, } #[derive(Clone)] diff --git a/crates/common/src/expr/functions/asynch.rs b/crates/common/src/expr/functions/asynch.rs index 55b36164..768c7b4e 100644 --- a/crates/common/src/expr/functions/asynch.rs +++ b/crates/common/src/expr/functions/asynch.rs @@ -2,7 +2,7 @@ use std::{cmp::Ordering, net::IpAddr, vec::IntoIter}; use directory::backend::RcptType; use mail_auth::IpLookupStrategy; -use store::{Deserialize, Rows, Value}; +use store::{dispatch::lookup::KeyValue, Deserialize, Rows, Value}; use trc::AddContext; use crate::Server; @@ -43,7 +43,7 @@ impl Server { let store = params.next_as_string(); let key = params.next_as_string(); - self.get_lookup_store(store.as_ref(), session_id) + self.get_in_memory_store(store.as_ref(), session_id) .key_get::(key.into_owned().into_bytes()) .await .map(|value| value.map(|v| v.into_inner()).unwrap_or_default()) @@ -53,7 +53,7 @@ impl Server { let store = params.next_as_string(); let key = params.next_as_string(); - self.get_lookup_store(store.as_ref(), session_id) + self.get_in_memory_store(store.as_ref(), session_id) .key_exists(key.into_owned().into_bytes()) .await .caused_by(trc::location!()) @@ -64,12 +64,11 @@ impl Server { let key = params.next_as_string(); let value = params.next_as_string(); - self.get_lookup_store(store.as_ref(), session_id) - .key_set( + self.get_in_memory_store(store.as_ref(), session_id) + .key_set(KeyValue::new( key.into_owned().into_bytes(), value.into_owned().into_bytes(), - None, - ) + )) .await .map(|_| true) .caused_by(trc::location!()) @@ -80,8 +79,8 @@ impl Server { let key = params.next_as_string(); let value = params.next_as_integer(); - self.get_lookup_store(store.as_ref(), session_id) - .counter_incr(key.into_owned().into_bytes(), value, None, true) + self.get_in_memory_store(store.as_ref(), session_id) + .counter_incr(KeyValue::new(key.into_owned(), value)) .await .map(Variable::Integer) .caused_by(trc::location!()) @@ -90,7 +89,7 @@ impl Server { let store = params.next_as_string(); let key = params.next_as_string(); - self.get_lookup_store(store.as_ref(), session_id) + self.get_in_memory_store(store.as_ref(), session_id) .counter_get(key.into_owned().into_bytes()) .await .map(Variable::Integer) @@ -107,7 +106,7 @@ impl Server { mut arguments: FncParams<'x>, session_id: u64, ) -> trc::Result> { - let store = self.get_lookup_store(arguments.next_as_string().as_ref(), session_id); + let store = self.get_data_store(arguments.next_as_string().as_ref(), session_id); let query = arguments.next_as_string(); if query.is_empty() { @@ -129,7 +128,7 @@ impl Server { .map_or(false, |q| q.eq_ignore_ascii_case(b"SELECT")) { let mut rows = store - .query::(&query, arguments) + .sql_query::(&query, arguments) .await .caused_by(trc::location!())?; Ok(match rows.rows.len().cmp(&1) { @@ -157,7 +156,7 @@ impl Server { }) } else { store - .query::(&query, arguments) + .sql_query::(&query, arguments) .await .caused_by(trc::location!()) .map(|v| v.into()) diff --git a/crates/common/src/ipc.rs b/crates/common/src/ipc.rs index 9b64a318..113a01fd 100644 --- a/crates/common/src/ipc.rs +++ b/crates/common/src/ipc.rs @@ -13,7 +13,7 @@ use mail_auth::{ mta_sts::TlsRpt, report::{tlsrpt::FailureDetails, Record}, }; -use store::{BlobStore, LookupStore, Store}; +use store::{BlobStore, InMemoryStore, Store}; use tokio::sync::{mpsc, oneshot}; use utils::{map::bitmap::Bitmap, BlobHash}; @@ -68,7 +68,7 @@ pub enum HousekeeperEvent { pub enum PurgeType { Data(Store), Blobs { store: Store, blob_store: BlobStore }, - Lookup(LookupStore), + Lookup(InMemoryStore), Account(Option), } diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 1f37b92e..5d5c7b01 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -34,7 +34,6 @@ use listener::{ }; use manager::webadmin::{Resource, WebAdminManager}; -use nlp::bayes::cache::BayesTokenCache; use parking_lot::{Mutex, RwLock}; use reqwest::Response; use rustls::sign::CertifiedKey; @@ -67,6 +66,28 @@ pub static DAEMON_NAME: &str = concat!("Stalwart Mail Server v", env!("CARGO_PKG pub const IPC_CHANNEL_BUFFER: usize = 1024; +pub const KV_ACME: u8 = 0; +pub const KV_OAUTH: u8 = 1; +pub const KV_RATE_LIMIT_RCPT: u8 = 2; +pub const KV_RATE_LIMIT_SCAN: u8 = 3; +pub const KV_RATE_LIMIT_LOITER: u8 = 4; +pub const KV_RATE_LIMIT_AUTH: u8 = 5; +pub const KV_RATE_LIMIT_HASH: u8 = 6; +pub const KV_RATE_LIMIT_CONTACT: u8 = 7; +pub const KV_RATE_LIMIT_JMAP: u8 = 8; +pub const KV_RATE_LIMIT_JMAP_AUTH: u8 = 9; +pub const KV_RATE_LIMIT_HTTP_ANONYM: u8 = 10; +pub const KV_RATE_LIMIT_IMAP: u8 = 11; +pub const KV_REPUTATION_IP: u8 = 12; +pub const KV_REPUTATION_FROM: u8 = 13; +pub const KV_REPUTATION_DOMAIN: u8 = 14; +pub const KV_REPUTATION_ASN: u8 = 15; +pub const KV_GREYLIST: u8 = 16; +pub const KV_BAYES_MODEL_GLOBAL: u8 = 17; +pub const KV_BAYES_MODEL_USER: u8 = 18; +pub const KV_TRUSTED_REPLY: u8 = 19; +pub const KV_LOCK_PURGE_ACCOUNT: u8 = 20; + #[derive(Clone)] pub struct Server { pub inner: Arc, @@ -92,7 +113,6 @@ pub struct Data { pub permissions: ADashMap>, pub permissions_version: AtomicU8, - pub bayes_cache: BayesTokenCache, pub remote_lists: RwLock>, pub asn_geo_data: AsnGeoLookupData, @@ -350,3 +370,27 @@ impl Default for Ipc { } } } + +pub fn ip_to_bytes(ip: &IpAddr) -> Vec { + match ip { + IpAddr::V4(ip) => ip.octets().to_vec(), + IpAddr::V6(ip) => ip.octets().to_vec(), + } +} + +pub fn ip_to_bytes_prefix(prefix: u8, ip: &IpAddr) -> Vec { + match ip { + IpAddr::V4(ip) => { + let mut buf = Vec::with_capacity(5); + buf.push(prefix); + buf.extend_from_slice(&ip.octets()); + buf + } + IpAddr::V6(ip) => { + let mut buf = Vec::with_capacity(17); + buf.push(prefix); + buf.extend_from_slice(&ip.octets()); + buf + } + } +} diff --git a/crates/common/src/listener/acme/order.rs b/crates/common/src/listener/acme/order.rs index 1417a916..2bfefb26 100644 --- a/crates/common/src/listener/acme/order.rs +++ b/crates/common/src/listener/acme/order.rs @@ -9,12 +9,13 @@ use rustls::sign::CertifiedKey; use rustls_pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer}; use std::sync::Arc; use std::time::{Duration, Instant}; +use store::dispatch::lookup::KeyValue; use trc::{AcmeEvent, EventType}; use x509_parser::parse_x509_certificate; use crate::listener::acme::directory::Identifier; use crate::listener::acme::ChallengeSettings; -use crate::Server; +use crate::{Server, KV_ACME}; use super::directory::{Account, AuthStatus, Directory, OrderStatus}; use super::AcmeProvider; @@ -207,20 +208,26 @@ impl Server { match &provider.challenge { ChallengeSettings::TlsAlpn01 => { - self.lookup_store() + self.in_memory_store() .key_set( - format!("acme:{domain}").into_bytes(), - account.tls_alpn_key(challenge, domain.clone())?, - 3600.into(), + KeyValue::with_prefix( + KV_ACME, + &domain, + account.tls_alpn_key(challenge, domain.clone())?, + ) + .expires(3600), ) .await?; } ChallengeSettings::Http01 => { - self.lookup_store() + self.in_memory_store() .key_set( - format!("acme:{}", challenge.token).into_bytes(), - account.http_proof(challenge)?, - 3600.into(), + KeyValue::with_prefix( + KV_ACME, + &challenge.token, + account.http_proof(challenge)?, + ) + .expires(3600), ) .await?; } diff --git a/crates/common/src/listener/acme/resolver.rs b/crates/common/src/listener/acme/resolver.rs index fec948fc..25d5e2d8 100644 --- a/crates/common/src/listener/acme/resolver.rs +++ b/crates/common/src/listener/acme/resolver.rs @@ -44,7 +44,7 @@ impl Server { pub(crate) async fn build_acme_certificate(&self, domain: &str) -> Option> { match self - .lookup_store() + .in_memory_store() .key_get::>(format!("acme:{domain}").into_bytes()) .await { diff --git a/crates/common/src/listener/blocked.rs b/crates/common/src/listener/blocked.rs index 643ff95c..26710280 100644 --- a/crates/common/src/listener/blocked.rs +++ b/crates/common/src/listener/blocked.rs @@ -16,7 +16,10 @@ use utils::{ glob::GlobPattern, }; -use crate::{manager::config::MatchType, Server}; +use crate::{ + ip_to_bytes, manager::config::MatchType, Server, KV_RATE_LIMIT_AUTH, KV_RATE_LIMIT_LOITER, + KV_RATE_LIMIT_RCPT, KV_RATE_LIMIT_SCAN, +}; #[derive(Debug, Clone)] pub struct Security { @@ -138,13 +141,13 @@ impl Server { if let Some(rate) = &self.core.network.security.rcpt_fail_rate { let is_allowed = self.is_ip_allowed(&ip) || (self - .lookup_store() - .is_rate_allowed(format!("r:{ip}").as_bytes(), rate, false) + .in_memory_store() + .is_rate_allowed(KV_RATE_LIMIT_RCPT, &ip_to_bytes(&ip), rate, false) .await? .is_none() && self - .lookup_store() - .is_rate_allowed(format!("r:{rcpt}").as_bytes(), rate, false) + .in_memory_store() + .is_rate_allowed(KV_RATE_LIMIT_RCPT, rcpt.as_bytes(), rate, false) .await? .is_none()); @@ -160,8 +163,8 @@ impl Server { if let Some(rate) = &self.core.network.security.scanner_fail_rate { let is_allowed = self.is_ip_allowed(&ip) || self - .lookup_store() - .is_rate_allowed(format!("h:{ip}").as_bytes(), rate, false) + .in_memory_store() + .is_rate_allowed(KV_RATE_LIMIT_SCAN, &ip_to_bytes(&ip), rate, false) .await? .is_none(); @@ -187,8 +190,8 @@ impl Server { if let Some(rate) = &self.core.network.security.loiter_fail_rate { let is_allowed = self.is_ip_allowed(&ip) || self - .lookup_store() - .is_rate_allowed(format!("l:{ip}").as_bytes(), rate, false) + .in_memory_store() + .is_rate_allowed(KV_RATE_LIMIT_LOITER, &ip_to_bytes(&ip), rate, false) .await? .is_none(); @@ -205,14 +208,14 @@ impl Server { let login = login.unwrap_or_default(); let is_allowed = self.is_ip_allowed(&ip) || (self - .lookup_store() - .is_rate_allowed(format!("b:{ip}").as_bytes(), rate, false) + .in_memory_store() + .is_rate_allowed(KV_RATE_LIMIT_AUTH, &ip_to_bytes(&ip), rate, false) .await? .is_none() && (login.is_empty() || self - .lookup_store() - .is_rate_allowed(format!("b:{login}").as_bytes(), rate, false) + .in_memory_store() + .is_rate_allowed(KV_RATE_LIMIT_AUTH, login.as_bytes(), rate, false) .await? .is_none())); if !is_allowed { diff --git a/crates/common/src/manager/reload.rs b/crates/common/src/manager/reload.rs index 143d4008..af9d2b03 100644 --- a/crates/common/src/manager/reload.rs +++ b/crates/common/src/manager/reload.rs @@ -53,10 +53,10 @@ impl Server { pub async fn reload_lookups(&self) -> trc::Result { let mut config = self.core.storage.config.build_config("lookup").await?; let mut stores = Stores::default(); - stores.parse_memory_stores(&mut config); + stores.parse_static_stores(&mut config); let mut core = self.core.as_ref().clone(); - for (id, store) in stores.lookup_stores { + for (id, store) in stores.in_memory_stores { core.storage.lookups.insert(id, store); } @@ -75,7 +75,7 @@ impl Server { stores: self.core.storage.stores.clone(), blob_stores: self.core.storage.blobs.clone(), fts_stores: self.core.storage.ftss.clone(), - lookup_stores: self.core.storage.lookups.clone(), + in_memory_stores: self.core.storage.lookups.clone(), purge_schedules: Default::default(), }; stores.parse_stores(&mut config).await; diff --git a/crates/common/src/scripts/plugins/lookup.rs b/crates/common/src/scripts/plugins/lookup.rs index 1567a130..398ec1de 100644 --- a/crates/common/src/scripts/plugins/lookup.rs +++ b/crates/common/src/scripts/plugins/lookup.rs @@ -5,7 +5,7 @@ */ use sieve::{runtime::Variable, FunctionMap}; -use store::{Deserialize, Value}; +use store::{dispatch::lookup::KeyValue, Deserialize, Value}; use crate::scripts::into_sieve_value; @@ -93,13 +93,15 @@ pub async fn exec_set(ctx: PluginContext<'_>) -> trc::Result { .details("Unknown store") })? .key_set( - ctx.arguments[1].to_string().into_owned().into_bytes(), - if !ctx.arguments[2].is_empty() { - bincode::serialize(&ctx.arguments[2]).unwrap_or_default() - } else { - vec![] - }, - expires, + KeyValue::new( + ctx.arguments[1].to_string().into_owned().into_bytes(), + if !ctx.arguments[2].is_empty() { + bincode::serialize(&ctx.arguments[2]).unwrap_or_default() + } else { + vec![] + }, + ) + .expires_opt(expires), ) .await .map(|_| true.into()) diff --git a/crates/common/src/scripts/plugins/query.rs b/crates/common/src/scripts/plugins/query.rs index 15092caf..f908645a 100644 --- a/crates/common/src/scripts/plugins/query.rs +++ b/crates/common/src/scripts/plugins/query.rs @@ -19,8 +19,8 @@ pub fn register(plugin_id: u32, fnc_map: &mut FunctionMap) { pub async fn exec(ctx: PluginContext<'_>) -> trc::Result { // Obtain store name let store = match &ctx.arguments[0] { - Variable::String(v) if !v.is_empty() => ctx.server.core.storage.lookups.get(v.as_ref()), - _ => Some(&ctx.server.core.storage.lookup), + Variable::String(v) if !v.is_empty() => ctx.server.core.storage.stores.get(v.as_ref()), + _ => Some(&ctx.server.core.storage.data), } .ok_or_else(|| { trc::SieveEvent::RuntimeError @@ -48,7 +48,7 @@ pub async fn exec(ctx: PluginContext<'_>) -> trc::Result { .get(..6) .map_or(false, |q| q.eq_ignore_ascii_case(b"SELECT")) { - let mut rows = store.query::(&query, arguments).await?; + let mut rows = store.sql_query::(&query, arguments).await?; Ok(match rows.rows.len().cmp(&1) { Ordering::Equal => { let mut row = rows.rows.pop().unwrap().values; @@ -82,6 +82,10 @@ pub async fn exec(ctx: PluginContext<'_>) -> trc::Result { .into(), }) } else { - Ok(store.query::(&query, arguments).await.is_ok().into()) + Ok(store + .sql_query::(&query, arguments) + .await + .is_ok() + .into()) } } diff --git a/crates/directory/src/backend/sql/config.rs b/crates/directory/src/backend/sql/config.rs index c224ab99..d1cdbb17 100644 --- a/crates/directory/src/backend/sql/config.rs +++ b/crates/directory/src/backend/sql/config.rs @@ -18,13 +18,14 @@ impl SqlDirectory { ) -> Option { let prefix = prefix.as_key(); let store_id = config.value_require((&prefix, "store"))?.to_string(); - let store = if let Some(store) = stores.lookup_stores.get(&store_id) { - store.clone() - } else { - let err = format!("Directory references a non-existent store {store_id:?}"); - config.new_build_error((&prefix, "store"), err); - return None; - }; + let sql_store = + if let Some(sql_store) = stores.stores.get(&store_id).filter(|store| store.is_sql()) { + sql_store.clone() + } else { + let err = format!("Directory references a non-existent store {store_id:?}"); + config.new_build_error((&prefix, "store"), err); + return None; + }; let mut mappings = SqlMappings { column_description: config @@ -64,7 +65,7 @@ impl SqlDirectory { } Some(SqlDirectory { - store, + sql_store, mappings, data_store, }) diff --git a/crates/directory/src/backend/sql/lookup.rs b/crates/directory/src/backend/sql/lookup.rs index 80888eca..141ed497 100644 --- a/crates/directory/src/backend/sql/lookup.rs +++ b/crates/directory/src/backend/sql/lookup.rs @@ -32,8 +32,11 @@ impl SqlDirectory { QueryBy::Name(username) => ( self.mappings .row_to_principal( - self.store - .query::(&self.mappings.query_name, vec![username.into()]) + self.sql_store + .sql_query::( + &self.mappings.query_name, + vec![username.into()], + ) .await .caused_by(trc::location!())?, ) @@ -51,8 +54,8 @@ impl SqlDirectory { ( self.mappings .row_to_principal( - self.store - .query::( + self.sql_store + .sql_query::( &self.mappings.query_name, vec![principal.name().into()], ) @@ -76,8 +79,11 @@ impl SqlDirectory { match self .mappings .row_to_principal( - self.store - .query::(&self.mappings.query_name, vec![username.into()]) + self.sql_store + .sql_query::( + &self.mappings.query_name, + vec![username.into()], + ) .await .caused_by(trc::location!())?, ) @@ -108,8 +114,8 @@ impl SqlDirectory { // Obtain members if return_member_of && !self.mappings.query_members.is_empty() { for row in self - .store - .query::( + .sql_store + .sql_query::( &self.mappings.query_members, vec![external_principal.name().into()], ) @@ -134,8 +140,8 @@ impl SqlDirectory { external_principal.set( PrincipalField::Emails, PrincipalValue::StringList( - self.store - .query::( + self.sql_store + .sql_query::( &self.mappings.query_emails, vec![external_principal.name().into()], ) @@ -151,8 +157,8 @@ impl SqlDirectory { external_principal.set( PrincipalField::Secrets, PrincipalValue::StringList( - self.store - .query::( + self.sql_store + .sql_query::( &self.mappings.query_secrets, vec![external_principal.name().into()], ) @@ -198,8 +204,8 @@ impl SqlDirectory { pub async fn email_to_id(&self, address: &str) -> trc::Result> { let names = self - .store - .query::(&self.mappings.query_recipients, vec![address.into()]) + .sql_store + .sql_query::(&self.mappings.query_recipients, vec![address.into()]) .await .caused_by(trc::location!())?; @@ -219,8 +225,8 @@ impl SqlDirectory { pub async fn rcpt(&self, address: &str) -> trc::Result { let result = self - .store - .query::( + .sql_store + .sql_query::( &self.mappings.query_recipients, vec![address.to_string().into()], ) diff --git a/crates/directory/src/backend/sql/mod.rs b/crates/directory/src/backend/sql/mod.rs index f518a211..19549f46 100644 --- a/crates/directory/src/backend/sql/mod.rs +++ b/crates/directory/src/backend/sql/mod.rs @@ -4,13 +4,13 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use store::{LookupStore, Store}; +use store::Store; pub mod config; pub mod lookup; pub struct SqlDirectory { - store: LookupStore, + sql_store: Store, mappings: SqlMappings, pub(crate) data_store: Store, } diff --git a/crates/directory/src/core/mod.rs b/crates/directory/src/core/mod.rs index 859ea631..2683eb5e 100644 --- a/crates/directory/src/core/mod.rs +++ b/crates/directory/src/core/mod.rs @@ -73,7 +73,7 @@ impl Permission { Permission::BlobFetch => "Retrieve arbitrary blobs", Permission::PurgeBlobStore => "Purge the blob storage", Permission::PurgeDataStore => "Purge the data storage", - Permission::PurgeLookupStore => "Purge the lookup storage", + Permission::PurgeInMemoryStore => "Purge the lookup storage", Permission::PurgeAccount => "Purge user accounts", Permission::FtsReindex => "Rebuild the full-text search index", Permission::Undelete => "Restore deleted items", diff --git a/crates/directory/src/lib.rs b/crates/directory/src/lib.rs index cdfda42e..a6c1280e 100644 --- a/crates/directory/src/lib.rs +++ b/crates/directory/src/lib.rs @@ -122,7 +122,7 @@ pub enum Permission { BlobFetch, PurgeBlobStore, PurgeDataStore, - PurgeLookupStore, + PurgeInMemoryStore, PurgeAccount, FtsReindex, Undelete, diff --git a/crates/imap/src/core/client.rs b/crates/imap/src/core/client.rs index a4b4b1b0..e6733c39 100644 --- a/crates/imap/src/core/client.rs +++ b/crates/imap/src/core/client.rs @@ -8,7 +8,7 @@ use std::{iter::Peekable, sync::Arc, vec::IntoIter}; use common::{ listener::{limiter::ConcurrencyLimiter, SessionResult, SessionStream}, - ConcurrencyLimiters, + ConcurrencyLimiters, KV_RATE_LIMIT_IMAP, }; use imap_proto::{ receiver::{self, Request}, @@ -292,7 +292,12 @@ impl Session { .core .storage .lookup - .is_rate_allowed(format!("ireq:{}", data.account_id).as_bytes(), rate, true) + .is_rate_allowed( + KV_RATE_LIMIT_IMAP, + &data.account_id.to_be_bytes(), + rate, + true, + ) .await? .is_some() { diff --git a/crates/jmap/src/api/form.rs b/crates/jmap/src/api/form.rs index 1ef718a5..29b9b7e1 100644 --- a/crates/jmap/src/api/form.rs +++ b/crates/jmap/src/api/form.rs @@ -9,8 +9,9 @@ use std::{borrow::Cow, fmt::Write, future::Future}; use chrono::Utc; use common::{ config::network::{ContactForm, FieldOrDefault}, + ip_to_bytes, ipc::{DeliveryResult, IngestMessage}, - psl, Server, + psl, Server, KV_RATE_LIMIT_CONTACT, }; use hyper::StatusCode; use mail_builder::{ @@ -61,7 +62,8 @@ impl FormHandler for Server { .storage .lookup .is_rate_allowed( - format!("contact:{}", session.remote_ip).as_bytes(), + KV_RATE_LIMIT_CONTACT, + &ip_to_bytes(&session.remote_ip), rate, false, ) diff --git a/crates/jmap/src/api/management/stores.rs b/crates/jmap/src/api/management/stores.rs index 6014f816..6e2a5648 100644 --- a/crates/jmap/src/api/management/stores.rs +++ b/crates/jmap/src/api/management/stores.rs @@ -123,7 +123,7 @@ impl ManageStore for Server { } (Some("purge"), Some("lookup"), id, &Method::GET) => { // Validate the access token - access_token.assert_has_permission(Permission::PurgeLookupStore)?; + access_token.assert_has_permission(Permission::PurgeInMemoryStore)?; let store = if let Some(id) = id { if let Some(store) = self.core.storage.lookups.get(id) { diff --git a/crates/jmap/src/auth/oauth/auth.rs b/crates/jmap/src/auth/oauth/auth.rs index db5cdd49..5b39c6c4 100644 --- a/crates/jmap/src/auth/oauth/auth.rs +++ b/crates/jmap/src/auth/oauth/auth.rs @@ -11,13 +11,14 @@ use common::{ oauth::{CLIENT_ID_MAX_LEN, DEVICE_CODE_LEN, USER_CODE_ALPHABET, USER_CODE_LEN}, AccessToken, }, - Server, + Server, KV_OAUTH, }; use rand::distributions::Standard; use serde::Deserialize; use serde_json::json; use std::future::Future; use store::{ + dispatch::lookup::KeyValue, rand::{distributions::Alphanumeric, thread_rng, Rng}, write::Bincode, Serialize, @@ -120,9 +121,8 @@ impl OAuthApiHandler for Server { .storage .lookup .key_set( - format!("oauth:{client_code}").into_bytes(), - value, - self.core.oauth.oauth_expiry_auth_code.into(), + KeyValue::with_prefix(KV_OAUTH, client_code.as_bytes(), value) + .expires(self.core.oauth.oauth_expiry_auth_code), ) .await?; @@ -147,7 +147,10 @@ impl OAuthApiHandler for Server { .core .storage .lookup - .key_get::>(format!("oauth:{code}").into_bytes()) + .key_get::>(KeyValue::<()>::build_key( + KV_OAUTH, + code.as_bytes(), + )) .await? { if auth_code.inner.status == OAuthStatus::Pending { @@ -160,7 +163,7 @@ impl OAuthApiHandler for Server { self.core .storage .lookup - .key_delete(format!("oauth:{code}").into_bytes()) + .key_delete(KeyValue::<()>::build_key(KV_OAUTH, code.as_bytes())) .await?; // Update device code status @@ -168,9 +171,12 @@ impl OAuthApiHandler for Server { .storage .lookup .key_set( - format!("oauth:{device_code}").into_bytes(), - auth_code.serialize(), - self.core.oauth.oauth_expiry_auth_code.into(), + KeyValue::with_prefix( + KV_OAUTH, + device_code.as_bytes(), + auth_code.serialize(), + ) + .expires(self.core.oauth.oauth_expiry_auth_code), ) .await?; } @@ -238,9 +244,8 @@ impl OAuthApiHandler for Server { .storage .lookup .key_set( - format!("oauth:{device_code}").into_bytes(), - oauth_code.clone(), - self.core.oauth.oauth_expiry_user_code.into(), + KeyValue::with_prefix(KV_OAUTH, device_code.as_bytes(), oauth_code.clone()) + .expires(self.core.oauth.oauth_expiry_user_code), ) .await?; @@ -249,9 +254,8 @@ impl OAuthApiHandler for Server { .storage .lookup .key_set( - format!("oauth:{user_code}").into_bytes(), - oauth_code, - self.core.oauth.oauth_expiry_user_code.into(), + KeyValue::with_prefix(KV_OAUTH, user_code.as_bytes(), oauth_code) + .expires(self.core.oauth.oauth_expiry_user_code), ) .await?; diff --git a/crates/jmap/src/auth/rate_limit.rs b/crates/jmap/src/auth/rate_limit.rs index 6d7be622..bfe62981 100644 --- a/crates/jmap/src/auth/rate_limit.rs +++ b/crates/jmap/src/auth/rate_limit.rs @@ -7,8 +7,10 @@ use std::{net::IpAddr, sync::Arc}; use common::{ + ip_to_bytes, listener::limiter::{ConcurrencyLimiter, InFlight}, - ConcurrencyLimiters, Server, + ConcurrencyLimiters, Server, KV_RATE_LIMIT_HTTP_ANONYM, KV_RATE_LIMIT_JMAP, + KV_RATE_LIMIT_JMAP_AUTH, }; use directory::Permission; use trc::AddContext; @@ -59,7 +61,8 @@ impl RateLimiter for Server { .storage .lookup .is_rate_allowed( - format!("j:{}", access_token.primary_id).as_bytes(), + KV_RATE_LIMIT_JMAP, + &access_token.primary_id.to_be_bytes(), rate, false, ) @@ -91,7 +94,7 @@ impl RateLimiter for Server { .core .storage .lookup - .is_rate_allowed(format!("jreq:{}", addr).as_bytes(), rate, false) + .is_rate_allowed(KV_RATE_LIMIT_HTTP_ANONYM, &ip_to_bytes(addr), rate, false) .await .caused_by(trc::location!())? .is_some() @@ -122,7 +125,7 @@ impl RateLimiter for Server { .core .storage .lookup - .is_rate_allowed(format!("jauth:{}", addr).as_bytes(), rate, true) + .is_rate_allowed(KV_RATE_LIMIT_JMAP_AUTH, &ip_to_bytes(addr), rate, true) .await .caused_by(trc::location!())? .is_some() @@ -139,7 +142,7 @@ impl RateLimiter for Server { .core .storage .lookup - .is_rate_allowed(format!("jauth:{}", addr).as_bytes(), rate, false) + .is_rate_allowed(KV_RATE_LIMIT_JMAP_AUTH, &ip_to_bytes(addr), rate, false) .await .caused_by(trc::location!())? .is_some() diff --git a/crates/jmap/src/email/delete.rs b/crates/jmap/src/email/delete.rs index 5d587aa3..3ad55545 100644 --- a/crates/jmap/src/email/delete.rs +++ b/crates/jmap/src/email/delete.rs @@ -6,7 +6,7 @@ use std::time::Duration; -use common::Server; +use common::{Server, KV_LOCK_PURGE_ACCOUNT}; use jmap_proto::types::{ collection::Collection, id::Id, keyword::Keyword, property::Property, state::StateChange, type_state::DataType, @@ -255,21 +255,12 @@ impl EmailDeletion for Server { .core .storage .lookup - .counter_incr( - format!("purge:{account_id}").into_bytes(), - 1, - Some(3600), - true, - ) + .try_lock(KV_LOCK_PURGE_ACCOUNT, &account_id.to_be_bytes(), 3600) .await { - Ok(1) => (), - Ok(count) => { - trc::event!( - Purge(trc::PurgeEvent::PurgeActive), - AccountId = account_id, - Total = count, - ); + Ok(true) => (), + Ok(false) => { + trc::event!(Purge(trc::PurgeEvent::PurgeActive), AccountId = account_id,); return; } Err(err) => { diff --git a/crates/jmap/src/services/housekeeper.rs b/crates/jmap/src/services/housekeeper.rs index fa65dddb..0be886e0 100644 --- a/crates/jmap/src/services/housekeeper.rs +++ b/crates/jmap/src/services/housekeeper.rs @@ -299,7 +299,7 @@ pub fn spawn_housekeeper(inner: Arc, mut rx: mpsc::Receiver, mut rx: mpsc::Receiver { ("blob", store.purge_blobs(blob_store).await) } - PurgeStore::Lookup(lookup_store) => { - ("lookup", lookup_store.purge_lookup_store().await) - } + PurgeStore::Lookup(in_memory_store) => ( + "lookup", + in_memory_store.purge_in_memory_store().await, + ), }; match result { diff --git a/crates/main/Cargo.toml b/crates/main/Cargo.toml index 32c58976..639c020c 100644 --- a/crates/main/Cargo.toml +++ b/crates/main/Cargo.toml @@ -22,6 +22,7 @@ jmap_proto = { path = "../jmap-proto" } smtp = { path = "../smtp" } imap = { path = "../imap" } pop3 = { path = "../pop3" } +spam-filter = { path = "../spam-filter" } managesieve = { path = "../managesieve" } common = { path = "../common" } directory = { path = "../directory" } @@ -45,4 +46,9 @@ elastic = ["store/elastic"] s3 = ["store/s3"] redis = ["store/redis"] azure = ["store/azure"] -enterprise = ["jmap/enterprise", "common/enterprise", "store/enterprise", "managesieve/enterprise", "directory/enterprise"] +enterprise = [ "jmap/enterprise", + "common/enterprise", + "store/enterprise", + "managesieve/enterprise", + "directory/enterprise", + "spam-filter/enterprise"] diff --git a/crates/managesieve/src/core/client.rs b/crates/managesieve/src/core/client.rs index 159669a0..0fe1b703 100644 --- a/crates/managesieve/src/core/client.rs +++ b/crates/managesieve/src/core/client.rs @@ -4,7 +4,10 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use common::listener::{SessionResult, SessionStream}; +use common::{ + listener::{SessionResult, SessionStream}, + KV_RATE_LIMIT_IMAP, +}; use imap_proto::receiver::{self, Request}; use jmap_proto::types::{collection::Collection, property::Property}; use store::query::Filter; @@ -186,7 +189,8 @@ impl Session { .storage .lookup .is_rate_allowed( - format!("ireq:{}", access_token.primary_id()).as_bytes(), + KV_RATE_LIMIT_IMAP, + &access_token.primary_id().to_be_bytes(), rate, true, ) diff --git a/crates/nlp/Cargo.toml b/crates/nlp/Cargo.toml index 865f40a0..66a3a734 100644 --- a/crates/nlp/Cargo.toml +++ b/crates/nlp/Cargo.toml @@ -11,15 +11,16 @@ siphasher = "1.0" serde = { version = "1.0", features = ["derive"]} bincode = "1.3.3" nohash = "0.2.0" -ahash = "0.8.3" +ahash = { version = "0.8.3", features = ["serde"] } whatlang = "0.16" # Language detection rust-stemmers = "1.2" # Stemmers -tinysegmenter = "0.1" # Japanese tokenizer jieba-rs = "0.7" # Chinese stemmer phf = { version = "0.11", features = ["macros"] } lru-cache = "0.1.2" parking_lot = "0.12.1" psl = "2" +radix_trie = "0.2.1" +maplit = "1.0.2" [features] test_mode = [] diff --git a/crates/nlp/src/bayes/cache.rs b/crates/nlp/src/bayes/cache.rs deleted file mode 100644 index af1b1234..00000000 --- a/crates/nlp/src/bayes/cache.rs +++ /dev/null @@ -1,112 +0,0 @@ -/* - * SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd - * - * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL - */ - -use std::{ - hash::BuildHasherDefault, - time::{Duration, Instant}, -}; - -use lru_cache::LruCache; -use nohash::NoHashHasher; -use parking_lot::Mutex; - -use super::{TokenHash, Weights}; - -#[derive(Debug)] -pub struct BayesTokenCache { - positive: Mutex>>>, - negative: Mutex>>>, - ttl_negative: Duration, - ttl_positive: Duration, -} - -#[derive(Debug, Clone)] -pub struct CacheItem { - item: Weights, - valid_until: Instant, -} - -impl BayesTokenCache { - pub fn new(capacity: usize, ttl_positive: Duration, ttl_negative: Duration) -> Self { - Self { - positive: Mutex::new(LruCache::with_hasher(capacity, Default::default())), - negative: Mutex::new(LruCache::with_hasher(capacity, Default::default())), - ttl_negative, - ttl_positive, - } - } - - pub fn get(&self, hash: &TokenHash) -> Option> { - { - let mut pos_cache = self.positive.lock(); - if let Some(entry) = pos_cache.get_mut(hash) { - return if entry.valid_until >= Instant::now() { - Some(Some(entry.item)) - } else { - pos_cache.remove(hash); - None - }; - } - } - { - let mut neg_cache = self.negative.lock(); - if let Some(entry) = neg_cache.get_mut(hash) { - return if *entry >= Instant::now() { - Some(None) - } else { - neg_cache.remove(hash); - None - }; - } - } - - None - } - - pub fn insert_positive(&self, hash: TokenHash, weights: Weights) { - self.positive.lock().insert( - hash, - CacheItem { - item: weights, - valid_until: Instant::now() + self.ttl_positive, - }, - ); - } - - pub fn insert_negative(&self, hash: TokenHash) { - self.negative - .lock() - .insert(hash, Instant::now() + self.ttl_negative); - } - - pub fn invalidate(&self, hash: &TokenHash) { - if self.positive.lock().remove(hash).is_none() { - self.negative.lock().remove(hash); - } - } -} - -impl Default for BayesTokenCache { - fn default() -> Self { - Self { - positive: Mutex::new(LruCache::with_hasher(1024, Default::default())), - negative: Mutex::new(LruCache::with_hasher(1024, Default::default())), - ttl_negative: Default::default(), - ttl_positive: Default::default(), - } - } -} - -impl Clone for BayesTokenCache { - fn clone(&self) -> Self { - Self { - positive: Mutex::new(self.positive.lock().clone()), - negative: Mutex::new(self.negative.lock().clone()), - ttl_negative: self.ttl_negative, - ttl_positive: self.ttl_positive, - } - } -} diff --git a/crates/nlp/src/bayes/mod.rs b/crates/nlp/src/bayes/mod.rs index 7516044e..da77bb5c 100644 --- a/crates/nlp/src/bayes/mod.rs +++ b/crates/nlp/src/bayes/mod.rs @@ -4,21 +4,19 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use std::{collections::HashMap, hash::BuildHasherDefault}; - -use nohash::NoHashHasher; +use ahash::AHashMap; +use radix_trie::TrieKey; use serde::{Deserialize, Serialize}; use crate::tokenizers::osb::Gram; -pub mod cache; pub mod classify; pub mod tokenize; pub mod train; #[derive(Debug, Serialize, Deserialize, Default)] pub struct BayesModel { - pub weights: HashMap>>, + pub weights: AHashMap, pub spam_learns: u32, pub ham_learns: u32, } @@ -29,13 +27,15 @@ pub struct BayesClassifier { pub min_tokens: u32, pub min_prob_strength: f64, pub min_learns: u32, + pub min_balance: f64, } -#[derive(Debug, Serialize, Deserialize, Default, Copy, Clone, PartialEq, Eq)] +#[derive(Debug, Serialize, Deserialize, Default, Copy, Clone, PartialEq, Eq, Hash)] pub struct TokenHash { - pub h1: u64, - pub h2: u64, + hash: [u8; HASH_LEN], + len: u8, } +const HASH_LEN: usize = std::mem::size_of::() * 2; #[derive(Debug, Serialize, Deserialize, Default, Copy, Clone, Hash, PartialEq, Eq)] pub struct Weights { @@ -50,6 +50,7 @@ impl BayesClassifier { min_tokens: 11, min_prob_strength: 0.05, min_learns: 200, + min_balance: 0.1, } } } @@ -62,33 +63,70 @@ impl Default for BayesClassifier { impl From> for TokenHash { fn from(value: Gram<'_>) -> Self { + let mut hash = TokenHash { + hash: [0; HASH_LEN], + len: 0, + }; + match value { - Gram::Uni { t1 } => TokenHash { - h1: xxhash_rust::xxh3::xxh3_64(t1.as_bytes()), - h2: farmhash::hash64(t1.as_bytes()), - }, + Gram::Uni { t1 } => { + if t1.len() <= HASH_LEN { + hash.hash[..t1.len()].copy_from_slice(t1); + hash.len = t1.len() as u8; + } else { + let h1 = xxhash_rust::xxh3::xxh3_64(t1).to_be_bytes(); + let h2 = farmhash::hash64(t1).to_be_bytes(); + hash.hash[..std::mem::size_of::()].copy_from_slice(&h1); + hash.hash[std::mem::size_of::()..].copy_from_slice(&h2); + hash.len = HASH_LEN as u8; + } + } Gram::Bi { t1, t2, .. } => { - let mut buf = Vec::with_capacity(t1.len() + t2.len() + 1); - buf.extend_from_slice(t1.as_bytes()); - buf.push(b' '); - buf.extend_from_slice(t2.as_bytes()); - TokenHash { - h1: xxhash_rust::xxh3::xxh3_64(&buf), - h2: farmhash::hash64(&buf), + let len = t1.len() + t2.len() + 1; + if len <= HASH_LEN { + for (h, b) in hash.hash.iter_mut().zip( + t1.iter() + .copied() + .chain([b' '].into_iter()) + .chain(t2.iter().copied()), + ) { + *h = b; + } + hash.len = len as u8; + } else if t1.len() <= std::mem::size_of::() { + for (h, b) in hash.hash.iter_mut().zip( + t1.iter() + .copied() + .chain(xxhash_rust::xxh3::xxh3_64(t2).to_be_bytes().into_iter()) + .chain(farmhash::hash64(t2).to_be_bytes().into_iter()), + ) { + *h = b; + } + hash.len = HASH_LEN as u8; + } else { + let mut buf = Vec::with_capacity(t1.len() + t2.len() + 1); + buf.extend_from_slice(t1); + buf.push(b' '); + buf.extend_from_slice(t2); + let h1 = xxhash_rust::xxh3::xxh3_64(&buf).to_be_bytes(); + let h2 = farmhash::fingerprint64(&buf).to_be_bytes(); + hash.hash[..std::mem::size_of::()].copy_from_slice(&h1); + hash.hash[std::mem::size_of::()..].copy_from_slice(&h2); + hash.len = HASH_LEN as u8; } } } + + hash } } -impl std::hash::Hash for TokenHash { - fn hash(&self, state: &mut H) { - state.write_u64(self.h1 ^ self.h2); +impl TrieKey for TokenHash { + fn encode_bytes(&self) -> Vec { + self.hash[..self.len as usize].to_vec() } } -impl nohash::IsEnabled for TokenHash {} - impl From for Weights { fn from(value: i64) -> Self { Weights { @@ -103,3 +141,31 @@ impl From for i64 { (value.ham as i64) << 32 | value.spam as i64 } } + +impl TokenHash { + pub fn serialize_index(prefix: u8, account_id: Option) -> Vec { + if let Some(account_id) = account_id { + let mut buf = Vec::with_capacity(std::mem::size_of::() + 1); + buf.push(prefix); + buf.extend_from_slice(&account_id.to_be_bytes()); + buf + } else { + vec![prefix] + } + } + + pub fn serialize(&self, prefix: u8, account_id: Option) -> Vec { + if let Some(account_id) = account_id { + let mut buf = Vec::with_capacity(std::mem::size_of::() + self.len as usize + 1); + buf.push(prefix); + buf.extend_from_slice(&account_id.to_be_bytes()); + buf.extend_from_slice(&self.hash[..self.len as usize]); + buf + } else { + let mut buf = Vec::with_capacity(self.len as usize + 1); + buf.push(prefix); + buf.extend_from_slice(&self.hash[..self.len as usize]); + buf + } + } +} diff --git a/crates/nlp/src/bayes/tokenize.rs b/crates/nlp/src/bayes/tokenize.rs index fd59ac90..0c4402fd 100644 --- a/crates/nlp/src/bayes/tokenize.rs +++ b/crates/nlp/src/bayes/tokenize.rs @@ -4,7 +4,7 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use std::borrow::Cow; +use std::{borrow::Cow, net::IpAddr}; use crate::{ language::{ @@ -13,18 +13,19 @@ use crate::{ stopwords::STOP_WORDS, Language, }, - tokenizers::{ - chinese::JIEBA, - types::{TokenType, TypesTokenizer}, - }, + tokenizers::{chinese::JIEBA, japanese, types::TokenType}, }; -pub struct BayesTokenizer<'x> { - text: &'x str, - tokenizer: TypesTokenizer<'x>, +pub struct BayesTokenizer> { + stream: T, stemmer: Stemmer, stop_words: Option<&'static phf::Set<&'static str>>, - tokens: Vec>, + tokens: Vec>, +} + +pub enum BayesInputToken { + Word(String), + Raw(Vec), } enum Stemmer { @@ -34,8 +35,8 @@ enum Stemmer { None, } -impl<'x> BayesTokenizer<'x> { - pub fn new(text: &'x str) -> Self { +impl> BayesTokenizer { + pub fn new(text: &str, stream: T) -> Self { // Detect language let (mut language, score) = LanguageDetector::detect_single(text).unwrap_or((Language::English, 1.0)); @@ -44,8 +45,7 @@ impl<'x> BayesTokenizer<'x> { } Self { - text, - tokenizer: TypesTokenizer::new(text), + stream, stemmer: match language { Language::Mandarin => Stemmer::Mandarin, Language::Japanese => Stemmer::Japanese, @@ -59,20 +59,17 @@ impl<'x> BayesTokenizer<'x> { } } -impl<'x> Iterator for BayesTokenizer<'x> { - type Item = Cow<'x, str>; +impl> Iterator for BayesTokenizer { + type Item = Vec; fn next(&mut self) -> Option { if let Some(prev_token) = self.tokens.pop() { return Some(prev_token); } - loop { - let token = self.tokenizer.next()?; - - let word: Cow = match token.word { - TokenType::Alphabetic(word) => { - let word = word.to_lowercase(); + for token in self.stream.by_ref() { + return match token { + BayesInputToken::Word(word) => { if self .stop_words .map_or(false, |sw| sw.contains(word.as_str())) @@ -81,8 +78,8 @@ impl<'x> Iterator for BayesTokenizer<'x> { } match &self.stemmer { Stemmer::IndoEuropean(stemmer) => match stemmer.stem(&word) { - Cow::Borrowed(_) => word.into(), - Cow::Owned(stemmed_word) => stemmed_word.into(), + Cow::Borrowed(_) => word.into_bytes(), + Cow::Owned(stemmed_word) => stemmed_word.into_bytes(), }, Stemmer::Mandarin => { let mut result = JIEBA.cut(&word, false).into_iter(); @@ -90,77 +87,90 @@ impl<'x> Iterator for BayesTokenizer<'x> { let stemmed_word = stemmed_word.to_string(); self.tokens = result .rev() - .map(|word| Cow::from(word.to_string())) + .map(|word| word.to_string().into_bytes()) .collect::>(); - stemmed_word.into() + stemmed_word.into_bytes() } else { // This shouldn't happen, but just in case continue; } } Stemmer::Japanese => { - let mut result = tinysegmenter::tokenize(&word).into_iter(); + let mut result = japanese::tokenize(&word).into_iter(); if let Some(stemmed_word) = result.next() { - self.tokens = result.rev().map(Cow::from).collect::>(); - stemmed_word.into() + self.tokens = + result.rev().map(|b| b.into_bytes()).collect::>(); + stemmed_word.into_bytes() } else { // This shouldn't happen, but just in case continue; } } - Stemmer::None => word.into(), + Stemmer::None => word.into_bytes(), } } + BayesInputToken::Raw(raw) => raw, + } + .into(); + } - TokenType::Url(word) => { - if let Some((_, host)) = word.split_once("://") { - host.split_once('/') - .map_or(host, |(h, _)| h) - .to_lowercase() - .into() - } else { - continue; - } - } - TokenType::IpAddr(word) => word.into(), - TokenType::UrlNoScheme(word) => word - .split_once('/') - .map_or(word, |(h, _)| h) - .to_lowercase() - .into(), - TokenType::Alphanumeric(word) - | TokenType::Email(word) - | TokenType::UrlNoHost(word) => word.to_lowercase().into(), - TokenType::Other(ch) => { - if SYMBOLS.contains(&ch) { - (&self.text[token.from..token.to]).into() - } else { - continue; - } - } - TokenType::Integer(word) => number_to_tag("INTEGER", word).into(), - TokenType::Float(word) => number_to_tag("FLOAT", word).into(), - TokenType::Punctuation(_) | TokenType::Space => { - continue; - } - }; + None + } +} - return Some(word); +impl> TokenType { + pub fn to_bayes_token(&self) -> Option { + match self { + TokenType::Alphabetic(word) => { + Some(BayesInputToken::Word(word.as_ref().to_lowercase())) + } + TokenType::Url(word) => { + let word = word.as_ref(); + word.split_once("://") + .map(|(_, host)| BayesInputToken::Raw(url_host_as_bytes(host))) + } + TokenType::IpAddr(word) => word.as_ref().parse::().ok().map(|ip| { + BayesInputToken::Raw(match ip { + IpAddr::V4(ip) => ip.octets().to_vec(), + IpAddr::V6(ip) => ip.octets().to_vec(), + }) + }), + TokenType::UrlNoScheme(word) => { + BayesInputToken::Raw(url_host_as_bytes(word.as_ref())).into() + } + TokenType::Alphanumeric(word) | TokenType::Email(word) | TokenType::UrlNoHost(word) => { + BayesInputToken::Raw(word.as_ref().to_lowercase().into_bytes()).into() + } + TokenType::Other(ch) => { + if SYMBOLS.contains(ch) { + Some(BayesInputToken::Raw(ch.to_string().into_bytes())) + } else { + None + } + } + TokenType::Integer(word) => number_to_tag(false, word.as_ref()).into(), + TokenType::Float(word) => number_to_tag(true, word.as_ref()).into(), + TokenType::Punctuation(_) | TokenType::Space => None, } } } -fn number_to_tag(prefix: &str, num: &str) -> String { - format!( - "{}_{}_{}", - prefix, - if prefix.starts_with('-') { - "NEG" - } else { - "POS" - }, - num.len() - ) +fn url_host_as_bytes(host: &str) -> Vec { + host.split_once('/') + .map_or(host, |(h, _)| h.rsplit_once(':').map_or(h, |(h, _)| h)) + .to_lowercase() + .into_bytes() +} + +fn number_to_tag(is_float: bool, num: &str) -> BayesInputToken { + let t = match (is_float, num.starts_with('-')) { + (true, true) => b'F', + (true, false) => b'f', + (false, true) => b'I', + (false, false) => b'i', + }; + + BayesInputToken::Raw([t, num.len() as u8].to_vec()) } pub static SYMBOLS: phf::Set = phf::phf_set! { @@ -1147,7 +1157,7 @@ pub static SYMBOLS: phf::Set = phf::phf_set! { mod tests { use std::borrow::Cow; - use crate::bayes::tokenize::BayesTokenizer; + use crate::{bayes::tokenize::BayesTokenizer, tokenizers::types::TypesTokenizer}; #[test] fn bayes_tokenizer() { @@ -1233,7 +1243,12 @@ mod tests { ]; for (input, expect) in inputs.iter() { - let input = BayesTokenizer::new(input).collect::>(); + let input = BayesTokenizer::new( + input, + TypesTokenizer::new(input).filter_map(|t| t.word.to_bayes_token()), + ) + .map(|word| String::from_utf8(word).unwrap()) + .collect::>(); let expect = expect.iter().copied().map(Cow::from).collect::>(); assert_eq!(input, expect,); diff --git a/crates/nlp/src/bayes/train.rs b/crates/nlp/src/bayes/train.rs index c2281f4e..c200d329 100644 --- a/crates/nlp/src/bayes/train.rs +++ b/crates/nlp/src/bayes/train.rs @@ -20,12 +20,16 @@ impl BayesModel { } for token in tokens { - let hs = self.weights.entry(token.inner).or_default(); - if is_spam { - hs.spam += 1; - } else { - hs.ham += 1; - } + self.train_token(token.inner, is_spam); + } + } + + pub fn train_token(&mut self, token: TokenHash, is_spam: bool) { + let hs = self.weights.entry(token).or_default(); + if is_spam { + hs.spam += 1; + } else { + hs.ham += 1; } } diff --git a/crates/nlp/src/lib.rs b/crates/nlp/src/lib.rs index 68a40005..760ca635 100644 --- a/crates/nlp/src/lib.rs +++ b/crates/nlp/src/lib.rs @@ -8,25 +8,37 @@ mod test { use crate::{ bayes::{tokenize::BayesTokenizer, BayesClassifier, BayesModel}, - tokenizers::osb::{OsbToken, OsbTokenizer}, + tokenizers::{ + osb::{OsbToken, OsbTokenizer}, + types::TypesTokenizer, + }, }; #[test] #[ignore] fn train() { - let db = - fs::read_to_string("/Users/me/code/mail-server/_ignore/spam_or_not_spam.csv").unwrap(); + let db = fs::read_to_string("/Users/me/code/mail-server/_ignore/old/spam_or_not_spam.csv") + .unwrap(); let mut bayes = BayesModel::default(); for line in db.lines() { let (text, is_spam) = line.rsplit_once(',').unwrap(); let is_spam = is_spam == "1"; - bayes.train(OsbTokenizer::new(BayesTokenizer::new(text), 5), is_spam); + bayes.train( + OsbTokenizer::new( + BayesTokenizer::new( + text, + TypesTokenizer::new(text).filter_map(|t| t.word.to_bayes_token()), + ), + 5, + ), + is_spam, + ); } println!("Ham: {} Spam: {}", bayes.ham_learns, bayes.spam_learns,); fs::write( - "/Users/me/code/mail-server/_ignore/spam_or_not_spam.bin", + "/Users/me/code/mail-server/_ignore/old/spam_or_not_spam.bin", bincode::serialize(&bayes).unwrap(), ) .unwrap(); @@ -36,27 +48,46 @@ mod test { #[ignore] fn classify() { let model: BayesModel = bincode::deserialize( - &fs::read("/Users/me/code/mail-server/_ignore/spam_or_not_spam.bin").unwrap(), + &fs::read("/Users/me/code/mail-server/_ignore/old/spam_or_not_spam.bin").unwrap(), ) .unwrap(); let bayes = BayesClassifier::new(); for text in [ - "i am attaching to this email a presentation to integrate the spreadsheet into our server", + concat!( + "i am attaching to this email a presentation to integrate the ", + "spreadsheet into our server and obtain the data from the database" + ), "buy this great product special offer sales", - "i m using simple dns from jhsoft we support only a few web sites and i d like to swap secondary services with someone in a similar position", + concat!( + "i m using simple dns from jhsoft we support only a few web sites ", + "and i d like to swap secondary services with someone in a similar position" + ), "viagra xenical vioxx zyban propecia we only offer the real viagra xenical ", ] { println!( - "{:?} -> {}", + "{:?} -> {:?}", text, bayes - .classify(OsbTokenizer::new(BayesTokenizer::new(text), 5).filter_map(|x| model.weights.get(&x.inner).map(|w| { - OsbToken { - idx: x.idx, - inner: *w, - } - })), model.ham_learns, model.spam_learns) + .classify( + OsbTokenizer::new( + BayesTokenizer::new( + text, + TypesTokenizer::new(text).filter_map(|t| t.word.to_bayes_token()) + ), + 5 + ) + .filter_map(|x| model.weights.get(&x.inner).map( + |w| { + OsbToken { + idx: x.idx, + inner: *w, + } + } + )), + model.ham_learns, + model.spam_learns + ) .unwrap() ); } diff --git a/crates/nlp/src/tokenizers/japanese.rs b/crates/nlp/src/tokenizers/japanese.rs index 451d4b10..5f9fc15d 100644 --- a/crates/nlp/src/tokenizers/japanese.rs +++ b/crates/nlp/src/tokenizers/japanese.rs @@ -4,9 +4,11 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use std::vec::IntoIter; - use super::{InnerToken, Token}; +use maplit::hashmap; +use std::collections::HashMap; +use std::vec::IntoIter; +use std::{hash::Hash, sync::LazyLock}; pub struct JapaneseTokenizer<'x, T, I> where @@ -47,7 +49,7 @@ where let token = self.tokenizer.next()?; if token.word.is_alphabetic_8bit() { let mut token_to = token.from; - self.tokens = tinysegmenter::tokenize(token.word.unwrap_alphabetic().as_ref()) + self.tokens = tokenize(token.word.unwrap_alphabetic().as_ref()) .into_iter() .map(|word| { let token_from = token_to; @@ -68,6 +70,249 @@ where } } +// Ported from https://github.com/woxtu/rust-tinysegmenter, MIT license + +const BIAS: i32 = -332; + +fn get_score(d: &HashMap, s: &T) -> i32 { + d.get(s).cloned().unwrap_or(0) +} + +fn get_ctype(c: char) -> char { + match c as u32 { + 0x4E00 | 0x4E8C | 0x4E09 | 0x56DB | 0x4E94 | 0x516D | 0x4E03 | 0x516B | 0x4E5D | 0x5341 => { + 'M' + } + 0x767E | 0x5343 | 0x4E07 | 0x5104 | 0x5146 => 'M', + 0x4E00..=0x9FA0 | 0x3005 | 0x3006 | 0x30F5 | 0x30F6 => 'H', + 0x3041..=0x3093 => 'I', + 0x30A1..=0x30F4 | 0x30FC | 0xFF71..=0xFF9D | 0xFF9E | 0xFF70 => 'K', + 0x61..=0x7A | 0x41..=0x5A | 0xFF41..=0xFF5A | 0xFF21..=0xFF3A => 'A', + 0x30..=0x3a | 0xFF10..=0xFF19 => 'N', + _ => 'O', + } +} + +pub fn tokenize(s: &str) -> Vec { + if s.is_empty() { + return Vec::new(); + } + + let mut result = Vec::with_capacity(s.chars().count()); + + let segments = [B3, B2, B1] + .into_iter() + .chain(s.chars()) + .chain([E1, E2, E3]) + .collect::>(); + + let ctypes = ['O'; 3] + .into_iter() + .chain(s.chars().map(get_ctype)) + .chain(['O'; 3]) + .collect::>(); + + let mut word = segments[3].to_string(); + let mut p = vec!['U'; 3]; + + for index in 4..segments.len() - 3 { + let mut score = BIAS; + let w = &segments[index - 3..index + 3]; + let c = &ctypes[index - 3..index + 3]; + + score += get_score(&*UP1, &p[0]); + score += get_score(&*UP2, &p[1]); + score += get_score(&*UP3, &p[2]); + score += get_score(&*BP1, &(p[0], p[1])); + score += get_score(&*BP2, &(p[1], p[2])); + score += get_score(&*UW1, &w[0]); + score += get_score(&*UW2, &w[1]); + score += get_score(&*UW3, &w[2]); + score += get_score(&*UW4, &w[3]); + score += get_score(&*UW5, &w[4]); + score += get_score(&*UW6, &w[5]); + score += get_score(&*BW1, &(w[1], w[2])); + score += get_score(&*BW2, &(w[2], w[3])); + score += get_score(&*BW3, &(w[3], w[4])); + score += get_score(&*TW1, &(w[0], w[1], w[2])); + score += get_score(&*TW2, &(w[1], w[2], w[3])); + score += get_score(&*TW3, &(w[2], w[3], w[4])); + score += get_score(&*TW4, &(w[3], w[4], w[5])); + score += get_score(&*UC1, &c[0]); + score += get_score(&*UC2, &c[1]); + score += get_score(&*UC3, &c[2]); + score += get_score(&*UC4, &c[3]); + score += get_score(&*UC5, &c[4]); + score += get_score(&*UC6, &c[5]); + score += get_score(&*BC1, &(c[1], c[2])); + score += get_score(&*BC2, &(c[2], c[3])); + score += get_score(&*BC3, &(c[3], c[4])); + score += get_score(&*TC1, &(c[0], c[1], c[2])); + score += get_score(&*TC2, &(c[1], c[2], c[3])); + score += get_score(&*TC3, &(c[2], c[3], c[4])); + score += get_score(&*TC4, &(c[3], c[4], c[5])); + score += get_score(&*UQ1, &(p[0], c[0])); + score += get_score(&*UQ2, &(p[1], c[1])); + score += get_score(&*UQ3, &(p[2], c[2])); + score += get_score(&*BQ1, &(p[1], c[1], c[2])); + score += get_score(&*BQ2, &(p[1], c[2], c[3])); + score += get_score(&*BQ3, &(p[2], c[1], c[2])); + score += get_score(&*BQ4, &(p[2], c[2], c[3])); + score += get_score(&*TQ1, &(p[1], c[0], c[1], c[2])); + score += get_score(&*TQ2, &(p[1], c[1], c[2], c[3])); + score += get_score(&*TQ3, &(p[2], c[0], c[1], c[2])); + score += get_score(&*TQ4, &(p[2], c[1], c[2], c[3])); + + p.remove(0); + p.push(if score < 0 { 'O' } else { 'B' }); + + if 0 < score { + result.push(word.clone()); + word.clear(); + } + word.push(segments[index]); + } + + result.push(word.clone()); + result +} + +const B1: char = '\u{F0000}'; +const B2: char = '\u{F0001}'; +const B3: char = '\u{F0002}'; +const E1: char = '\u{F0003}'; +const E2: char = '\u{F0004}'; +const E3: char = '\u{F0005}'; + +static BC1: LazyLock> = LazyLock::new(|| { + hashmap! { ('H', 'H') => 6, ('I', 'I') => 2461, ('K', 'H') => 406, ('O', 'H') => -1378, } +}); +static BC2: LazyLock> = LazyLock::new(|| { + hashmap! { ('A', 'A') => -3267, ('A', 'I') => 2744, ('A', 'N') => -878, ('H', 'H') => -4070, ('H', 'M') => -1711, ('H', 'N') => 4012, ('H', 'O') => 3761, ('I', 'A') => 1327, ('I', 'H') => -1184, ('I', 'I') => -1332, ('I', 'K') => 1721, ('I', 'O') => 5492, ('K', 'I') => 3831, ('K', 'K') => -8741, ('M', 'H') => -3132, ('M', 'K') => 3334, ('O', 'O') => -2920, } +}); +static BC3: LazyLock> = LazyLock::new(|| { + hashmap! { ('H', 'H') => 996, ('H', 'I') => 626, ('H', 'K') => -721, ('H', 'N') => -1307, ('H', 'O') => -836, ('I', 'H') => -301, ('K', 'K') => 2762, ('M', 'K') => 1079, ('M', 'M') => 4034, ('O', 'A') => -1652, ('O', 'H') => 266, } +}); +static BP1: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'B') => 295, ('O', 'B') => 304, ('O', 'O') => -125, ('U', 'B') => 352, } +}); +static BP2: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'O') => 60, ('O', 'O') => -1762, } +}); +static BQ1: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'H', 'H') => 1150, ('B', 'H', 'M') => 1521, ('B', 'I', 'I') => -1158, ('B', 'I', 'M') => 886, ('B', 'M', 'H') => 1208, ('B', 'N', 'H') => 449, ('B', 'O', 'H') => -91, ('B', 'O', 'O') => -2597, ('O', 'H', 'I') => 451, ('O', 'I', 'H') => -296, ('O', 'K', 'A') => 1851, ('O', 'K', 'H') => -1020, ('O', 'K', 'K') => 904, ('O', 'O', 'O') => 2965, } +}); +static BQ2: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'H', 'H') => 118, ('B', 'H', 'I') => -1159, ('B', 'H', 'M') => 466, ('B', 'I', 'H') => -919, ('B', 'K', 'K') => -1720, ('B', 'K', 'O') => 864, ('O', 'H', 'H') => -1139, ('O', 'H', 'M') => -181, ('O', 'I', 'H') => 153, ('U', 'H', 'I') => -1146, } +}); +static BQ3: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'H', 'H') => -792, ('B', 'H', 'I') => 2664, ('B', 'I', 'I') => -299, ('B', 'K', 'I') => 419, ('B', 'M', 'H') => 937, ('B', 'M', 'M') => 8335, ('B', 'N', 'N') => 998, ('B', 'O', 'H') => 775, ('O', 'H', 'H') => 2174, ('O', 'H', 'M') => 439, ('O', 'I', 'I') => 280, ('O', 'K', 'H') => 1798, ('O', 'K', 'I') => -793, ('O', 'K', 'O') => -2242, ('O', 'M', 'H') => -2402, ('O', 'O', 'O') => 11699, } +}); +static BQ4: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'H', 'H') => -3895, ('B', 'I', 'H') => 3761, ('B', 'I', 'I') => -4654, ('B', 'I', 'K') => 1348, ('B', 'K', 'K') => -1806, ('B', 'M', 'I') => -3385, ('B', 'O', 'O') => -12396, ('O', 'A', 'H') => 926, ('O', 'H', 'H') => 266, ('O', 'H', 'K') => -2036, ('O', 'N', 'N') => -973, } +}); +static BW1: LazyLock> = LazyLock::new(|| { + hashmap! { (',', 'と') => 660, (',', '同') => 727, (B1, 'あ') => 1404, (B1, '同') => 542, ('、', 'と') => 660, ('、', '同') => 727, ('」', 'と') => 1682, ('あ', 'っ') => 1505, ('い', 'う') => 1743, ('い', 'っ') => -2055, ('い', 'る') => 672, ('う', 'し') => -4817, ('う', 'ん') => 665, ('か', 'ら') => 3472, ('が', 'ら') => 600, ('こ', 'う') => -790, ('こ', 'と') => 2083, ('こ', 'ん') => -1262, ('さ', 'ら') => -4143, ('さ', 'ん') => 4573, ('し', 'た') => 2641, ('し', 'て') => 1104, ('す', 'で') => -3399, ('そ', 'こ') => 1977, ('そ', 'れ') => -871, ('た', 'ち') => 1122, ('た', 'め') => 601, ('っ', 'た') => 3463, ('つ', 'い') => -802, ('て', 'い') => 805, ('て', 'き') => 1249, ('で', 'き') => 1127, ('で', 'す') => 3445, ('で', 'は') => 844, ('と', 'い') => -4915, ('と', 'み') => 1922, ('ど', 'こ') => 3887, ('な', 'い') => 5713, ('な', 'っ') => 3015, ('な', 'ど') => 7379, ('な', 'ん') => -1113, ('に', 'し') => 2468, ('に', 'は') => 1498, ('に', 'も') => 1671, ('に', '対') => -912, ('の', '一') => -501, ('の', '中') => 741, ('ま', 'せ') => 2448, ('ま', 'で') => 1711, ('ま', 'ま') => 2600, ('ま', 'る') => -2155, ('や', 'む') => -1947, ('よ', 'っ') => -2565, ('れ', 'た') => 2369, ('れ', 'で') => -913, ('を', 'し') => 1860, ('を', '見') => 731, ('亡', 'く') => -1886, ('京', '都') => 2558, ('取', 'り') => -2784, ('大', 'き') => -2604, ('大', '阪') => 1497, ('平', '方') => -2314, ('引', 'き') => -1336, ('日', '本') => -195, ('本', '当') => -2423, ('毎', '日') => -2113, ('目', '指') => -724, ('｣', 'と') => 1682, } +}); +static BW2: LazyLock> = LazyLock::new(|| { + hashmap! { ('.', '.') => -11822, ('1', '1') => -669, ('―', '―') => -5730, ('−', '−') => -13175, ('い', 'う') => -1609, ('う', 'か') => 2490, ('か', 'し') => -1350, ('か', 'も') => -602, ('か', 'ら') => -7194, ('か', 'れ') => 4612, ('が', 'い') => 853, ('が', 'ら') => -3198, ('き', 'た') => 1941, ('く', 'な') => -1597, ('こ', 'と') => -8392, ('こ', 'の') => -4193, ('さ', 'せ') => 4533, ('さ', 'れ') => 13168, ('さ', 'ん') => -3977, ('し', 'い') => -1819, ('し', 'か') => -545, ('し', 'た') => 5078, ('し', 'て') => 972, ('し', 'な') => 939, ('そ', 'の') => -3744, ('た', 'い') => -1253, ('た', 'た') => -662, ('た', 'だ') => -3857, ('た', 'ち') => -786, ('た', 'と') => 1224, ('た', 'は') => -939, ('っ', 'た') => 4589, ('っ', 'て') => 1647, ('っ', 'と') => -2094, ('て', 'い') => 6144, ('て', 'き') => 3640, ('て', 'く') => 2551, ('て', 'は') => -3110, ('て', 'も') => -3065, ('で', 'い') => 2666, ('で', 'き') => -1528, ('で', 'し') => -3828, ('で', 'す') => -4761, ('で', 'も') => -4203, ('と', 'い') => 1890, ('と', 'こ') => -1746, ('と', 'と') => -2279, ('と', 'の') => 720, ('と', 'み') => 5168, ('と', 'も') => -3941, ('な', 'い') => -2488, ('な', 'が') => -1313, ('な', 'ど') => -6509, ('な', 'の') => 2614, ('な', 'ん') => 3099, ('に', 'お') => -1615, ('に', 'し') => 2748, ('に', 'な') => 2454, ('に', 'よ') => -7236, ('に', '対') => -14943, ('に', '従') => -4688, ('に', '関') => -11388, ('の', 'か') => 2093, ('の', 'で') => -7059, ('の', 'に') => -6041, ('の', 'の') => -6125, ('は', 'い') => 1073, ('は', 'が') => -1033, ('は', 'ず') => -2532, ('ば', 'れ') => 1813, ('ま', 'し') => -1316, ('ま', 'で') => -6621, ('ま', 'れ') => 5409, ('め', 'て') => -3153, ('も', 'い') => 2230, ('も', 'の') => -10713, ('ら', 'か') => -944, ('ら', 'し') => -1611, ('ら', 'に') => -1897, ('り', 'し') => 651, ('り', 'ま') => 1620, ('れ', 'た') => 4270, ('れ', 'て') => 849, ('れ', 'ば') => 4114, ('ろ', 'う') => 6067, ('わ', 'れ') => 7901, ('を', '通') => -11877, ('ん', 'だ') => 728, ('ん', 'な') => -4115, ('一', '人') => 602, ('一', '方') => -1375, ('一', '日') => 970, ('一', '部') => -1051, ('上', 'が') => -4479, ('会', '社') => -1116, ('出', 'て') => 2163, ('分', 'の') => -7758, ('同', '党') => 970, ('同', '日') => -913, ('大', '阪') => -2471, ('委', '員') => -1250, ('少', 'な') => -1050, ('年', '度') => -8669, ('年', '間') => -1626, ('府', '県') => -2363, ('手', '権') => -1982, ('新', '聞') => -4066, ('日', '新') => -722, ('日', '本') => -7068, ('日', '米') => 3372, ('曜', '日') => -601, ('朝', '鮮') => -2355, ('本', '人') => -2697, ('東', '京') => -1543, ('然', 'と') => -1384, ('社', '会') => -1276, ('立', 'て') => -990, ('第', 'に') => -1612, ('米', '国') => -4268, ('１', '１') => -669, ('ｸ', 'ﾞ') => 1319,} +}); +static BW3: LazyLock> = LazyLock::new(|| { + hashmap! { ('あ', 'た') => -2194, ('あ', 'り') => 719, ('あ', 'る') => 3846, ('い', '.') => -1185, ('い', '。') => -1185, ('い', 'い') => 5308, ('い', 'え') => 2079, ('い', 'く') => 3029, ('い', 'た') => 2056, ('い', 'っ') => 1883, ('い', 'る') => 5600, ('い', 'わ') => 1527, ('う', 'ち') => 1117, ('う', 'と') => 4798, ('え', 'と') => 1454, ('か', '.') => 2857, ('か', '。') => 2857, ('か', 'け') => -743, ('か', 'っ') => -4098, ('か', 'に') => -669, ('か', 'ら') => 6520, ('か', 'り') => -2670, ('が', ',') => 1816, ('が', '、') => 1816, ('が', 'き') => -4855, ('が', 'け') => -1127, ('が', 'っ') => -913, ('が', 'ら') => -4977, ('が', 'り') => -2064, ('き', 'た') => 1645, ('け', 'ど') => 1374, ('こ', 'と') => 7397, ('こ', 'の') => 1542, ('こ', 'ろ') => -2757, ('さ', 'い') => -714, ('さ', 'を') => 976, ('し', ',') => 1557, ('し', '、') => 1557, ('し', 'い') => -3714, ('し', 'た') => 3562, ('し', 'て') => 1449, ('し', 'な') => 2608, ('し', 'ま') => 1200, ('す', '.') => -1310, ('す', '。') => -1310, ('す', 'る') => 6521, ('ず', ',') => 3426, ('ず', '、') => 3426, ('ず', 'に') => 841, ('そ', 'う') => 428, ('た', '.') => 8875, ('た', '。') => 8875, ('た', 'い') => -594, ('た', 'の') => 812, ('た', 'り') => -1183, ('た', 'る') => -853, ('だ', '.') => 4098, ('だ', '。') => 4098, ('だ', 'っ') => 1004, ('っ', 'た') => -4748, ('っ', 'て') => 300, ('て', 'い') => 6240, ('て', 'お') => 855, ('て', 'も') => 302, ('で', 'す') => 1437, ('で', 'に') => -1482, ('で', 'は') => 2295, ('と', 'う') => -1387, ('と', 'し') => 2266, ('と', 'の') => 541, ('と', 'も') => -3543, ('ど', 'う') => 4664, ('な', 'い') => 1796, ('な', 'く') => -903, ('な', 'ど') => 2135, ('に', ',') => -1021, ('に', '、') => -1021, ('に', 'し') => 1771, ('に', 'な') => 1906, ('に', 'は') => 2644, ('の', ',') => -724, ('の', '、') => -724, ('の', '子') => -1000, ('は', ',') => 1337, ('は', '、') => 1337, ('べ', 'き') => 2181, ('ま', 'し') => 1113, ('ま', 'す') => 6943, ('ま', 'っ') => -1549, ('ま', 'で') => 6154, ('ま', 'れ') => -793, ('ら', 'し') => 1479, ('ら', 'れ') => 6820, ('る', 'る') => 3818, ('れ', ',') => 854, ('れ', '、') => 854, ('れ', 'た') => 1850, ('れ', 'て') => 1375, ('れ', 'ば') => -3246, ('れ', 'る') => 1091, ('わ', 'れ') => -605, ('ん', 'だ') => 606, ('ん', 'で') => 798, ('カ', '月') => 990, ('会', '議') => 860, ('入', 'り') => 1232, ('大', '会') => 2217, ('始', 'め') => 1681, ('市', ' ') => 965, ('新', '聞') => -5055, ('日', ',') => 974, ('日', '、') => 974, ('社', '会') => 2024, ('ｶ', '月') => 990, } +}); + +static TC1: LazyLock> = LazyLock::new(|| { + hashmap! { ('A', 'A', 'A') => 1093, ('H', 'H', 'H') => 1029, ('H', 'H', 'M') => 580, ('H', 'I', 'I') => 998, ('H', 'O', 'H') => -390, ('H', 'O', 'M') => -331, ('I', 'H', 'I') => 1169, ('I', 'O', 'H') => -142, ('I', 'O', 'I') => -1015, ('I', 'O', 'M') => 467, ('M', 'M', 'H') => 187, ('O', 'O', 'I') => -1832, } +}); +static TC2: LazyLock> = LazyLock::new(|| { + hashmap! { ('H', 'H', 'O') => 2088, ('H', 'I', 'I') => -1023, ('H', 'M', 'M') => -1154, ('I', 'H', 'I') => -1965, ('K', 'K', 'H') => 703, ('O', 'I', 'I') => -2649, } +}); +static TC3: LazyLock> = LazyLock::new(|| { + hashmap! { ('A', 'A', 'A') => -294, ('H', 'H', 'H') => 346, ('H', 'H', 'I') => -341, ('H', 'I', 'I') => -1088, ('H', 'I', 'K') => 731, ('H', 'O', 'H') => -1486, ('I', 'H', 'H') => 128, ('I', 'H', 'I') => -3041, ('I', 'H', 'O') => -1935, ('I', 'I', 'H') => -825, ('I', 'I', 'M') => -1035, ('I', 'O', 'I') => -542, ('K', 'H', 'H') => -1216, ('K', 'K', 'A') => 491, ('K', 'K', 'H') => -1217, ('K', 'O', 'K') => -1009, ('M', 'H', 'H') => -2694, ('M', 'H', 'M') => -457, ('M', 'H', 'O') => 123, ('M', 'M', 'H') => -471, ('N', 'N', 'H') => -1689, ('N', 'N', 'O') => 662, ('O', 'H', 'O') => -3393, } +}); +static TC4: LazyLock> = LazyLock::new(|| { + hashmap! { ('H', 'H', 'H') => -203, ('H', 'H', 'I') => 1344, ('H', 'H', 'K') => 365, ('H', 'H', 'M') => -122, ('H', 'H', 'N') => 182, ('H', 'H', 'O') => 669, ('H', 'I', 'H') => 804, ('H', 'I', 'I') => 679, ('H', 'O', 'H') => 446, ('I', 'H', 'H') => 695, ('I', 'H', 'O') => -2324, ('I', 'I', 'H') => 321, ('I', 'I', 'I') => 1497, ('I', 'I', 'O') => 656, ('I', 'O', 'O') => 54, ('K', 'A', 'K') => 4845, ('K', 'K', 'A') => 3386, ('K', 'K', 'K') => 3065, ('M', 'H', 'H') => -405, ('M', 'H', 'I') => 201, ('M', 'M', 'H') => -241, ('M', 'M', 'M') => 661, ('M', 'O', 'M') => 841, } +}); +static TQ1: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'H', 'H', 'H') => -227, ('B', 'H', 'H', 'I') => 316, ('B', 'H', 'I', 'H') => -132, ('B', 'I', 'H', 'H') => 60, ('B', 'I', 'I', 'I') => 1595, ('B', 'N', 'H', 'H') => -744, ('B', 'O', 'H', 'H') => 225, ('B', 'O', 'O', 'O') => -908, ('O', 'A', 'K', 'K') => 482, ('O', 'H', 'H', 'H') => 281, ('O', 'H', 'I', 'H') => 249, ('O', 'I', 'H', 'I') => 200, ('O', 'I', 'I', 'H') => -68, } +}); +static TQ2: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'I', 'H', 'H') => -1401, ('B', 'I', 'I', 'I') => -1033, ('B', 'K', 'A', 'K') => -543, ('B', 'O', 'O', 'O') => -5591, } +}); +static TQ3: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'H', 'H', 'H') => 478, ('B', 'H', 'H', 'M') => -1073, ('B', 'H', 'I', 'H') => 222, ('B', 'H', 'I', 'I') => -504, ('B', 'I', 'I', 'H') => -116, ('B', 'I', 'I', 'I') => -105, ('B', 'M', 'H', 'I') => -863, ('B', 'M', 'H', 'M') => -464, ('B', 'O', 'M', 'H') => 620, ('O', 'H', 'H', 'H') => 346, ('O', 'H', 'H', 'I') => 1729, ('O', 'H', 'I', 'I') => 997, ('O', 'H', 'M', 'H') => 481, ('O', 'I', 'H', 'H') => 623, ('O', 'I', 'I', 'H') => 1344, ('O', 'K', 'A', 'K') => 2792, ('O', 'K', 'H', 'H') => 587, ('O', 'K', 'K', 'A') => 679, ('O', 'O', 'H', 'H') => 110, ('O', 'O', 'I', 'I') => -685, } +}); +static TQ4: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'H', 'H', 'H') => -721, ('B', 'H', 'H', 'M') => -3604, ('B', 'H', 'I', 'I') => -966, ('B', 'I', 'I', 'H') => -607, ('B', 'I', 'I', 'I') => -2181, ('O', 'A', 'A', 'A') => -2763, ('O', 'A', 'K', 'K') => 180, ('O', 'H', 'H', 'H') => -294, ('O', 'H', 'H', 'I') => 2446, ('O', 'H', 'H', 'O') => 480, ('O', 'H', 'I', 'H') => -1573, ('O', 'I', 'H', 'H') => 1935, ('O', 'I', 'H', 'I') => -493, ('O', 'I', 'I', 'H') => 626, ('O', 'I', 'I', 'I') => -4007, ('O', 'K', 'A', 'K') => -8156, } +}); +static TW1: LazyLock> = LazyLock::new(|| { + hashmap! { ('に', 'つ', 'い') => -4681, ('東', '京', '都') => 2026, } +}); +static TW2: LazyLock> = LazyLock::new(|| { + hashmap! { ('あ', 'る', '程') => -2049, ('い', 'っ', 'た') => -1256, ('こ', 'ろ', 'が') => -2434, ('し', 'ょ', 'う') => 3873, ('そ', 'の', '後') => -4430, ('だ', 'っ', 'て') => -1049, ('て', 'い', 'た') => 1833, ('と', 'し', 'て') => -4657, ('と', 'も', 'に') => -4517, ('も', 'の', 'で') => 1882, ('一', '気', 'に') => -792, ('初', 'め', 'て') => -1512, ('同', '時', 'に') => -8097, ('大', 'き', 'な') => -1255, ('対', 'し', 'て') => -2721, ('社', '会', '党') => -3216, } +}); +static TW3: LazyLock> = LazyLock::new(|| { + hashmap! { ('い', 'た', 'だ') => -1734, ('し', 'て', 'い') => 1314, ('と', 'し', 'て') => -4314, ('に', 'つ', 'い') => -5483, ('に', 'と', 'っ') => -5989, ('に', '当', 'た') => -6247, ('の', 'で', ',') => -727, ('の', 'で', '、') => -727, ('の', 'も', 'の') => -600, ('れ', 'か', 'ら') => -3752, ('十', '二', '月') => -2287, } +}); +static TW4: LazyLock> = LazyLock::new(|| { + hashmap! { ('い', 'う', '.') => 8576, ('い', 'う', '。') => 8576, ('か', 'ら', 'な') => -2348, ('し', 'て', 'い') => 2958, ('た', 'が', ',') => 1516, ('た', 'が', '、') => 1516, ('て', 'い', 'る') => 1538, ('と', 'い', 'う') => 1349, ('ま', 'し', 'た') => 5543, ('ま', 'せ', 'ん') => 1097, ('よ', 'う', 'と') => -4258, ('よ', 'る', 'と') => 5865, } +}); + +static UC1: LazyLock> = LazyLock::new(|| { + hashmap! { 'A' => 484, 'K' => 93, 'M' => 645, 'O' => -505, } +}); +static UC2: LazyLock> = LazyLock::new(|| { + hashmap! { 'A' => 819, 'H' => 1059, 'I' => 409, 'M' => 3987, 'N' => 5775, 'O' => 646, } +}); +static UC3: LazyLock> = LazyLock::new(|| { + hashmap! { 'A' => -1370, 'I' => 2311, } +}); +static UC4: LazyLock> = LazyLock::new(|| { + hashmap! { 'A' => -2643, 'H' => 1809, 'I' => -1032, 'K' => -3450, 'M' => 3565, 'N' => 3876, 'O' => 6646, } +}); +static UC5: LazyLock> = LazyLock::new(|| { + hashmap! { 'H' => 313, 'I' => -1238, 'K' => -799, 'M' => 539, 'O' => -831, } +}); +static UC6: LazyLock> = LazyLock::new(|| { + hashmap! { 'H' => -506, 'I' => -253, 'K' => 87, 'M' => 247, 'O' => -387, } +}); +static UP1: LazyLock> = LazyLock::new(|| { + hashmap! { 'O' => -214, } +}); +static UP2: LazyLock> = LazyLock::new(|| { + hashmap! { 'B' => 69, 'O' => 935, } +}); +static UP3: LazyLock> = LazyLock::new(|| { + hashmap! { 'B' => 189, } +}); +static UQ1: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'H') => 21, ('B', 'I') => -12, ('B', 'K') => -99, ('B', 'N') => 142, ('B', 'O') => -56, ('O', 'H') => -95, ('O', 'I') => 477, ('O', 'K') => 410, ('O', 'O') => -2422, } +}); +static UQ2: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'H') => 216, ('B', 'I') => 113, ('O', 'K') => 1759, } +}); +static UQ3: LazyLock> = LazyLock::new(|| { + hashmap! { ('B', 'A') => -479, ('B', 'H') => 42, ('B', 'I') => 1913, ('B', 'K') => -7198, ('B', 'M') => 3160, ('B', 'N') => 6427, ('B', 'O') => 14761, ('O', 'I') => -827, ('O', 'N') => -3212, } +}); +static UW1: LazyLock> = LazyLock::new(|| { + hashmap! { ',' => 156, '、' => 156, '「' => -463, 'あ' => -941, 'う' => -127, 'が' => -553, 'き' => 121, 'こ' => 505, 'で' => -201, 'と' => -547, 'ど' => -123, 'に' => -789, 'の' => -185, 'は' => -847, 'も' => -466, 'や' => -470, 'よ' => 182, 'ら' => -292, 'り' => 208, 'れ' => 169, 'を' => -446, 'ん' => -137, '・' => -135, '主' => -402, '京' => -268, '区' => -912, '午' => 871, '国' => -460, '大' => 561, '委' => 729, '市' => -411, '日' => -141, '理' => 361, '生' => -408, '県' => -386, '都' => -718, '｢' => -463, '･' => -135, } +}); +static UW2: LazyLock> = LazyLock::new(|| { + hashmap! { ',' => -829, '、' => -829, '〇' => 892, '「' => -645, '」' => 3145, 'あ' => -538, 'い' => 505, 'う' => 134, 'お' => -502, 'か' => 1454, 'が' => -856, 'く' => -412, 'こ' => 1141, 'さ' => 878, 'ざ' => 540, 'し' => 1529, 'す' => -675, 'せ' => 300, 'そ' => -1011, 'た' => 188, 'だ' => 1837, 'つ' => -949, 'て' => -291, 'で' => -268, 'と' => -981, 'ど' => 1273, 'な' => 1063, 'に' => -1764, 'の' => 130, 'は' => -409, 'ひ' => -1273, 'べ' => 1261, 'ま' => 600, 'も' => -1263, 'や' => -402, 'よ' => 1639, 'り' => -579, 'る' => -694, 'れ' => 571, 'を' => -2516, 'ん' => 2095, 'ア' => -587, 'カ' => 306, 'キ' => 568, 'ッ' => 831, '三' => -758, '不' => -2150, '世' => -302, '中' => -968, '主' => -861, '事' => 492, '人' => -123, '会' => 978, '保' => 362, '入' => 548, '初' => -3025, '副' => -1566, '北' => -3414, '区' => -422, '大' => -1769, '天' => -865, '太' => -483, '子' => -1519, '学' => 760, '実' => 1023, '小' => -2009, '市' => -813, '年' => -1060, '強' => 1067, '手' => -1519, '揺' => -1033, '政' => 1522, '文' => -1355, '新' => -1682, '日' => -1815, '明' => -1462, '最' => -630, '朝' => -1843, '本' => -1650, '東' => -931, '果' => -665, '次' => -2378, '民' => -180, '気' => -1740, '理' => 752, '発' => 529, '目' => -1584, '相' => -242, '県' => -1165, '立' => -763, '第' => 810, '米' => 509, '自' => -1353, '行' => 838, '西' => -744, '見' => -3874, '調' => 1010, '議' => 1198, '込' => 3041, '開' => 1758, '間' => -1257, '｢' => -645, '｣' => 3145, 'ｯ' => 831, 'ｱ' => -587, 'ｶ' => 306, 'ｷ' => 568, } +}); +static UW3: LazyLock> = LazyLock::new(|| { + hashmap! { ',' => 4889, '1' => -800, '−' => -1723, '、' => 4889, '々' => -2311, '〇' => 5827, '」' => 2670, '〓' => -3573, 'あ' => -2696, 'い' => 1006, 'う' => 2342, 'え' => 1983, 'お' => -4864, 'か' => -1163, 'が' => 3271, 'く' => 1004, 'け' => 388, 'げ' => 401, 'こ' => -3552, 'ご' => -3116, 'さ' => -1058, 'し' => -395, 'す' => 584, 'せ' => 3685, 'そ' => -5228, 'た' => 842, 'ち' => -521, 'っ' => -1444, 'つ' => -1081, 'て' => 6167, 'で' => 2318, 'と' => 1691, 'ど' => -899, 'な' => -2788, 'に' => 2745, 'の' => 4056, 'は' => 4555, 'ひ' => -2171, 'ふ' => -1798, 'へ' => 1199, 'ほ' => -5516, 'ま' => -4384, 'み' => -120, 'め' => 1205, 'も' => 2323, 'や' => -788, 'よ' => -202, 'ら' => 727, 'り' => 649, 'る' => 5905, 'れ' => 2773, 'わ' => -1207, 'を' => 6620, 'ん' => -518, 'ア' => 551, 'グ' => 1319, 'ス' => 874, 'ッ' => -1350, 'ト' => 521, 'ム' => 1109, 'ル' => 1591, 'ロ' => 2201, 'ン' => 278, '・' => -3794, '一' => -1619, '下' => -1759, '世' => -2087, '両' => 3815, '中' => 653, '主' => -758, '予' => -1193, '二' => 974, '人' => 2742, '今' => 792, '他' => 1889, '以' => -1368, '低' => 811, '何' => 4265, '作' => -361, '保' => -2439, '元' => 4858, '党' => 3593, '全' => 1574, '公' => -3030, '六' => 755, '共' => -1880, '円' => 5807, '再' => 3095, '分' => 457, '初' => 2475, '別' => 1129, '前' => 2286, '副' => 4437, '力' => 365, '動' => -949, '務' => -1872, '化' => 1327, '北' => -1038, '区' => 4646, '千' => -2309, '午' => -783, '協' => -1006, '口' => 483, '右' => 1233, '各' => 3588, '合' => -241, '同' => 3906, '和' => -837, '員' => 4513, '国' => 642, '型' => 1389, '場' => 1219, '外' => -241, '妻' => 2016, '学' => -1356, '安' => -423, '実' => -1008, '家' => 1078, '小' => -513, '少' => -3102, '州' => 1155, '市' => 3197, '平' => -1804, '年' => 2416, '広' => -1030, '府' => 1605, '度' => 1452, '建' => -2352, '当' => -3885, '得' => 1905, '思' => -1291, '性' => 1822, '戸' => -488, '指' => -3973, '政' => -2013, '教' => -1479, '数' => 3222, '文' => -1489, '新' => 1764, '日' => 2099, '旧' => 5792, '昨' => -661, '時' => -1248, '曜' => -951, '最' => -937, '月' => 4125, '期' => 360, '李' => 3094, '村' => 364, '東' => -805, '核' => 5156, '森' => 2438, '業' => 484, '氏' => 2613, '民' => -1694, '決' => -1073, '法' => 1868, '海' => -495, '無' => 979, '物' => 461, '特' => -3850, '生' => -273, '用' => 914, '町' => 1215, '的' => 7313, '直' => -1835, '省' => 792, '県' => 6293, '知' => -1528, '私' => 4231, '税' => 401, '立' => -960, '第' => 1201, '米' => 7767, '系' => 3066, '約' => 3663, '級' => 1384, '統' => -4229, '総' => 1163, '線' => 1255, '者' => 6457, '能' => 725, '自' => -2869, '英' => 785, '見' => 1044, '調' => -562, '財' => -733, '費' => 1777, '車' => 1835, '軍' => 1375, '込' => -1504, '通' => -1136, '選' => -681, '郎' => 1026, '郡' => 4404, '部' => 1200, '金' => 2163, '長' => 421, '開' => -1432, '間' => 1302, '関' => -1282, '雨' => 2009, '電' => -1045, '非' => 2066, '駅' => 1620, '１' => -800, '｣' => 2670, '･' => -3794, 'ｯ' => -1350, 'ｱ' => 551, 'ｽ' => 874, 'ﾄ' => 521, 'ﾑ' => 1109, 'ﾙ' => 1591, 'ﾛ' => 2201, 'ﾝ' => 278, } +}); +static UW4: LazyLock> = LazyLock::new(|| { + hashmap! { ',' => 3930, '.' => 3508, '―' => -4841, '、' => 3930, '。' => 3508, '〇' => 4999, '「' => 1895, '」' => 3798, '〓' => -5156, 'あ' => 4752, 'い' => -3435, 'う' => -640, 'え' => -2514, 'お' => 2405, 'か' => 530, 'が' => 6006, 'き' => -4482, 'ぎ' => -3821, 'く' => -3788, 'け' => -4376, 'げ' => -4734, 'こ' => 2255, 'ご' => 1979, 'さ' => 2864, 'し' => -843, 'じ' => -2506, 'す' => -731, 'ず' => 1251, 'せ' => 181, 'そ' => 4091, 'た' => 5034, 'だ' => 5408, 'ち' => -3654, 'っ' => -5882, 'つ' => -1659, 'て' => 3994, 'で' => 7410, 'と' => 4547, 'な' => 5433, 'に' => 6499, 'ぬ' => 1853, 'ね' => 1413, 'の' => 7396, 'は' => 8578, 'ば' => 1940, 'ひ' => 4249, 'び' => -4134, 'ふ' => 1345, 'へ' => 6665, 'べ' => -744, 'ほ' => 1464, 'ま' => 1051, 'み' => -2082, 'む' => -882, 'め' => -5046, 'も' => 4169, 'ゃ' => -2666, 'や' => 2795, 'ょ' => -1544, 'よ' => 3351, 'ら' => -2922, 'り' => -9726, 'る' => -14896, 'れ' => -2613, 'ろ' => -4570, 'わ' => -1783, 'を' => 13150, 'ん' => -2352, 'カ' => 2145, 'コ' => 1789, 'セ' => 1287, 'ッ' => -724, 'ト' => -403, 'メ' => -1635, 'ラ' => -881, 'リ' => -541, 'ル' => -856, 'ン' => -3637, '・' => -4371, 'ー' => -11870, '一' => -2069, '中' => 2210, '予' => 782, '事' => -190, '井' => -1768, '人' => 1036, '以' => 544, '会' => 950, '体' => -1286, '作' => 530, '側' => 4292, '先' => 601, '党' => -2006, '共' => -1212, '内' => 584, '円' => 788, '初' => 1347, '前' => 1623, '副' => 3879, '力' => -302, '動' => -740, '務' => -2715, '化' => 776, '区' => 4517, '協' => 1013, '参' => 1555, '合' => -1834, '和' => -681, '員' => -910, '器' => -851, '回' => 1500, '国' => -619, '園' => -1200, '地' => 866, '場' => -1410, '塁' => -2094, '士' => -1413, '多' => 1067, '大' => 571, '子' => -4802, '学' => -1397, '定' => -1057, '寺' => -809, '小' => 1910, '屋' => -1328, '山' => -1500, '島' => -2056, '川' => -2667, '市' => 2771, '年' => 374, '庁' => -4556, '後' => 456, '性' => 553, '感' => 916, '所' => -1566, '支' => 856, '改' => 787, '政' => 2182, '教' => 704, '文' => 522, '方' => -856, '日' => 1798, '時' => 1829, '最' => 845, '月' => -9066, '木' => -485, '来' => -442, '校' => -360, '業' => -1043, '氏' => 5388, '民' => -2716, '気' => -910, '沢' => -939, '済' => -543, '物' => -735, '率' => 672, '球' => -1267, '生' => -1286, '産' => -1101, '田' => -2900, '町' => 1826, '的' => 2586, '目' => 922, '省' => -3485, '県' => 2997, '空' => -867, '立' => -2112, '第' => 788, '米' => 2937, '系' => 786, '約' => 2171, '経' => 1146, '統' => -1169, '総' => 940, '線' => -994, '署' => 749, '者' => 2145, '能' => -730, '般' => -852, '行' => -792, '規' => 792, '警' => -1184, '議' => -244, '谷' => -1000, '賞' => 730, '車' => -1481, '軍' => 1158, '輪' => -1433, '込' => -3370, '近' => 929, '道' => -1291, '選' => 2596, '郎' => -4866, '都' => 1192, '野' => -1100, '銀' => -2213, '長' => 357, '間' => -2344, '院' => -2297, '際' => -2604, '電' => -878, '領' => -1659, '題' => -792, '館' => -1984, '首' => 1749, '高' => 2120, '｢' => 1895, '｣' => 3798, '･' => -4371, 'ｯ' => -724, 'ｰ' => -11870, 'ｶ' => 2145, 'ｺ' => 1789, 'ｾ' => 1287, 'ﾄ' => -403, 'ﾒ' => -1635, 'ﾗ' => -881, 'ﾘ' => -541, 'ﾙ' => -856, 'ﾝ' => -3637, } +}); +static UW5: LazyLock> = LazyLock::new(|| { + hashmap! { ',' => 465, '.' => -299, '1' => -514, E2 => -32768, ']' => -2762, '、' => 465, '。' => -299, '「' => 363, 'あ' => 1655, 'い' => 331, 'う' => -503, 'え' => 1199, 'お' => 527, 'か' => 647, 'が' => -421, 'き' => 1624, 'ぎ' => 1971, 'く' => 312, 'げ' => -983, 'さ' => -1537, 'し' => -1371, 'す' => -852, 'だ' => -1186, 'ち' => 1093, 'っ' => 52, 'つ' => 921, 'て' => -18, 'で' => -850, 'と' => -127, 'ど' => 1682, 'な' => -787, 'に' => -1224, 'の' => -635, 'は' => -578, 'べ' => 1001, 'み' => 502, 'め' => 865, 'ゃ' => 3350, 'ょ' => 854, 'り' => -208, 'る' => 429, 'れ' => 504, 'わ' => 419, 'を' => -1264, 'ん' => 327, 'イ' => 241, 'ル' => 451, 'ン' => -343, '中' => -871, '京' => 722, '会' => -1153, '党' => -654, '務' => 3519, '区' => -901, '告' => 848, '員' => 2104, '大' => -1296, '学' => -548, '定' => 1785, '嵐' => -1304, '市' => -2991, '席' => 921, '年' => 1763, '思' => 872, '所' => -814, '挙' => 1618, '新' => -1682, '日' => 218, '月' => -4353, '査' => 932, '格' => 1356, '機' => -1508, '氏' => -1347, '田' => 240, '町' => -3912, '的' => -3149, '相' => 1319, '省' => -1052, '県' => -4003, '研' => -997, '社' => -278, '空' => -813, '統' => 1955, '者' => -2233, '表' => 663, '語' => -1073, '議' => 1219, '選' => -1018, '郎' => -368, '長' => 786, '間' => 1191, '題' => 2368, '館' => -689, '１' => -514, '｢' => 363, 'ｲ' => 241, 'ﾙ' => 451, 'ﾝ' => -343, } +}); +static UW6: LazyLock> = LazyLock::new(|| { + hashmap! { ',' => 227, '.' => 808, '1' => -270, E1 => 306, '、' => 227, '。' => 808, 'あ' => -307, 'う' => 189, 'か' => 241, 'が' => -73, 'く' => -121, 'こ' => -200, 'じ' => 1782, 'す' => 383, 'た' => -428, 'っ' => 573, 'て' => -1014, 'で' => 101, 'と' => -105, 'な' => -253, 'に' => -149, 'の' => -417, 'は' => -236, 'も' => -206, 'り' => 187, 'る' => -135, 'を' => 195, 'ル' => -673, 'ン' => -496, '一' => -277, '中' => 201, '件' => -800, '会' => 624, '前' => 302, '区' => 1792, '員' => -1212, '委' => 798, '学' => -960, '市' => 887, '広' => -695, '後' => 535, '業' => -697, '相' => 753, '社' => -507, '福' => 974, '空' => -822, '者' => 1811, '連' => 463, '郎' => 1082, '１' => -270, 'ﾙ' => -673, 'ﾝ' => -496, } +}); + #[cfg(test)] mod tests { use crate::tokenizers::{japanese::JapaneseTokenizer, word::WordTokenizer, Token}; diff --git a/crates/nlp/src/tokenizers/osb.rs b/crates/nlp/src/tokenizers/osb.rs index e6cba5f8..cd91fa3d 100644 --- a/crates/nlp/src/tokenizers/osb.rs +++ b/crates/nlp/src/tokenizers/osb.rs @@ -4,7 +4,7 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use std::{borrow::Cow, iter::Peekable}; +use std::iter::Peekable; #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct OsbToken { @@ -14,26 +14,26 @@ pub struct OsbToken { #[derive(Debug, Clone, PartialEq, Eq)] pub enum Gram<'x> { - Uni { t1: &'x str }, - Bi { t1: &'x str, t2: &'x str }, + Uni { t1: &'x [u8] }, + Bi { t1: &'x [u8], t2: &'x [u8] }, } -pub struct OsbTokenizer<'x, I, R> +pub struct OsbTokenizer where - I: Iterator>, + I: Iterator>, R: for<'y> From> + 'static, { iter: Peekable, - buf: Vec>>, + buf: Vec>>, window_size: usize, window_pos: usize, window_idx: usize, phantom: std::marker::PhantomData, } -impl<'x, I, R> OsbTokenizer<'x, I, R> +impl OsbTokenizer where - I: Iterator>, + I: Iterator>, R: for<'y> From> + 'static, { pub fn new(iter: I, window_size: usize) -> Self { @@ -48,9 +48,9 @@ where } } -impl<'x, I, R> Iterator for OsbTokenizer<'x, I, R> +impl Iterator for OsbTokenizer where - I: Iterator>, + I: Iterator>, R: for<'y> From> + 'static, { type Item = OsbToken; @@ -91,15 +91,17 @@ where #[cfg(test)] mod test { - use std::borrow::Cow; - use crate::tokenizers::osb::{Gram, OsbToken}; impl From> for String { fn from(value: Gram<'_>) -> Self { match value { - Gram::Uni { t1 } => t1.to_string(), - Gram::Bi { t1, t2 } => format!("{t1} {t2}"), + Gram::Uni { t1 } => std::str::from_utf8(t1).unwrap().to_string(), + Gram::Bi { t1, t2 } => format!( + "{} {}", + std::str::from_utf8(t1).unwrap(), + std::str::from_utf8(t2).unwrap() + ), } } } @@ -110,7 +112,7 @@ mod test { super::OsbTokenizer::new( "The quick brown fox jumps over the lazy dog and the lazy cat" .split_ascii_whitespace() - .map(Cow::from), + .map(|b| b.as_bytes().to_vec()), 5, ) .collect::>(), diff --git a/crates/pop3/src/client.rs b/crates/pop3/src/client.rs index b1648c9a..ad866f57 100644 --- a/crates/pop3/src/client.rs +++ b/crates/pop3/src/client.rs @@ -4,7 +4,10 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use common::listener::{SessionResult, SessionStream}; +use common::{ + listener::{SessionResult, SessionStream}, + KV_RATE_LIMIT_IMAP, +}; use mail_send::Credentials; use trc::{AddContext, SecurityEvent}; @@ -239,7 +242,8 @@ impl Session { .storage .lookup .is_rate_allowed( - format!("ireq:{}", mailbox.account_id).as_bytes(), + KV_RATE_LIMIT_IMAP, + &mailbox.account_id.to_be_bytes(), rate, true, ) diff --git a/crates/smtp/src/core/throttle.rs b/crates/smtp/src/core/throttle.rs index 9abed15c..5dc27be3 100644 --- a/crates/smtp/src/core/throttle.rs +++ b/crates/smtp/src/core/throttle.rs @@ -8,7 +8,7 @@ use common::{ config::smtp::{queue::QueueQuota, *}, expr::{functions::ResolveVariable, *}, listener::{limiter::ConcurrencyLimiter, SessionStream}, - ThrottleKey, + ThrottleKey, KV_RATE_LIMIT_HASH, }; use dashmap::mapref::entry::Entry; use trc::SmtpEvent; @@ -212,27 +212,33 @@ impl Session { // Check rate if let Some(rate) = &t.rate { - if self + match self .server .core .storage .lookup - .is_rate_allowed(key.hash.as_slice(), rate, false) + .is_rate_allowed(KV_RATE_LIMIT_HASH, key.hash.as_slice(), rate, false) .await - .unwrap_or_default() - .is_some() { - trc::event!( - Smtp(SmtpEvent::RateLimitExceeded), - SpanId = self.data.session_id, - Id = t.id.clone(), - Limit = vec![ - trc::Value::from(rate.requests), - trc::Value::from(rate.period) - ], - ); + Ok(Some(_)) => { + trc::event!( + Smtp(SmtpEvent::RateLimitExceeded), + SpanId = self.data.session_id, + Id = t.id.clone(), + Limit = vec![ + trc::Value::from(rate.requests), + trc::Value::from(rate.period) + ], + ); - return false; + return false; + } + Err(err) => { + trc::error!(err + .span_id(self.data.session_id) + .caused_by(trc::location!())); + } + _ => (), } } } @@ -248,13 +254,27 @@ impl Session { hasher.update(&rate.period.as_secs().to_ne_bytes()[..]); hasher.update(&rate.requests.to_ne_bytes()[..]); - self.server + match self + .server .core .storage .lookup - .is_rate_allowed(hasher.finalize().as_bytes(), rate, false) + .is_rate_allowed( + KV_RATE_LIMIT_HASH, + hasher.finalize().as_bytes(), + rate, + false, + ) .await - .unwrap_or_default() - .is_none() + { + Ok(None) => true, + Ok(Some(_)) => false, + Err(err) => { + trc::error!(err + .span_id(self.data.session_id) + .caused_by(trc::location!())); + true + } + } } } diff --git a/crates/smtp/src/inbound/rcpt.rs b/crates/smtp/src/inbound/rcpt.rs index 9c18ec4a..d98a3b45 100644 --- a/crates/smtp/src/inbound/rcpt.rs +++ b/crates/smtp/src/inbound/rcpt.rs @@ -4,11 +4,14 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use common::{config::smtp::session::Stage, listener::SessionStream, scripts::ScriptModification}; +use common::{ + config::smtp::session::Stage, listener::SessionStream, scripts::ScriptModification, KV_GREYLIST, +}; use directory::backend::RcptType; use smtp_proto::{ RcptTo, RCPT_NOTIFY_DELAY, RCPT_NOTIFY_FAILURE, RCPT_NOTIFY_NEVER, RCPT_NOTIFY_SUCCESS, }; +use store::dispatch::lookup::KeyValue; use trc::{SecurityEvent, SmtpEvent}; use crate::{ @@ -299,24 +302,31 @@ impl Session { .greylist_duration .filter(|_| self.data.authenticated_as.is_none()) { - let key = format!( - "g:{}:{}:{}", - self.data.remote_ip_str, - self.data.mail_from.as_ref().unwrap().address_lcase, - self.data.rcpt_to.last().unwrap().address_lcase + let mut key = Vec::with_capacity(64); + key.push(KV_GREYLIST); + match self.data.remote_ip { + std::net::IpAddr::V4(ipv4_addr) => key.extend_from_slice(&ipv4_addr.octets()), + std::net::IpAddr::V6(ipv6_addr) => key.extend_from_slice(&ipv6_addr.octets()), + }; + key.extend_from_slice( + self.data + .mail_from + .as_ref() + .unwrap() + .address_lcase + .as_bytes(), ); - match self - .server - .lookup_store() - .key_exists(key.clone().into_bytes()) - .await - { + key.extend_from_slice(self.data.rcpt_to.last().unwrap().address_lcase.as_bytes()); + + match self.server.in_memory_store().key_exists(key.clone()).await { Ok(true) => (), Ok(false) => { match self .server - .lookup_store() - .key_set(key.into_bytes(), vec![], greylist_duration.as_secs().into()) + .in_memory_store() + .key_set( + KeyValue::new(key, vec![]).expires(greylist_duration.as_secs()), + ) .await { Ok(_) => { diff --git a/crates/smtp/src/queue/throttle.rs b/crates/smtp/src/queue/throttle.rs index e80c3a44..7aef6253 100644 --- a/crates/smtp/src/queue/throttle.rs +++ b/crates/smtp/src/queue/throttle.rs @@ -10,7 +10,7 @@ use common::{ config::smtp::Throttle, expr::functions::ResolveVariable, listener::limiter::{ConcurrencyLimiter, InFlight}, - Server, + Server, KV_RATE_LIMIT_HASH, }; use dashmap::mapref::entry::Entry; use store::write::now; @@ -52,26 +52,32 @@ impl IsAllowed for Server { let key = throttle.new_key(envelope); if let Some(rate) = &throttle.rate { - if let Ok(Some(next_refill)) = self + match self .core .storage .lookup - .is_rate_allowed(key.as_ref(), rate, false) + .is_rate_allowed(KV_RATE_LIMIT_HASH, key.as_ref(), rate, false) .await { - trc::event!( - Queue(trc::QueueEvent::RateLimitExceeded), - SpanId = session_id, - Id = throttle.id.clone(), - Limit = vec![ - trc::Value::from(rate.requests), - trc::Value::from(rate.period) - ], - ); + Ok(Some(next_refill)) => { + trc::event!( + Queue(trc::QueueEvent::RateLimitExceeded), + SpanId = session_id, + Id = throttle.id.clone(), + Limit = vec![ + trc::Value::from(rate.requests), + trc::Value::from(rate.period) + ], + ); - return Err(Error::Rate { - retry_at: now() + next_refill, - }); + return Err(Error::Rate { + retry_at: now() + next_refill, + }); + } + Err(err) => { + trc::error!(err.span_id(session_id).caused_by(trc::location!())); + } + _ => (), } } diff --git a/crates/spam-filter/src/analysis/bayes.rs b/crates/spam-filter/src/analysis/bayes.rs new file mode 100644 index 00000000..bdfec28d --- /dev/null +++ b/crates/spam-filter/src/analysis/bayes.rs @@ -0,0 +1,52 @@ +use std::future::Future; + +use common::Server; + +use crate::{ + modules::bayes::{bayes_classify, bayes_train_if_balanced}, + SpamFilterContext, +}; + +pub trait SpamFilterAnalyzeBayes: Sync + Send { + fn spam_filter_analyze_bayes_classify( + &self, + ctx: &mut SpamFilterContext<'_>, + ) -> impl Future + Send; + + fn spam_filter_analyze_spam_trap( + &self, + ctx: &mut SpamFilterContext<'_>, + ) -> impl Future + Send; +} + +impl SpamFilterAnalyzeBayes for Server { + async fn spam_filter_analyze_bayes_classify(&self, ctx: &mut SpamFilterContext<'_>) { + if let Some(config) = &self.core.spam.bayes { + if !ctx.result.has_tag("SPAM_TRAP") && !ctx.result.has_tag("TRUSTED_REPLY") { + match bayes_classify(self, ctx).await { + Ok(score) => { + if score > config.score_spam { + ctx.result.add_tag("BAYES_SPAM"); + } else if score < config.score_ham { + ctx.result.add_tag("BAYES_HAM"); + } + } + Err(err) => { + trc::error!(err.span_id(ctx.input.span_id).caused_by(trc::location!())); + } + } + } + } + } + + async fn spam_filter_analyze_spam_trap(&self, ctx: &mut SpamFilterContext<'_>) { + if ctx + .output + .env_to_addr + .iter() + .any(|addr| self.core.spam.list_spamtraps.contains(&addr.address)) + { + ctx.result.add_tag("SPAM_TRAP"); + } + } +} diff --git a/crates/spam-filter/src/analysis/init.rs b/crates/spam-filter/src/analysis/init.rs index 8d49419f..410ff61c 100644 --- a/crates/spam-filter/src/analysis/init.rs +++ b/crates/spam-filter/src/analysis/init.rs @@ -220,27 +220,3 @@ impl SpamFilterInit for Server { } } } - -/* - -use std::future::Future; - -use common::Server; - -use crate::SpamFilterContext; - -pub trait SpamFilterAnalyze!: Sync + Send { - fn spam_filter_analyze_*( - &self, - ctx: &mut SpamFilterContext<'_>, - ) -> impl Future + Send; -} - -impl SpamFilterAnalyze! for Server { - async fn spam_filter_analyze_*(&self, ctx: &mut SpamFilterContext<'_>) { - todo!() - } -} - - -*/ diff --git a/crates/spam-filter/src/analysis/llm.rs b/crates/spam-filter/src/analysis/llm.rs new file mode 100644 index 00000000..593bf484 --- /dev/null +++ b/crates/spam-filter/src/analysis/llm.rs @@ -0,0 +1,90 @@ +// SPDX-SnippetBegin +// SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd +// SPDX-License-Identifier: LicenseRef-SEL + +use std::{future::Future, time::Instant}; + +use common::Server; +use trc::AiEvent; + +use crate::SpamFilterContext; + +pub trait SpamFilterAnalyzeLlm: Sync + Send { + fn spam_filter_analyze_llm( + &self, + ctx: &mut SpamFilterContext<'_>, + ) -> impl Future + Send; +} + +impl SpamFilterAnalyzeLlm for Server { + async fn spam_filter_analyze_llm(&self, ctx: &mut SpamFilterContext<'_>) { + if let Some(config) = self + .core + .enterprise + .as_ref() + .and_then(|c| c.spam_filter_llm.as_ref()) + { + let time = Instant::now(); + let prompt = config.prompt.clone(); + match config + .model + .send_request(prompt, config.temperature.into()) + .await + { + Ok(response) => { + trc::event!( + Ai(AiEvent::LlmResponse), + Id = config.model.id.clone(), + Details = response.clone(), + Elapsed = time.elapsed(), + SpanId = ctx.input.span_id, + ); + + let mut category = None; + let mut confidence = None; + let mut explanation = None; + + for (idx, value) in response.split(config.separator).enumerate() { + let value = value.trim(); + if !value.is_empty() { + if idx == config.index_category { + let value = value.to_uppercase(); + if config.categories.contains(value.as_str()) { + category = Some(value); + } + } else if config.index_confidence.map_or(false, |i| i == idx) { + let value = value.to_uppercase(); + if config.confidence.contains(value.as_str()) { + confidence = Some(value); + } + } else if config.index_explanation.map_or(false, |i| i == idx) { + explanation = Some(value); + } + } + } + + let category = match (category, confidence) { + (Some(category), Some(confidence)) => { + ctx.result.add_tag(format!("LLM_{category}_{confidence}")); + category + } + (Some(category), None) => { + ctx.result.add_tag(format!("LLM_{category}")); + category + } + _ => return, + }; + + if let Some(explanation) = explanation { + ctx.result.llm_header = + format!("X-Spam-Llm-Explanation: {category} ({explanation})\r\n",) + .into(); + } + } + Err(err) => { + trc::error!(err.span_id(ctx.input.span_id)); + } + } + } + } +} diff --git a/crates/spam-filter/src/analysis/mod.rs b/crates/spam-filter/src/analysis/mod.rs index a2eab40b..263976f7 100644 --- a/crates/spam-filter/src/analysis/mod.rs +++ b/crates/spam-filter/src/analysis/mod.rs @@ -12,6 +12,7 @@ use mail_parser::{parsers::MessageStream, Header}; use crate::{Recipient, SpamFilterContext, SpamFilterInput, SpamFilterOutput, SpamFilterResult}; +pub mod bayes; pub mod bounce; pub mod date; pub mod dmarc; @@ -22,6 +23,8 @@ pub mod headers; pub mod html; pub mod init; pub mod ip; +#[cfg(feature = "enterprise")] +pub mod llm; pub mod messageid; pub mod mime; pub mod pyzor; @@ -29,7 +32,9 @@ pub mod received; pub mod recipient; pub mod replyto; pub mod reputation; +pub mod score; pub mod subject; +pub mod trusted_reply; pub mod url; impl SpamFilterInput<'_> { @@ -57,6 +62,10 @@ impl SpamFilterResult { pub fn add_tag(&mut self, tag: impl Into) { self.tags.insert(tag.into()); } + + pub fn has_tag(&self, tag: impl AsRef) -> bool { + self.tags.contains(tag.as_ref()) + } } pub(crate) struct SpamFilterResolver<'x, T: ResolveVariable> { diff --git a/crates/spam-filter/src/analysis/reputation.rs b/crates/spam-filter/src/analysis/reputation.rs index 88cb20d9..2318aa90 100644 --- a/crates/spam-filter/src/analysis/reputation.rs +++ b/crates/spam-filter/src/analysis/reputation.rs @@ -1,8 +1,11 @@ -use std::future::Future; +use std::{borrow::Cow, future::Future}; -use common::Server; +use common::{ + ip_to_bytes, Server, KV_REPUTATION_ASN, KV_REPUTATION_DOMAIN, KV_REPUTATION_FROM, + KV_REPUTATION_IP, +}; use mail_auth::DmarcResult; -use store::{Deserialize, Serialize}; +use store::{dispatch::lookup::KeyValue, Deserialize, Serialize}; use crate::{ modules::{key_get, key_set}, @@ -39,25 +42,32 @@ impl SpamFilterAnalyzeReputation for Server { }; // Do not penalize forged domains - let prefix = if matches!(ctx.input.dmarc_result, DmarcResult::Pass) { - "" - } else { - "_" - }; + let is_dmarc_pass = matches!(ctx.input.dmarc_result, DmarcResult::Pass); let mut types = vec![ - (Type::Ip, format!("i:{}", ctx.input.remote_ip)), - (Type::From, format!("f:{}{}", prefix, sender.address)), + (Type::Ip, Cow::Owned(ip_to_bytes(&ctx.input.remote_ip))), + ( + Type::From, + if is_dmarc_pass { + Cow::Borrowed(sender.address.as_bytes()) + } else { + Cow::Owned(format!("_{}", sender.domain_part.sld_or_default()).into_bytes()) + }, + ), ( Type::Domain, - format!("d:{}{}", prefix, sender.domain_part.sld_or_default()), + if is_dmarc_pass { + Cow::Borrowed(sender.domain_part.sld_or_default().as_bytes()) + } else { + Cow::Owned(format!("_{}", sender.domain_part.sld_or_default()).into_bytes()) + }, ), ]; // Add ASN if let Some(asn_id) = &ctx.input.asn { ctx.result.add_tag(format!("SOURCE_ASN_{asn_id}")); - types.push((Type::Asn, format!("a:{asn_id}"))); + types.push((Type::Asn, Cow::Owned(asn_id.to_be_bytes().to_vec()))); } if let Some(country) = &ctx.input.country { @@ -68,8 +78,6 @@ impl SpamFilterAnalyzeReputation for Server { let mut reputation = 0.0; for (rep_type, key) in types { - let key = key.into_bytes(); - let mut token = match key_get::(self, ctx.input.span_id, key.clone()).await { Ok(Some(token)) => token, @@ -77,13 +85,16 @@ impl SpamFilterAnalyzeReputation for Server { key_set( self, ctx.input.span_id, - key, - Reputation { - count: 1, - score: ctx.result.score, - } - .serialize(), - config.expiry.into(), + KeyValue::with_prefix( + rep_type.prefix(), + key.as_ref(), + Reputation { + count: 1, + score: ctx.result.score, + } + .serialize(), + ) + .expires(config.expiry), ) .await; continue; @@ -100,9 +111,8 @@ impl SpamFilterAnalyzeReputation for Server { key_set( self, ctx.input.span_id, - key, - token.serialize(), - config.expiry.into(), + KeyValue::with_prefix(rep_type.prefix(), key.as_ref(), token.serialize()) + .expires(config.expiry), ) .await; } @@ -126,6 +136,17 @@ impl SpamFilterAnalyzeReputation for Server { } } +impl Type { + pub fn prefix(&self) -> u8 { + match self { + Type::Ip => KV_REPUTATION_IP, + Type::From => KV_REPUTATION_FROM, + Type::Domain => KV_REPUTATION_DOMAIN, + Type::Asn => KV_REPUTATION_ASN, + } + } +} + impl Serialize for &Reputation { fn serialize(self) -> Vec { let mut buf = Vec::with_capacity(12); diff --git a/crates/spam-filter/src/analysis/score.rs b/crates/spam-filter/src/analysis/score.rs new file mode 100644 index 00000000..cea0cd25 --- /dev/null +++ b/crates/spam-filter/src/analysis/score.rs @@ -0,0 +1,97 @@ +use common::{config::spamfilter::SpamFilterAction, Server}; +use std::{fmt::Write, future::Future, vec}; + +use crate::{modules::bayes::bayes_train_if_balanced, SpamFilterContext}; + +pub trait SpamFilterAnalyzeScore: Sync + Send { + fn spam_filter_score( + &self, + ctx: &mut SpamFilterContext<'_>, + ) -> impl Future> + Send; + + fn spam_filter_finalize( + &self, + ctx: &mut SpamFilterContext<'_>, + header: String, + ) -> impl Future> + Send; +} + +impl SpamFilterAnalyzeScore for Server { + async fn spam_filter_score(&self, ctx: &mut SpamFilterContext<'_>) -> SpamFilterAction { + let mut results = vec![]; + let mut header_len = 60; + + for tag in &ctx.result.tags { + let score = match self.core.spam.list_scores.get(tag) { + Some(SpamFilterAction::Allow(score)) => *score, + Some(SpamFilterAction::Discard) => { + return SpamFilterAction::Discard; + } + Some(SpamFilterAction::Reject) => { + return SpamFilterAction::Reject; + } + None => 0.0, + }; + ctx.result.score += score; + header_len += tag.len() + 10; + results.push((tag.as_str(), score)); + } + + // Sort by score + let mut header = String::with_capacity(header_len); + results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap().then_with(|| a.0.cmp(b.0))); + header.push_str("X-Spam-Result: "); + for (idx, (tag, score)) in results.into_iter().enumerate() { + if idx > 0 { + header.push_str(",\r\n\t"); + } + let _ = write!(&mut header, "{} ({:.2})", tag, score); + } + header.push_str("\r\n"); + + SpamFilterAction::Allow(header) + } + + async fn spam_filter_finalize( + &self, + ctx: &mut SpamFilterContext<'_>, + mut header: String, + ) -> SpamFilterAction { + // Train Bayes classifier + if let Some(config) = self.core.spam.bayes.as_ref().filter(|c| c.auto_learn) { + let was_classified = + ctx.result.has_tag("BAYES_SPAM") || ctx.result.has_tag("BAYES_HAM"); + if ctx.result.has_tag("SPAM_TRAP") + || (ctx.result.score >= config.auto_learn_spam_threshold && !was_classified) + { + bayes_train_if_balanced(self, ctx, true).await; + } else if ctx.result.has_tag("TRUSTED_REPLY") + || (ctx.result.score <= config.auto_learn_ham_threshold && !was_classified) + { + bayes_train_if_balanced(self, ctx, false).await; + } + } + + if self.core.spam.score_reject_threshold > 0.0 + && ctx.result.score >= self.core.spam.score_reject_threshold + { + SpamFilterAction::Reject + } else if self.core.spam.score_discard_threshold > 0.0 + && ctx.result.score >= self.core.spam.score_discard_threshold + { + SpamFilterAction::Discard + } else { + let _ = write!( + &mut header, + "X-Spam-Status: {}, score={:.2}\r\n", + if ctx.result.score >= self.core.spam.score_spam_threshold { + "Yes" + } else { + "No" + }, + ctx.result.score + ); + SpamFilterAction::Allow(header) + } + } +} diff --git a/crates/spam-filter/src/analysis/trusted_reply.rs b/crates/spam-filter/src/analysis/trusted_reply.rs new file mode 100644 index 00000000..b44bd864 --- /dev/null +++ b/crates/spam-filter/src/analysis/trusted_reply.rs @@ -0,0 +1,83 @@ +use std::future::Future; + +use common::{Server, KV_TRUSTED_REPLY}; +use mail_parser::{HeaderName, HeaderValue}; +use store::dispatch::lookup::KeyValue; + +use crate::{modules::bayes::bayes_train_if_balanced, SpamFilterContext}; + +pub trait SpamFilterAnalyzeTrustedReply: Sync + Send { + fn spam_filter_analyze_reply_in( + &self, + ctx: &mut SpamFilterContext<'_>, + ) -> impl Future + Send; + + fn spam_filter_analyze_reply_out( + &self, + ctx: &mut SpamFilterContext<'_>, + ) -> impl Future + Send; +} + +impl SpamFilterAnalyzeTrustedReply for Server { + async fn spam_filter_analyze_reply_in(&self, ctx: &mut SpamFilterContext<'_>) { + if self.core.spam.trusted_reply.is_some() { + for header in ctx.input.message.headers() { + if let HeaderName::InReplyTo | HeaderName::References = &header.name { + let ids: Box + Send> = match &header.value { + HeaderValue::Text(cow) => Box::new(std::iter::once(cow.as_ref())), + HeaderValue::TextList(vec) => Box::new(vec.iter().map(|s| s.as_ref())), + _ => { + continue; + } + }; + + for id in ids { + match self + .in_memory_store() + .key_exists(KeyValue::<()>::build_key(KV_TRUSTED_REPLY, id.as_bytes())) + .await + { + Ok(true) => { + ctx.result.add_tag("TRUSTED_REPLY"); + return; + } + Err(err) => { + trc::error!(err + .span_id(ctx.input.span_id) + .caused_by(trc::location!())); + } + _ => {} + } + } + } + } + } + } + + async fn spam_filter_analyze_reply_out(&self, ctx: &mut SpamFilterContext<'_>) { + if let (Some(hold_time), Some(message_id)) = + (self.core.spam.trusted_reply, ctx.input.message.message_id()) + { + if let Err(err) = self + .in_memory_store() + .key_set( + KeyValue::with_prefix(KV_TRUSTED_REPLY, message_id.as_bytes(), vec![]) + .expires(hold_time), + ) + .await + { + trc::error!(err.span_id(ctx.input.span_id).caused_by(trc::location!())); + } + } + + if self + .core + .spam + .bayes + .as_ref() + .map_or(false, |config| config.auto_learn_reply_ham) + { + bayes_train_if_balanced(self, ctx, false).await; + } + } +} diff --git a/crates/spam-filter/src/lib.rs b/crates/spam-filter/src/lib.rs index 302ef9ab..23c5203a 100644 --- a/crates/spam-filter/src/lib.rs +++ b/crates/spam-filter/src/lib.rs @@ -40,6 +40,7 @@ pub struct SpamFilterInput<'x> { pub env_from_flags: u64, pub env_rcpt_to: &'x [&'x str], + pub account_id: Option, pub is_test: bool, } @@ -84,6 +85,7 @@ pub struct SpamFilterResult { pub rbl_domain_checks: usize, pub rbl_url_checks: usize, pub rbl_email_checks: usize, + pub llm_header: Option, } pub struct SpamFilterContext<'x> { diff --git a/crates/spam-filter/src/modules/bayes.rs b/crates/spam-filter/src/modules/bayes.rs index 814f5a97..823d2351 100644 --- a/crates/spam-filter/src/modules/bayes.rs +++ b/crates/spam-filter/src/modules/bayes.rs @@ -4,68 +4,83 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ +use std::collections::HashSet; + +use common::{ip_to_bytes, Server, KV_BAYES_MODEL_GLOBAL, KV_BAYES_MODEL_USER}; +use mail_auth::DmarcResult; use nlp::{ bayes::{ - cache::BayesTokenCache, tokenize::BayesTokenizer, BayesClassifier, BayesModel, TokenHash, - Weights, + tokenize::{BayesInputToken, BayesTokenizer}, + BayesModel, TokenHash, Weights, + }, + tokenizers::{ + osb::{Gram, OsbToken, OsbTokenizer}, + types::TokenType, }, - tokenizers::osb::{OsbToken, OsbTokenizer}, }; -use sieve::{runtime::Variable, FunctionMap}; -use store::{write::key::KeySerializer, LookupStore, U64_LEN}; +use store::dispatch::lookup::KeyValue; use trc::AddContext; -use super::PluginContext; - -pub fn register_train(plugin_id: u32, fnc_map: &mut FunctionMap) { - fnc_map.set_external_function("bayes_train", plugin_id, 3); -} - -pub fn register_untrain(plugin_id: u32, fnc_map: &mut FunctionMap) { - fnc_map.set_external_function("bayes_untrain", plugin_id, 3); -} - -pub fn register_classify(plugin_id: u32, fnc_map: &mut FunctionMap) { - fnc_map.set_external_function("bayes_classify", plugin_id, 3); -} - -pub fn register_is_balanced(plugin_id: u32, fnc_map: &mut FunctionMap) { - fnc_map.set_external_function("bayes_is_balanced", plugin_id, 3); -} - -pub async fn exec_train(ctx: PluginContext<'_>) -> trc::Result { - train(ctx, true).await -} - -pub async fn exec_untrain(ctx: PluginContext<'_>) -> trc::Result { - train(ctx, false).await -} - -async fn train(ctx: PluginContext<'_>, is_train: bool) -> trc::Result { - let store = match &ctx.arguments[0] { - Variable::String(v) if !v.is_empty() => ctx.server.core.storage.lookups.get(v.as_ref()), - _ => Some(&ctx.server.core.storage.lookup), - } - .ok_or_else(|| { - trc::SieveEvent::RuntimeError - .ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned()) - .details("Unknown store") - })?; - - let text = ctx.arguments[1].to_string(); - let is_spam = ctx.arguments[2].to_bool(); - if text.is_empty() { - trc::bail!(trc::SpamEvent::TrainError - .into_err() - .reason("Empty message")); - } +use crate::{SpamFilterContext, TextPart}; +pub(crate) async fn bayes_train( + server: &Server, + ctx: &SpamFilterContext<'_>, + is_spam: bool, + is_train: bool, +) -> trc::Result<()> { // Train the model let mut model = BayesModel::default(); + + // Train metadata tokens + for token in ctx.spam_tokens() { + model.train_token(TokenHash::from(Gram::Uni { t1: &token }), is_spam); + } + + // Train the subject model.train( - OsbTokenizer::new(BayesTokenizer::new(text.as_ref()), 5), + OsbTokenizer::new( + BayesTokenizer::new( + &ctx.output.subject_thread, + ctx.output.subject_tokens.iter().filter_map(to_bayes_token), + ), + 5, + ), is_spam, ); + + // Train the body + match ctx + .input + .message + .html_body + .first() + .or_else(|| ctx.input.message.text_body.first()) + .and_then(|idx| ctx.output.text_parts.get(*idx)) + { + Some(TextPart::Html { + text_body, tokens, .. + }) => { + model.train( + OsbTokenizer::new( + BayesTokenizer::new(text_body, tokens.iter().filter_map(to_bayes_token_owned)), + 5, + ), + is_spam, + ); + } + Some(TextPart::Plain { text_body, tokens }) => { + model.train( + OsbTokenizer::new( + BayesTokenizer::new(text_body, tokens.iter().filter_map(to_bayes_token)), + 5, + ), + is_spam, + ); + } + _ => {} + } + if model.weights.is_empty() { trc::bail!(trc::SpamEvent::TrainError .into_err() @@ -74,29 +89,27 @@ async fn train(ctx: PluginContext<'_>, is_train: bool) -> trc::Result trc::event!( Spam(trc::SpamEvent::Train), - SpanId = ctx.session_id, + SpanId = ctx.input.span_id, Details = is_spam, Total = model.weights.len(), ); // Update weight and invalidate cache - let bayes_cache = &ctx.server.inner.data.bayes_cache; + let prefix = if ctx.input.account_id.is_some() { + KV_BAYES_MODEL_GLOBAL + } else { + KV_BAYES_MODEL_USER + }; if is_train { for (hash, weights) in model.weights { - store - .counter_incr( - KeySerializer::new(U64_LEN) - .write(hash.h1) - .write(hash.h2) - .finalize(), - weights.into(), - None, - false, - ) + server + .in_memory_store() + .counter_incr(KeyValue::new( + hash.serialize(prefix, ctx.input.account_id), + i64::from(weights), + )) .await .caused_by(trc::location!())?; - - bayes_cache.invalidate(&hash); } // Update training counts @@ -105,97 +118,148 @@ async fn train(ctx: PluginContext<'_>, is_train: bool) -> trc::Result } else { Weights { spam: 0, ham: 1 } }; - store - .counter_incr( - KeySerializer::new(U64_LEN) - .write(0u64) - .write(0u64) - .finalize(), - weights.into(), - None, - false, - ) + server + .in_memory_store() + .counter_incr(KeyValue::new( + TokenHash::serialize_index(prefix, ctx.input.account_id), + i64::from(weights), + )) .await - .caused_by(trc::location!())?; + .caused_by(trc::location!()) + .map(|_| ()) } else { //TODO: Implement untrain - return Ok(false.into()); + Ok(()) } - - bayes_cache.invalidate(&TokenHash::default()); - - Ok(true.into()) } -pub async fn exec_classify(ctx: PluginContext<'_>) -> trc::Result { - let store = match &ctx.arguments[0] { - Variable::String(v) if !v.is_empty() => ctx.server.core.storage.lookups.get(v.as_ref()), - _ => Some(&ctx.server.core.storage.lookup), - } - .ok_or_else(|| { - trc::SieveEvent::RuntimeError - .ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned()) - .details("Unknown store") - })?; - let text = ctx.arguments[1].to_string(); - if text.is_empty() { - trc::bail!(trc::SpamEvent::ClassifyError - .into_err() - .reason("Empty message")); - } - - // Create classifier from defaults - let mut classifier = BayesClassifier::default(); - if let Some(params) = ctx.arguments[2].as_array() { - if let Some(Variable::Integer(value)) = params.first() { - classifier.min_token_hits = *value as u32; - } - if let Some(Variable::Integer(value)) = params.get(1) { - classifier.min_tokens = *value as u32; - } - if let Some(Variable::Float(value)) = params.get(2) { - classifier.min_prob_strength = *value; - } - if let Some(Variable::Integer(value)) = params.get(3) { - classifier.min_learns = *value as u32; - } - } +pub(crate) async fn bayes_classify( + server: &Server, + ctx: &SpamFilterContext<'_>, +) -> trc::Result { + let classifier = if let Some(config) = &server.core.spam.bayes { + &config.classifier + } else { + return Ok(0.0); + }; // Obtain training counts - let bayes_cache = &ctx.server.inner.data.bayes_cache; - let (spam_learns, ham_learns) = bayes_cache - .get_or_update(TokenHash::default(), store) + let prefix = if ctx.input.account_id.is_some() { + KV_BAYES_MODEL_GLOBAL + } else { + KV_BAYES_MODEL_USER + }; + let (spam_learns, ham_learns) = server + .in_memory_store() + .counter_get(TokenHash::serialize_index(prefix, ctx.input.account_id)) .await - .map(|w| (w.spam, w.ham))?; + .map(|w| { + let w = Weights::from(w); + (w.spam, w.ham) + })?; // Make sure we have enough training data if spam_learns < classifier.min_learns || ham_learns < classifier.min_learns { trc::event!( Spam(trc::SpamEvent::NotEnoughTrainingData), - SpanId = ctx.session_id, + SpanId = ctx.input.span_id, Details = vec![ trc::Value::from(spam_learns), trc::Value::from(ham_learns), trc::Value::from(classifier.min_learns) ], ); - return Ok(Variable::default()); + return Ok(0.0); } // Classify the text - let mut tokens = Vec::new(); - for token in OsbTokenizer::<_, TokenHash>::new(BayesTokenizer::new(text.as_ref()), 5) { - let weights = bayes_cache.get_or_update(token.inner, store).await?; - tokens.push(OsbToken { + let mut osb_tokens = Vec::new(); + + // Classify metadata tokens + for token in ctx.spam_tokens() { + let weights = server + .in_memory_store() + .counter_get( + TokenHash::from(Gram::Uni { t1: &token }).serialize(prefix, ctx.input.account_id), + ) + .await + .map(Weights::from)?; + osb_tokens.push(OsbToken { + inner: weights, + idx: 1, + }); + } + + // Classify the subject + for token in OsbTokenizer::<_, TokenHash>::new( + BayesTokenizer::new( + &ctx.output.subject_thread, + ctx.output.subject_tokens.iter().filter_map(to_bayes_token), + ), + 5, + ) { + let weights = server + .in_memory_store() + .counter_get(token.inner.serialize(prefix, ctx.input.account_id)) + .await + .map(Weights::from)?; + osb_tokens.push(OsbToken { inner: weights, idx: token.idx, }); } - let result = classifier.classify(tokens.into_iter(), ham_learns, spam_learns); + + // Classify the body + match ctx + .input + .message + .html_body + .first() + .or_else(|| ctx.input.message.text_body.first()) + .and_then(|idx| ctx.output.text_parts.get(*idx)) + { + Some(TextPart::Html { + text_body, tokens, .. + }) => { + for token in OsbTokenizer::<_, TokenHash>::new( + BayesTokenizer::new(text_body, tokens.iter().filter_map(to_bayes_token_owned)), + 5, + ) { + let weights = server + .in_memory_store() + .counter_get(token.inner.serialize(prefix, ctx.input.account_id)) + .await + .map(Weights::from)?; + osb_tokens.push(OsbToken { + inner: weights, + idx: token.idx, + }); + } + } + Some(TextPart::Plain { text_body, tokens }) => { + for token in OsbTokenizer::<_, TokenHash>::new( + BayesTokenizer::new(text_body, tokens.iter().filter_map(to_bayes_token)), + 5, + ) { + let weights = server + .in_memory_store() + .counter_get(token.inner.serialize(prefix, ctx.input.account_id)) + .await + .map(Weights::from)?; + osb_tokens.push(OsbToken { + inner: weights, + idx: token.idx, + }); + } + } + _ => {} + } + + let result = classifier.classify(osb_tokens.into_iter(), ham_learns, spam_learns); trc::event!( Spam(trc::SpamEvent::Classify), - SpanId = ctx.session_id, + SpanId = ctx.input.span_id, Details = vec![ trc::Value::from(spam_learns), trc::Value::from(ham_learns), @@ -204,38 +268,39 @@ pub async fn exec_classify(ctx: PluginContext<'_>) -> trc::Result { Result = result.unwrap_or_default() ); - Ok(result.map(Variable::from).unwrap_or_default()) + Ok(result.unwrap_or_default()) } -pub async fn exec_is_balanced(ctx: PluginContext<'_>) -> trc::Result { - let min_balance = match &ctx.arguments[2] { - Variable::Float(n) => *n, - Variable::Integer(n) => *n as f64, - _ => 0.0, - }; +pub(crate) async fn bayes_is_balanced( + server: &Server, + ctx: &SpamFilterContext<'_>, + learn_spam: bool, +) -> trc::Result { + let min_balance = server + .core + .spam + .bayes + .as_ref() + .map_or(0.0, |c| c.classifier.min_balance); if min_balance == 0.0 { - return Ok(true.into()); + return Ok(true); } - let store = match &ctx.arguments[0] { - Variable::String(v) if !v.is_empty() => ctx.server.core.storage.lookups.get(v.as_ref()), - _ => Some(&ctx.server.core.storage.lookup), - } - .ok_or_else(|| { - trc::SieveEvent::RuntimeError - .ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned()) - .details("Unknown store") - })?; - - let learn_spam = ctx.arguments[1].to_bool(); - // Obtain training counts - let bayes_cache = &ctx.server.inner.data.bayes_cache; - let (spam_learns, ham_learns) = bayes_cache - .get_or_update(TokenHash::default(), store) + let prefix = if ctx.input.account_id.is_some() { + KV_BAYES_MODEL_GLOBAL + } else { + KV_BAYES_MODEL_USER + }; + let (spam_learns, ham_learns) = server + .in_memory_store() + .counter_get(TokenHash::serialize_index(prefix, ctx.input.account_id)) .await - .map(|w| (w.spam as f64, w.ham as f64))?; + .map(|w| { + let w = Weights::from(w); + (w.spam as f64, w.ham as f64) + })?; let result = if spam_learns > 0.0 || ham_learns > 0.0 { if learn_spam { @@ -249,7 +314,7 @@ pub async fn exec_is_balanced(ctx: PluginContext<'_>) -> trc::Result { trc::event!( Spam(trc::SpamEvent::TrainBalance), - SpanId = ctx.session_id, + SpanId = ctx.input.span_id, Details = vec![ trc::Value::from(learn_spam), trc::Value::from(min_balance), @@ -259,40 +324,74 @@ pub async fn exec_is_balanced(ctx: PluginContext<'_>) -> trc::Result { Result = result ); - Ok(result.into()) + Ok(result) } -trait LookupOrInsert { - async fn get_or_update(&self, hash: TokenHash, get_token: &LookupStore) - -> trc::Result; -} - -impl LookupOrInsert for BayesTokenCache { - async fn get_or_update( - &self, - hash: TokenHash, - get_token: &LookupStore, - ) -> trc::Result { - if let Some(weights) = self.get(&hash) { - Ok(weights.unwrap_or_default()) - } else { - let num = get_token - .counter_get( - KeySerializer::new(U64_LEN) - .write(hash.h1) - .write(hash.h2) - .finalize(), - ) - .await - .caused_by(trc::location!())?; - Ok(if num != 0 { - let weights = Weights::from(num); - self.insert_positive(hash, weights); - weights - } else { - self.insert_negative(hash); - Weights::default() - }) +pub(crate) async fn bayes_train_if_balanced( + server: &Server, + ctx: &SpamFilterContext<'_>, + learn_spam: bool, +) { + let err = match bayes_is_balanced(server, ctx, learn_spam).await { + Ok(true) => match bayes_train(server, ctx, learn_spam, true).await { + Ok(_) => { + return; + } + Err(err) => err, + }, + Ok(false) => { + return; } + Err(err) => err, + }; + + trc::error!(err.span_id(ctx.input.span_id).caused_by(trc::location!())); +} + +const P_FROM_NAME: u8 = 0; +const P_FROM_EMAIL: u8 = 1; +const P_FROM_DOMAIN: u8 = 2; +const P_ASN: u8 = 3; +const P_REMOTE_IP: u8 = 4; + +impl SpamFilterContext<'_> { + pub fn spam_tokens(&self) -> HashSet> { + let mut tokens = HashSet::new(); + if matches!(self.input.dmarc_result, DmarcResult::Pass) { + for addr in [&self.output.env_from_addr, &self.output.from.email] { + if !addr.address.is_empty() { + tokens.insert(add_prefix(P_FROM_EMAIL, addr.address.as_bytes())); + tokens.insert(add_prefix( + P_FROM_DOMAIN, + addr.domain_part.sld_or_default().as_bytes(), + )); + } + } + if let Some(name) = &self.output.from.name { + for name_part in name.split_whitespace() { + tokens.insert(add_prefix(P_FROM_NAME, name_part.to_lowercase().as_bytes())); + } + } + } + if let Some(asn) = self.input.asn { + tokens.insert(add_prefix(P_ASN, &asn.to_be_bytes())); + } + tokens.insert(add_prefix(P_REMOTE_IP, &ip_to_bytes(&self.input.remote_ip))); + tokens } } + +fn add_prefix(prefix: u8, key: &[u8]) -> Vec { + let mut buf = Vec::with_capacity(key.len() + 1); + buf.extend_from_slice(key); + buf.push(prefix); + buf +} + +fn to_bayes_token(token: &TokenType<&str>) -> Option { + token.to_bayes_token() +} + +fn to_bayes_token_owned(token: &TokenType) -> Option { + token.to_bayes_token() +} diff --git a/crates/spam-filter/src/modules/mod.rs b/crates/spam-filter/src/modules/mod.rs index 5576c7f9..b1480539 100644 --- a/crates/spam-filter/src/modules/mod.rs +++ b/crates/spam-filter/src/modules/mod.rs @@ -1,6 +1,7 @@ use common::Server; -use store::{Deserialize, Value}; +use store::{dispatch::lookup::KeyValue, Deserialize, Value}; +pub mod bayes; pub mod dnsbl; pub mod html; pub mod pyzor; @@ -13,7 +14,7 @@ pub(crate) async fn key_get> + std::fmt::De key: impl Into>, ) -> Result, ()> { server - .lookup_store() + .in_memory_store() .key_get(key.into()) .await .map_err(|err| { @@ -21,14 +22,8 @@ pub(crate) async fn key_get> + std::fmt::De }) } -pub(crate) async fn key_set( - server: &Server, - span_id: u64, - key: Vec, - value: Vec, - expires: Option, -) { - if let Err(err) = server.lookup_store().key_set(key, value, expires).await { +pub(crate) async fn key_set(server: &Server, span_id: u64, kv: KeyValue>) { + if let Err(err) = server.in_memory_store().key_set(kv).await { trc::error!(err.span_id(span_id).caused_by(trc::location!())); } } diff --git a/crates/store/src/backend/memory/mod.rs b/crates/store/src/backend/memory/mod.rs index 5426df6d..d0453bd8 100644 --- a/crates/store/src/backend/memory/mod.rs +++ b/crates/store/src/backend/memory/mod.rs @@ -7,12 +7,12 @@ use ahash::AHashMap; use utils::{config::Config, glob::GlobMap}; -use crate::{LookupStore, Stores, Value}; +use crate::{InMemoryStore, Stores, Value}; -pub type MemoryStore = GlobMap>; +pub type StaticMemoryStore = GlobMap>; impl Stores { - pub fn parse_memory_stores(&mut self, config: &mut Config) { + pub fn parse_static_stores(&mut self, config: &mut Config) { let mut lookups = AHashMap::new(); let mut errors = Vec::new(); @@ -68,7 +68,7 @@ impl Stores { // Add entry lookups .entry(id.to_string()) - .or_insert_with(MemoryStore::default) + .or_insert_with(StaticMemoryStore::default) .insert(key, value); } else { errors.push(key.to_string()); @@ -80,8 +80,8 @@ impl Stores { } for (id, store) in lookups { - self.lookup_stores - .insert(id, LookupStore::Memory(store.into())); + self.in_memory_stores + .insert(id, InMemoryStore::Static(store.into())); } } } diff --git a/crates/store/src/config.rs b/crates/store/src/config.rs index 26620727..4f206248 100644 --- a/crates/store/src/config.rs +++ b/crates/store/src/config.rs @@ -4,14 +4,12 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use std::sync::Arc; - use utils::config::{cron::SimpleCron, utils::ParseValue, Config}; use crate::{ backend::fs::FsStore, write::purge::{PurgeSchedule, PurgeStore}, - BlobStore, CompressionAlgo, LookupStore, QueryStore, Store, Stores, + BlobStore, CompressionAlgo, InMemoryStore, Store, Stores, }; #[cfg(feature = "s3")] @@ -106,7 +104,7 @@ impl Stores { store_id.clone(), BlobStore::from(db.clone()).with_compression(compression_algo), ); - self.lookup_stores.insert(store_id, db.into()); + self.in_memory_stores.insert(store_id, db.into()); } } #[cfg(feature = "foundation")] @@ -128,7 +126,7 @@ impl Stores { store_id.clone(), BlobStore::from(db.clone()).with_compression(compression_algo), ); - self.lookup_stores.insert(store_id, db.into()); + self.in_memory_stores.insert(store_id, db.into()); } } #[cfg(feature = "postgres")] @@ -144,7 +142,7 @@ impl Stores { store_id.clone(), BlobStore::from(db.clone()).with_compression(compression_algo), ); - self.lookup_stores.insert(store_id.clone(), db.into()); + self.in_memory_stores.insert(store_id.clone(), db.into()); } } #[cfg(feature = "mysql")] @@ -159,7 +157,7 @@ impl Stores { store_id.clone(), BlobStore::from(db.clone()).with_compression(compression_algo), ); - self.lookup_stores.insert(store_id.clone(), db.into()); + self.in_memory_stores.insert(store_id.clone(), db.into()); } } #[cfg(feature = "sqlite")] @@ -181,7 +179,7 @@ impl Stores { store_id.clone(), BlobStore::from(db.clone()).with_compression(compression_algo), ); - self.lookup_stores.insert(store_id.clone(), db.into()); + self.in_memory_stores.insert(store_id.clone(), db.into()); } } "fs" => { @@ -210,9 +208,9 @@ impl Stores { "redis" => { if let Some(db) = RedisStore::open(config, prefix) .await - .map(LookupStore::from) + .map(InMemoryStore::from) { - self.lookup_stores.insert(store_id, db); + self.in_memory_stores.insert(store_id, db); } } #[cfg(feature = "enterprise")] @@ -262,7 +260,7 @@ impl Stores { id.to_string(), BlobStore::from(db.clone()).with_compression(compression), ); - self.lookup_stores.insert(id.to_string(), db.into()); + self.in_memory_stores.insert(id.to_string(), db.into()); } } "distributed-blob" => { @@ -285,44 +283,7 @@ impl Stores { pub async fn parse_lookups(&mut self, config: &mut Config) { // Parse memory stores - self.parse_memory_stores(config); - - // Add SQL queries as lookup stores - for (store_id, lookup_store) in self.stores.iter().filter_map(|(id, store)| { - if store.is_sql() { - Some((id.clone(), LookupStore::from(store.clone()))) - } else { - None - } - }) { - // Add queries as lookup stores - for lookup_id in config.sub_keys(("store", store_id.as_str(), "query"), "") { - if let Some(query) = config.value(("store", store_id.as_str(), "query", lookup_id)) - { - self.lookup_stores.insert( - format!("{store_id}/{lookup_id}"), - LookupStore::Query(Arc::new(QueryStore { - store: lookup_store.clone(), - query: query.to_string(), - })), - ); - } - } - - // Run init queries on database - for query in config - .values(("store", store_id.as_str(), "init.execute")) - .map(|(_, s)| s.to_string()) - .collect::>() - { - if let Err(err) = lookup_store.query::(&query, Vec::new()).await { - config.new_build_error( - ("store", store_id.as_str()), - format!("Failed to initialize store: {err}"), - ); - } - } - } + self.parse_static_stores(config); // Parse purge schedules if let Some(store) = config @@ -361,8 +322,8 @@ impl Stores { }); } } - for (store_id, store) in &self.lookup_stores { - if matches!(store, LookupStore::Store(_)) { + for (store_id, store) in &self.in_memory_stores { + if matches!(store, InMemoryStore::Store(_)) { self.purge_schedules.push(PurgeSchedule { cron: config .property_or_default::( diff --git a/crates/store/src/dispatch/lookup.rs b/crates/store/src/dispatch/lookup.rs index 904581ef..5c32896a 100644 --- a/crates/store/src/dispatch/lookup.rs +++ b/crates/store/src/dispatch/lookup.rs @@ -7,59 +7,33 @@ use trc::AddContext; use utils::config::Rate; -use crate::{write::LookupClass, Row}; +use crate::write::LookupClass; #[allow(unused_imports)] use crate::{ write::{ key::{DeserializeBigEndian, KeySerializer}, now, BatchBuilder, Operation, ValueClass, ValueOp, }, - Deserialize, IterateParams, LookupStore, QueryResult, Store, Value, ValueKey, U64_LEN, + Deserialize, InMemoryStore, IterateParams, QueryResult, Store, Value, ValueKey, U64_LEN, }; -impl LookupStore { - #[allow(unreachable_patterns)] - #[allow(unused_variables)] - pub async fn query( - &self, - query: &str, - params: Vec>, - ) -> trc::Result { - let result = match self { - #[cfg(feature = "sqlite")] - LookupStore::Store(Store::SQLite(store)) => store.query(query, ¶ms).await, - #[cfg(feature = "postgres")] - LookupStore::Store(Store::PostgreSQL(store)) => store.query(query, ¶ms).await, - #[cfg(feature = "mysql")] - LookupStore::Store(Store::MySQL(store)) => store.query(query, ¶ms).await, - _ => Err(trc::StoreEvent::NotSupported.into_err()), - }; +pub struct KeyValue { + key: Vec, + value: T, + expires: Option, +} - trc::event!( - Store(trc::StoreEvent::SqlQuery), - Details = query.to_string(), - Value = params.as_slice(), - Result = &result, - ); - - result.caused_by(trc::location!()) - } - - pub async fn key_set( - &self, - key: Vec, - value: Vec, - expires: Option, - ) -> trc::Result<()> { +impl InMemoryStore { + pub async fn key_set(&self, kv: KeyValue>) -> trc::Result<()> { match self { - LookupStore::Store(store) => { + InMemoryStore::Store(store) => { let mut batch = BatchBuilder::new(); batch.ops.push(Operation::Value { - class: ValueClass::Lookup(LookupClass::Key(key)), + class: ValueClass::Lookup(LookupClass::Key(kv.key)), op: ValueOp::Set( - KeySerializer::new(value.len() + U64_LEN) - .write(expires.map_or(u64::MAX, |expires| now() + expires)) - .write(value.as_slice()) + KeySerializer::new(kv.value.len() + U64_LEN) + .write(kv.expires.map_or(u64::MAX, |expires| now() + expires)) + .write(kv.value.as_slice()) .finalize() .into(), ), @@ -67,34 +41,20 @@ impl LookupStore { store.write(batch.build()).await.map(|_| ()) } #[cfg(feature = "redis")] - LookupStore::Redis(store) => store.key_set(key, value, expires).await, - LookupStore::Query(lookup) => lookup - .store - .query::( - &lookup.query, - vec![String::from_utf8(key).unwrap_or_default().into()], - ) - .await - .map(|_| ()), - LookupStore::Memory(_) => Err(trc::StoreEvent::NotSupported.into_err()), + InMemoryStore::Redis(store) => store.key_set(kv.key, kv.value, kv.expires).await, + InMemoryStore::Static(_) => Err(trc::StoreEvent::NotSupported.into_err()), } .caused_by(trc::location!()) } - pub async fn counter_incr( - &self, - key: Vec, - value: i64, - expires: Option, - return_value: bool, - ) -> trc::Result { + pub async fn counter_incr(&self, kv: KeyValue) -> trc::Result { match self { - LookupStore::Store(store) => { + InMemoryStore::Store(store) => { let mut batch = BatchBuilder::new(); - if let Some(expires) = expires { + if let Some(expires) = kv.expires { batch.ops.push(Operation::Value { - class: ValueClass::Lookup(LookupClass::Key(key.clone())), + class: ValueClass::Lookup(LookupClass::Key(kv.key.clone())), op: ValueOp::Set( KeySerializer::new(U64_LEN * 2) .write(0u64) @@ -106,34 +66,25 @@ impl LookupStore { } batch.ops.push(Operation::Value { - class: ValueClass::Lookup(LookupClass::Counter(key)), - op: if return_value { - ValueOp::AddAndGet(value) - } else { - ValueOp::AtomicAdd(value) - }, + class: ValueClass::Lookup(LookupClass::Counter(kv.key)), + op: ValueOp::AddAndGet(kv.value), }); - store.write(batch.build()).await.and_then(|r| { - if return_value { - r.last_counter_id() - } else { - Ok(0) - } - }) + store + .write(batch.build()) + .await + .and_then(|r| r.last_counter_id()) } #[cfg(feature = "redis")] - LookupStore::Redis(store) => store.key_incr(key, value, expires).await, - LookupStore::Query(_) | LookupStore::Memory(_) => { - Err(trc::StoreEvent::NotSupported.into_err()) - } + InMemoryStore::Redis(store) => store.key_incr(kv.key, kv.value, kv.expires).await, + InMemoryStore::Static(_) => Err(trc::StoreEvent::NotSupported.into_err()), } .caused_by(trc::location!()) } pub async fn key_delete(&self, key: Vec) -> trc::Result<()> { match self { - LookupStore::Store(store) => { + InMemoryStore::Store(store) => { let mut batch = BatchBuilder::new(); batch.ops.push(Operation::Value { class: ValueClass::Lookup(LookupClass::Key(key)), @@ -142,17 +93,15 @@ impl LookupStore { store.write(batch.build()).await.map(|_| ()) } #[cfg(feature = "redis")] - LookupStore::Redis(store) => store.key_delete(key).await, - LookupStore::Query(_) | LookupStore::Memory(_) => { - Err(trc::StoreEvent::NotSupported.into_err()) - } + InMemoryStore::Redis(store) => store.key_delete(key).await, + InMemoryStore::Static(_) => Err(trc::StoreEvent::NotSupported.into_err()), } .caused_by(trc::location!()) } pub async fn counter_delete(&self, key: Vec) -> trc::Result<()> { match self { - LookupStore::Store(store) => { + InMemoryStore::Store(store) => { let mut batch = BatchBuilder::new(); batch.ops.push(Operation::Value { class: ValueClass::Lookup(LookupClass::Counter(key)), @@ -161,10 +110,8 @@ impl LookupStore { store.write(batch.build()).await.map(|_| ()) } #[cfg(feature = "redis")] - LookupStore::Redis(store) => store.key_delete(key).await, - LookupStore::Query(_) | LookupStore::Memory(_) => { - Err(trc::StoreEvent::NotSupported.into_err()) - } + InMemoryStore::Redis(store) => store.key_delete(key).await, + InMemoryStore::Static(_) => Err(trc::StoreEvent::NotSupported.into_err()), } .caused_by(trc::location!()) } @@ -174,26 +121,15 @@ impl LookupStore { key: Vec, ) -> trc::Result> { match self { - LookupStore::Store(store) => store + InMemoryStore::Store(store) => store .get_value::>(ValueKey::from(ValueClass::Lookup(LookupClass::Key( key, )))) .await .map(|value| value.and_then(|v| v.into())), #[cfg(feature = "redis")] - LookupStore::Redis(store) => store.key_get(key).await, - LookupStore::Query(lookup) => lookup - .store - .query::>( - &lookup.query, - vec![String::from_utf8(key).unwrap_or_default().into()], - ) - .await - .map(|row| { - row.and_then(|row| row.values.into_iter().next()) - .map(|value| T::from(value)) - }), - LookupStore::Memory(store) => Ok(store + InMemoryStore::Redis(store) => store.key_get(key).await, + InMemoryStore::Static(store) => Ok(store .get(std::str::from_utf8(&key).unwrap_or_default()) .map(|value| T::from(value.clone()))), } @@ -202,7 +138,7 @@ impl LookupStore { pub async fn counter_get(&self, key: Vec) -> trc::Result { match self { - LookupStore::Store(store) => { + InMemoryStore::Store(store) => { store .get_counter(ValueKey::from(ValueClass::Lookup(LookupClass::Counter( key, @@ -210,33 +146,23 @@ impl LookupStore { .await } #[cfg(feature = "redis")] - LookupStore::Redis(store) => store.counter_get(key).await, - LookupStore::Query(_) | LookupStore::Memory(_) => { - Err(trc::StoreEvent::NotSupported.into_err()) - } + InMemoryStore::Redis(store) => store.counter_get(key).await, + InMemoryStore::Static(_) => Err(trc::StoreEvent::NotSupported.into_err()), } .caused_by(trc::location!()) } pub async fn key_exists(&self, key: Vec) -> trc::Result { match self { - LookupStore::Store(store) => store + InMemoryStore::Store(store) => store .get_value::>(ValueKey::from(ValueClass::Lookup(LookupClass::Key( key, )))) .await .map(|value| matches!(value, Some(LookupValue::Value(())))), #[cfg(feature = "redis")] - LookupStore::Redis(store) => store.key_exists(key).await, - LookupStore::Query(lookup) => lookup - .store - .query::>( - &lookup.query, - vec![String::from_utf8(key).unwrap_or_default().into()], - ) - .await - .map(|row| row.is_some()), - LookupStore::Memory(store) => Ok(store + InMemoryStore::Redis(store) => store.key_exists(key).await, + InMemoryStore::Static(store) => Ok(store .get(std::str::from_utf8(&key).unwrap_or_default()) .is_some()), } @@ -245,6 +171,7 @@ impl LookupStore { pub async fn is_rate_allowed( &self, + prefix: u8, key: &[u8], rate: &Rate, soft_check: bool, @@ -254,12 +181,13 @@ impl LookupStore { let range_end = (range_start * rate.period.as_secs()) + rate.period.as_secs(); let expires_in = range_end - now; - let mut bucket = Vec::with_capacity(key.len() + U64_LEN); + let mut bucket = Vec::with_capacity(key.len() + U64_LEN + 1); + bucket.push(prefix); bucket.extend_from_slice(key); bucket.extend_from_slice(range_start.to_be_bytes().as_slice()); let requests = if !soft_check { - self.counter_incr(bucket, 1, expires_in.into(), true) + self.counter_incr(KeyValue::new(bucket, 1).expires(expires_in)) .await .caused_by(trc::location!())? } else { @@ -273,9 +201,15 @@ impl LookupStore { } } - pub async fn purge_lookup_store(&self) -> trc::Result<()> { + pub async fn try_lock(&self, prefix: u8, key: &[u8], duration: u64) -> trc::Result { + self.counter_incr(KeyValue::with_prefix(prefix, key, 1).expires(duration)) + .await + .map(|count| count == 1) + } + + pub async fn purge_in_memory_store(&self) -> trc::Result<()> { match self { - LookupStore::Store(store) => { + InMemoryStore::Store(store) => { // Delete expired keys and counters let from_key = ValueKey::from(ValueClass::Lookup(LookupClass::Key(vec![0u8]))); let to_key = @@ -354,8 +288,8 @@ impl LookupStore { } } #[cfg(feature = "redis")] - LookupStore::Redis(_) => {} - LookupStore::Query(_) | LookupStore::Memory(_) => {} + InMemoryStore::Redis(_) => {} + InMemoryStore::Static(_) => {} } Ok(()) @@ -363,12 +297,48 @@ impl LookupStore { pub fn is_sql(&self) -> bool { match self { - LookupStore::Store(store) => store.is_sql(), + InMemoryStore::Store(store) => store.is_sql(), _ => false, } } } +impl KeyValue { + pub fn build_key(prefix: u8, key: impl AsRef<[u8]>) -> Vec { + let key_ = key.as_ref(); + let mut key = Vec::with_capacity(key_.len() + 1); + key.push(prefix); + key.extend_from_slice(key_); + key + } + + pub fn with_prefix(prefix: u8, key: impl AsRef<[u8]>, value: T) -> Self { + Self { + key: Self::build_key(prefix, key), + value, + expires: None, + } + } + + pub fn new(key: impl Into>, value: T) -> Self { + Self { + key: key.into(), + value, + expires: None, + } + } + + pub fn expires(mut self, expires: u64) -> Self { + self.expires = expires.into(); + self + } + + pub fn expires_opt(mut self, expires: Option) -> Self { + self.expires = expires; + self + } +} + enum LookupValue { Value(T), None, diff --git a/crates/store/src/dispatch/store.rs b/crates/store/src/dispatch/store.rs index 872dc71f..3931ded2 100644 --- a/crates/store/src/dispatch/store.rs +++ b/crates/store/src/dispatch/store.rs @@ -18,8 +18,9 @@ use crate::{ now, AnyClass, AnyKey, AssignedIds, Batch, BatchBuilder, BitmapClass, BitmapHash, Operation, ReportClass, ValueClass, ValueOp, }, - BitmapKey, Deserialize, IterateParams, Key, Store, ValueKey, SUBSPACE_BITMAP_ID, - SUBSPACE_BITMAP_TAG, SUBSPACE_BITMAP_TEXT, SUBSPACE_INDEXES, SUBSPACE_LOGS, U32_LEN, + BitmapKey, Deserialize, IterateParams, Key, QueryResult, Store, Value, ValueKey, + SUBSPACE_BITMAP_ID, SUBSPACE_BITMAP_TAG, SUBSPACE_BITMAP_TEXT, SUBSPACE_INDEXES, SUBSPACE_LOGS, + U32_LEN, }; use super::DocumentSet; @@ -154,6 +155,33 @@ impl Store { .caused_by(trc::location!()) } + #[allow(unreachable_patterns)] + #[allow(unused_variables)] + pub async fn sql_query( + &self, + query: &str, + params: Vec>, + ) -> trc::Result { + let result = match self { + #[cfg(feature = "sqlite")] + Self::SQLite(store) => store.query(query, ¶ms).await, + #[cfg(feature = "postgres")] + Self::PostgreSQL(store) => store.query(query, ¶ms).await, + #[cfg(feature = "mysql")] + Self::MySQL(store) => store.query(query, ¶ms).await, + _ => Err(trc::StoreEvent::NotSupported.into_err()), + }; + + trc::event!( + Store(trc::StoreEvent::SqlQuery), + Details = query.to_string(), + Value = params.as_slice(), + Result = &result, + ); + + result.caused_by(trc::location!()) + } + pub async fn write(&self, batch: Batch) -> trc::Result { #[cfg(feature = "test_mode")] if std::env::var("PARANOID_WRITE").map_or(false, |v| v == "1") { diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs index 431ccfa4..43df6482 100644 --- a/crates/store/src/lib.rs +++ b/crates/store/src/lib.rs @@ -15,7 +15,7 @@ pub mod write; pub use ahash; use ahash::AHashMap; -use backend::{fs::FsStore, memory::MemoryStore}; +use backend::{fs::FsStore, memory::StaticMemoryStore}; pub use blake3; pub use parking_lot; pub use rand; @@ -171,7 +171,7 @@ pub struct Stores { pub stores: AHashMap, pub blob_stores: AHashMap, pub fts_stores: AHashMap, - pub lookup_stores: AHashMap, + pub in_memory_stores: AHashMap, pub purge_schedules: Vec, } @@ -225,18 +225,11 @@ pub enum FtsStore { } #[derive(Clone, Debug)] -pub enum LookupStore { +pub enum InMemoryStore { Store(Store), - Query(Arc), #[cfg(feature = "redis")] Redis(Arc), - Memory(Arc), -} - -#[derive(Debug)] -pub struct QueryStore { - pub store: LookupStore, - pub query: String, + Static(Arc), } #[cfg(feature = "sqlite")] @@ -311,7 +304,7 @@ impl From for FtsStore { } #[cfg(feature = "redis")] -impl From for LookupStore { +impl From for InMemoryStore { fn from(store: RedisStore) -> Self { Self::Redis(Arc::new(store)) } @@ -332,7 +325,7 @@ impl From for BlobStore { } } -impl From for LookupStore { +impl From for InMemoryStore { fn from(store: Store) -> Self { Self::Store(store) } @@ -347,7 +340,7 @@ impl Default for BlobStore { } } -impl Default for LookupStore { +impl Default for InMemoryStore { fn default() -> Self { Self::Store(Store::None) } diff --git a/crates/store/src/write/purge.rs b/crates/store/src/write/purge.rs index 79886b2e..626c9116 100644 --- a/crates/store/src/write/purge.rs +++ b/crates/store/src/write/purge.rs @@ -10,13 +10,13 @@ use tokio::sync::watch; use trc::PurgeEvent; use utils::config::cron::SimpleCron; -use crate::{BlobStore, LookupStore, Store}; +use crate::{BlobStore, InMemoryStore, Store}; #[derive(Clone)] pub enum PurgeStore { Data(Store), Blobs { store: Store, blob_store: BlobStore }, - Lookup(LookupStore), + Lookup(InMemoryStore), } #[derive(Clone)] @@ -59,7 +59,7 @@ impl PurgeSchedule { PurgeStore::Blobs { store, blob_store } => { store.purge_blobs(blob_store.clone()).await } - PurgeStore::Lookup(store) => store.purge_lookup_store().await, + PurgeStore::Lookup(store) => store.purge_in_memory_store().await, }; if let Err(err) = result { diff --git a/resources/config/spamfilter/scripts/bayes_classify.sieve b/resources/config/spamfilter/scripts/bayes_classify.sieve deleted file mode 100644 index ce5835df..00000000 --- a/resources/config/spamfilter/scripts/bayes_classify.sieve +++ /dev/null @@ -1,17 +0,0 @@ -if eval "!t.SPAM_TRAP && !t.TRUSTED_REPLY" { - - # Classification parameters - # min_token_hits: 2 - # min_tokens: 11 - # min_prob_strength: 0.05 - # min_learns: 200 - - let "bayes_result" "bayes_classify(SPAM_DB, body_and_subject, [2, 11, 0.05, 200])"; - if eval "!is_empty(bayes_result)" { - if eval "bayes_result > 0.7" { - let "t.BAYES_SPAM" "1"; - } elsif eval "bayes_result < 0.5" { - let "t.BAYES_HAM" "1"; - } - } -} diff --git a/resources/config/spamfilter/scripts/epilogue.sieve b/resources/config/spamfilter/scripts/epilogue.sieve deleted file mode 100644 index 4eb1fd33..00000000 --- a/resources/config/spamfilter/scripts/epilogue.sieve +++ /dev/null @@ -1,28 +0,0 @@ - -# Train the bayes classifier automatically -if eval "AUTOLEARN_ENABLE && (score >= AUTOLEARN_SPAM_THRESHOLD || score <= AUTOLEARN_HAM_THRESHOLD)" { - let "is_spam" "score >= AUTOLEARN_SPAM_THRESHOLD"; - eval "bayes_is_balanced(SPAM_DB, is_spam, AUTOLEARN_SPAM_HAM_BALANCE) && - bayes_train(SPAM_DB, body_and_subject, is_spam)"; -} - -# Process score actions -if eval "SCORE_REJECT_THRESHOLD && score >= SCORE_REJECT_THRESHOLD" { - reject "Your message has been rejected because it has an excessive spam score. If you feel this is an error, please contact the postmaster."; - stop; -} elsif eval "SCORE_DISCARD_THRESHOLD && score >= SCORE_DISCARD_THRESHOLD" { - discard; - stop; -} elsif eval "ADD_HEADER_SPAM" { - let "spam_status" ""; - if eval "score >= SCORE_SPAM_THRESHOLD" { - let "spam_status" "'Yes, score=' + score"; - } else { - let "spam_status" "'No, score=' + score"; - } - eval "add_header('X-Spam-Status', spam_status)"; - if eval "!is_empty(spam_result)" { - eval "add_header('X-Spam-Result', spam_result)"; - } -} - diff --git a/resources/config/spamfilter/scripts/llm.sieve b/resources/config/spamfilter/scripts/llm.sieve deleted file mode 100644 index c1e16d5e..00000000 --- a/resources/config/spamfilter/scripts/llm.sieve +++ /dev/null @@ -1,41 +0,0 @@ -if eval "LLM_MODEL && LLM_PROMPT_TEXT" { - let "llm_result" "trim(split_n(llm_prompt(LLM_MODEL, LLM_PROMPT_TEXT + '\n\nSubject: ' + subject_clean + '\n\n' + text_body, 0.5), ',', 3))"; - - if eval "eq_ignore_case(llm_result[0], 'Unsolicited')" { - if eval "eq_ignore_case(llm_result[1], 'High')" { - let "t.LLM_UNSOLICITED_HIGH" "1"; - } elsif eval "eq_ignore_case(llm_result[1], 'Medium')" { - let "t.LLM_UNSOLICITED_MEDIUM" "1"; - } else { - let "t.LLM_UNSOLICITED_LOW" "1"; - } - } elsif eval "eq_ignore_case(llm_result[0], 'Commercial')" { - if eval "eq_ignore_case(llm_result[1], 'High')" { - let "t.LLM_COMMERCIAL_HIGH" "1"; - } elsif eval "eq_ignore_case(llm_result[1], 'Medium')" { - let "t.LLM_COMMERCIAL_MEDIUM" "1"; - } else { - let "t.LLM_COMMERCIAL_LOW" "1"; - } - } elsif eval "eq_ignore_case(llm_result[0], 'Harmful')" { - if eval "eq_ignore_case(llm_result[1], 'High')" { - let "t.LLM_HARMFUL_HIGH" "1"; - } elsif eval "eq_ignore_case(llm_result[1], 'Medium')" { - let "t.LLM_HARMFUL_MEDIUM" "1"; - } else { - let "t.LLM_HARMFUL_LOW" "1"; - } - } elsif eval "eq_ignore_case(llm_result[0], 'Legitimate')" { - if eval "eq_ignore_case(llm_result[1], 'High')" { - let "t.LLM_LEGITIMATE_HIGH" "1"; - } elsif eval "eq_ignore_case(llm_result[1], 'Medium')" { - let "t.LLM_LEGITIMATE_MEDIUM" "1"; - } else { - let "t.LLM_LEGITIMATE_LOW" "1"; - } - } - - if eval "ADD_HEADER_LLM && count(llm_result) > 2" { - eval "add_header('X-Spam-Llm-Result', 'Category=' + llm_result[0] + '; Confidence=' + llm_result[1] + '; Explanation=' + llm_result[2])"; - } -} diff --git a/resources/config/spamfilter/scripts/replies_in.sieve b/resources/config/spamfilter/scripts/replies_in.sieve deleted file mode 100644 index 4f485bde..00000000 --- a/resources/config/spamfilter/scripts/replies_in.sieve +++ /dev/null @@ -1,12 +0,0 @@ - -let "message_ids" "header.In-Reply-To:References"; - -let "i" "count(message_ids)"; -while "i > 0" { - let "i" "i - 1"; - - if eval "key_exists(SPAM_DB, 'm:' + message_ids[i])" { - let "t.TRUSTED_REPLY" "1"; - break; - } -} diff --git a/resources/config/spamfilter/scripts/replies_out.sieve b/resources/config/spamfilter/scripts/replies_out.sieve deleted file mode 100644 index 76c21875..00000000 --- a/resources/config/spamfilter/scripts/replies_out.sieve +++ /dev/null @@ -1,12 +0,0 @@ - -# This script should be used on authenticated SMTP sessions only -let "message_id" "header.Message-ID"; - -if eval "!is_empty(message_id)" { - # Store the message ID for 30 days - eval "key_set(SPAM_DB, 'm:' + message_id, '', 2592000)"; - - if eval "AUTOLEARN_ENABLE && AUTOLEARN_REPLIES_HAM && bayes_is_balanced(SPAM_DB, false, AUTOLEARN_SPAM_HAM_BALANCE)" { - eval "bayes_train(SPAM_DB, thread_name(header.subject) + ' ' + body.to_text, false)"; - } -} diff --git a/resources/config/spamfilter/scripts/scores.sieve b/resources/config/spamfilter/scripts/scores.sieve deleted file mode 100644 index f300c08c..00000000 --- a/resources/config/spamfilter/scripts/scores.sieve +++ /dev/null @@ -1,27 +0,0 @@ -# Add scores -let "tags" "var_names()"; -let "i" "count(tags)"; -let "spam_result" ""; -while "i > 0" { - let "i" "i - 1"; - let "tag" "tags[i]"; - let "tag_score" "key_get('spam-scores', tag)"; - - if eval "is_number(tag_score)" { - let "score" "score + tag_score"; - if eval "ADD_HEADER_SPAM_RESULT" { - if eval "!is_empty(spam_result)" { - let "spam_result" "spam_result + ',\r\n\t' + tag + ' (' + tag_score + ')'"; - } else { - let "spam_result" "spam_result + tag + ' (' + tag_score + ')'"; - } - } - } elsif eval "tag_score == 'reject'" { - let "SCORE_REJECT_THRESHOLD" "1"; - let "score" "2"; - break; - } elsif eval "tag_score == 'discard'" { - discard; - stop; - } -} diff --git a/resources/config/spamfilter/scripts/spamtrap.sieve b/resources/config/spamfilter/scripts/spamtrap.sieve deleted file mode 100644 index 201a1700..00000000 --- a/resources/config/spamfilter/scripts/spamtrap.sieve +++ /dev/null @@ -1,9 +0,0 @@ - -# Check if the message was sent to a spam trap address -if eval "AUTOLEARN_ENABLE && key_exists('spam-trap', envelope.to)" { - eval "bayes_is_balanced(SPAM_DB, false, AUTOLEARN_SPAM_HAM_BALANCE) && bayes_train(SPAM_DB, body_and_subject, true)"; - let "t.SPAM_TRAP" "1"; - - # Disable autolearn so the classifier is not trained twice - let "AUTOLEARN_ENABLE" "0"; -} diff --git a/tests/Cargo.toml b/tests/Cargo.toml index f14069c7..e877b636 100644 --- a/tests/Cargo.toml +++ b/tests/Cargo.toml @@ -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" } diff --git a/tests/src/directory/mod.rs b/tests/src/directory/mod.rs index 67aca0f8..b253c058 100644 --- a/tests/src/directory/mod.rs +++ b/tests/src/directory/mod.rs @@ -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 { diff --git a/tests/src/directory/sql.rs b/tests/src/directory/sql.rs index 72b9156a..0eb5e1eb 100644 --- a/tests/src/directory/sql.rs +++ b/tests/src/directory/sql.rs @@ -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::(&format!("DROP TABLE IF EXISTS {table}"), vec![]) + .sql_query::(&format!("DROP TABLE IF EXISTS {table}"), vec![]) .await .unwrap(); } @@ -383,7 +383,7 @@ impl DirectoryStore { }; self.store - .query::(&query, vec![]) + .sql_query::(&query, vec![]) .await .unwrap_or_else(|_| panic!("failed for {query}")); } @@ -396,7 +396,7 @@ impl DirectoryStore { "individual" }; self.store - .query::( + .sql_query::( 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::( + .sql_query::( 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::( + .sql_query::( 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::( + .sql_query::( 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::( + .sql_query::( 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::( + .sql_query::( 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::( + .sql_query::( 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"))] { diff --git a/tests/src/jmap/auth_oauth.rs b/tests/src/jmap/auth_oauth.rs index 83b43d53..fb6f5fe8 100644 --- a/tests/src/jmap/auth_oauth.rs +++ b/tests/src/jmap/auth_oauth.rs @@ -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); diff --git a/tests/src/jmap/enterprise.rs b/tests/src/jmap/enterprise.rs index 63264221..78cf4cfc 100644 --- a/tests/src/jmap/enterprise.rs +++ b/tests/src/jmap/enterprise.rs @@ -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 diff --git a/tests/src/smtp/inbound/antispam.rs b/tests/src/smtp/inbound/antispam.rs index 90a89a2b..9e4beb2a 100644 --- a/tests/src/smtp/inbound/antispam.rs +++ b/tests/src/smtp/inbound/antispam.rs @@ -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); diff --git a/tests/src/smtp/lookup/sql.rs b/tests/src/smtp/lookup/sql.rs index 45e0ab03..bca69908 100644 --- a/tests/src/smtp/lookup/sql.rs +++ b/tests/src/smtp/lookup/sql.rs @@ -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::(query, Vec::new()) + .sql_query::(query, Vec::new()) .await .unwrap(); } diff --git a/tests/src/store/lookup.rs b/tests/src/store/lookup.rs index 09e45b06..6ab529fe 100644 --- a/tests/src/store/lookup.rs +++ b/tests/src/store/lookup.rs @@ -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::(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::(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; } }