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

Spam filter performance and accuracy improvements (part 8)

Browse Source
This commit is contained in:
mdecimus
2025-12-09 22:46:12 +01:00
parent 6c4d28a877
commit efd1d86ff3
21 changed files with 1780 additions and 661 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/common/src/config/inner.rs
View File

@@ -7,9 +7,12 @@
use super::server::tls::{build_self_signed_cert, parse_certificates};
use crate::{
CacheSwap, Caches, Data, DavResource, DavResources, MailboxCache, MessageStoreCache,
MessageUidCache, SpamClassifier, TlsConnectors,
MessageUidCache, TlsConnectors,
auth::{AccessToken, roles::RolePermissions},
config::smtp::resolver::{Policy, Tlsa},
config::{
smtp::resolver::{Policy, Tlsa},
spamfilter::SpamClassifier,
},
listener::blocked::BlockedIps,
manager::webadmin::WebAdminManager,
};

114
crates/common/src/config/spamfilter.rs
View File

@@ -7,7 +7,7 @@
use super::{Variable, functions::ResolveVariable, if_block::IfBlock, tokenizer::TokenMap};
use ahash::AHashSet;
use mail_auth::common::resolver::ToReverseName;
use nlp::classifier::sgd::TextClassifier;
use nlp::classifier::model::{CcfhClassifier, FhClassifier};
use std::{
net::{IpAddr, SocketAddr},
time::Duration,
@@ -20,12 +20,17 @@ use utils::{
};
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
pub struct SpamClassifierModel {
pub classifier: TextClassifier,
pub ham_count: u64,
pub spam_count: u64,
pub last_sample_expiry: u64,
pub last_trained_at: u64,
pub enum SpamClassifier {
FhClassifier {
classifier: FhClassifier,
last_trained_at: u64,
},
CcfhClassifier {
classifier: CcfhClassifier,
last_trained_at: u64,
},
#[default]
Disabled,
}
#[derive(Debug, Clone, Default)]
@@ -75,10 +80,10 @@ pub enum SpamFilterAction<T> {
#[derive(Debug, Clone, Default)]
pub struct ClassifierConfig {
pub epochs: usize,
pub feature_hash_size: usize,
pub alpha: f32,
pub train_batch_size: usize,
pub w_params: FtrlParameters,
pub i_params: Option<FtrlParameters>,
pub num_epochs: usize,
pub reservoir_capacity: usize,
pub min_ham_samples: u64,
pub min_spam_samples: u64,
pub auto_learn_reply_ham: bool,
@@ -87,6 +92,15 @@ pub struct ClassifierConfig {
pub train_frequency: Option<u64>,
}
#[derive(Debug, Clone, Default)]
pub struct FtrlParameters {
pub feature_hash_size: usize,
pub alpha: f64,
pub beta: f64,
pub l1_ratio: f64,
pub l2_ratio: f64,
}
#[derive(Debug, Clone)]
pub struct PyzorConfig {
pub address: SocketAddr,
@@ -443,25 +457,29 @@ impl ClassifierConfig {
return None;
}
let feature_hash_size: usize = config
.property_or_default("spam-filter.classifier.parameters.features", "1048576")
.unwrap_or(1048576);
if !feature_hash_size.is_power_of_two() {
config.new_build_error(
"spam-filter.classifier.parameters.features",
"Feature size must be a power of two.",
);
}
let w_params = FtrlParameters::parse(config, "spam-filter.classifier.parameters", 20);
let i_params = if config
.property_or_default("spam-filter.classifier.ccfh.enable", "false")
.unwrap_or(false)
{
Some(FtrlParameters::parse(
config,
"spam-filter.classifier.ccfh.parameters",
w_params.feature_hash_size - 2,
))
} else {
None
};
ClassifierConfig {
feature_hash_size,
epochs: config
.property_or_default("spam-filter.classifier.parameters.epochs", "1000")
.unwrap_or(1000),
alpha: config
.property_or_default("spam-filter.classifier.parameters.alpha", "0.00001")
.unwrap_or(0.00001),
w_params,
i_params,
num_epochs: config
.property_or_default("spam-filter.classifier.training.epochs", "3")
.unwrap_or(3),
reservoir_capacity: config
.property_or_default("spam-filter.classifier.reservoir-capacity", "1024")
.unwrap_or(1024),
auto_learn_card_is_ham: config
.property_or_default("spam-filter.card-is-ham.learn", "true")
.unwrap_or(true),
@@ -478,9 +496,6 @@ impl ClassifierConfig {
min_spam_samples: config
.property_or_default("spam-filter.classifier.samples.min-spam", "10")
.unwrap_or(10),
train_batch_size: config
.property_or_default("spam-filter.classifier.training.batch-size", "100")
.unwrap_or(100),
train_frequency: config
.property_or_default::<Option<Duration>>(
"spam-filter.classifier.training.frequency",
@@ -493,6 +508,43 @@ impl ClassifierConfig {
}
}
impl FtrlParameters {
pub fn parse(config: &mut Config, prefix: &str, default_features: usize) -> Self {
let feature_hash_size: usize = config
.property((prefix, "features"))
.unwrap_or(default_features);
if !(16..=28).contains(&feature_hash_size) {
config.new_build_error(
(prefix, "features"),
"Feature size must be between 2^16 and 2^28.",
);
}
FtrlParameters {
feature_hash_size: 1 << feature_hash_size,
alpha: config
.property_or_default((prefix, "alpha"), "2.0")
.unwrap_or(2.0),
beta: config
.property_or_default((prefix, "beta"), "1.0")
.unwrap_or(1.0),
l1_ratio: config
.property_or_default((prefix, "l1"), "0.001")
.unwrap_or(0.001),
l2_ratio: config
.property_or_default((prefix, "l2"), "0.0001")
.unwrap_or(0.0001),
}
}
}
impl SpamClassifier {
pub fn is_active(&self) -> bool {
!matches!(self, SpamClassifier::Disabled)
}
}
impl SpamFilterScoreConfig {
pub fn parse(config: &mut Config) -> Self {
SpamFilterScoreConfig {

42
crates/common/src/core.rs
View File

@@ -5,7 +5,7 @@
*/
use crate::{
Inner, Server, SpamClassifier,
Inner, Server,
auth::{AccessToken, ResourceToken, TenantInfo},
config::{
smtp::{
@@ -15,9 +15,10 @@ use crate::{
QueueStrategy, RequireOptional, RoutingStrategy, TlsStrategy, VirtualQueue,
},
},
spamfilter::SpamClassifierModel,
spamfilter::SpamClassifier,
},
ipc::{BroadcastEvent, PushEvent, PushNotification},
manager::SPAM_CLASSIFIER_KEY,
};
use directory::{Directory, QueryParams, Type, backend::internal::manage::ManageDirectory};
use mail_auth::IpLookupStrategy;
@@ -41,7 +42,7 @@ use types::{
blob::{BlobClass, BlobId},
blob_hash::BlobHash,
collection::{Collection, SyncCollection},
field::{Field, PrincipalField},
field::Field,
type_state::{DataType, StateChange},
};
use utils::{map::bitmap::Bitmap, snowflake::SnowflakeIdGenerator};
@@ -1044,41 +1045,20 @@ impl Server {
}
pub async fn spam_model_reload(&self) -> trc::Result<()> {
if let Some(config) = &self.core.spam.classifier {
if self.core.spam.classifier.is_some() {
if let Some(model) = self
.store()
.get_value::<Archive<AlignedBytes>>(ValueKey::property(
u32::MAX,
Collection::Principal,
u32::MAX,
PrincipalField::SpamModel,
))
.blob_store()
.get_blob(SPAM_CLASSIFIER_KEY, 0..usize::MAX)
.await
.and_then(|archive| match archive {
Some(archive) => archive.deserialize::<SpamClassifierModel>().map(Some),
Some(archive) => <Archive<AlignedBytes> as Deserialize>::deserialize(&archive)
.and_then(|archive| archive.deserialize_untrusted::<SpamClassifier>())
.map(Some),
None => Ok(None),
})
.caused_by(trc::location!())?
{
if model.ham_count >= config.min_ham_samples
&& model.spam_count >= config.min_spam_samples
{
self.inner
.data
.spam_classifier
.store(Arc::new(SpamClassifier {
model: model.classifier,
last_trained_at: model.last_trained_at,
}));
} else {
trc::event!(
Spam(SpamEvent::ModelNotReady),
Details = vec![
trc::Value::from(model.ham_count),
trc::Value::from(model.spam_count)
],
);
}
self.inner.data.spam_classifier.store(Arc::new(model));
} else {
trc::event!(Spam(SpamEvent::ModelNotFound));
}

9
crates/common/src/lib.rs
View File

@@ -28,7 +28,6 @@ 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::classifier::sgd::TextClassifier;
use parking_lot::{Mutex, RwLock};
use rustls::sign::CertifiedKey;
use std::{
@@ -73,6 +72,8 @@ pub mod enterprise;
pub use psl;
use crate::config::spamfilter::SpamClassifier;
pub static VERSION_PRIVATE: &str = env!("CARGO_PKG_VERSION");
pub static VERSION_PUBLIC: &str = "1.0.0";
@@ -131,12 +132,6 @@ pub struct Inner {
pub ipc: Ipc,
}
#[derive(Default)]
pub struct SpamClassifier {
pub model: TextClassifier,
pub last_trained_at: u64,
}
pub struct Data {
pub spam_classifier: ArcSwap<SpamClassifier>,

2
crates/common/src/manager/mod.rs
View File

@@ -21,6 +21,8 @@ pub mod webadmin;
const DEFAULT_SPAMFILTER_URL: &str =
"https://github.com/stalwartlabs/spam-filter/releases/latest/download/spam-filter.toml";
pub const WEBADMIN_KEY: &[u8] = "STALWART_WEBADMIN".as_bytes();
pub const SPAM_TRAINER_KEY: &[u8] = "STALWART_SPAM_TRAIN_DATA.lz4".as_bytes();
pub const SPAM_CLASSIFIER_KEY: &[u8] = "STALWART_SPAM_CLASSIFIER_MODEL.lz4".as_bytes();
// SPDX-SnippetBegin
// SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>

111
crates/nlp/src/classifier/adam.rs Normal file
View File

@@ -0,0 +1,111 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::classifier::{Optimizer, model::FhClassifier};
pub struct Adam {
parameters: Vec<f32>,
bias: f32,
learning_rate: f32,
beta1: f32,
beta2: f32,
epsilon: f32,
t: f32,
m0: Vec<f32>,
v0: Vec<f32>,
m_bias: f32,
v_bias: f32,
// Step info
bias2_sqrt: f32,
alpha_t: f32,
}
impl Adam {
pub fn new(n_parameters: usize, learning_rate: f32) -> Self {
Adam {
parameters: vec![0.0; n_parameters],
learning_rate,
beta1: 0.9,
beta2: 0.999,
epsilon: 1e-8,
t: 0.0,
m0: vec![0.0; n_parameters],
v0: vec![0.0; n_parameters],
m_bias: 0.0,
v_bias: 0.0,
bias: 0.0,
bias2_sqrt: 0.0,
alpha_t: 0.0,
}
}
pub fn with_hyperparams(mut self, beta1: f32, beta2: f32, epsilon: f32) -> Self {
self.beta1 = beta1;
self.beta2 = beta2;
self.epsilon = epsilon;
self
}
pub fn with_initial_weights(self, value: f32) -> Self {
Adam {
parameters: vec![value; self.parameters.len()],
..self
}
}
}
impl Optimizer for Adam {
#[inline(always)]
fn step(&mut self) {
self.t += 1.0;
let bias1 = 1.0 - self.beta1.powf(self.t);
self.bias2_sqrt = (1.0 - self.beta2.powf(self.t)).sqrt();
self.alpha_t = self.learning_rate / bias1;
}
#[inline(always)]
fn update_param(&mut self, i: usize, g: f32) {
self.m0[i] = self.beta1 * self.m0[i] + (1.0 - self.beta1) * g;
self.v0[i] = self.beta2 * self.v0[i] + (1.0 - self.beta2) * g * g;
self.parameters[i] -=
self.alpha_t * self.m0[i] / (self.v0[i].sqrt() / self.bias2_sqrt + self.epsilon);
}
#[inline(always)]
fn update_bias(&mut self, g: f32) {
self.m_bias = self.beta1 * self.m_bias + (1.0 - self.beta1) * g;
self.v_bias = self.beta2 * self.v_bias + (1.0 - self.beta2) * g * g;
self.bias -=
self.alpha_t * self.m_bias / (self.v_bias.sqrt() / self.bias2_sqrt + self.epsilon);
}
#[inline(always)]
fn get_param(&self, idx: usize) -> f32 {
self.parameters[idx]
}
#[inline(always)]
fn get_bias(&self) -> f32 {
self.bias
}
#[inline(always)]
fn get_param_mut(&mut self, idx: usize) -> &mut f32 {
&mut self.parameters[idx]
}
fn build_classifier(&self) -> FhClassifier {
FhClassifier {
parameters: self.parameters.clone(),
bias: self.bias,
}
}
fn num_parameters(&self) -> usize {
self.parameters.len()
}
}

162
crates/nlp/src/classifier/feature.rs
View File

@@ -4,82 +4,160 @@
* 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,
#[derive(Debug)]
pub struct Sample<T> {
pub features: Vec<T>,
pub class: f32,
}
pub struct Features(pub(super) HashMap<u32, f32, BuildNoHashHasher<u32>>);
pub struct FeatureBuilder {
pub(super) features_mask: u32,
pub struct FhFeatureBuilder {
pub(super) weight_mask: u64,
}
pub trait Feature {
#[derive(Debug)]
pub struct FhFeature {
pub idx: usize,
pub weight: f32,
}
#[derive(Debug)]
pub struct CcfhFeature {
pub idx_w1: usize,
pub idx_w2: usize,
pub idx_i: usize,
pub weight: f32,
}
pub struct CcfhFeatureBuilder {
pub(super) weight_mask: u64,
pub(super) indicator_mask: u64,
}
pub trait FeatureWeight {
fn idx(&self) -> usize;
fn weight(&self) -> f32;
fn weight_mut(&mut self) -> &mut f32;
}
pub trait UnprocessedFeature {
fn prefix(&self) -> u16;
fn value(&self) -> &[u8];
}
impl FeatureBuilder {
pub fn scale<I: Feature>(&self, features: &mut HashMap<I, f32>) {
impl FeatureWeight for FhFeature {
fn weight(&self) -> f32 {
self.weight
}
fn weight_mut(&mut self) -> &mut f32 {
&mut self.weight
}
fn idx(&self) -> usize {
self.idx
}
}
impl FeatureWeight for CcfhFeature {
fn weight(&self) -> f32 {
self.weight
}
fn weight_mut(&mut self) -> &mut f32 {
&mut self.weight
}
fn idx(&self) -> usize {
self.idx_w1
}
}
impl FeatureBuilder for FhFeatureBuilder {
type Feature = FhFeature;
fn build_feature(&self, bytes: &[u8], weight: f32) -> FhFeature {
let hash1 = xxh3_64_with_seed(bytes, 0);
let sign = if hash1 & (1 << 63) == 0 { 1.0 } else { -1.0 };
FhFeature {
idx: (hash1 & self.weight_mask) as usize,
weight: sign * weight,
}
}
}
impl FeatureBuilder for CcfhFeatureBuilder {
type Feature = CcfhFeature;
fn build_feature(&self, bytes: &[u8], weight: f32) -> CcfhFeature {
let hash1 = xxh3_64_with_seed(bytes, 0);
let hash2 = xxh3_64_with_seed(bytes, 0x9E3779B97F4A7C15);
let hash3 = xxh3_64_with_seed(bytes, 0x517CC1B727220A95);
let sign = if hash3 & (1 << 63) == 0 { 1.0 } else { -1.0 };
CcfhFeature {
idx_w1: (hash1 & self.weight_mask) as usize,
idx_w2: (hash2 & self.weight_mask) as usize,
idx_i: (hash3 & self.indicator_mask) as usize,
weight: sign * weight,
}
}
}
pub trait FeatureBuilder {
// Feature type associated type
type Feature: FeatureWeight;
fn build_feature(&self, bytes: &[u8], weight: f32) -> Self::Feature;
fn scale<I: UnprocessedFeature>(&self, features: &mut HashMap<I, f32>) {
// Log frequency scaling
for x in features.values_mut() {
*x = x.ln_1p();
}
}
pub fn build<I: Feature>(
fn build<I: UnprocessedFeature>(
&self,
features: &HashMap<I, f32>,
features_in: &HashMap<I, f32>,
account_id: Option<u32>,
) -> Features {
// Do the "hash trick"
let mut features_map =
HashMap::with_capacity_and_hasher(features.len(), BuildNoHashHasher::default());
) -> Vec<Self::Feature> {
let mut features_out = Vec::with_capacity(features_in.len());
let mut buf = Vec::with_capacity(2 + 4 + 63);
for (feature, count) in features {
for (feature, count) in features_in {
buf.extend_from_slice(&feature.prefix().to_be_bytes());
buf.extend_from_slice(feature.value());
let big_hash = xxh3_64_with_seed(&buf, 0);
let hash = big_hash as u32 & self.features_mask;
let sign = if big_hash & (1 << 63) == 0 { 1.0 } else { -1.0 };
*features_map.entry(hash).or_default() += sign * count;
features_out.push(self.build_feature(&buf, *count));
if let Some(account_id) = account_id {
buf.extend_from_slice(&account_id.to_be_bytes());
let big_hash = xxh3_64_with_seed(&buf, 0);
let hash = big_hash as u32 & self.features_mask;
let sign = if big_hash & (1 << 63) == 0 { 1.0 } else { -1.0 };
*features_map.entry(hash).or_default() += sign * count;
features_out.push(self.build_feature(&buf, *count));
}
buf.clear();
}
// L2 normalization
let sum_of_squares = features_map
.values()
.map(|&x| x as f64 * x as f64)
let sum_of_squares = features_out
.iter()
.map(|f| f.weight() as f64 * f.weight() as f64)
.sum::<f64>();
if sum_of_squares > 0.0 {
let norm = sum_of_squares.sqrt() as f32;
for x in features_map.values_mut() {
*x /= norm;
for feature in &mut features_out {
*feature.weight_mut() /= norm;
}
}
Features(features_map)
features_out
}
}
impl Sample {
pub fn new(features: Features, class: bool) -> Self {
impl<T> Sample<T> {
pub fn new(features: Vec<T>, class: bool) -> Self {
Self {
features,
class: if class { 1.0 } else { 0.0 },
@@ -87,14 +165,8 @@ impl Sample {
}
}
impl AsRef<Sample> for Sample {
fn as_ref(&self) -> &Sample {
self
}
}
impl AsRef<Features> for Features {
fn as_ref(&self) -> &Features {
impl<T> AsRef<Sample<T>> for Sample<T> {
fn as_ref(&self) -> &Sample<T> {
self
}
}

132
crates/nlp/src/classifier/ftrl.rs Normal file
View File

@@ -0,0 +1,132 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::classifier::{Optimizer, model::FhClassifier};
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
pub struct Ftrl {
alpha: f64,
beta: f64,
l1_ratio: f64,
l2_ratio: f64,
zn: Vec<Zn>,
zn_bias: Zn,
}
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Clone, Copy, Debug, Default)]
pub struct Zn {
z: f32,
n: f64,
}
impl Ftrl {
pub fn new(n_features: usize) -> Self {
Ftrl {
alpha: 2.0,
beta: 1.0,
l1_ratio: 0.001,
l2_ratio: 0.0001,
zn: vec![Zn::default(); n_features],
zn_bias: Zn::default(),
}
}
pub fn with_hyperparams(mut self, alpha: f64, beta: f64, l1_ratio: f64, l2_ratio: f64) -> Self {
self.alpha = alpha;
self.beta = beta;
self.l1_ratio = l1_ratio;
self.l2_ratio = l2_ratio;
self
}
pub fn set_hyperparams(&mut self, alpha: f64, beta: f64, l1_ratio: f64, l2_ratio: f64) {
self.alpha = alpha;
self.beta = beta;
self.l1_ratio = l1_ratio;
self.l2_ratio = l2_ratio;
}
pub fn with_initial_weights(self, value: f32) -> Self {
Ftrl {
zn: vec![Zn { z: value, n: 0.0 }; self.zn.len()],
..self
}
}
}
impl Optimizer for Ftrl {
#[inline(always)]
fn update_param(&mut self, idx: usize, grad: f32) {
let zn = &mut self.zn[idx];
let current_w = if zn.z.abs() as f64 <= self.l1_ratio {
0.0
} else {
-(zn.z - zn.z.signum() * self.l1_ratio as f32)
/ (self.l2_ratio + (self.beta + zn.n.sqrt()) / self.alpha) as f32
};
let grad = grad as f64;
let grad_sq = grad * grad;
let sigma = ((zn.n + grad_sq).sqrt() - zn.n.sqrt()) / self.alpha;
zn.z += (grad - sigma * current_w as f64) as f32;
zn.n += grad_sq;
}
#[inline(always)]
fn update_bias(&mut self, grad: f32) {
let current_bias = -self.zn_bias.z
/ ((self.zn_bias.n.sqrt() + self.beta) / self.alpha + self.l2_ratio) as f32;
let grad = grad as f64;
let grad_sq = grad * grad;
let sigma = ((self.zn_bias.n + grad_sq).sqrt() - self.zn_bias.n.sqrt()) / self.alpha;
self.zn_bias.z += (grad - sigma * current_bias as f64) as f32;
self.zn_bias.n += grad_sq;
}
#[inline(always)]
fn get_param(&self, idx: usize) -> f32 {
let zn = self.zn[idx];
if zn.z.abs() as f64 <= self.l1_ratio {
0.0
} else {
-(zn.z - zn.z.signum() * self.l1_ratio as f32)
/ (self.l2_ratio + (self.beta + zn.n.sqrt()) / self.alpha) as f32
}
}
#[inline(always)]
fn get_bias(&self) -> f32 {
-self.zn_bias.z / ((self.zn_bias.n.sqrt() + self.beta) / self.alpha + self.l2_ratio) as f32
}
fn step(&mut self) {}
#[inline(always)]
fn get_param_mut(&mut self, idx: usize) -> &mut f32 {
&mut self.zn[idx].z
}
fn build_classifier(&self) -> FhClassifier {
FhClassifier {
parameters: self
.zn
.iter()
.map(|zn| {
if zn.z.abs() as f64 <= self.l1_ratio {
0.0
} else {
-(zn.z - zn.z.signum() * self.l1_ratio as f32)
/ (self.l2_ratio + (self.beta + zn.n.sqrt()) / self.alpha) as f32
}
})
.collect(),
bias: self.get_bias(),
}
}
fn num_parameters(&self) -> usize {
self.zn.len()
}
}

41
crates/nlp/src/classifier/mod.rs
View File

@@ -4,5 +4,46 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::classifier::model::FhClassifier;
pub mod adam;
pub mod feature;
pub mod ftrl;
pub mod model;
pub mod reservoir;
pub mod sgd;
pub mod train;
const MAX_DLOSS: f32 = 1e4;
pub trait Optimizer {
fn step(&mut self);
fn update_param(&mut self, i: usize, g: f32);
fn update_bias(&mut self, g: f32);
fn get_param(&self, idx: usize) -> f32;
fn get_param_mut(&mut self, idx: usize) -> &mut f32;
fn get_bias(&self) -> f32;
fn build_classifier(&self) -> FhClassifier;
fn num_parameters(&self) -> usize;
}
#[inline(always)]
fn sigmoid(z: f32) -> f32 {
let z = z.clamp(-35.0, 35.0);
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 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
}
}

109
crates/nlp/src/classifier/model.rs Normal file
View File

@@ -0,0 +1,109 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::classifier::{
feature::{CcfhFeature, CcfhFeatureBuilder, FhFeature, FhFeatureBuilder},
sigmoid,
};
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
pub struct FhClassifier {
pub(crate) parameters: Vec<f32>,
pub(crate) bias: f32,
}
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
pub struct CcfhClassifier {
pub(crate) parameters: Vec<f32>,
pub(crate) indicators: Vec<f32>,
pub(crate) bias: f32,
}
impl FhClassifier {
pub fn predict_proba_sample(&self, features: &[FhFeature]) -> f32 {
let mut z: f32 = 0.0;
for f in features {
z += self.parameters[f.idx] * f.weight;
}
sigmoid(z + self.bias)
}
pub fn predict(&self, features: &[FhFeature]) -> f32 {
if self.predict_proba_sample(features) > 0.7 {
1.0
} else {
0.0
}
}
pub fn predict_batch<I>(&self, test: I) -> Vec<f32>
where
I: IntoIterator,
I::Item: AsRef<Vec<FhFeature>>,
{
test.into_iter()
.map(|features| self.predict(features.as_ref()))
.collect()
}
pub fn feature_builder(&self) -> FhFeatureBuilder {
FhFeatureBuilder {
weight_mask: (self.parameters.len() - 1) as u64,
}
}
pub fn parameters(&self) -> &[f32] {
&self.parameters
}
pub fn bias(&self) -> f32 {
self.bias
}
}
impl CcfhClassifier {
pub fn predict_proba_sample(&self, features: &[CcfhFeature]) -> f32 {
let mut z: f32 = 0.0;
for f in features {
let q = self.indicators[f.idx_i];
let v1 = self.parameters[f.idx_w1];
let v2 = self.parameters[f.idx_w2];
z += (q * v1 + (1.0 - q) * v2) * f.weight;
}
sigmoid(z + self.bias)
}
pub fn predict(&self, features: &[CcfhFeature]) -> f32 {
if self.predict_proba_sample(features) >= 0.5 {
1.0
} else {
0.0
}
}
pub fn predict_batch<I>(&self, test: I) -> Vec<f32>
where
I: IntoIterator,
I::Item: AsRef<Vec<CcfhFeature>>,
{
test.into_iter()
.map(|features| self.predict(features.as_ref()))
.collect()
}
pub fn feature_builder(&self) -> CcfhFeatureBuilder {
CcfhFeatureBuilder {
weight_mask: (self.parameters.len() - 1) as u64,
indicator_mask: (self.indicators.len() - 1) as u64,
}
}
pub fn is_active(&self) -> bool {
!self.parameters.is_empty()
}
}

81
crates/nlp/src/classifier/reservoir.rs Normal file
View File

@@ -0,0 +1,81 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use rand::{Rng, seq::IndexedRandom};
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
pub struct SampleReservoir<T> {
pub spam: SampleReservoirClass<T>,
pub ham: SampleReservoirClass<T>,
}
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
pub struct SampleReservoirClass<T> {
pub buffer: Vec<T>,
pub total_seen: u64,
}
impl<T: Clone + Eq> SampleReservoir<T> {
pub fn update_reservoir(&mut self, item: &T, is_spam: bool, capacity: usize) {
let class = if is_spam {
&mut self.spam
} else {
&mut self.ham
};
class.total_seen += 1;
if class.buffer.len() < capacity {
class.buffer.push(item.clone());
} else if let Some(buf) = class
.buffer
.get_mut(rand::rng().random_range(0..class.total_seen as usize))
{
*buf = item.clone();
}
}
pub fn replay_samples(
&mut self,
count_needed: usize,
is_spam: bool,
) -> impl Iterator<Item = &T> {
(if is_spam {
&mut self.spam
} else {
&mut self.ham
})
.buffer
.choose_multiple(&mut rand::rng(), count_needed)
}
pub fn remove_sample(&mut self, item: &T, is_spam: bool) {
let class = if is_spam {
&mut self.spam
} else {
&mut self.ham
};
if let Some(pos) = class.buffer.iter().position(|x| x == item) {
class.buffer.swap_remove(pos);
}
}
}
impl<T> Default for SampleReservoir<T> {
fn default() -> Self {
SampleReservoir {
spam: SampleReservoirClass {
buffer: Vec::new(),
total_seen: 0,
},
ham: SampleReservoirClass {
buffer: Vec::new(),
total_seen: 0,
},
}
}
}

512
crates/nlp/src/classifier/sgd.rs
View File

@@ -4,163 +4,124 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::classifier::feature::{FeatureBuilder, Features, Sample};
use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
use crate::classifier::{Optimizer, gradient, model::FhClassifier};
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
pub struct TextClassifier {
weights: Vec<f32>,
intercept: f32,
pub struct Sgd {
parameters: Vec<f32>,
bias: f32,
alpha: f64,
l1_ratio: f64,
l2_ratio: f64,
t: f64,
w_scale: f32,
optimal_init: f64,
eta: f32,
u: f32,
q: Vec<f32>,
}
const MAX_DLOSS: f32 = 1e4;
impl TextClassifier {
pub fn new(n_features: usize) -> Self {
TextClassifier {
weights: vec![0.0; n_features],
intercept: 0.0,
}
}
pub fn fit(&mut self, samples: &mut [impl AsRef<Sample>], n_epochs: usize, alpha: f32) {
let mut rng = StdRng::seed_from_u64(42);
let mut t = 1;
let mut w_scale = 1.0;
// Heuristic to initialize 'optimal' learning rate
impl Sgd {
pub fn new(n_features: usize, alpha: f64, l1_ratio: f64, l2_ratio: f64) -> Self {
let typw = (1.0 / alpha.sqrt()).sqrt();
let initial_eta0 = typw / 1.0_f32.max(gradient(1.0, -typw));
let initial_eta0 = typw / 1.0_f64.max(gradient(1.0, -typw as f32) as f64);
let optimal_init = 1.0 / (initial_eta0 * alpha);
for _ in 0..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 / (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 * 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;
}
Sgd {
parameters: vec![0.0; n_features],
bias: 0.0,
alpha,
l1_ratio,
l2_ratio,
t: 0.0,
w_scale: 1.0,
optimal_init,
eta: initial_eta0 as f32,
u: 0.0,
q: vec![0.0; n_features],
}
}
if w_scale != 1.0 {
for w in &mut self.weights {
*w *= w_scale;
pub fn with_initial_parameters(self, value: f32) -> Self {
Sgd {
parameters: vec![value; self.parameters.len()],
..self
}
}
fn maybe_rescale(&mut self) {
if !(1e-6..=1e6).contains(&self.w_scale) {
for w in &mut self.parameters {
*w *= self.w_scale;
}
self.w_scale = 1.0;
}
}
#[inline(always)]
fn apply_l1_penalty(&mut self) {
if self.l1_ratio > 0.0 {
for (z, q) in self.parameters.iter_mut().zip(self.q.iter_mut()) {
let z_orig = *z;
let scaled_z = *z * self.w_scale;
if scaled_z > 0.0 {
*z = (*z - (self.u + *q) / self.w_scale).max(0.0);
} else if scaled_z < 0.0 {
*z = (*z + (self.u - *q) / self.w_scale).min(0.0);
}
*q += self.w_scale * (z_orig - *z);
}
}
}
}
pub 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)
impl Optimizer for Sgd {
fn step(&mut self) {
self.t += 1.0;
self.eta = (1.0 / ((self.alpha) * (self.optimal_init + self.t - 1.0))) as f32;
self.w_scale *= 1.0 - ((1.0 - self.l1_ratio) as f32 * self.eta * self.l2_ratio as f32);
self.u += self.eta * self.l1_ratio as f32 * self.alpha as f32;
}
pub fn predict(&self, features: &Features) -> f32 {
let proba = self.predict_proba_sample(features);
if proba >= 0.5 { 1.0 } else { 0.0 }
fn update_param(&mut self, i: usize, g: f32) {
self.parameters[i] += (-self.eta * g) / self.w_scale;
}
pub fn predict_batch<I>(&self, test: I) -> Vec<f32>
where
I: IntoIterator,
I::Item: AsRef<Features>,
{
test.into_iter()
.map(|features| self.predict(features.as_ref()))
.collect()
fn update_bias(&mut self, g: f32) {
self.bias += -self.eta * g;
self.maybe_rescale();
self.apply_l1_penalty();
}
pub fn feature_builder(&self) -> FeatureBuilder {
FeatureBuilder {
features_mask: (self.weights.len() - 1) as u32,
#[inline(always)]
fn get_param(&self, idx: usize) -> f32 {
self.parameters[idx] * self.w_scale
}
#[inline(always)]
fn get_bias(&self) -> f32 {
self.bias
}
#[inline(always)]
fn get_param_mut(&mut self, idx: usize) -> &mut f32 {
&mut self.parameters[idx]
}
fn build_classifier(&self) -> FhClassifier {
FhClassifier {
parameters: self.parameters.iter().map(|w| w * self.w_scale).collect(),
bias: self.bias,
}
}
pub fn is_active(&self) -> bool {
!self.weights.is_empty()
fn num_parameters(&self) -> usize {
self.parameters.len()
}
}
#[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::{Feature, Sample},
sgd::TextClassifier,
};
use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
use std::{
collections::HashMap,
@@ -169,6 +130,178 @@ pub mod tests {
time::Instant,
};
use crate::classifier::{
Optimizer,
adam::Adam,
feature::{
CcfhFeature, CcfhFeatureBuilder, FeatureBuilder, FhFeature, FhFeatureBuilder, Sample,
UnprocessedFeature,
},
ftrl::Ftrl,
train::{CcfhTrainer, FhTrainer},
};
#[test]
fn text_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 time = Instant::now();
for line in reader.lines().skip(1) {
let line = line.unwrap();
let (text, class) = line.trim().rsplit_once(',').unwrap();
//let (class, text) = line.trim().split_once(',').unwrap();
let text = text.trim_start_matches('"').trim_end_matches('"');
samples.push((text.to_string(), class == "1"));
}
println!("Loaded {} samples in {:?}", samples.len(), time.elapsed());
samples.shuffle(&mut StdRng::seed_from_u64(42));
let (train_samples, test_samples) = train_test_split(&samples, 0.2);
println!(
"Training samples: {}, Testing samples: {}",
train_samples.len(),
test_samples.len()
);
const FH_SIZE: usize = 16;
const CCFH_SIZE: usize = FH_SIZE - 2;
let mut rng = StdRng::seed_from_u64(42);
let fh_builder = FhFeatureBuilder {
weight_mask: (1 << FH_SIZE) - 1,
};
let mut fh_train_samples = build_fh_samples(train_samples.as_slice(), &fh_builder);
fh_train_samples.shuffle(&mut rng);
let fh_test_samples = build_fh_samples(test_samples.as_slice(), &fh_builder);
let ccfh_builder = CcfhFeatureBuilder {
weight_mask: (1 << FH_SIZE) - 1,
indicator_mask: (1 << CCFH_SIZE) - 1,
};
let mut ccfh_train_samples = build_ccfh_samples(train_samples.as_slice(), &ccfh_builder);
ccfh_train_samples.shuffle(&mut rng);
let ccfh_test_samples = build_ccfh_samples(test_samples.as_slice(), &ccfh_builder);
fh_model_stats(
"FTRL",
FhTrainer::new(Ftrl::new(1 << FH_SIZE)),
&fh_train_samples,
&fh_test_samples,
);
ccfh_model_stats(
"FTRL + FTRL",
CcfhTrainer::new(
Ftrl::new(1 << FH_SIZE),
Ftrl::new(1 << CCFH_SIZE).with_initial_weights(0.5),
),
&ccfh_train_samples,
&ccfh_test_samples,
);
fh_model_stats(
"Adam",
FhTrainer::new(Adam::new(1 << FH_SIZE, 0.01)),
&fh_train_samples,
&fh_test_samples,
);
ccfh_model_stats(
"Adam + Adam",
CcfhTrainer::new(
Adam::new(1 << FH_SIZE, 0.01),
Adam::new(1 << CCFH_SIZE, 0.01).with_initial_weights(0.5),
),
&ccfh_train_samples,
&ccfh_test_samples,
);
/*fh_model_stats(
"SGD",
FhTrainer::new(Sgd::new(1 << FH_SIZE, 0.0001, 0.0, 0.0001)),
&fh_train_samples,
&fh_test_samples,
);
ccfh_model_stats(
"FTRL + SGD",
CcfhTrainer::new(
Ftrl::new(1 << FH_SIZE),
Sgd::new(1 << CCFH_SIZE, 0.0001, 0.0, 0.0001).with_initial_parameters(0.5),
),
&ccfh_train_samples,
&ccfh_test_samples,
);*/
}
fn fh_model_stats(
name: &str,
mut model: FhTrainer<impl Optimizer>,
train_samples: &[Sample<FhFeature>],
test_samples: &[Sample<FhFeature>],
) {
print!("⏳ Training {}... ", name);
let time = Instant::now();
let mut batch = Vec::new();
for sample in train_samples {
batch.push(sample);
if batch.len() == 128 {
model.fit(&mut batch, 5);
batch.clear();
}
}
if !batch.is_empty() {
model.fit(&mut batch, 5);
}
println!(" trained in {:?}", time.elapsed());
let y_pred = model
.build_classifier()
.predict_batch(test_samples.iter().map(|s| &s.features));
let y_train: Vec<f32> = 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));
}
fn ccfh_model_stats(
name: &str,
mut model: CcfhTrainer<impl Optimizer, impl Optimizer>,
train_samples: &[Sample<CcfhFeature>],
test_samples: &[Sample<CcfhFeature>],
) {
print!("⏳ Training {}... ", name);
let time = Instant::now();
let mut batch = Vec::new();
for sample in train_samples {
batch.push(sample);
if batch.len() == 128 {
model.fit(&mut batch, 5);
batch.clear();
}
}
if !batch.is_empty() {
model.fit(&mut batch, 5);
}
println!(" trained in {:?}", time.elapsed());
let y_pred = model
.build_classifier()
.predict_batch(test_samples.iter().map(|s| &s.features));
let y_train: Vec<f32> = 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));
}
fn accuracy_score(y_true: &[f32], y_pred: &[f32]) -> f32 {
y_true
.iter()
@@ -231,15 +364,19 @@ pub mod tests {
}
}
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();
#[allow(clippy::type_complexity)]
pub fn train_test_split(
data: &[(String, bool)],
test_size: f32,
) -> (Vec<(&String, bool)>, Vec<(&String, bool)>) {
let mut class_0: Vec<(&String, bool)> = Vec::new();
let mut class_1: Vec<(&String, bool)> = Vec::new();
for sample in data {
if sample.class == 0.0 {
class_0.push(sample);
for (sample, class) in data {
if !*class {
class_0.push((sample, *class));
} else {
class_1.push(sample);
class_1.push((sample, *class));
}
}
@@ -260,7 +397,49 @@ pub mod tests {
(train, test)
}
impl Feature for String {
pub fn build_fh_samples(
data: &[(&String, bool)],
builder: &FhFeatureBuilder,
) -> Vec<Sample<FhFeature>> {
let mut samples = Vec::with_capacity(data.len());
for (text, class) in data {
let mut sample: HashMap<String, f32> = HashMap::new();
for word in text.split_whitespace() {
*sample.entry(word.to_string()).or_default() += 1.0;
}
builder.scale(&mut sample);
samples.push(Sample {
features: builder.build(&sample, 12345.into()),
class: if *class { 1.0 } else { 0.0 },
});
}
samples
}
pub fn build_ccfh_samples(
data: &[(&String, bool)],
builder: &CcfhFeatureBuilder,
) -> Vec<Sample<CcfhFeature>> {
let mut samples = Vec::with_capacity(data.len());
for (text, class) in data {
let mut sample: HashMap<String, f32> = HashMap::new();
for word in text.split_whitespace() {
*sample.entry(word.to_string()).or_default() += 1.0;
}
builder.scale(&mut sample);
samples.push(Sample {
features: builder.build(&sample, 12345.into()),
class: if *class { 1.0 } else { 0.0 },
});
}
samples
}
impl UnprocessedFeature for String {
fn prefix(&self) -> u16 {
0
}
@@ -269,61 +448,4 @@ pub mod tests {
self.as_bytes()
}
}
#[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 = TextClassifier::new(1 << 20);
let builder = model.feature_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('"');
let mut sample: HashMap<String, f32> = HashMap::new();
for word in text.split_whitespace() {
*sample.entry(word.to_string()).or_default() += 1.0;
}
builder.scale(&mut sample);
samples.push(Sample {
features: builder.build(&sample, None),
class: 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, 1000, 0.0001);
println!("SGD Classifier trained in {:?}", time.elapsed());
let y_pred = model.predict_batch(test_samples.iter().map(|s| &s.features));
let y_train: Vec<f32> = 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));
}
}

157
crates/nlp/src/classifier/train.rs Normal file
View File

@@ -0,0 +1,157 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::classifier::{
MAX_DLOSS, Optimizer,
feature::{CcfhFeature, CcfhFeatureBuilder, FhFeature, FhFeatureBuilder, Sample},
gradient,
model::{CcfhClassifier, FhClassifier},
};
use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
pub struct FhTrainer<T: Optimizer> {
pub optimizer: T,
}
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
pub struct CcfhTrainer<W: Optimizer, I: Optimizer> {
pub w_optimizer: W,
pub i_optimizer: I,
}
impl<T: Optimizer> FhTrainer<T> {
pub fn new(optimizer: T) -> Self {
FhTrainer { optimizer }
}
pub fn fit(&mut self, samples: &mut [impl AsRef<Sample<FhFeature>>], num_epochs: usize) {
for _ in 0..num_epochs {
samples.shuffle(&mut StdRng::seed_from_u64(42));
for sample in samples.iter() {
let sample = sample.as_ref();
let mut dot: f32 = 0.0;
for f in &sample.features {
dot += self.optimizer.get_param(f.idx) * f.weight;
}
let p = dot + self.optimizer.get_bias();
let dloss = gradient(sample.class, p).clamp(-MAX_DLOSS, MAX_DLOSS);
self.optimizer.step();
for f in &sample.features {
self.optimizer.update_param(f.idx, dloss * f.weight);
}
self.optimizer.update_bias(dloss);
}
}
}
pub fn feature_builder(&self) -> FhFeatureBuilder {
FhFeatureBuilder {
weight_mask: (self.optimizer.num_parameters() - 1) as u64,
}
}
pub fn build_classifier(&self) -> FhClassifier {
self.optimizer.build_classifier()
}
pub fn optimizer(&self) -> &T {
&self.optimizer
}
pub fn optimizer_mut(&mut self) -> &mut T {
&mut self.optimizer
}
}
impl<W: Optimizer, I: Optimizer> CcfhTrainer<W, I> {
pub fn new(w_optimizer: W, i_optimizer: I) -> Self {
CcfhTrainer {
w_optimizer,
i_optimizer,
}
}
pub fn fit(&mut self, samples: &mut [impl AsRef<Sample<CcfhFeature>>], num_epochs: usize) {
for _ in 0..num_epochs {
samples.shuffle(&mut StdRng::seed_from_u64(42));
for sample in samples.iter() {
let sample = sample.as_ref();
let mut dot: f32 = 0.0;
for f in &sample.features {
let q = self.i_optimizer.get_param(f.idx_i);
let v1 = self.w_optimizer.get_param(f.idx_w1);
let v2 = self.w_optimizer.get_param(f.idx_w2);
dot += (q * v1 + (1.0 - q) * v2) * f.weight;
}
let p = dot + self.w_optimizer.get_bias();
let dloss = gradient(sample.class, p).clamp(-MAX_DLOSS, MAX_DLOSS);
self.w_optimizer.step();
self.i_optimizer.step();
for f in &sample.features {
let q = self.i_optimizer.get_param(f.idx_i);
let v1 = self.w_optimizer.get_param(f.idx_w1);
let v2 = self.w_optimizer.get_param(f.idx_w2);
// Update weights
let d_v1 = f.weight * q;
let d_v2 = f.weight * (1.0 - q);
self.w_optimizer.update_param(f.idx_w1, dloss * d_v1);
self.w_optimizer.update_param(f.idx_w2, dloss * d_v2);
// Update indicator
let d_q = (v1 - v2) * f.weight;
self.i_optimizer.update_param(f.idx_i, dloss * d_q);
let fi = self.i_optimizer.get_param_mut(f.idx_i);
*fi = fi.clamp(0.0, 1.0);
}
self.w_optimizer.update_bias(dloss);
}
}
}
pub fn feature_builder(&self) -> CcfhFeatureBuilder {
CcfhFeatureBuilder {
weight_mask: (self.w_optimizer.num_parameters() - 1) as u64,
indicator_mask: (self.i_optimizer.num_parameters() - 1) as u64,
}
}
pub fn build_classifier(&self) -> CcfhClassifier {
let w_classifier = self.w_optimizer.build_classifier();
let i_classifier = self.i_optimizer.build_classifier();
CcfhClassifier {
parameters: w_classifier.parameters,
indicators: i_classifier.parameters,
bias: w_classifier.bias,
}
}
pub fn w_optimizer(&self) -> &W {
&self.w_optimizer
}
pub fn w_optimizer_mut(&mut self) -> &mut W {
&mut self.w_optimizer
}
pub fn i_optimizer(&self) -> &I {
&self.i_optimizer
}
pub fn i_optimizer_mut(&mut self) -> &mut I {
&mut self.i_optimizer
}
}

15
crates/services/src/housekeeper/mod.rs
View File

@@ -6,7 +6,7 @@
use common::{
Inner, KV_LOCK_HOUSEKEEPER, LONG_1D_SLUMBER, Server,
config::telemetry::OtelMetrics,
config::{spamfilter, telemetry::OtelMetrics},
core::BuildServer,
ipc::{BroadcastEvent, HousekeeperEvent, PurgeType},
};
@@ -109,11 +109,14 @@ pub fn spawn_housekeeper(inner: Arc<Inner>, mut rx: mpsc::Receiver<HousekeeperEv
.as_ref()
.and_then(|c| c.train_frequency)
{
let last_trained_at = server.inner.data.spam_classifier.load().last_trained_at;
let next_train = if last_trained_at > 0 {
now().saturating_sub(last_trained_at).min(train_frequency)
} else {
train_frequency
let next_train = match server.inner.data.spam_classifier.load().as_ref() {
spamfilter::SpamClassifier::FhClassifier {
last_trained_at, ..
}
| spamfilter::SpamClassifier::CcfhClassifier {
last_trained_at, ..
} => now().saturating_sub(*last_trained_at).min(train_frequency),
spamfilter::SpamClassifier::Disabled => train_frequency,
};
queue.schedule(

14
crates/spam-filter/src/analysis/score.rs
View File

@@ -250,19 +250,17 @@ impl SpamFilterAnalyzeScore for Server {
pub trait ConfidenceStore {
fn spam_tag(&self) -> &'static str;
fn is_certain(&self) -> Option<bool>;
}
impl ConfidenceStore for f32 {
fn spam_tag(&self) -> &'static str {
match *self {
p if p < 0.10 => "PROB_HAM_HIGH",
p if p < 0.15 => "PROB_HAM_HIGH",
p if p < 0.25 => "PROB_HAM_MEDIUM",
p if p < 0.40 => "PROB_HAM_LOW",
p if p < 0.50 => "PROB_HAM_UNCERTAIN",
p if p < 0.60 => "PROB_SPAM_UNCERTAIN",
p if p < 0.75 => "PROB_SPAM_LOW",
p if p < 0.90 => "PROB_SPAM_MEDIUM",
p if p < 0.85 => "PROB_SPAM_MEDIUM",
p => {
if p.is_finite() {
"PROB_SPAM_HIGH"
@@ -272,12 +270,4 @@ impl ConfidenceStore for f32 {
}
}
}
fn is_certain(&self) -> Option<bool> {
match *self {
p if p < 0.40 => Some(false), // certain ham
p if p > 0.60 => Some(true), // certain spam
_ => None, // uncertain
}
}
}

637
crates/spam-filter/src/modules/classifier.rs
View File

@@ -11,16 +11,20 @@ use crate::analysis::url::SpamFilterAnalyzeUrl;
use crate::modules::html::{A, ALT, HREF, HtmlToken, IMG, SRC, TITLE};
use crate::{Email, SpamFilterContext, TextPart};
use crate::{Hostname, SpamFilterInput};
use common::config::spamfilter::SpamClassifierModel;
use common::config::spamfilter;
use common::manager::{SPAM_CLASSIFIER_KEY, SPAM_TRAINER_KEY};
use common::{Server, config::spamfilter::Location, ipc::BroadcastEvent};
use mail_auth::DmarcResult;
use mail_parser::{MessageParser, MimeHeaders};
use nlp::classifier::feature::Sample;
use nlp::tokenizers::types::TypesTokenizer;
use nlp::{
classifier::{feature::Feature, sgd::TextClassifier},
tokenizers::{stream::WordStemTokenizer, types::TokenType},
use nlp::classifier::feature::{
CcfhFeature, CcfhFeatureBuilder, FeatureBuilder, FhFeature, FhFeatureBuilder, Sample,
UnprocessedFeature,
};
use nlp::classifier::ftrl::Ftrl;
use nlp::classifier::reservoir::SampleReservoir;
use nlp::classifier::train::{CcfhTrainer, FhTrainer};
use nlp::tokenizers::types::TypesTokenizer;
use nlp::tokenizers::{stream::WordStemTokenizer, types::TokenType};
use std::time::Instant;
use std::{
borrow::Cow,
@@ -28,9 +32,12 @@ use std::{
hash::{Hash, RandomState},
sync::Arc,
};
use store::rand::SeedableRng;
use store::rand::rngs::StdRng;
use store::rand::seq::SliceRandom;
use store::write::{BlobLink, now};
use store::{
IterateParams, Serialize, U32_LEN, U64_LEN, ValueKey,
Deserialize, IterateParams, Serialize, U32_LEN, U64_LEN, ValueKey,
write::{
AlignedBytes, Archive, Archiver, BatchBuilder, BlobOp, ValueClass,
key::DeserializeBigEndian,
@@ -38,7 +45,7 @@ use store::{
};
use tokio::sync::{mpsc, oneshot};
use trc::{AddContext, SpamEvent};
use types::{blob_hash::BlobHash, collection::Collection, field::PrincipalField};
use types::blob_hash::BlobHash;
use unicode_general_category::{GeneralCategory, get_general_category};
use unicode_normalization::UnicodeNormalization;
use unicode_security::mixed_script::AugmentedScriptSet;
@@ -57,13 +64,32 @@ pub trait SpamClassifier {
) -> impl Future<Output = Tokens<'x>> + Send;
}
struct TrainingSample {
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Clone, PartialEq, Eq, Debug)]
pub struct TrainingSample {
hash: BlobHash,
account_id: u32,
}
struct TrainingTask {
sample: TrainingSample,
is_spam: bool,
is_replay: bool,
remove: Option<u64>,
}
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
pub struct SpamTrainer {
pub trainer: SpamTrainerClass,
pub reservoir: SampleReservoir<TrainingSample>,
pub last_sample_expiry: u64,
}
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
pub enum SpamTrainerClass {
FtrlFh(Box<FhTrainer<Ftrl>>),
FtrlCfh(Box<CcfhTrainer<Ftrl, Ftrl>>),
}
impl SpamClassifier for Server {
async fn spam_train(&self, retrain: bool) -> trc::Result<()> {
let Some(config) = &self.core.spam.classifier else {
@@ -73,34 +99,64 @@ impl SpamClassifier for Server {
let started = Instant::now();
trc::event!(Spam(SpamEvent::TrainStarted));
// Fetch model
let mut model = if !retrain
&& let Some(model) = self
.store()
.get_value::<Archive<AlignedBytes>>(ValueKey::property(
u32::MAX,
Collection::Principal,
u32::MAX,
PrincipalField::SpamModel,
))
// Fetch or build trainer
let mut trainer = if !retrain
&& let Some(trainer) = self
.blob_store()
.get_blob(SPAM_TRAINER_KEY, 0..usize::MAX)
.await
.and_then(|archive| match archive {
Some(archive) => archive.deserialize::<SpamClassifierModel>().map(Some),
Some(archive) => <Archive<AlignedBytes> as Deserialize>::deserialize(&archive)
.and_then(|archive| archive.deserialize_untrusted::<SpamTrainer>())
.map(Some),
None => Ok(None),
})
.caused_by(trc::location!())?
{
model
trainer
} else {
SpamClassifierModel {
classifier: TextClassifier::new(config.feature_hash_size),
ham_count: 0,
spam_count: 0,
SpamTrainer {
trainer: match &config.i_params {
Some(i_params) => SpamTrainerClass::FtrlCfh(Box::new(CcfhTrainer::new(
Ftrl::new(config.w_params.feature_hash_size),
Ftrl::new(i_params.feature_hash_size).with_initial_weights(0.5),
))),
None => SpamTrainerClass::FtrlFh(Box::new(FhTrainer::new(Ftrl::new(
config.w_params.feature_hash_size,
)))),
},
reservoir: SampleReservoir::default(),
last_sample_expiry: 0,
last_trained_at: 0,
}
};
// Update hyperparameters
match (&mut trainer.trainer, &config.i_params) {
(SpamTrainerClass::FtrlFh(trainer), None) => {
trainer.optimizer_mut().set_hyperparams(
config.w_params.alpha,
config.w_params.beta,
config.w_params.l1_ratio,
config.w_params.l2_ratio,
);
}
(SpamTrainerClass::FtrlCfh(trainer), Some(i_params)) => {
trainer.w_optimizer_mut().set_hyperparams(
config.w_params.alpha,
config.w_params.beta,
config.w_params.l1_ratio,
config.w_params.l2_ratio,
);
trainer.i_optimizer_mut().set_hyperparams(
i_params.alpha,
i_params.beta,
i_params.l1_ratio,
i_params.l2_ratio,
);
}
_ => {}
}
// Fetch blob hashes for samples
let mut samples = Vec::new();
let mut remove_entries = false;
@@ -110,7 +166,7 @@ impl SpamClassifier for Server {
document_id: 0,
class: ValueClass::Blob(BlobOp::SpamSample {
hash: BlobHash::default(),
until: model.last_sample_expiry + 1,
until: trainer.last_sample_expiry + 1,
}),
};
let to_key = ValueKey {
@@ -122,6 +178,8 @@ impl SpamClassifier for Server {
until: u64::MAX,
}),
};
let mut spam_count = 0;
let mut ham_count = 0;
self.store()
.iterate(
IterateParams::new(from_key, to_key).ascending(),
@@ -144,21 +202,32 @@ impl SpamClassifier for Server {
let do_remove = *hold == 0;
let is_spam = *is_spam == 1;
samples.push(TrainingSample {
hash,
account_id,
let sample = TrainingSample { hash, account_id };
// Add to reservoir
if !do_remove {
trainer.reservoir.update_reservoir(
&sample,
is_spam,
config.reservoir_capacity,
);
}
samples.push(TrainingTask {
sample,
is_spam,
is_replay: false,
remove: do_remove.then_some(until),
});
remove_entries |= do_remove;
// Update model stats
model.last_sample_expiry = until;
// Update trainer stats
trainer.last_sample_expiry = until;
if is_spam {
model.spam_count += 1;
spam_count += 1;
} else {
model.ham_count += 1;
ham_count += 1;
}
Ok(true)
@@ -176,57 +245,80 @@ impl SpamClassifier for Server {
return Ok(());
}
// Balance classes if needed
if spam_count > ham_count {
// We have too much spam today. We need to replay old HAM.
samples.extend(
trainer
.reservoir
.replay_samples(spam_count - ham_count, false)
.map(|sample| TrainingTask {
sample: sample.clone(),
is_spam: false,
is_replay: true,
remove: None,
}),
);
} else if ham_count > spam_count {
// We have too much ham today. We need to replay old SPAM.
samples.extend(
trainer
.reservoir
.replay_samples(ham_count - spam_count, true)
.map(|sample| TrainingTask {
sample: sample.clone(),
is_spam: true,
is_replay: true,
remove: None,
}),
);
}
let num_samples = samples.len();
samples.shuffle(&mut StdRng::seed_from_u64(42));
// Spawn training task
struct TrainJob {
samples: Vec<Sample>,
done: oneshot::Sender<()>,
}
let builder = model.classifier.feature_builder();
let n_epochs = config.epochs;
let alpha = config.alpha;
let (batch_tx, mut batch_rx) = mpsc::channel::<TrainJob>(1);
let (model_tx, model_rx) = oneshot::channel();
let task = trainer.trainer.spawn(config.num_epochs)?;
let is_fh = matches!(task, TrainTask::Fh { .. });
std::thread::Builder::new()
.name("SGD Train Task".into())
.spawn(move || {
while let Some(mut job) = batch_rx.blocking_recv() {
model.classifier.fit(&mut job.samples, n_epochs, alpha);
let _ = job.done.send(());
}
// Send model back when done
let _ = model_tx.send(model);
})
.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
.reason(err)
.details("Failed to spawn spam train task")
.caused_by(trc::location!())
})?;
// Train
for chunk in samples.chunks(128) {
let mut fh_samples = if is_fh {
Vec::with_capacity(chunk.len())
} else {
Vec::new()
};
let mut ccfh_samples = if !is_fh {
Vec::with_capacity(chunk.len())
} else {
Vec::new()
};
// Train model
for chunk in samples.chunks(config.train_batch_size.max(10)) {
let mut samples = Vec::with_capacity(chunk.len());
for sample in chunk {
let account_id = if sample.account_id != u32::MAX {
Some(sample.account_id)
let account_id = if sample.sample.account_id != u32::MAX {
Some(sample.sample.account_id)
} else {
None
};
let Some(raw_message) = self
.blob_store()
.get_blob(sample.hash.as_slice(), 0..usize::MAX)
.get_blob(sample.sample.hash.as_slice(), 0..usize::MAX)
.await
.caused_by(trc::location!())?
else {
trc::event!(
Spam(SpamEvent::TrainSampleNotFound),
Reason = "Blob not found",
AccountId = account_id,
BlobId = sample.hash.to_hex(),
);
if sample.is_replay {
trainer
.reservoir
.remove_sample(&sample.sample, sample.is_spam);
} else {
trc::event!(
Spam(SpamEvent::TrainSampleNotFound),
Reason = "Blob not found",
AccountId = account_id,
BlobId = sample.sample.hash.to_hex(),
);
}
continue;
};
@@ -237,26 +329,57 @@ impl SpamClassifier for Server {
self.spam_filter_analyze_domain(&mut ctx).await;
self.spam_filter_analyze_url(&mut ctx).await;
let mut tokens = self.spam_build_tokens(&ctx).await.0;
builder.scale(&mut tokens);
let features = builder.build(&tokens, account_id);
samples.push(Sample::new(features, sample.is_spam));
match &task {
TrainTask::Fh { builder, .. } => {
builder.scale(&mut tokens);
fh_samples.push(Sample::new(
builder.build(&tokens, account_id),
sample.is_spam,
));
}
TrainTask::Ccfh { builder, .. } => {
builder.scale(&mut tokens);
ccfh_samples.push(Sample::new(
builder.build(&tokens, account_id),
sample.is_spam,
));
}
}
}
// Send batch for training
let (done_tx, done_rx) = oneshot::channel();
batch_tx
.send(TrainJob {
samples,
done: done_tx,
})
.await
.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
.reason(err)
.details("Spam train task failed")
.caused_by(trc::location!())
})?;
let (done_tx, done_rx) = oneshot::channel::<()>();
match &task {
TrainTask::Fh { batch_tx, .. } => {
batch_tx
.send(FhTrainJob {
samples: fh_samples,
done: done_tx,
})
.await
.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
.reason(err)
.details("Spam train task failed")
.caused_by(trc::location!())
})?;
}
TrainTask::Ccfh { batch_tx, .. } => {
batch_tx
.send(CcfhTrainJob {
samples: ccfh_samples,
done: done_tx,
})
.await
.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
.reason(err)
.details("Spam train task failed")
.caused_by(trc::location!())
})?;
}
}
done_rx.await.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
@@ -266,53 +389,89 @@ impl SpamClassifier for Server {
})?;
}
// Take ownership of model
drop(batch_tx);
let mut model = model_rx.await.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
.reason(err)
.details("Spam train task failed")
.caused_by(trc::location!())
})?;
// Take ownership of trainer
trainer.trainer = match task {
TrainTask::Fh {
batch_tx,
trainer_rx,
..
} => {
drop(batch_tx);
SpamTrainerClass::FtrlFh(trainer_rx.await.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
.reason(err)
.details("Spam train task failed")
.caused_by(trc::location!())
})?)
}
TrainTask::Ccfh {
batch_tx,
trainer_rx,
..
} => {
drop(batch_tx);
SpamTrainerClass::FtrlCfh(trainer_rx.await.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
.reason(err)
.details("Spam train task failed")
.caused_by(trc::location!())
})?)
}
};
// Store updated model
model.last_trained_at = now();
let archiver = Archiver::new(model);
let mut batch = BatchBuilder::new();
batch
.with_account_id(u32::MAX)
.with_collection(Collection::Principal)
.with_document(u32::MAX)
.set(
ValueClass::Property(PrincipalField::SpamModel.into()),
archiver.serialize().caused_by(trc::location!())?,
);
self.store()
.write(batch.build_all())
// Store updated trainer and classifier
let ham_count = trainer.reservoir.ham.total_seen;
let spam_count = trainer.reservoir.spam.total_seen;
let classifier = Archiver::new(match &trainer.trainer {
SpamTrainerClass::FtrlFh(fh_trainer) => spamfilter::SpamClassifier::FhClassifier {
classifier: fh_trainer.build_classifier(),
last_trained_at: now(),
},
SpamTrainerClass::FtrlCfh(ccfh_trainer) => spamfilter::SpamClassifier::CcfhClassifier {
classifier: ccfh_trainer.build_classifier(),
last_trained_at: now(),
},
});
self.blob_store()
.put_blob(
SPAM_TRAINER_KEY,
&Archiver::new(trainer)
.serialize()
.caused_by(trc::location!())?,
)
.await
.caused_by(trc::location!())?;
if ham_count >= config.min_ham_samples && spam_count >= config.min_spam_samples {
self.blob_store()
.put_blob(
SPAM_CLASSIFIER_KEY,
&classifier.serialize().caused_by(trc::location!())?,
)
.await
.caused_by(trc::location!())?;
// Reload model
let model = archiver.inner;
if model.ham_count >= config.min_ham_samples && model.spam_count >= config.min_spam_samples
{
self.inner
.data
.spam_classifier
.store(Arc::new(common::SpamClassifier {
model: model.classifier,
last_trained_at: model.last_trained_at,
}));
.store(Arc::new(classifier.inner));
self.cluster_broadcast(BroadcastEvent::ReloadSpamFilter)
.await;
} else {
self.blob_store()
.delete_blob(SPAM_CLASSIFIER_KEY)
.await
.caused_by(trc::location!())?;
trc::event!(
Spam(SpamEvent::ModelNotReady),
Details = vec![trc::Value::from(ham_count), trc::Value::from(spam_count)],
);
}
trc::event!(
Spam(SpamEvent::TrainCompleted),
Total = num_samples,
Details = vec![
trc::Value::from(model.ham_count),
trc::Value::from(model.spam_count)
],
Details = vec![trc::Value::from(ham_count), trc::Value::from(spam_count)],
Elapsed = started.elapsed()
);
@@ -322,13 +481,13 @@ impl SpamClassifier for Server {
for sample in samples {
if let Some(until) = sample.remove {
batch
.with_account_id(sample.account_id)
.with_account_id(sample.sample.account_id)
.clear(BlobOp::Link {
hash: sample.hash.clone(),
hash: sample.sample.hash.clone(),
to: BlobLink::Temporary { until },
})
.clear(BlobOp::SpamSample {
hash: sample.hash,
hash: sample.sample.hash,
until,
});
if batch.is_large_batch() {
@@ -353,59 +512,102 @@ impl SpamClassifier for Server {
async fn spam_classify(&self, ctx: &mut SpamFilterContext<'_>) -> trc::Result<()> {
let classifier = self.inner.data.spam_classifier.load_full();
let model = &classifier.model;
if model.is_active() {
let started = Instant::now();
let mut classifier_confidence = Vec::with_capacity(ctx.input.env_rcpt_to.len());
let mut has_prediction = false;
let mut tokens = self.spam_build_tokens(ctx).await.0;
let feature_builder = model.feature_builder();
feature_builder.scale(&mut tokens);
let started = Instant::now();
match classifier.as_ref() {
spamfilter::SpamClassifier::FhClassifier { classifier, .. } => {
let mut classifier_confidence = Vec::with_capacity(ctx.input.env_rcpt_to.len());
let mut has_prediction = false;
let mut tokens = self.spam_build_tokens(ctx).await.0;
let feature_builder = classifier.feature_builder();
feature_builder.scale(&mut tokens);
for rcpt in &ctx.input.env_rcpt_to {
let prediction = if let Some(account_id) = self
.directory()
.email_to_id(rcpt)
.await
.caused_by(trc::location!())?
{
has_prediction = true;
model
.predict_proba_sample(&feature_builder.build(&tokens, account_id.into()))
.into()
for rcpt in &ctx.input.env_rcpt_to {
let prediction = if let Some(account_id) = self
.directory()
.email_to_id(rcpt)
.await
.caused_by(trc::location!())?
{
has_prediction = true;
classifier
.predict_proba_sample(
&feature_builder.build(&tokens, account_id.into()),
)
.into()
} else {
None
};
classifier_confidence.push(prediction);
}
if has_prediction {
ctx.result.classifier_confidence = classifier_confidence;
} else {
None
};
classifier_confidence.push(prediction);
// None of the recipients are local, default to global model prediction
let prediction =
classifier.predict_proba_sample(&feature_builder.build(&tokens, None));
ctx.result.classifier_confidence =
vec![prediction.into(); ctx.input.env_rcpt_to.len()];
}
}
spamfilter::SpamClassifier::CcfhClassifier { classifier, .. } => {
let mut classifier_confidence = Vec::with_capacity(ctx.input.env_rcpt_to.len());
let mut has_prediction = false;
let mut tokens = self.spam_build_tokens(ctx).await.0;
let feature_builder = classifier.feature_builder();
feature_builder.scale(&mut tokens);
if has_prediction {
ctx.result.classifier_confidence = classifier_confidence;
} else {
// None of the recipients are local, default to global model prediction
let prediction = model.predict_proba_sample(&feature_builder.build(&tokens, None));
ctx.result.classifier_confidence =
vec![prediction.into(); ctx.input.env_rcpt_to.len()];
for rcpt in &ctx.input.env_rcpt_to {
let prediction = if let Some(account_id) = self
.directory()
.email_to_id(rcpt)
.await
.caused_by(trc::location!())?
{
has_prediction = true;
classifier
.predict_proba_sample(
&feature_builder.build(&tokens, account_id.into()),
)
.into()
} else {
None
};
classifier_confidence.push(prediction);
}
if has_prediction {
ctx.result.classifier_confidence = classifier_confidence;
} else {
// None of the recipients are local, default to global model prediction
let prediction =
classifier.predict_proba_sample(&feature_builder.build(&tokens, None));
ctx.result.classifier_confidence =
vec![prediction.into(); ctx.input.env_rcpt_to.len()];
}
}
spamfilter::SpamClassifier::Disabled => {
return Ok(());
}
trc::event!(
Spam(SpamEvent::Classify),
Result = ctx
.result
.classifier_confidence
.iter()
.zip(ctx.input.env_rcpt_to.iter())
.map(|(v, rcpt)| trc::Value::Array(vec![
trc::Value::from(rcpt.to_string()),
trc::Value::from(*v)
]))
.collect::<Vec<_>>(),
SpanId = ctx.input.span_id,
Elapsed = started.elapsed()
);
}
trc::event!(
Spam(SpamEvent::Classify),
Result = ctx
.result
.classifier_confidence
.iter()
.zip(ctx.input.env_rcpt_to.iter())
.map(|(v, rcpt)| trc::Value::Array(vec![
trc::Value::from(rcpt.to_string()),
trc::Value::from(*v)
]))
.collect::<Vec<_>>(),
SpanId = ctx.input.span_id,
Elapsed = started.elapsed()
);
Ok(())
}
@@ -597,6 +799,92 @@ impl SpamClassifier for Server {
}
}
struct FhTrainJob {
samples: Vec<Sample<FhFeature>>,
done: oneshot::Sender<()>,
}
struct CcfhTrainJob {
samples: Vec<Sample<CcfhFeature>>,
done: oneshot::Sender<()>,
}
enum TrainTask {
Fh {
batch_tx: mpsc::Sender<FhTrainJob>,
trainer_rx: oneshot::Receiver<Box<FhTrainer<Ftrl>>>,
builder: FhFeatureBuilder,
},
Ccfh {
batch_tx: mpsc::Sender<CcfhTrainJob>,
trainer_rx: oneshot::Receiver<Box<CcfhTrainer<Ftrl, Ftrl>>>,
builder: CcfhFeatureBuilder,
},
}
impl SpamTrainerClass {
fn spawn(self, num_epochs: usize) -> trc::Result<TrainTask> {
match self {
SpamTrainerClass::FtrlFh(mut trainer) => {
let builder = trainer.feature_builder();
let (batch_tx, mut batch_rx) = mpsc::channel::<FhTrainJob>(1);
let (trainer_tx, trainer_rx) = oneshot::channel();
std::thread::Builder::new()
.name("FTRL Train Task".into())
.spawn(move || {
while let Some(mut job) = batch_rx.blocking_recv() {
trainer.fit(&mut job.samples, num_epochs);
let _ = job.done.send(());
}
// Send trainer back when done
let _ = trainer_tx.send(trainer);
})
.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
.reason(err)
.details("Failed to spawn spam train task")
.caused_by(trc::location!())
})?;
Ok(TrainTask::Fh {
batch_tx,
trainer_rx,
builder,
})
}
SpamTrainerClass::FtrlCfh(mut trainer) => {
let builder = trainer.feature_builder();
let (batch_tx, mut batch_rx) = mpsc::channel::<CcfhTrainJob>(1);
let (trainer_tx, trainer_rx) = oneshot::channel();
std::thread::Builder::new()
.name("FTRL Train Task".into())
.spawn(move || {
while let Some(mut job) = batch_rx.blocking_recv() {
trainer.fit(&mut job.samples, num_epochs);
let _ = job.done.send(());
}
// Send trainer back when done
let _ = trainer_tx.send(trainer);
})
.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
.reason(err)
.details("Failed to spawn spam train task")
.caused_by(trc::location!())
})?;
Ok(TrainTask::Ccfh {
batch_tx,
trainer_rx,
builder,
})
}
}
}
}
const MAX_TOKEN_LENGTH: usize = 16;
#[derive(
@@ -744,15 +1032,10 @@ impl<'x> Tokens<'x> {
});
}
if is_body {
if is_body && word.len() == upper_count && word.len() > 3 {
self.insert(Token::Word {
value: "_word".into(),
value: "_allcaps".into(),
});
if word.len() == upper_count && word.len() > 3 {
self.insert(Token::Word {
value: "_allcaps".into(),
});
}
}
}
TokenType::Alphanumeric(word) => {
@@ -1024,7 +1307,7 @@ fn truncate_word(word: &str, max_len: usize) -> &str {
}
}
impl Feature for Token<'_> {
impl UnprocessedFeature for Token<'_> {
fn prefix(&self) -> u16 {
match self {
Token::Word { .. } => 0,

2
crates/types/src/field.rs
View File

@@ -83,7 +83,6 @@ pub enum PrincipalField {
DefaultAddressBookId,
ActiveScriptId,
PushSubscriptions,
SpamModel,
}
impl From<ContactField> for u8 {
@@ -163,7 +162,6 @@ impl From<PrincipalField> for u8 {
PrincipalField::DefaultAddressBookId => 48,
PrincipalField::ActiveScriptId => 49,
PrincipalField::PushSubscriptions => 44,
PrincipalField::SpamModel => 52,
PrincipalField::Archive => ARCHIVE_FIELD,
}
}

6
tests/resources/smtp/antispam/classifier.test
View File

@@ -1,5 +1,5 @@
envelope_to hello@world.com
expect PROB_SPAM_HIGH
expect PROB_SPAM_MEDIUM
Subject: save up to NUMBER on life insurance
@@ -7,14 +7,14 @@ why spend more than you have to life quote savings ensuring your family s financ
<!-- NEXT TEST -->
envelope_to hello@world.com
expect PROB_HAM_HIGH
expect PROB_HAM_MEDIUM
Subject: can someone explain
what type of operating system solaris is as ive never seen or used it i dont know wheather to get a server from sun or from dell i would prefer a linux based server and sun seems to be the one for that but im not sure if solaris is a distro of linux or a completely different operating system can someone explain kiall mac innes irish linux users group ilug URL URL for un subscription information list maintainer listmaster URL
<!-- NEXT TEST -->
envelope_to hello@world.com
expect PROB_HAM_LOW
expect PROB_SPAM_LOW
Subject: Lorem ipsum dolor sit amet, consectetur adipiscing elit

280
tests/src/imap/antispam.rs
View File

@@ -6,14 +6,14 @@
use super::{IMAPTest, ImapConnection};
use crate::{imap::Type, jmap::mail::delivery::SmtpConnection, smtp::session::VerifyResponse};
use common::{Server, config::spamfilter::SpamClassifierModel};
use common::{Server, manager::SPAM_TRAINER_KEY};
use imap_proto::ResponseType;
use spam_filter::modules::classifier::SpamClassifier;
use spam_filter::modules::classifier::{SpamClassifier, SpamTrainer};
use store::{
IterateParams, U32_LEN, U64_LEN, ValueKey,
Deserialize, IterateParams, U32_LEN, U64_LEN, ValueKey,
write::{AlignedBytes, Archive, BlobOp, ValueClass, key::DeserializeBigEndian},
};
use types::{blob_hash::BlobHash, collection::Collection, field::PrincipalField};
use types::blob_hash::BlobHash;
pub async fn test(handle: &IMAPTest) {
println!("Running Spam classifier tests...");
@@ -71,23 +71,14 @@ pub async fn test(handle: &IMAPTest) {
// Train the classifier
handle.server.spam_train(false).await.unwrap();
let model = spam_classifier_model(&handle.server).await;
assert_eq!(model.ham_count, 10);
assert_eq!(model.spam_count, 10);
assert_eq!(model.reservoir.ham.total_seen, 10);
assert_eq!(model.reservoir.spam.total_seen, 10);
assert_eq!(
model.last_sample_expiry,
samples.samples.iter().map(|s| s.until).max().unwrap()
);
assert_eq!(spam_training_samples(&handle.server).await.total_count, 20);
assert!(
handle
.server
.inner
.data
.spam_classifier
.load()
.model
.is_active()
);
assert!(handle.server.inner.data.spam_classifier.load().is_active());
// Send 3 test emails
for message in TEST {
@@ -110,7 +101,7 @@ pub async fn test(handle: &IMAPTest) {
.assert_not_contains("FLAGS ($Junk")
.assert_contains("Subject: classifier test")
.assert_contains("X-Spam-Status: No")
.assert_contains("PROB_HAM_MEDIUM");
.assert_contains("PROB_SPAM_UNCERTAIN");
imap.send_ok("SELECT \"Junk Mail\"").await;
imap.send("FETCH 10 (FLAGS RFC822.TEXT)").await;
imap.assert_read(Type::Tagged, ResponseType::Ok)
@@ -152,18 +143,15 @@ pub struct TrainingSample {
pub until: u64,
}
pub async fn spam_classifier_model(server: &Server) -> SpamClassifierModel {
pub async fn spam_classifier_model(server: &Server) -> SpamTrainer {
server
.store()
.get_value::<Archive<AlignedBytes>>(ValueKey::property(
u32::MAX,
Collection::Principal,
u32::MAX,
PrincipalField::SpamModel,
))
.blob_store()
.get_blob(SPAM_TRAINER_KEY, 0..usize::MAX)
.await
.and_then(|archive| match archive {
Some(archive) => archive.deserialize::<SpamClassifierModel>().map(Some),
Some(archive) => <Archive<AlignedBytes> as Deserialize>::deserialize(&archive)
.and_then(|archive| archive.deserialize_untrusted::<SpamTrainer>())
.map(Some),
None => Ok(None),
})
.unwrap()
@@ -276,7 +264,7 @@ impl ImapConnection {
pub const SPAM: [&str; 10] = [
concat!(
"Subject: save up to NUMBER on life insurance\r\n\r\n wh",
"Subject: save up to = on life insurance\r\n\r\n wh",
"y spend more than you have to life quote savings e",
"nsuring your family s financial security is very i",
"mportant life quote savings makes buying life insu",
@@ -287,8 +275,8 @@ pub const SPAM: [&str; 10] = [
"in the country on new coverage you can save hundre",
"ds or even thousands of dollars by requesting a fr",
"ee quote from lifequote savings our service will t",
"ake you less than NUMBER minutes to complete shop ",
"and compare save up to NUMBER on all types of life",
"ake you less than = minutes to complete shop ",
"and compare save up to = on all types of life",
" insurance hyperlink click here for your free quot",
"e protecting your family is the best investment yo",
"u ll ever make if you are in receipt of this email",
@@ -304,38 +292,38 @@ pub const SPAM: [&str; 10] = [
"ayers are on this one for once be where the player",
"s are this is your private invitation experts are ",
"calling this the fastest way to huge cash flow eve",
"r conceived leverage NUMBER NUMBER into NUMBER NUM",
"r conceived leverage = = into = NUM",
"BER over and over again the question here is you e",
"ither want to be wealthy or you don t which one ar",
"e you i am tossing you a financial lifeline and fo",
"r your sake i hope you grab onto it and hold on ti",
"ght for the ride of your life testimonials hear wh",
"at average people are doing their first few days w",
"e ve received NUMBER NUMBER in NUMBER day and we a",
"e ve received = = in = day and we a",
"re doing that over and over again q s in al i m a ",
"single mother in fl and i ve received NUMBER NUMBE",
"R in the last NUMBER days d s in fl i was not sure",
" about this when i sent off my NUMBER NUMBER pledg",
"e but i got back NUMBER NUMBER the very next day l",
"single mother in fl and i ve received = NUMBE",
"R in the last = days d s in fl i was not sure",
" about this when i sent off my = = pledg",
"e but i got back = = the very next day l",
" l in ky i didn t have the money so i found myself",
" a partner to work this with we have received NUMB",
"ER NUMBER over the last NUMBER days i think i made",
"ER = over the last = days i think i made",
" the right decision don t you k c in fl i pick up ",
"NUMBER NUMBER my first day and i they gave me free",
"= = my first day and i they gave me free",
" leads and all the training you can too j w in ca ",
"announcing we will close your sales for you and he",
"lp you get a fax blast immediately upon your entry",
" you make the money free leads training don t wait",
" call now fax back to NUMBER NUMBER NUMBER NUMBER ",
"or call NUMBER NUMBER NUMBER NUMBER name__________",
" call now fax back to = = = = ",
"or call = = = = name__________",
"________________________phone_____________________",
"______________________ fax________________________",
"_____________email________________________________",
"____________ best time to call____________________",
"_____time zone____________________________________",
"____ this message is sent in compliance of the new",
" e mail bill per section NUMBER paragraph a NUMBER",
" c of s NUMBER further transmissions by the sender",
" e mail bill per section = paragraph a =",
" c of s = further transmissions by the sender",
" of this email may be stopped at no cost to you by",
" sending a reply to this email address with the wo",
"rd remove in the subject line errors omissions and",
@@ -345,70 +333,70 @@ pub const SPAM: [&str; 10] = [
" for the sole purpose of these communications your",
" continued inclusion is only by your gracious perm",
"ission if you wish to not receive this mail from m",
"e please send an email to tesrewinter URL with rem",
"e please send an email to tesrewinter with rem",
"ove in the subject and you will be deleted immedia",
"tely\r\n\r\n"
),
concat!(
"Subject: help wanted \r\n\r\nwe are a NUMBER year old f",
"ortune NUMBER company that is growing at a tremend",
"Subject: help wanted \r\n\r\nwe are a = year old f",
"ortune = company that is growing at a tremend",
"ous rate we are looking for individuals who want t",
"o work from home this is an opportunity to make an",
" excellent income no experience is required we wil",
"l train you so if you are looking to be employed f",
"rom home with a career that has vast opportunities",
" then go URL we are looking for energetic and self",
" then go we are looking for energetic and self",
" motivated people if that is you than click on the",
" link and fill out the form and one of our employe",
"ment specialist will contact you to be removed fro",
"m our link simple go to URL \r\n\r\n"
"m our link simple go to \r\n\r\n"
),
concat!(
"Subject: tired of the bull out there\r\n\r\n want to st",
"op losing money want a real money maker receive NU",
"MBER NUMBER NUMBER NUMBER today experts are callin",
"MBER = = = today experts are callin",
"g this the fastest way to huge cash flow ever conc",
"eived a powerhouse gifting program you don t want ",
"to miss we work as a team this is your private inv",
"itation get in with the founders this is where the",
" big boys play the major players are on this one f",
"or once be where the players are this is a system ",
"that will drive NUMBER NUMBER s to your doorstep i",
"n a short period of time leverage NUMBER NUMBER in",
"to NUMBER NUMBER over and over again the question ",
"that will drive = = s to your doorstep i",
"n a short period of time leverage = = in",
"to = = over and over again the question ",
"here is you either want to be wealthy or you don t",
" which one are you i am tossing you a financial li",
"feline and for your sake i hope you grab onto it a",
"nd hold on tight for the ride of your life testimo",
"nials hear what average people are doing their fir",
"st few days we ve received NUMBER NUMBER in NUMBER",
"st few days we ve received = = in =",
" day and we are doing that over and over again q s",
" in al i m a single mother in fl and i ve received",
" NUMBER NUMBER in the last NUMBER days d s in fl i",
" was not sure about this when i sent off my NUMBER",
" NUMBER pledge but i got back NUMBER NUMBER the ve",
" = = in the last = days d s in fl i",
" was not sure about this when i sent off my =",
" = pledge but i got back = = the ve",
"ry next day l l in ky i didn t have the money so i",
" found myself a partner to work this with we have ",
"received NUMBER NUMBER over the last NUMBER days i",
"received = = over the last = days i",
" think i made the right decision don t you k c in ",
"fl i pick up NUMBER NUMBER my first day and i they",
"fl i pick up = = my first day and i they",
" gave me free leads and all the training you can t",
"oo j w in ca this will be the most important call ",
"you make this year free leads training announcing ",
"we will close your sales for you and help you get ",
"a fax blast immediately upon your entry you make t",
"he money free leads training don t wait call now N",
"UMBER NUMBER NUMBER NUMBER print and fax to NUMBER",
" NUMBER NUMBER NUMBER or send an email requesting ",
"more information to successleads URL please includ",
"e your name and telephone number receive NUMBER NU",
"MBER free leads just for responding a NUMBER NUMBE",
"UMBER = = = print and fax to =",
" = = = or send an email requesting ",
"more information to successleads please includ",
"e your name and telephone number receive = NU",
"MBER free leads just for responding a = NUMBE",
"R value name___________________________________ ph",
"one___________________________________ fax________",
"_____________________________ email_______________",
"____________________ this message is sent in compl",
"iance of the new e mail bill per section NUMBER pa",
"ragraph a NUMBER c of s NUMBER further transmissio",
"iance of the new e mail bill per section = pa",
"ragraph a = c of s = further transmissio",
"ns by the sender of this email may be stopped at n",
"o cost to you by sending a reply to this email add",
"ress with the word remove in the subject line erro",
@@ -419,35 +407,35 @@ pub const SPAM: [&str; 10] = [
"munications your continued inclusion is only by yo",
"ur gracious permission if you wish to not receive ",
"this mail from me please send an email to tesrewin",
"ter URL with remove in the subject and you will be",
"ter with remove in the subject and you will be",
" deleted immediately\r\n\r\n"
),
concat!(
"Subject: cellular phone accessories \r\n\r\n all at bel",
"ow wholesale prices http NUMBER NUMBER NUMBER NUMB",
"ER NUMBER sites merchant sales hands free ear buds",
" NUMBER NUMBER phone holsters NUMBER NUMBER booste",
"r antennas only NUMBER NUMBER phone cases NUMBER N",
"UMBER car chargers NUMBER NUMBER face plates as lo",
"w as NUMBER NUMBER lithium ion batteries as low as",
" NUMBER NUMBER http NUMBER NUMBER NUMBER NUMBER NU",
"ow wholesale prices http = = = NUMB",
"ER = sites merchant sales hands free ear buds",
" = = phone holsters = = booste",
"r antennas only = = phone cases = N",
"UMBER car chargers = = face plates as lo",
"w as = = lithium ion batteries as low as",
" = = http = = = = NU",
"MBER sites merchant sales click below for accessor",
"ies on all nokia motorola lg nextel samsung qualco",
"mm ericsson audiovox phones at below wholesale pri",
"ces http NUMBER NUMBER NUMBER NUMBER NUMBER sites ",
"ces http = = = = = sites ",
"merchant sales if you need assistance please call ",
"us NUMBER NUMBER NUMBER to be removed from future ",
"us = = = to be removed from future ",
"mailings please send your remove request to remove",
" me now NUMBER URL thank you and have a super day\r\n",
" me now = thank you and have a super day\r\n",
" \r\n"
),
concat!(
"Subject: conferencing made easy\r\n\r\n only NUMBER cen",
"Subject: conferencing made easy\r\n\r\n only = cen",
"ts per minute including long distance no setup fee",
"s no contracts or monthly fees call anytime from a",
"nywhere to anywhere connects up to NUMBER particip",
"nywhere to anywhere connects up to = particip",
"ants simplicity in set up and administration opera",
"tor help available NUMBER NUMBER the highest quali",
"tor help available = = the highest quali",
"ty service for the lowest rate in the industry fil",
"l out the form below to find out how you can lower",
" your phone bill every month required input field ",
@@ -483,8 +471,8 @@ pub const SPAM: [&str; 10] = [
"ated by the french government and as such i had to",
" change my identity so that my investment will not",
" be traced and confiscated i have deposited the su",
"m eighteen million united state dollars us NUMBER ",
"NUMBER NUMBER NUMBER with a security company for s",
"m eighteen million united state dollars us = ",
"= = = with a security company for s",
"afekeeping the funds are security coded to prevent",
" them from knowing the content what i want you to ",
"do is to indicate your interest that you will assi",
@@ -506,35 +494,35 @@ pub const SPAM: [&str; 10] = [
" remunerationfor your services for this reason kin",
"dly furnish us your contact information that is yo",
"ur personal telephone and fax number for confident",
"ial URL regards mrs m sese seko\r\n\r\n"
"ial regards mrs m sese seko\r\n\r\n"
),
concat!(
"Subject: lowest rates available for term life insu",
"rance\r\n\r\n take a moment and fill out our online for",
"m to see the low rate you qualify for save up to N",
"UMBER from regular rates smokers accepted URL repr",
"UMBER from regular rates smokers accepted repr",
"esenting quality nationwide carriers act now to ea",
"sily remove your address from the list go to URL p",
"lease allow NUMBER NUMBER hours for removal\r\n\r\n"
"sily remove your address from the list go to p",
"lease allow = = hours for removal\r\n\r\n"
),
concat!(
"Subject: central bank of nigeria foreign remittanc",
"e \r\n\r\n dept tinubu square lagos nigeria email smith",
"_j URL NUMBERth of august NUMBER attn president ce",
"_j =th of august = attn president ce",
"o strictly private business proposal i am mr johns",
"on s abu the bills and exchange director at the fo",
"reignremittance department of the central bank of ",
"nigeria i am writingyou this letter to ask for you",
"r support and cooperation to carrying thisbusiness",
" opportunity in my department we discovered abando",
"ned the sumof us NUMBER NUMBER NUMBER NUMBER thirt",
"ned the sumof us = = = = thirt",
"y seven million four hundred thousand unitedstates",
" dollars in an account that belong to one of our f",
"oreign customers an american late engr john creek ",
"junior an oil merchant with the federal government",
" of nigeria who died along with his entire family ",
"of a wifeand two children in kenya airbus aNUMBER ",
"NUMBER flight kqNUMBER in novemberNUMBER since we ",
"of a wifeand two children in kenya airbus a= ",
"= flight kq= in november= since we ",
"heard of his death we have been expecting his next",
" of kin tocome over and put claims for his money a",
"s the heir because we cannotrelease the fund from ",
@@ -550,11 +538,11 @@ pub const SPAM: [&str; 10] = [
"r bank other wisethe fund will be returned to the ",
"bank treasury as unclaimed fund we have agreed tha",
"t our ratio of sharing will be as stated thus NUMB",
"ER for you as foreign partner and NUMBER for us th",
"ER for you as foreign partner and = for us th",
"e officials in my department upon the successful c",
"ompletion of this transfer my colleague and i will",
"come to your country and mind our share it is from",
" our NUMBER we intendto import computer accessorie",
" our = we intendto import computer accessorie",
"s into my country as way of recycling thefund to c",
"ommence this transaction we require you to immedia",
"tely indicateyour interest by calling me or sendin",
@@ -567,23 +555,23 @@ pub const SPAM: [&str; 10] = [
"t be kept strictly confidential becauseof its natu",
"re nb please remember to give me your phone and fa",
"x no mr johnson smith abu irish linux users group ",
"ilug URL URL for un subscription information list ",
"maintainer listmaster URL\r\n\r\n"
"ilug for un subscription information list ",
"maintainer listmaster \r\n\r\n"
),
concat!(
"Subject: dear stuart\r\n\r\n are you tired of searching",
" for love in all the wrong places find love now at",
" URL URL browse through thousands of personals in ",
"your area join for free URL search e mail chat use",
" URL to meet cool guys and hot girls go NUMBER on ",
"NUMBER or use our private chat rooms click on the ",
"link to get started URL find love now you have rec",
" browse through thousands of personals in ",
"your area join for free search e mail chat use",
" to meet cool guys and hot girls go = on ",
"= or use our private chat rooms click on the ",
"link to get started find love now you have rec",
"eived this email because you have registerd with e",
"mailrewardz or subscribed through one of our marke",
"ting partners if you have received this message in",
" error or wish to stop receiving these great offer",
"s please click the remove link above to unsubscrib",
"e from these mailings please click here URL\r\n\r\n"
"e from these mailings please click here \r\n\r\n"
),
];
@@ -600,20 +588,20 @@ pub const HAM: [&str; 10] = [
"ering i would prefer not to have to write a script",
" myself but will appreciate any suggestions this U",
"RL email is sponsored by osdn tired of that same o",
"ld cell phone get a new here for free URL ________",
"ld cell phone get a new here for free ________",
"_______________________________________ spamassass",
"in talk mailing list spamassassin talk URL URL\r\n\r\n"
"in talk mailing list spamassassin talk \r\n\r\n"
),
concat!(
"Message-ID: mid2@foobar.org\r\nSubject: hello\r\n\r\nhave y",
"ou seen and discussed this article and his approac",
"h thank you URL hell there are no rules here we re",
"h thank you hell there are no rules here we re",
" trying to accomplish something thomas alva edison",
" this URL email is sponsored by osdn tired of that",
" same old cell phone get a new here for free URL _",
" this email is sponsored by osdn tired of that",
" same old cell phone get a new here for free _",
"______________________________________________ spa",
"massassin devel mailing list spamassassin devel UR",
"L URL \r\n\r\n"
"L \r\n\r\n"
),
concat!(
"Message-ID: <mid3@foobar.org>\r\nSubject: hi all apol",
@@ -625,9 +613,9 @@ pub const HAM: [&str; 10] = [
"internet wild i e machines with static real ips an",
"y help pointers would be helpful cheers rgrds bern",
"ard bernard tyers national centre for sensor resea",
"rch p NUMBER NUMBER NUMBER NUMBER e bernard tyers ",
"URL w URL l nNUMBER ______________________________",
"_________________ iiu mailing list iiu URL URL \r\n\r\n"
"rch p = = = = e bernard tyers ",
" w l n= ______________________________",
"_________________ iiu mailing list iiu \r\n\r\n"
),
concat!(
"Message-ID: <mid4@foobar.org>\r\nSubject: can someone",
@@ -638,8 +626,8 @@ pub const HAM: [&str; 10] = [
"ne for that but im not sure if solaris is a distro",
" of linux or a completely different operating syst",
"em can someone explain kiall mac innes irish linux",
" users group ilug URL URL for un subscription info",
"rmation list maintainer listmaster URL \r\n\r\n"
" users group ilug for un subscription info",
"rmation list maintainer listmaster \r\n\r\n"
),
concat!(
"Message-ID: <mid5@foobar.org>\r\nSubject: folks my fi",
@@ -647,13 +635,13 @@ pub const HAM: [&str; 10] = [
"t am new to linux just got a new pc at home dell b",
"ox with windows xp added a second hard disk for li",
"nux partitioned the disk and have installed suse N",
"UMBER NUMBER from cd which went fine except it did",
"UMBER = from cd which went fine except it did",
"n t pick up my monitor i have a dell branded eNUMB",
"ERfpp NUMBER lcd flat panel monitor and a nvidia g",
"eforceNUMBER tiNUMBER video card both of which are",
"ERfpp = lcd flat panel monitor and a nvidia g",
"eforce= ti= video card both of which are",
" probably too new to feature in suse s default set",
" i downloaded a driver from the nvidia website and",
" installed it using rpm then i ran saxNUMBER as wa",
" installed it using rpm then i ran sax= as wa",
"s recommended in some postings i found on the net ",
"but it still doesn t feature my video card in the ",
"available list what next another problem i have a ",
@@ -665,9 +653,9 @@ pub const HAM: [&str; 10] = [
"ful i ve searched the net but have run out of idea",
"s or should i be going for a different version of ",
"linux such as redhat opinions welcome thanks a lot",
" peter irish linux users group ilug URL URL for un",
" peter irish linux users group ilug for un",
" subscription information list maintainer listmast",
"er URL\r\n\r\n"
"er \r\n\r\n"
),
concat!(
"Message-ID: <mid6@foobar.org>\r\nSubject: has anyone\r\n",
@@ -675,23 +663,23 @@ pub const HAM: [&str; 10] = [
"random person go to a webpage create a mailing lis",
"t then administer that list also of course let ppl",
" sign up for the lists and manage their subscripti",
"ons similar to the old URL but i d like to have it",
" running on my server not someone elses chris URL ",
"ons similar to the old but i d like to have it",
" running on my server not someone elses chris ",
"\r\n\r\n"
),
concat!(
"Message-ID: <mid7@foobar.org>\r\nSubject: hi thank yo",
"u for the useful replies\r\n\r\ni have found some intere",
"sting tutorials in the ibm developer connection UR",
"L and URL registration is needed i will post the s",
"L and registration is needed i will post the s",
"ame message on the web application security list a",
"s suggested by someone for now i thing i will use ",
"mdNUMBER for password checking i will use the appr",
"md= for password checking i will use the appr",
"oach described in secure programmin fo linux and u",
"nix how to i will separate the authentication modu",
"le so i can change its implementation at anytime t",
"hank you again mario torre please avoid sending me",
" word or powerpoint attachments see URL \r\n\r\n"
" word or powerpoint attachments see \r\n\r\n"
),
concat!(
"Message-ID: <mid8@foobar.org>\r\nSubject: hehe sorry\r\n",
@@ -701,21 +689,21 @@ pub const HAM: [&str; 10] = [
"edhat dell provide some computers pre loaded with ",
"red hat i dont know for sure tho so get someone el",
"ses opnion as well as mine original message from i",
"lug admin URL mailto ilug admin URL on behalf of p",
"eter staunton sent NUMBER august NUMBER NUMBER NUM",
"BER to ilug URL subject ilug newbie seeks advice s",
"use NUMBER NUMBER folks my first time posting have",
"lug admin mailto ilug admin on behalf of p",
"eter staunton sent = august = = NUM",
"BER to ilug subject ilug newbie seeks advice s",
"use = = folks my first time posting have",
" a bit of unix experience but am new to linux just",
" got a new pc at home dell box with windows xp add",
"ed a second hard disk for linux partitioned the di",
"sk and have installed suse NUMBER NUMBER from cd w",
"sk and have installed suse = = from cd w",
"hich went fine except it didn t pick up my monitor",
" i have a dell branded eNUMBERfpp NUMBER lcd flat ",
"panel monitor and a nvidia geforceNUMBER tiNUMBER ",
" i have a dell branded e=fpp = lcd flat ",
"panel monitor and a nvidia geforce= ti= ",
"video card both of which are probably too new to f",
"eature in suse s default set i downloaded a driver",
" from the nvidia website and installed it using rp",
"m then i ran saxNUMBER as was recommended in some ",
"m then i ran sax= as was recommended in some ",
"postings i found on the net but it still doesn t f",
"eature my video card in the available list what ne",
"xt another problem i have a dell branded keyboard ",
@@ -727,32 +715,32 @@ pub const HAM: [&str; 10] = [
"net but have run out of ideas or should i be going",
" for a different version of linux such as redhat o",
"pinions welcome thanks a lot peter irish linux use",
"rs group ilug URL URL for un subscription informat",
"ion list maintainer listmaster URL irish linux use",
"rs group ilug URL URL for un subscription informat",
"ion list maintainer listmaster URL\r\n\r\n"
"rs group ilug for un subscription informat",
"ion list maintainer listmaster irish linux use",
"rs group ilug for un subscription informat",
"ion list maintainer listmaster \r\n\r\n"
),
concat!(
"Message-ID: <mid9@foobar.org>\r\nSubject: it will fun",
"ction as a router\r\n\r\nif that is what you wish it eve",
"n looks like the modem s embedded os is some kind ",
"of linux being that it has interesting interfaces ",
"like ethNUMBER i don t use it as a router though i",
"like eth= i don t use it as a router though i",
" just have it do the absolute minimum dsl stuff an",
"d do all the really fun stuff like pppoe on my lin",
"ux box also the manual tells you what the default ",
"password is don t forget to run pppoe over the alc",
"atel speedtouch NUMBERi as in my case you have to ",
"atel speedtouch =i as in my case you have to ",
"have a bridge configured in the router modem s sof",
"tware this lists your vci values etc also does any",
"one know if the high end speedtouch with NUMBER et",
"one know if the high end speedtouch with = et",
"hernet ports can act as a full router or do i stil",
"l need to run a pppoe stack on the linux box regar",
"ds vin irish linux users group ilug URL URL for un",
"ds vin irish linux users group ilug for un",
" subscription information list maintainer listmast",
"er URL irish linux users group ilug URL URL for un",
"er irish linux users group ilug for un",
" subscription information list maintainer listmast",
"er URL \r\n\r\n"
"er \r\n\r\n"
),
concat!(
"Message-ID: <mid10@foobar.org>\r\nSubject: all is it ",
@@ -764,23 +752,23 @@ pub const HAM: [&str; 10] = [
"sia and elsewhere coupled with the false emails i ",
"received myself it s really starting to annoy me a",
"m i the only one seeing an increase in recent week",
"s martin martin whelan déise design URL tel NUMBE",
"R NUMBER our core product déiseditor allows organ",
"s martin martin whelan déise design tel NUMBE",
"R = our core product déiseditor allows organ",
"isations to publish information to their web site ",
"in a fast and cost effective manner there is no ne",
"ed for a full time web developer as the site can b",
"e easily updated by the organisations own staff in",
"stant updates to keep site information fresh sites",
" which are updated regularly bring users back visi",
"t URL for a demonstration déiseditor managing you",
"t for a demonstration déiseditor managing you",
"r information ____________________________________",
"___________ iiu mailing list iiu URL URL ,0\r\n"
"___________ iiu mailing list iiu ,0\r\n"
),
];
const TEST: [&str; 3] = [
concat!(
"Subject: save up to NUMBER on life insurance\r\n\r\nwhy ",
"Subject: save up to = on life insurance\r\n\r\nwhy ",
"spend more than you have to life quote savings ens",
"uring your family s financial security is very imp",
"ortant life quote savings makes buying life insura",
@@ -791,8 +779,8 @@ const TEST: [&str; 3] = [
" the country on new coverage you can save hundreds",
" or even thousands of dollars by requesting a free",
" quote from lifequote savings our service will tak",
"e you less than NUMBER minutes to complete shop an",
"d compare save up to NUMBER on all types of life i",
"e you less than = minutes to complete shop an",
"d compare save up to = on all types of life i",
"nsurance hyperlink click here for your free quote ",
"protecting your family is the best investment you ",
"ll ever make if you are in receipt of this email i",
@@ -809,9 +797,9 @@ const TEST: [&str; 3] = [
"un seems to be the one for that but im not sure if",
" solaris is a distro of linux or a completely diff",
"erent operating system can someone explain kiall m",
"ac innes irish linux users group ilug URL URL for ",
"ac innes irish linux users group ilug for ",
"un subscription information list maintainer listma",
"ster URL \r\n"
"ster \r\n"
),
concat!(
"Subject: classifier test\r\n\r\nthis is a novel text tha",

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

@@ -667,7 +667,7 @@ wait = "1ms"
enable = true
[spam-filter.list]
scores = {"PROB_SPAM_HIGH" = "10.0"}
scores = {"PROB_SPAM_LOW" = "10.0", "PROB_SPAM_HIGH" = "10.0"}
[lookup]
"spam-traps" = {"spamtrap@*"}

6
tests/src/smtp/inbound/antispam.rs
View File

@@ -331,8 +331,8 @@ async fn antispam() {
"dmarc",
"rbl",
"spamtrap",
//"classifier_html",
//"classifier_features",
"classifier_html",
"classifier_features",
"classifier",
"pyzor",
"llm",
@@ -372,7 +372,7 @@ async fn antispam() {
.put_temporary_blob(u32::MAX, sample.as_bytes(), 60)
.await
.unwrap();
server.add_spam_sample(&mut batch, hash, class == "spam", false, 0);
server.add_spam_sample(&mut batch, hash, class == "spam", true, 0);
batch.clear(blob_hold);
}
}