From 05ef5a7c1040d25aa1817ef3f69869326ea8ce14 Mon Sep 17 00:00:00 2001 From: mdecimus <11444311+mdecimus@users.noreply.github.com> Date: Fri, 28 Nov 2025 19:00:04 +0100 Subject: [PATCH] Spam filter performance and accuracy improvements (part 1) --- crates/common/src/config/inner.rs | 17 +- crates/common/src/config/spamfilter.rs | 200 ++++------ crates/common/src/lib.rs | 17 +- crates/nlp/Cargo.toml | 3 +- crates/nlp/src/bayes/classify.rs | 145 -------- crates/nlp/src/bayes/mod.rs | 185 ---------- crates/nlp/src/bayes/train.rs | 55 --- crates/nlp/src/classifier/feature.rs | 56 +++ crates/nlp/src/classifier/mod.rs | 8 + crates/nlp/src/classifier/sgd.rs | 317 ++++++++++++++++ crates/nlp/src/lib.rs | 97 +---- crates/nlp/src/tokenizers/mod.rs | 2 +- crates/nlp/src/tokenizers/osb.rs | 343 ------------------ .../tokenize.rs => tokenizers/stream.rs} | 106 +++--- crates/spam-filter/src/analysis/mod.rs | 1 - crates/spam-filter/src/analysis/reputation.rs | 243 ++++--------- crates/spam-filter/src/analysis/score.rs | 10 +- .../spam-filter/src/analysis/trusted_reply.rs | 88 ----- crates/spam-filter/src/modules/bayes.rs | 2 +- crates/utils/src/config/utils.rs | 8 + 20 files changed, 604 insertions(+), 1299 deletions(-) delete mode 100644 crates/nlp/src/bayes/classify.rs delete mode 100644 crates/nlp/src/bayes/mod.rs delete mode 100644 crates/nlp/src/bayes/train.rs create mode 100644 crates/nlp/src/classifier/feature.rs create mode 100644 crates/nlp/src/classifier/mod.rs create mode 100644 crates/nlp/src/classifier/sgd.rs delete mode 100644 crates/nlp/src/tokenizers/osb.rs rename crates/nlp/src/{bayes/tokenize.rs => tokenizers/stream.rs} (98%) delete mode 100644 crates/spam-filter/src/analysis/trusted_reply.rs diff --git a/crates/common/src/config/inner.rs b/crates/common/src/config/inner.rs index 13386389..cb642b3a 100644 --- a/crates/common/src/config/inner.rs +++ b/crates/common/src/config/inner.rs @@ -9,7 +9,10 @@ use crate::{ CacheSwap, Caches, Data, DavResource, DavResources, MailboxCache, MessageStoreCache, MessageUidCache, TlsConnectors, auth::{AccessToken, roles::RolePermissions}, - config::smtp::resolver::{Policy, Tlsa}, + config::{ + smtp::resolver::{Policy, Tlsa}, + spamfilter::Reputation, + }, listener::blocked::BlockedIps, manager::webadmin::WebAdminManager, }; @@ -17,7 +20,7 @@ use ahash::{AHashMap, AHashSet}; use arc_swap::ArcSwap; use mail_auth::{MX, Parameters, Txt}; use mail_send::smtp::tls::build_tls_connector; -use nlp::bayes::{TokenHash, Weights}; +use nlp::classifier::sgd::SGDClassifier; use parking_lot::RwLock; use std::{ net::{IpAddr, Ipv4Addr, Ipv6Addr}, @@ -49,6 +52,8 @@ impl Data { } Data { + spam_classifier: ArcSwap::from_pointee(SGDClassifier::default()), + spam_reputation: ArcSwap::from_pointee(Reputation::default()), tls_certificates: ArcSwap::from_pointee(certificates), tls_self_signed_cert: build_self_signed_cert( subject_names.into_iter().collect::>(), @@ -138,12 +143,6 @@ impl Caches { (std::mem::size_of::() + (500 * std::mem::size_of::())) as u64, ), - bayes: CacheWithTtl::from_config( - config, - "bayes", - MB_10, - (std::mem::size_of::() + std::mem::size_of::()) as u64, - ), dns_txt: CacheWithTtl::from_config( config, "dns.txt", @@ -223,6 +222,8 @@ impl Caches { impl Default for Data { fn default() -> Self { Self { + spam_classifier: Default::default(), + spam_reputation: Default::default(), tls_certificates: Default::default(), tls_self_signed_cert: Default::default(), blocked_ips: Default::default(), diff --git a/crates/common/src/config/spamfilter.rs b/crates/common/src/config/spamfilter.rs index 95584f03..67cffc21 100644 --- a/crates/common/src/config/spamfilter.rs +++ b/crates/common/src/config/spamfilter.rs @@ -4,14 +4,14 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ +use super::{Variable, functions::ResolveVariable, if_block::IfBlock, tokenizer::TokenMap}; +use ahash::{AHashMap, AHashSet}; +use compact_str::CompactString; +use mail_auth::common::resolver::ToReverseName; use std::{ net::{IpAddr, SocketAddr}, time::Duration, }; - -use ahash::AHashSet; -use mail_auth::common::resolver::ToReverseName; -use nlp::bayes::BayesClassifier; use tokio::net::lookup_host; use utils::{ cache::CacheItemWeight, @@ -19,20 +19,38 @@ use utils::{ glob::GlobMap, }; -use super::{Variable, functions::ResolveVariable, if_block::IfBlock, tokenizer::TokenMap}; +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum ReputationType { + Domain(CompactString), + Asn(u32), + Ip(IpAddr), +} + +#[derive(Debug, Clone, Copy, Default)] +pub struct ReputationCount { + pub ham: u32, + pub spam: u32, +} + +#[derive(Debug, Clone, Default)] +pub struct Reputation { + pub items: AHashMap, + pub total: ReputationCount, + pub last_fetch: u64, +} #[derive(Debug, Clone, Default)] pub struct SpamFilterConfig { pub enabled: bool, pub card_is_ham: bool, + pub grey_list_expiry: Option, + pub dnsbl: DnsBlConfig, pub rules: SpamFilterRules, pub lists: SpamFilterLists, pub pyzor: Option, - pub reputation: Option, - pub bayes: Option, + pub classifier: Option, pub scores: SpamFilterScoreConfig, - pub expiry: SpamFilterExpiryConfig, pub headers: SpamFilterHeaderConfig, } @@ -40,7 +58,6 @@ pub struct SpamFilterConfig { pub struct SpamFilterHeaderConfig { pub status: Option, pub result: Option, - pub bayes_result: Option, pub llm: Option, } @@ -51,12 +68,6 @@ pub struct SpamFilterScoreConfig { pub spam_threshold: f64, } -#[derive(Debug, Clone, Default)] -pub struct SpamFilterExpiryConfig { - pub grey_list: Option, - pub trusted_reply: Option, -} - #[derive(Debug, Clone, Default)] pub struct DnsBlConfig { pub max_ip_checks: usize, @@ -80,29 +91,14 @@ pub enum SpamFilterAction { } #[derive(Debug, Clone, Default)] -pub struct BayesConfig { - pub classifier: BayesClassifier, - pub auto_learn: bool, +pub struct ClassifierConfig { + pub epochs: usize, + pub feature_hash_size: usize, + pub alpha: f32, pub auto_learn_reply_ham: bool, - pub auto_learn_spam_threshold: f64, - pub auto_learn_ham_threshold: f64, pub auto_learn_card_is_ham: bool, pub score_spam: f64, pub score_ham: f64, - pub account_score_spam: f64, - pub account_score_ham: f64, - pub account_classify: bool, -} - -#[derive(Debug, Clone, Default)] -pub struct ReputationConfig { - pub expiry: u64, - pub token_score: f64, - pub factor: f64, - pub ip_weight: f64, - pub domain_weight: f64, - pub asn_weight: f64, - pub sender_weight: f64, } #[derive(Debug, Clone)] @@ -187,11 +183,13 @@ impl SpamFilterConfig { rules: SpamFilterRules::parse(config), lists: SpamFilterLists::parse(config), pyzor: PyzorConfig::parse(config).await, - reputation: ReputationConfig::parse(config), - bayes: BayesConfig::parse(config), + classifier: ClassifierConfig::parse(config), scores: SpamFilterScoreConfig::parse(config), - expiry: SpamFilterExpiryConfig::parse(config), headers: SpamFilterHeaderConfig::parse(config), + grey_list_expiry: config + .property::>("spam-filter.grey-list.duration") + .unwrap_or_default() + .map(|d| d.as_secs()), } } } @@ -326,7 +324,6 @@ impl SpamFilterHeaderConfig { ("status", &mut header.status), ("result", &mut header.result), ("llm", &mut header.llm), - ("bayes", &mut header.bayes_result), ] { if config .property_or_default(("spam-filter.header", typ, "enable"), "true") @@ -474,99 +471,45 @@ impl PyzorConfig { } } -impl ReputationConfig { +impl ClassifierConfig { pub fn parse(config: &mut Config) -> Option { if !config - .property_or_default("spam-filter.reputation.enable", "false") - .unwrap_or(false) - { - return None; - } - - ReputationConfig { - expiry: config - .property_or_default::("spam-filter.reputation.expiry", "30d") - .map(|d| d.as_secs()) - .unwrap_or(2592000), - token_score: config - .property_or_default("spam-filter.reputation.score", "0.98") - .unwrap_or(0.98), - factor: config - .property_or_default("spam-filter.reputation.factor", "0.5") - .unwrap_or(0.5), - ip_weight: config - .property_or_default("spam-filter.reputation.weight.ip", "0.2") - .unwrap_or(0.2), - domain_weight: config - .property_or_default("spam-filter.reputation.weight.domain", "0.2") - .unwrap_or(0.2), - asn_weight: config - .property_or_default("spam-filter.reputation.weight.asn", "0.1") - .unwrap_or(0.1), - sender_weight: config - .property_or_default("spam-filter.reputation.weight.sender", "0.5") - .unwrap_or(0.5), - } - .into() - } -} - -impl BayesConfig { - pub fn parse(config: &mut Config) -> Option { - if !config - .property_or_default("spam-filter.bayes.enable", "true") + .property_or_default("spam-filter.classifier.enable", "true") .unwrap_or(true) { return None; } - BayesConfig { - classifier: BayesClassifier { - min_token_hits: config - .property_or_default("spam-filter.bayes.classify.tokens.hits", "2") - .unwrap_or(2), - min_tokens: config - .property_or_default("spam-filter.bayes.classify.tokens.min", "11") - .unwrap_or(11), - min_prob_strength: config - .property_or_default("spam-filter.bayes.classify.strength", "0.05") - .unwrap_or(0.05), - min_learns: config - .property_or_default("spam-filter.bayes.classify.learns", "200") - .unwrap_or(200), - min_balance: config - .property_or_default("spam-filter.bayes.classify.balance", "0.9") - .unwrap_or(0.9), - }, - auto_learn: config - .property_or_default("spam-filter.bayes.auto-learn.enable", "true") - .unwrap_or(true), - auto_learn_reply_ham: config - .property_or_default("spam-filter.bayes.auto-learn.trusted-reply", "true") - .unwrap_or(true), - auto_learn_spam_threshold: config - .property_or_default("spam-filter.bayes.auto-learn.threshold.spam", "6.0") - .unwrap_or(6.0), - auto_learn_ham_threshold: config - .property_or_default("spam-filter.bayes.auto-learn.threshold.ham", "-1.0") - .unwrap_or(-2.0), + let feature_hash_size: usize = config + .property_or_default("spam-filter.classifier.feature-hash-size", "1048576") + .unwrap_or(1048576); + + if !feature_hash_size.is_power_of_two() { + config.new_build_error( + "spam-filter.classifier.feature-hash-size", + "Feature hash size must be a power of two.", + ); + } + + ClassifierConfig { + feature_hash_size, + epochs: config + .property_or_default("spam-filter.classifier.epochs", "1000") + .unwrap_or(1000), + alpha: config + .property_or_default("spam-filter.classifier.alpha", "0.0001") + .unwrap_or(0.0001), score_spam: config - .property_or_default("spam-filter.bayes.score.spam", "0.7") + .property_or_default("spam-filter.classifier.score.spam", "0.7") .unwrap_or(0.7), score_ham: config - .property_or_default("spam-filter.bayes.score.ham", "0.5") - .unwrap_or(0.5), - account_classify: config - .property_or_default("spam-filter.bayes.account.enable", "false") - .unwrap_or(false), - account_score_spam: config - .property_or_default("spam-filter.bayes.account.score.spam", "0.7") - .unwrap_or(0.7), - account_score_ham: config - .property_or_default("spam-filter.bayes.account.score.ham", "0.5") + .property_or_default("spam-filter.classifier.score.ham", "0.5") .unwrap_or(0.5), auto_learn_card_is_ham: config - .property_or_default("spam-filter.bayes.auto-learn.card-is-ham", "true") + .property_or_default("spam-filter.classifier.auto-learn.card-is-ham", "true") + .unwrap_or(true), + auto_learn_reply_ham: config + .property_or_default("spam-filter.classifier.auto-learn.trusted-reply", "true") .unwrap_or(true), } .into() @@ -589,24 +532,6 @@ impl SpamFilterScoreConfig { } } -impl SpamFilterExpiryConfig { - pub fn parse(config: &mut Config) -> Self { - SpamFilterExpiryConfig { - grey_list: config - .property::>("spam-filter.grey-list.duration") - .unwrap_or_default() - .map(|d| d.as_secs()), - trusted_reply: config - .property_or_default::>( - "spam-filter.trusted-reply.duration", - "30d", - ) - .unwrap_or_default() - .map(|d| d.as_secs()), - } - } -} - impl ParseValue for Element { fn parse_value(value: &str) -> utils::config::Result { match value { @@ -651,7 +576,6 @@ impl Default for SpamFilterHeaderConfig { SpamFilterHeaderConfig { status: "X-Spam-Status".to_string().into(), result: "X-Spam-Result".to_string().into(), - bayes_result: "X-Spam-Bayes".to_string().into(), llm: "X-Spam-LLM".to_string().into(), } } diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index efa7ecd1..1d3abc28 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -28,7 +28,7 @@ use ipc::{BroadcastEvent, HousekeeperEvent, PushEvent, QueueEvent, ReportingEven use listener::{asn::AsnGeoLookupData, blocked::Security, tls::AcmeProviders}; use mail_auth::{MX, Txt}; use manager::webadmin::{Resource, WebAdminManager}; -use nlp::bayes::{TokenHash, Weights}; +use nlp::classifier::sgd::SGDClassifier; use parking_lot::{Mutex, RwLock}; use rustls::sign::CertifiedKey; use std::{ @@ -73,6 +73,8 @@ pub mod enterprise; pub use psl; +use crate::config::spamfilter::Reputation; + pub static VERSION_PRIVATE: &str = env!("CARGO_PKG_VERSION"); pub static VERSION_PUBLIC: &str = "1.0.0"; @@ -109,14 +111,7 @@ pub const KV_RATE_LIMIT_CONTACT: u8 = 7; pub const KV_RATE_LIMIT_HTTP_AUTHENTICATED: u8 = 8; pub const KV_RATE_LIMIT_HTTP_ANONYMOUS: u8 = 9; pub const KV_RATE_LIMIT_IMAP: u8 = 10; -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; pub const KV_LOCK_QUEUE_MESSAGE: u8 = 21; pub const KV_LOCK_QUEUE_REPORT: u8 = 22; @@ -139,6 +134,9 @@ pub struct Inner { } pub struct Data { + pub spam_classifier: ArcSwap, + pub spam_reputation: ArcSwap, + pub tls_certificates: ArcSwap>>, pub tls_self_signed_cert: Option>, @@ -168,8 +166,6 @@ pub struct Caches { pub events: Cache>, pub scheduling: Cache>, - pub bayes: CacheWithTtl, - pub dns_txt: CacheWithTtl, pub dns_mx: CacheWithTtl>>, pub dns_ptr: CacheWithTtl>>, @@ -488,7 +484,6 @@ impl Default for Caches { contacts: Cache::new(1024, 10 * 1024 * 1024), events: Cache::new(1024, 10 * 1024 * 1024), scheduling: Cache::new(1024, 10 * 1024 * 1024), - bayes: CacheWithTtl::new(1024, 10 * 1024 * 1024), dns_rbl: CacheWithTtl::new(1024, 10 * 1024 * 1024), dns_txt: CacheWithTtl::new(1024, 10 * 1024 * 1024), dns_mx: CacheWithTtl::new(1024, 10 * 1024 * 1024), diff --git a/crates/nlp/Cargo.toml b/crates/nlp/Cargo.toml index 8bee0e40..db732c5a 100644 --- a/crates/nlp/Cargo.toml +++ b/crates/nlp/Cargo.toml @@ -6,8 +6,6 @@ edition = "2024" [dependencies] utils = { path = "../utils" } xxhash-rust = { version = "0.8.5", features = ["xxh3"] } -farmhash = "1.1.5" -siphasher = "1.0" serde = { version = "1.0", features = ["derive"]} nohash = "0.2.0" ahash = { version = "0.8.3", features = ["serde"] } @@ -20,6 +18,7 @@ psl = "2" radix_trie = "0.3" maplit = "1.0.2" hashify = "0.2.1" +rand = "0.9.2" [features] test_mode = [] diff --git a/crates/nlp/src/bayes/classify.rs b/crates/nlp/src/bayes/classify.rs deleted file mode 100644 index a1a07c71..00000000 --- a/crates/nlp/src/bayes/classify.rs +++ /dev/null @@ -1,145 +0,0 @@ -/* - * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC - * - * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL - */ - -use crate::tokenizers::osb::OsbToken; - -use super::{BayesClassifier, Weights}; - -// Position 0 represents Unigram weights -const FEATURE_WEIGHT: [f64; 8] = [1.0, 3125.0, 256.0, 27.0, 1.0, 0.0, 0.0, 0.0]; - -// Credits: ported from RSpamd -impl BayesClassifier { - pub fn classify(&self, tokens: T, ham_learns: u32, spam_learns: u32) -> Option - where - T: Iterator>, - { - if self.min_learns > 0 && (spam_learns < self.min_learns || ham_learns < self.min_learns) { - return None; - } - - let mut processed_tokens = 0; - let mut total_spam_prob = 0.0; - let mut total_ham_prob = 0.0; - - for token in tokens { - let weights = token.inner; - let total_count = weights.spam + weights.ham; - - if total_count >= self.min_token_hits { - let total_count = total_count as f64; - let spam_freq = weights.spam as f64 / f64::max(1.0, spam_learns as f64); - let ham_freq = weights.ham as f64 / f64::max(1.0, ham_learns as f64); - let spam_prob = spam_freq / (spam_freq + ham_freq); - let ham_prob = ham_freq / (spam_freq + ham_freq); - - let fw = FEATURE_WEIGHT[token.idx]; - let w = (fw * total_count) / (1.0 + fw * total_count); - let bayes_spam_prob = prob_combine(spam_prob, total_count, w, 0.5); - - if !((bayes_spam_prob > 0.5 && bayes_spam_prob < 0.5 + self.min_prob_strength) - || (bayes_spam_prob < 0.5 && bayes_spam_prob > 0.5 - self.min_prob_strength)) - { - let bayes_ham_prob = prob_combine(ham_prob, total_count, w, 0.5); - total_spam_prob += bayes_spam_prob.ln(); - total_ham_prob += bayes_ham_prob.ln(); - processed_tokens += 1; - } - } - } - - if processed_tokens == 0 || self.min_tokens > 0 && processed_tokens < self.min_tokens { - return None; - } - - let (h, s) = if total_spam_prob > -300.0 && total_ham_prob > -300.0 { - /* Fisher value is low enough to apply inv_chi_square */ - ( - 1.0 - inv_chi_square(total_spam_prob, processed_tokens), - 1.0 - inv_chi_square(total_ham_prob, processed_tokens), - ) - } else { - /* Use naive method */ - if total_spam_prob < total_ham_prob { - let h = (1.0 - (total_spam_prob - total_ham_prob).exp()) - / (1.0 + (total_spam_prob - total_ham_prob).exp()); - (h, 1.0 - h) - } else { - let s = (1.0 - (total_ham_prob - total_spam_prob).exp()) - / (1.0 + (total_ham_prob - total_spam_prob).exp()); - (1.0 - s, s) - } - }; - - let final_prob = if h.is_finite() && s.is_finite() { - (s + 1.0 - h) / 2.0 - } else { - /* - * We have some overflow, hence we need to check which class - * is NaN - */ - - if h.is_finite() { - 1.0 - } else if s.is_finite() { - 0.0 - } else { - 0.5 - } - }; - - if processed_tokens > 0 && (final_prob - 0.5).abs() > 0.05 { - Some(final_prob) - } else { - None - } - } -} - -/** - * Returns probability of chisquare > value with specified number of freedom - * degrees - */ -#[inline(always)] -fn inv_chi_square(value: f64, freedom_deg: u32) -> f64 { - let mut prob = value.exp(); - - if prob.is_finite() { - /* - * m is our confidence in class - * prob is e ^ x (small value since x is normally less than zero - * So we integrate over degrees of freedom and produce the total result - * from 1.0 (no confidence) to 0.0 (full confidence) - */ - - let mut sum = prob; - let m = -value; - - for i in 1..freedom_deg { - prob *= m / i as f64; - sum += prob; - } - - f64::min(1.0, sum) - } else { - /* - * e^x where x is large *NEGATIVE* number is OK, so we have a very strong - * confidence that inv-chi-square is close to zero - */ - - if value < 0.0 { 0.0 } else { 1.0 } - } -} - -/*#[inline(always)] -fn normalize_probability(x: f64, bias: f64) -> f64 { - ((x - bias) * 2.0).powi(8) -}*/ - -#[inline(always)] -fn prob_combine(prob: f64, cnt: f64, weight: f64, assumed: f64) -> f64 { - ((weight) * (assumed) + (cnt) * (prob)) / ((weight) + (cnt)) -} diff --git a/crates/nlp/src/bayes/mod.rs b/crates/nlp/src/bayes/mod.rs deleted file mode 100644 index f9bad387..00000000 --- a/crates/nlp/src/bayes/mod.rs +++ /dev/null @@ -1,185 +0,0 @@ -/* - * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC - * - * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL - */ - -use ahash::AHashMap; -use radix_trie::TrieKey; -use serde::{Deserialize, Serialize}; -use utils::cache::CacheItemWeight; - -use crate::tokenizers::osb::Gram; - -pub mod classify; -pub mod tokenize; -pub mod train; - -#[derive(Debug, Serialize, Deserialize, Default)] -pub struct BayesModel { - pub weights: AHashMap, - pub spam_learns: u32, - pub ham_learns: u32, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct BayesClassifier { - pub min_token_hits: u32, - 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, Hash)] -pub struct TokenHash { - 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 { - pub spam: u32, - pub ham: u32, -} - -impl BayesClassifier { - pub fn new() -> Self { - BayesClassifier { - min_token_hits: 2, - min_tokens: 11, - min_prob_strength: 0.05, - min_learns: 200, - min_balance: 0.1, - } - } -} - -impl Default for BayesClassifier { - fn default() -> Self { - Self::new() - } -} - -impl From> for TokenHash { - fn from(value: Gram<'_>) -> Self { - let mut hash = TokenHash { - hash: [0; HASH_LEN], - len: 0, - }; - - match value { - 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::fingerprint64(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 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::fingerprint64(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 TrieKey for TokenHash { - fn encode_bytes(&self) -> Vec { - self.hash[..self.len as usize].to_vec() - } -} - -impl From for Weights { - fn from(value: i64) -> Self { - Weights { - spam: value as u32, - ham: (value >> 32) as u32, - } - } -} - -impl From for i64 { - fn from(value: Weights) -> Self { - ((value.ham as i64) << 32) | value.spam as i64 - } -} - -impl CacheItemWeight for Weights { - fn weight(&self) -> u64 { - std::mem::size_of::() as u64 - } -} - -impl CacheItemWeight for TokenHash { - fn weight(&self) -> u64 { - std::mem::size_of::() as u64 - } -} - -impl TokenHash { - pub fn serialize(&self, prefix: u8, account_id: Option) -> Vec { - if let Some(account_id) = account_id { - self.serialize_account(prefix, account_id) - } else { - self.serialize_global(prefix) - } - } - - pub fn serialize_global(&self, prefix: u8) -> Vec { - let mut buf = Vec::with_capacity(self.len as usize + 1); - buf.push(prefix); - if self.len > 0 { - buf.extend_from_slice(&self.hash[..self.len as usize]); - } - buf - } - - pub fn serialize_account(&self, prefix: u8, account_id: u32) -> Vec { - 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()); - if self.len > 0 { - buf.extend_from_slice(&self.hash[..self.len as usize]); - } - buf - } -} diff --git a/crates/nlp/src/bayes/train.rs b/crates/nlp/src/bayes/train.rs deleted file mode 100644 index 608c367e..00000000 --- a/crates/nlp/src/bayes/train.rs +++ /dev/null @@ -1,55 +0,0 @@ -/* - * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC - * - * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL - */ - -use crate::tokenizers::osb::OsbToken; - -use super::{BayesModel, TokenHash}; - -impl BayesModel { - pub fn train(&mut self, tokens: T, is_spam: bool) - where - T: IntoIterator>, - { - if is_spam { - self.spam_learns += 1; - } else { - self.ham_learns += 1; - } - - for token in tokens { - 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; - } - } - - pub fn untrain(&mut self, tokens: T, is_spam: bool) - where - T: IntoIterator>, - { - if is_spam { - self.spam_learns -= 1; - } else { - self.ham_learns -= 1; - } - - for token in tokens { - let hs = self.weights.entry(token.inner).or_default(); - if is_spam { - hs.spam -= 1; - } else { - hs.ham -= 1; - } - } - } -} diff --git a/crates/nlp/src/classifier/feature.rs b/crates/nlp/src/classifier/feature.rs new file mode 100644 index 00000000..9683ac44 --- /dev/null +++ b/crates/nlp/src/classifier/feature.rs @@ -0,0 +1,56 @@ +/* + * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC + * + * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + */ + +use nohash::BuildNoHashHasher; +use std::collections::HashMap; +use xxhash_rust::xxh3::xxh3_64_with_seed; + +pub struct Sample { + pub(super) features: Features, + pub(super) class: f32, +} + +pub struct Features(pub(super) HashMap>); + +pub struct SampleBuilder { + pub(super) features_mask: u32, +} + +impl SampleBuilder { + pub fn build(&self, features: I, class: f32) -> Sample + where + I: IntoIterator, + I::Item: AsRef<[u8]>, + { + let mut features_map = HashMap::with_capacity_and_hasher(128, BuildNoHashHasher::default()); + + for feature in features { + let feature = feature.as_ref(); + let hash = xxh3_64_with_seed(feature, 0) as u32; + let hash_sign = xxh3_64_with_seed(feature, 1); + + *features_map.entry(hash & self.features_mask).or_default() += + if hash_sign & 1 == 0 { 1.0 } else { -1.0 }; + } + + Sample { + features: Features(features_map), + class, + } + } +} + +impl AsRef for Sample { + fn as_ref(&self) -> &Sample { + self + } +} + +impl AsRef for Features { + fn as_ref(&self) -> &Features { + self + } +} diff --git a/crates/nlp/src/classifier/mod.rs b/crates/nlp/src/classifier/mod.rs new file mode 100644 index 00000000..a11b8000 --- /dev/null +++ b/crates/nlp/src/classifier/mod.rs @@ -0,0 +1,8 @@ +/* + * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC + * + * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + */ + +pub mod feature; +pub mod sgd; diff --git a/crates/nlp/src/classifier/sgd.rs b/crates/nlp/src/classifier/sgd.rs new file mode 100644 index 00000000..27b52d03 --- /dev/null +++ b/crates/nlp/src/classifier/sgd.rs @@ -0,0 +1,317 @@ +/* + * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC + * + * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + */ + +use crate::classifier::feature::{Features, Sample, SampleBuilder}; +use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom}; + +#[derive(Default)] +pub struct SGDClassifier { + weights: Vec, + intercept: f32, + n_epochs: usize, + alpha: f32, + random_state: u64, +} + +const MAX_DLOSS: f32 = 1e4; + +impl SGDClassifier { + pub fn new(n_features: usize, n_epochs: usize, alpha: f32, random_state: u64) -> Self { + SGDClassifier { + weights: vec![0.0; n_features], + n_epochs, + random_state, + alpha, + intercept: 0.0, + } + } + + pub fn fit(&mut self, samples: &mut [impl AsRef]) { + let mut rng = StdRng::seed_from_u64(self.random_state); + let mut t = 1; + let mut w_scale = 1.0; + + // Heuristic to initialize 'optimal' learning rate + let typw = (1.0 / self.alpha.sqrt()).sqrt(); + let initial_eta0 = typw / 1.0_f32.max(gradient(1.0, -typw)); + let optimal_init = 1.0 / (initial_eta0 * self.alpha); + + for _ in 0..self.n_epochs { + samples.shuffle(&mut rng); + + for sample in samples.iter() { + // Prediction + let sample = sample.as_ref(); + let mut dot: f32 = 0.0; + for (idx, feature) in &sample.features.0 { + dot += self.weights[*idx as usize] * *feature; + } + let p = (dot * w_scale) + self.intercept; + let eta = 1.0 / (self.alpha * (optimal_init + (t as f32) - 1.0)); + + // Compute Loss & Gradient + let dloss = gradient(sample.class, p).clamp(-MAX_DLOSS, MAX_DLOSS); + + // Lazy weight decay + w_scale *= 1.0 - (eta * self.alpha); + + // Update weights + let update = -eta * dloss; + if update != 0.0 { + let scaled_update = update / w_scale; + + for (idx, feature) in &sample.features.0 { + self.weights[*idx as usize] += scaled_update * *feature; + } + + self.intercept += update; + } + + // Rescale weights if w_scale is too small or too large + if !(1e-6..=1e6).contains(&w_scale) { + for w in &mut self.weights { + *w *= w_scale; + } + w_scale = 1.0; + } + + t += 1; + } + } + + if w_scale != 1.0 { + for w in &mut self.weights { + *w *= w_scale; + } + } + } + + fn predict_proba_sample(&self, features: &Features) -> f32 { + let mut z: f32 = 0.0; + for (idx, feature) in &features.0 { + z += self.weights[*idx as usize] * *feature; + } + z += self.intercept; + + sigmoid(z) + } + + pub fn predict(&self, features: &Features) -> f32 { + let proba = self.predict_proba_sample(features); + if proba >= 0.5 { 1.0 } else { 0.0 } + } + + pub fn predict_batch(&self, test: I) -> Vec + where + I: IntoIterator, + I::Item: AsRef, + { + test.into_iter() + .map(|features| self.predict(features.as_ref())) + .collect() + } + + pub fn sample_builder(&self) -> SampleBuilder { + SampleBuilder { + features_mask: (self.weights.len() - 1) as u32, + } + } + + pub fn is_active(&self) -> bool { + !self.weights.is_empty() + } +} + +#[inline(always)] +fn gradient(y: f32, p: f32) -> f32 { + if p > -16.0 { + let exp_tmp = (-p).exp(); + ((1.0 - y) - y * exp_tmp) / (1.0 + exp_tmp) + } else { + p.exp() - y + } +} + +#[inline(always)] +fn sigmoid(z: f32) -> f32 { + if z >= 0.0 { + 1.0 / (1.0 + (-z).exp()) + } else { + let exp_z = z.exp(); + exp_z / (1.0 + exp_z) + } +} + +/*#[inline(always)] +fn loss(y: f32, p: f32) -> f32 { + log1pexp(p) - y * p +} + +#[inline(always)] +fn log1pexp(x: f32) -> f32 { + if x <= -16.0 { + x.exp() + } else if x <= 16.0 { + (1.0 + x.exp()).ln() + } else { + x + } +}*/ + +#[cfg(test)] +pub mod tests { + use crate::classifier::{feature::Sample, sgd::SGDClassifier}; + use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom}; + use std::{ + fs::File, + io::{BufRead, BufReader}, + time::Instant, + }; + + fn accuracy_score(y_true: &[f32], y_pred: &[f32]) -> f32 { + y_true + .iter() + .zip(y_pred.iter()) + .filter(|(true_val, pred_val)| **true_val == **pred_val) + .count() as f32 + / y_true.len() as f32 + } + + fn precision_score(y_true: &[f32], y_pred: &[f32], positive_class: f32) -> f32 { + let true_positives = y_true + .iter() + .zip(y_pred.iter()) + .filter(|(true_val, pred_val)| { + **pred_val == positive_class && **true_val == positive_class + }) + .count() as f32; + + let predicted_positives = y_pred + .iter() + .filter(|pred_val| **pred_val == positive_class) + .count() as f32; + + if predicted_positives == 0.0 { + 0.0 + } else { + true_positives / predicted_positives + } + } + + fn recall_score(y_true: &[f32], y_pred: &[f32], positive_class: f32) -> f32 { + let true_positives = y_true + .iter() + .zip(y_pred.iter()) + .filter(|(true_val, pred_val)| { + **pred_val == positive_class && **true_val == positive_class + }) + .count() as f32; + + let actual_positives = y_true + .iter() + .filter(|true_val| **true_val == positive_class) + .count() as f32; + + if actual_positives == 0.0 { + 0.0 + } else { + true_positives / actual_positives + } + } + + fn f1_score(y_true: &[f32], y_pred: &[f32], positive_class: f32) -> f32 { + let precision = precision_score(y_true, y_pred, positive_class); + let recall = recall_score(y_true, y_pred, positive_class); + + if precision + recall == 0.0 { + 0.0 + } else { + 2.0 * (precision * recall) / (precision + recall) + } + } + + fn train_test_split(data: &[Sample], test_size: f32) -> (Vec<&Sample>, Vec<&Sample>) { + let mut class_0: Vec<&Sample> = Vec::new(); + let mut class_1: Vec<&Sample> = Vec::new(); + + for sample in data { + if sample.class == 0.0 { + class_0.push(sample); + } else { + class_1.push(sample); + } + } + + let test_count_0 = (class_0.len() as f32 * test_size).round() as usize; + let test_count_1 = (class_1.len() as f32 * test_size).round() as usize; + + let (test_0, train_0) = class_0.split_at(test_count_0); + let (test_1, train_1) = class_1.split_at(test_count_1); + + let mut train = Vec::new(); + let mut test = Vec::new(); + + train.extend_from_slice(train_0); + train.extend_from_slice(train_1); + test.extend_from_slice(test_0); + test.extend_from_slice(test_1); + + (train, test) + } + + #[test] + fn sgd_classifier() { + let reader = BufReader::new( + File::open("/Users/me/code/playground/phishing_email.csv") + .expect("Could not open file"), + ); + let mut samples = Vec::with_capacity(1024); + + let mut model = SGDClassifier::new(1 << 20, 1000, 0.0001, 42); + let builder = model.sample_builder(); + + let time = Instant::now(); + + for line in reader.lines().skip(1) { + let line = line.unwrap(); + let (text, class) = line.trim().rsplit_once(',').unwrap(); + let text = text.trim_start_matches('"').trim_end_matches('"'); + samples.push( + builder.build( + text.split_whitespace(), + class + .parse() + .unwrap_or_else(|_| panic!("Invalid class value: {line}")), + ), + ); + } + + println!("Loaded {} samples in {:?}", samples.len(), time.elapsed()); + + samples.shuffle(&mut StdRng::seed_from_u64(42)); + + let (mut train_samples, test_samples) = train_test_split(&samples, 0.2); + + println!( + "Training samples: {}, Testing samples: {}", + train_samples.len(), + test_samples.len() + ); + + println!("Training SGD Classifier..."); + let time = Instant::now(); + model.fit(&mut train_samples); + println!("SGD Classifier trained in {:?}", time.elapsed()); + + let y_pred = model.predict_batch(test_samples.iter().map(|s| &s.features)); + let y_train: Vec = test_samples.iter().map(|s| s.class).collect(); + + println!("Accuracy: {:.4}", accuracy_score(&y_train, &y_pred)); + println!("Precision: {:.4}", precision_score(&y_train, &y_pred, 1.0)); + println!("Recall: {:.4}", recall_score(&y_train, &y_pred, 1.0)); + println!("F1 Score: {:.4}", f1_score(&y_train, &y_pred, 1.0)); + } +} diff --git a/crates/nlp/src/lib.rs b/crates/nlp/src/lib.rs index b09b08b7..a92aeda6 100644 --- a/crates/nlp/src/lib.rs +++ b/crates/nlp/src/lib.rs @@ -4,101 +4,6 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -pub mod bayes; +pub mod classifier; pub mod language; pub mod tokenizers; - -#[cfg(test)] -mod test { - use std::fs; - - use crate::{ - bayes::{ - BayesClassifier, BayesModel, - tokenize::{BayesTokenizer, tests::ToBayesToken}, - }, - tokenizers::{ - osb::{OsbToken, OsbTokenizer}, - types::TypesTokenizer, - }, - }; - - #[test] - #[ignore] - fn train() { - let db = - fs::read_to_string("/Users/me/code/stalwart/_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, - 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/stalwart/_ignore/old/spam_or_not_spam.bin", - bincode::serialize(&bayes).unwrap(), - ) - .unwrap(); - } - - #[test] - #[ignore] - fn classify() { - let model: BayesModel = bincode::deserialize( - &fs::read("/Users/me/code/stalwart/_ignore/old/spam_or_not_spam.bin").unwrap(), - ) - .unwrap(); - let bayes = BayesClassifier::new(); - - for text in [ - 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", - 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, - 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/mod.rs b/crates/nlp/src/tokenizers/mod.rs index ae861f04..9682eea9 100644 --- a/crates/nlp/src/tokenizers/mod.rs +++ b/crates/nlp/src/tokenizers/mod.rs @@ -6,8 +6,8 @@ pub mod chinese; pub mod japanese; -pub mod osb; pub mod space; +pub mod stream; pub mod types; pub mod word; diff --git a/crates/nlp/src/tokenizers/osb.rs b/crates/nlp/src/tokenizers/osb.rs deleted file mode 100644 index 47df92f1..00000000 --- a/crates/nlp/src/tokenizers/osb.rs +++ /dev/null @@ -1,343 +0,0 @@ -/* - * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC - * - * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL - */ - -use std::iter::Peekable; - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct OsbToken { - pub inner: T, - pub idx: usize, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum Gram<'x> { - Uni { t1: &'x [u8] }, - Bi { t1: &'x [u8], t2: &'x [u8] }, -} - -pub struct OsbTokenizer -where - I: Iterator>, - R: for<'y> From> + 'static, -{ - iter: Peekable, - buf: Vec>>, - window_size: usize, - window_pos: usize, - window_idx: usize, - phantom: std::marker::PhantomData, -} - -impl OsbTokenizer -where - I: Iterator>, - R: for<'y> From> + 'static, -{ - pub fn new(iter: I, window_size: usize) -> Self { - Self { - iter: iter.peekable(), - buf: vec![None; window_size], - window_pos: 0, - window_idx: 0, - window_size, - phantom: std::marker::PhantomData, - } - } -} - -impl Iterator for OsbTokenizer -where - I: Iterator>, - R: for<'y> From> + 'static, -{ - type Item = OsbToken; - - fn next(&mut self) -> Option { - let end_pos = (self.window_pos + self.window_idx) % self.window_size; - if self.buf[end_pos].is_none() { - self.buf[end_pos] = self.iter.next(); - } - - let t1 = self.buf[self.window_pos % self.window_size].as_deref()?; - let token = OsbToken { - inner: R::from(if self.window_idx != 0 { - Gram::Bi { - t1, - t2: self.buf[end_pos].as_deref()?, - } - } else { - Gram::Uni { t1 } - }), - idx: self.window_idx, - }; - - // Increment window index - self.window_idx += 1; - if self.window_idx == self.window_size - || (self.iter.peek().is_none() - && self.buf[(self.window_pos + self.window_idx) % self.window_size].is_none()) - { - self.buf[self.window_pos % self.window_size] = None; - self.window_idx = 0; - self.window_pos += 1; - } - - Some(token) - } -} - -#[cfg(test)] -mod test { - use crate::tokenizers::osb::{Gram, OsbToken}; - - impl From> for String { - fn from(value: Gram<'_>) -> Self { - match value { - 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() - ), - } - } - } - - #[test] - fn osb_tokenizer() { - assert_eq!( - super::OsbTokenizer::new( - "The quick brown fox jumps over the lazy dog and the lazy cat" - .split_ascii_whitespace() - .map(|b| b.as_bytes().to_vec()), - 5, - ) - .collect::>(), - vec![ - OsbToken { - inner: "The".to_string(), - idx: 0 - }, - OsbToken { - inner: "The quick".to_string(), - idx: 1 - }, - OsbToken { - inner: "The brown".to_string(), - idx: 2 - }, - OsbToken { - inner: "The fox".to_string(), - idx: 3 - }, - OsbToken { - inner: "The jumps".to_string(), - idx: 4 - }, - OsbToken { - inner: "quick".to_string(), - idx: 0 - }, - OsbToken { - inner: "quick brown".to_string(), - idx: 1 - }, - OsbToken { - inner: "quick fox".to_string(), - idx: 2 - }, - OsbToken { - inner: "quick jumps".to_string(), - idx: 3 - }, - OsbToken { - inner: "quick over".to_string(), - idx: 4 - }, - OsbToken { - inner: "brown".to_string(), - idx: 0 - }, - OsbToken { - inner: "brown fox".to_string(), - idx: 1 - }, - OsbToken { - inner: "brown jumps".to_string(), - idx: 2 - }, - OsbToken { - inner: "brown over".to_string(), - idx: 3 - }, - OsbToken { - inner: "brown the".to_string(), - idx: 4 - }, - OsbToken { - inner: "fox".to_string(), - idx: 0 - }, - OsbToken { - inner: "fox jumps".to_string(), - idx: 1 - }, - OsbToken { - inner: "fox over".to_string(), - idx: 2 - }, - OsbToken { - inner: "fox the".to_string(), - idx: 3 - }, - OsbToken { - inner: "fox lazy".to_string(), - idx: 4 - }, - OsbToken { - inner: "jumps".to_string(), - idx: 0 - }, - OsbToken { - inner: "jumps over".to_string(), - idx: 1 - }, - OsbToken { - inner: "jumps the".to_string(), - idx: 2 - }, - OsbToken { - inner: "jumps lazy".to_string(), - idx: 3 - }, - OsbToken { - inner: "jumps dog".to_string(), - idx: 4 - }, - OsbToken { - inner: "over".to_string(), - idx: 0 - }, - OsbToken { - inner: "over the".to_string(), - idx: 1 - }, - OsbToken { - inner: "over lazy".to_string(), - idx: 2 - }, - OsbToken { - inner: "over dog".to_string(), - idx: 3 - }, - OsbToken { - inner: "over and".to_string(), - idx: 4 - }, - OsbToken { - inner: "the".to_string(), - idx: 0 - }, - OsbToken { - inner: "the lazy".to_string(), - idx: 1 - }, - OsbToken { - inner: "the dog".to_string(), - idx: 2 - }, - OsbToken { - inner: "the and".to_string(), - idx: 3 - }, - OsbToken { - inner: "the the".to_string(), - idx: 4 - }, - OsbToken { - inner: "lazy".to_string(), - idx: 0 - }, - OsbToken { - inner: "lazy dog".to_string(), - idx: 1 - }, - OsbToken { - inner: "lazy and".to_string(), - idx: 2 - }, - OsbToken { - inner: "lazy the".to_string(), - idx: 3 - }, - OsbToken { - inner: "lazy lazy".to_string(), - idx: 4 - }, - OsbToken { - inner: "dog".to_string(), - idx: 0 - }, - OsbToken { - inner: "dog and".to_string(), - idx: 1 - }, - OsbToken { - inner: "dog the".to_string(), - idx: 2 - }, - OsbToken { - inner: "dog lazy".to_string(), - idx: 3 - }, - OsbToken { - inner: "dog cat".to_string(), - idx: 4 - }, - OsbToken { - inner: "and".to_string(), - idx: 0 - }, - OsbToken { - inner: "and the".to_string(), - idx: 1 - }, - OsbToken { - inner: "and lazy".to_string(), - idx: 2 - }, - OsbToken { - inner: "and cat".to_string(), - idx: 3 - }, - OsbToken { - inner: "the".to_string(), - idx: 0 - }, - OsbToken { - inner: "the lazy".to_string(), - idx: 1 - }, - OsbToken { - inner: "the cat".to_string(), - idx: 2 - }, - OsbToken { - inner: "lazy".to_string(), - idx: 0 - }, - OsbToken { - inner: "lazy cat".to_string(), - idx: 1 - }, - OsbToken { - inner: "cat".to_string(), - idx: 0 - } - ] - ); - } -} diff --git a/crates/nlp/src/bayes/tokenize.rs b/crates/nlp/src/tokenizers/stream.rs similarity index 98% rename from crates/nlp/src/bayes/tokenize.rs rename to crates/nlp/src/tokenizers/stream.rs index 9547f04d..56056315 100644 --- a/crates/nlp/src/bayes/tokenize.rs +++ b/crates/nlp/src/tokenizers/stream.rs @@ -4,8 +4,6 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use std::borrow::Cow; - use crate::{ language::{ Language, @@ -15,17 +13,23 @@ use crate::{ }, tokenizers::{chinese::JIEBA, japanese}, }; +use std::borrow::Cow; -pub struct BayesTokenizer> { +pub struct StreamTokenizer>, I: StreamInputTokenTrait> { stream: T, stemmer: Stemmer, stop_words: Option, - tokens: Vec>, + tokens: Vec, } -pub enum BayesInputToken { +pub enum StreamInputToken { Word(String), - Raw(Vec), + Other(T), +} + +pub trait StreamInputTokenTrait { + fn from_owned(word: String) -> Self; + fn from_borrowed(word: &str) -> Self; } enum Stemmer { @@ -35,7 +39,7 @@ enum Stemmer { None, } -impl> BayesTokenizer { +impl>, I: StreamInputTokenTrait> StreamTokenizer { pub fn new(text: &str, stream: T) -> Self { // Detect language let (mut language, score) = @@ -59,8 +63,10 @@ impl> BayesTokenizer { } } -impl> Iterator for BayesTokenizer { - type Item = Vec; +impl>, I: StreamInputTokenTrait> Iterator + for StreamTokenizer +{ + type Item = I; fn next(&mut self) -> Option { if let Some(prev_token) = self.tokens.pop() { @@ -69,24 +75,24 @@ impl> Iterator for BayesTokenizer { for token in self.stream.by_ref() { return match token { - BayesInputToken::Word(word) => { + StreamInputToken::Word(word) => { if self.stop_words.is_some_and(|sw| sw(word.as_str())) { continue; } match &self.stemmer { Stemmer::IndoEuropean(stemmer) => match stemmer.stem(&word) { - Cow::Borrowed(_) => word.into_bytes(), - Cow::Owned(stemmed_word) => stemmed_word.into_bytes(), + Cow::Borrowed(word) => I::from_borrowed(word), + Cow::Owned(stemmed_word) => I::from_owned(stemmed_word), }, Stemmer::Mandarin => { let mut result = JIEBA.cut(&word, false).into_iter(); if let Some(stemmed_word) = result.next() { - let stemmed_word = stemmed_word.to_string(); + let stemmed_word = I::from_borrowed(stemmed_word); self.tokens = result .rev() - .map(|word| word.to_string().into_bytes()) + .map(|word| I::from_borrowed(word)) .collect::>(); - stemmed_word.into_bytes() + stemmed_word } else { // This shouldn't happen, but just in case continue; @@ -96,17 +102,17 @@ impl> Iterator for BayesTokenizer { let mut result = japanese::tokenize(&word).into_iter(); if let Some(stemmed_word) = result.next() { self.tokens = - result.rev().map(|b| b.into_bytes()).collect::>(); - stemmed_word.into_bytes() + result.rev().map(|b| I::from_owned(b)).collect::>(); + I::from_owned(stemmed_word) } else { // This shouldn't happen, but just in case continue; } } - Stemmer::None => word.into_bytes(), + Stemmer::None => I::from_owned(word), } } - BayesInputToken::Raw(raw) => raw, + StreamInputToken::Other(raw) => raw, } .into(); } @@ -7866,53 +7872,61 @@ pub fn symbols(input: &str) -> bool { ) } -#[cfg(test)] -pub mod tests { - use std::{borrow::Cow, net::IpAddr}; - - use crate::{ - bayes::tokenize::BayesTokenizer, - tokenizers::types::{TokenType, TypesTokenizer}, - }; - - use super::{BayesInputToken, symbols}; - - pub trait ToBayesToken { - fn to_bayes_token(&self) -> Option; +impl StreamInputTokenTrait for Vec { + fn from_owned(word: String) -> Self { + word.into_bytes() } - impl, E: AsRef, U: AsRef, I: AsRef> ToBayesToken + fn from_borrowed(word: &str) -> Self { + word.as_bytes().to_vec() + } +} + +#[cfg(test)] +pub mod tests { + use super::{StreamInputToken, symbols}; + use crate::tokenizers::{ + stream::StreamTokenizer, + types::{TokenType, TypesTokenizer}, + }; + use std::{borrow::Cow, net::IpAddr}; + + pub trait ToStreamToken { + fn to_stream_token(&self) -> Option>>; + } + + impl, E: AsRef, U: AsRef, I: AsRef> ToStreamToken for TokenType { - fn to_bayes_token(&self) -> Option { + fn to_stream_token(&self) -> Option>> { match self { TokenType::Alphabetic(word) => { - Some(BayesInputToken::Word(word.as_ref().to_lowercase())) + Some(StreamInputToken::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))) + .map(|(_, host)| StreamInputToken::Other(url_host_as_bytes(host))) } TokenType::IpAddr(word) => word.as_ref().parse::().ok().map(|ip| { - BayesInputToken::Raw(match ip { + StreamInputToken::Other(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() + StreamInputToken::Other(url_host_as_bytes(word.as_ref())).into() } TokenType::Alphanumeric(word) | TokenType::UrlNoHost(word) => { - BayesInputToken::Raw(word.as_ref().to_lowercase().into_bytes()).into() + StreamInputToken::Other(word.as_ref().to_lowercase().into_bytes()).into() } TokenType::Email(word) => { - BayesInputToken::Raw(word.as_ref().to_lowercase().into_bytes()).into() + StreamInputToken::Other(word.as_ref().to_lowercase().into_bytes()).into() } TokenType::Other(ch) => { let ch = ch.to_string(); if symbols(&ch) { - Some(BayesInputToken::Raw(ch.into_bytes())) + Some(StreamInputToken::Other(ch.into_bytes())) } else { None } @@ -7931,7 +7945,7 @@ pub mod tests { .into_bytes() } - fn number_to_tag(is_float: bool, num: &str) -> BayesInputToken { + fn number_to_tag(is_float: bool, num: &str) -> StreamInputToken> { let t = match (is_float, num.starts_with('-')) { (true, true) => b'F', (true, false) => b'f', @@ -7939,11 +7953,11 @@ pub mod tests { (false, false) => b'i', }; - BayesInputToken::Raw([t, num.len() as u8].to_vec()) + StreamInputToken::Other([t, num.len() as u8].to_vec()) } #[test] - fn bayes_tokenizer() { + fn stream_tokenizer() { let inputs = [ ( "The quick brown fox jumps over the lazy dog", @@ -8026,9 +8040,9 @@ pub mod tests { ]; for (input, expect) in inputs.iter() { - let input = BayesTokenizer::new( + let input = StreamTokenizer::new( input, - TypesTokenizer::new(input).filter_map(|t| t.word.to_bayes_token()), + TypesTokenizer::new(input).filter_map(|t| t.word.to_stream_token()), ) .map(|word| String::from_utf8(word).unwrap()) .collect::>(); diff --git a/crates/spam-filter/src/analysis/mod.rs b/crates/spam-filter/src/analysis/mod.rs index abde7882..730ee365 100644 --- a/crates/spam-filter/src/analysis/mod.rs +++ b/crates/spam-filter/src/analysis/mod.rs @@ -37,7 +37,6 @@ pub mod reputation; pub mod rules; pub mod score; pub mod subject; -pub mod trusted_reply; pub mod url; // SPDX-SnippetBegin diff --git a/crates/spam-filter/src/analysis/reputation.rs b/crates/spam-filter/src/analysis/reputation.rs index 86710238..226bf8c6 100644 --- a/crates/spam-filter/src/analysis/reputation.rs +++ b/crates/spam-filter/src/analysis/reputation.rs @@ -4,19 +4,13 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use std::{borrow::Cow, future::Future}; - +use crate::SpamFilterContext; use common::{ - KV_REPUTATION_ASN, KV_REPUTATION_DOMAIN, KV_REPUTATION_FROM, KV_REPUTATION_IP, Server, - ip_to_bytes, + Server, + config::spamfilter::{ReputationCount, ReputationType}, }; use mail_auth::DmarcResult; -use store::{Deserialize, Serialize, dispatch::lookup::KeyValue}; - -use crate::{ - SpamFilterContext, - modules::{key_get, key_set}, -}; +use std::future::Future; pub trait SpamFilterAnalyzeReputation: Sync + Send { fn spam_filter_analyze_reputation( @@ -25,180 +19,85 @@ pub trait SpamFilterAnalyzeReputation: Sync + Send { ) -> impl Future + Send; } -#[derive(Debug)] -enum Type { - Ip, - From, - Domain, - Asn, -} - -#[derive(Debug)] -struct Reputation { - count: u32, - score: f64, -} - impl SpamFilterAnalyzeReputation for Server { async fn spam_filter_analyze_reputation(&self, ctx: &mut SpamFilterContext<'_>) { - // Obtain sender address - let sender = if !ctx.output.env_from_addr.address.is_empty() { - &ctx.output.env_from_addr - } else { - &ctx.output.from.email - }; - // Do not penalize forged domains - let is_dmarc_pass = matches!(ctx.input.dmarc_result, Some(DmarcResult::Pass)); - - let mut types = vec![ - (Type::Ip, Cow::Owned(ip_to_bytes(&ctx.input.remote_ip))), - ( - Type::From, - if is_dmarc_pass { - Cow::Borrowed(sender.address.as_bytes()) + let reputation = self.inner.data.spam_reputation.load(); + if reputation.total.ham > 0 + && reputation.total.spam > 0 + && reputation.total.ham + reputation.total.spam < 100 + { + if matches!(ctx.input.dmarc_result, Some(DmarcResult::Pass)) { + // Obtain sender address + let sender = if !ctx.output.env_from_addr.address.is_empty() { + &ctx.output.env_from_addr } else { - Cow::Owned(format!("_{}", sender.domain_part.sld_or_default()).into_bytes()) - }, - ), - ( - Type::Domain, - 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()) - }, - ), - ]; + &ctx.output.from.email + }; - // Add ASN - if let Some(asn_id) = &ctx.input.asn { - ctx.result.add_tag(format!("SOURCE_ASN_{asn_id}")); - types.push((Type::Asn, Cow::Owned(asn_id.to_be_bytes().to_vec()))); + if let Some(count) = reputation.items.get(&ReputationType::Domain( + sender.domain_part.sld_or_default().into(), + )) && let Some(tag) = reputation_tag("REP_DOMAIN", count, &reputation.total) + { + ctx.result.add_tag(tag); + } + } + + // Add ASN + if let Some(asn_id) = &ctx.input.asn { + ctx.result.add_tag(format!("SOURCE_ASN_{asn_id}")); + if let Some(count) = reputation.items.get(&ReputationType::Asn(*asn_id)) + && let Some(tag) = reputation_tag("REP_ASN", count, &reputation.total) + { + ctx.result.add_tag(tag); + } + } + + // Add IP + if let Some(count) = reputation + .items + .get(&ReputationType::Ip(ctx.input.remote_ip)) + && let Some(tag) = reputation_tag("REP_IP", count, &reputation.total) + { + ctx.result.add_tag(tag); + } + } else { + // Add ASN + if let Some(asn_id) = &ctx.input.asn { + ctx.result.add_tag(format!("SOURCE_ASN_{asn_id}")); + } } if let Some(country) = &ctx.input.country { ctx.result.add_tag(format!("SOURCE_COUNTRY_{country}")); } - - if let Some(config) = &self.core.spam.reputation { - let mut reputation = 0.0; - - for (rep_type, key) in types { - let token = match key_get::( - self, - ctx.input.span_id, - KeyValue::<()>::build_key(rep_type.prefix(), key.as_ref()), - ) - .await - { - Ok(Some(token)) => token, - Ok(None) if !ctx.input.is_test => { - key_set( - self, - ctx.input.span_id, - KeyValue::with_prefix( - rep_type.prefix(), - key.as_ref(), - Reputation { - count: 1, - score: ctx.result.score, - } - .serialize() - .unwrap(), - ) - .expires(config.expiry), - ) - .await; - continue; - } - Ok(None) | Err(_) => continue, - }; - - // Update reputation - let updated_score = (token.count + 1) as f64 - * (ctx.result.score + config.token_score * token.score) - / (config.token_score * token.count as f64 + 1.0); - let updated_count = token.count + 1; - - if !ctx.input.is_test { - key_set( - self, - ctx.input.span_id, - KeyValue::with_prefix( - rep_type.prefix(), - key.as_ref(), - Reputation { - count: updated_count, - score: updated_score, - } - .serialize() - .unwrap(), - ) - .expires(config.expiry), - ) - .await; - } - - // Assign weight - let weight = match rep_type { - Type::Ip => config.ip_weight, - Type::From => config.sender_weight, - Type::Domain => config.domain_weight, - Type::Asn => config.asn_weight, - }; - - reputation += token.score / token.count as f64 * weight; - } - - // Adjust score - if reputation > 0.0 { - ctx.result.score += (reputation - ctx.result.score) * config.factor; - } - } } } -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) -> trc::Result> { - let mut buf = Vec::with_capacity(12); - buf.extend_from_slice(&self.count.to_be_bytes()); - buf.extend_from_slice(&self.score.to_be_bytes()); - Ok(buf) - } -} - -impl Deserialize for Reputation { - fn deserialize(bytes: &[u8]) -> trc::Result { - if bytes.len() == 12 { - Ok(Reputation { - count: u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]), - score: f64::from_be_bytes([ - bytes[4], bytes[5], bytes[6], bytes[7], bytes[8], bytes[9], bytes[10], - bytes[11], - ]), - }) - } else { - Err(trc::StoreEvent::DataCorruption - .caused_by(trc::location!()) - .ctx(trc::Key::Value, bytes)) - } - } -} - -impl From> for Reputation { - fn from(_: store::Value<'_>) -> Self { - unimplemented!() +fn reputation_tag( + prefix: &str, + token_count: &ReputationCount, + total_count: &ReputationCount, +) -> Option { + let total_token_occurrences = token_count.spam + token_count.ham; + if total_token_occurrences < 10 { + return None; + } + let prob_token_given_spam = token_count.spam as f64 / total_count.spam as f64; + let prob_token_given_ham = token_count.ham as f64 / total_count.ham as f64; + if prob_token_given_spam + prob_token_given_ham == 0.0 { + return None; + } + let spam_probability = prob_token_given_spam / (prob_token_given_spam + prob_token_given_ham); + if spam_probability >= 0.90 { + Some(format!("{prefix}_VERY_BAD",)) + } else if spam_probability >= 0.75 { + Some(format!("{prefix}_BAD",)) + } else if spam_probability <= 0.10 { + Some(format!("{prefix}_VERY_GOOD",)) + } else if spam_probability <= 0.25 { + Some(format!("{prefix}_GOOD",)) + } else { + None } } diff --git a/crates/spam-filter/src/analysis/score.rs b/crates/spam-filter/src/analysis/score.rs index b14e6841..61f6f974 100644 --- a/crates/spam-filter/src/analysis/score.rs +++ b/crates/spam-filter/src/analysis/score.rs @@ -14,8 +14,7 @@ use crate::{ pyzor::SpamFilterAnalyzePyzor, received::SpamFilterAnalyzeReceived, recipient::SpamFilterAnalyzeRecipient, replyto::SpamFilterAnalyzeReplyTo, reputation::SpamFilterAnalyzeReputation, rules::SpamFilterAnalyzeRules, - subject::SpamFilterAnalyzeSubject, trusted_reply::SpamFilterAnalyzeTrustedReply, - url::SpamFilterAnalyzeUrl, + subject::SpamFilterAnalyzeSubject, url::SpamFilterAnalyzeUrl, }, modules::bayes::BayesClassifier, }; @@ -203,8 +202,8 @@ impl SpamFilterAnalyzeScore for Server { // SPDX-SnippetEnd - // Trusted reply analysis - self.spam_filter_analyze_reply_in(ctx).await; + // Reputation tracking and adjust score + self.spam_filter_analyze_reputation(ctx).await; // Spam trap self.spam_filter_analyze_spam_trap(ctx).await; @@ -225,9 +224,6 @@ impl SpamFilterAnalyzeScore for Server { SpamFilterAction::Reject => return SpamFilterAction::Reject, } - // Reputation tracking and adjust score - self.spam_filter_analyze_reputation(ctx).await; - // Final score calculation self.spam_filter_finalize(ctx).await } diff --git a/crates/spam-filter/src/analysis/trusted_reply.rs b/crates/spam-filter/src/analysis/trusted_reply.rs deleted file mode 100644 index a6fe1f99..00000000 --- a/crates/spam-filter/src/analysis/trusted_reply.rs +++ /dev/null @@ -1,88 +0,0 @@ -/* - * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC - * - * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL - */ - -use std::future::Future; - -use common::{KV_TRUSTED_REPLY, Server}; -use mail_parser::{HeaderName, HeaderValue}; -use store::dispatch::lookup::KeyValue; - -use crate::{SpamFilterContext, modules::bayes::BayesClassifier}; - -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.expiry.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.expiry.trusted_reply, - ctx.input.message.message_id(), - ) && 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() - .is_some_and(|config| config.auto_learn_reply_ham) - { - self.bayes_train_if_balanced(ctx, false).await; - } - } -} diff --git a/crates/spam-filter/src/modules/bayes.rs b/crates/spam-filter/src/modules/bayes.rs index 4068ed70..3819d159 100644 --- a/crates/spam-filter/src/modules/bayes.rs +++ b/crates/spam-filter/src/modules/bayes.rs @@ -6,7 +6,7 @@ use std::{borrow::Cow, collections::HashSet, future::Future, time::Duration}; -use common::{KV_BAYES_MODEL_GLOBAL, KV_BAYES_MODEL_USER, Server, ip_to_bytes}; +use common::{Server, ip_to_bytes}; use mail_auth::DmarcResult; use nlp::{ bayes::{ diff --git a/crates/utils/src/config/utils.rs b/crates/utils/src/config/utils.rs index e7a1a0ce..b7611502 100644 --- a/crates/utils/src/config/utils.rs +++ b/crates/utils/src/config/utils.rs @@ -463,6 +463,14 @@ impl ParseValue for i32 { } } +impl ParseValue for f32 { + fn parse_value(value: &str) -> super::Result { + value + .parse() + .map_err(|_| format!("Invalid floating point value {:?}.", value)) + } +} + impl ParseValue for IpAddr { fn parse_value(value: &str) -> super::Result { value