Port Spam filter to Rust - all tests passing

This commit is contained in:
mdecimus
2024-12-20 16:59:26 +01:00
parent e7298df610
commit 6416ff12d6
52 changed files with 755 additions and 705 deletions

View File

@@ -1,16 +1,20 @@
use std::{
borrow::Cow,
fs,
path::PathBuf,
sync::Arc,
time::{Duration, Instant},
};
use ahash::AHashSet;
use ahash::{AHashMap, AHashSet};
use common::{
auth::AccessToken,
enterprise::llm::{
AiApiConfig, ChatCompletionChoice, ChatCompletionRequest, ChatCompletionResponse, Message,
config::spamfilter::SpamFilterAction,
enterprise::{
llm::{
AiApiConfig, ChatCompletionChoice, ChatCompletionRequest, ChatCompletionResponse,
Message,
},
SpamFilterLlmConfig,
},
Core,
};
@@ -21,7 +25,6 @@ use mail_auth::{
IprevResult, SpfOutput, SpfResult, MX,
};
use mail_parser::MessageParser;
use sieve::runtime::Variable;
use smtp::core::{Session, SessionAddress};
use smtp_proto::{MAIL_BODY_8BITMIME, MAIL_SMTPUTF8};
use spam_filter::{
@@ -182,10 +185,12 @@ async fn antispam() {
let mut core = Core::parse(&mut config, stores, Default::default())
.await
.enable_enterprise();
core.enterprise.as_mut().unwrap().ai_apis.insert(
let ai_apis = AHashMap::from_iter([(
"dummy".to_string(),
AiApiConfig::parse(&mut config, "dummy").unwrap().into(),
);
)]);
core.enterprise.as_mut().unwrap().spam_filter_llm =
SpamFilterLlmConfig::parse(&mut config, &ai_apis);
crate::AssertConfig::assert_no_errors(config);
// Add mock DNS entries
@@ -279,7 +284,8 @@ async fn antispam() {
.join("smtp")
.join("antispam");
for test_name in [
/*"ip",
"combined",
"ip",
"helo",
"received",
"messageid",
@@ -298,11 +304,10 @@ async fn antispam() {
"replies_out",
"replies_in",
"spamtrap",
"bayes_classify",*/
"bayes_classify",
"reputation",
"pyzor",
"llm",
"combined",
] {
/*if test_name != "combined" {
continue;
@@ -324,7 +329,7 @@ async fn antispam() {
let mut dmarc_result = None;
let mut dmarc_policy = None;
let mut expected_tags = AHashSet::new();
let mut score_expect = 0.0;
let mut expect_headers = String::new();
let mut score_set = 0.0;
let mut score_final = 0.0;
let mut body_params = 0;
@@ -420,8 +425,14 @@ async fn antispam() {
expected_tags
.extend(value.split_ascii_whitespace().map(|v| v.to_uppercase()));
}
"expect_score" => {
score_expect = value.parse::<f64>().unwrap();
"expect_header" => {
let value = value.trim();
if !value.is_empty() {
if !expect_headers.is_empty() {
expect_headers.push(' ');
}
expect_headers.push_str(value);
}
}
"score" => {
score_set = value.parse::<f64>().unwrap();
@@ -478,6 +489,39 @@ async fn antispam() {
}
}
let parsed_message = MessageParser::new().parse(&message).unwrap();
// Combined tests
if test_name == "combined" {
match session
.spam_classify(
&parsed_message,
&dkim_domains,
arc_result.as_ref(),
dmarc_result.as_ref(),
dmarc_policy.as_ref(),
)
.await
{
SpamFilterAction::Allow(header) => {
let mut last_ch = 'x';
let mut result = String::with_capacity(header.len());
for ch in header.chars() {
if !ch.is_whitespace() {
if last_ch.is_whitespace() {
result.push(' ');
}
result.push(ch);
}
last_ch = ch;
}
assert_eq!(result, expect_headers);
}
other => panic!("Unexpected action {other:?}"),
}
continue;
}
// Initialize filter
let mut spam_input = session.build_spam_input(
&parsed_message,
&dkim_domains,
@@ -486,8 +530,6 @@ async fn antispam() {
dmarc_policy.as_ref(),
);
spam_input.is_tls = is_tls;
// Initialize filter
let mut spam_ctx = server.spam_filter_init(spam_input);
match test_name {
"html" => {
@@ -564,9 +606,7 @@ async fn antispam() {
}
"spamtrap" => {
server.spam_filter_analyze_spam_trap(&mut spam_ctx).await;
server
.spam_filter_finalize(&mut spam_ctx, String::new())
.await;
server.spam_filter_finalize(&mut spam_ctx).await;
}
"bayes_classify" => {
server
@@ -584,9 +624,6 @@ async fn antispam() {
"llm" => {
server.spam_filter_analyze_llm(&mut spam_ctx).await;
}
"combined" => {
todo!("combined");
}
_ => panic!("Invalid test {test_name:?}"),
}

View File

@@ -20,14 +20,14 @@ use smtp::core::{Session, State};
const CONFIG: &str = r#"
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
directory = "local"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/queue.db"
[directory."local"]

View File

@@ -20,16 +20,19 @@ use smtp::core::Session;
const CONFIG: &str = r#"
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
directory = "local"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/queue.db"
[spam-filter]
enable = false
[directory."local"]
type = "memory"
@@ -125,12 +128,7 @@ async fn data() {
// Send broken message
session
.send_message(
"john@doe.org",
&["bill@foobar.org"],
"From: john",
"550 5.7.7",
)
.send_message("john@doe.org", &["bill@foobar.org"], "invalid", "550 5.7.7")
.await;
// Naive Loop detection

View File

@@ -27,13 +27,13 @@ use smtp::core::Session;
const CONFIG: &str = r#"
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/queue.db"
[directory."local"]

View File

@@ -21,13 +21,13 @@ use crate::smtp::{
const CONFIG: &str = r#"
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/data.db"
[session.ehlo]

View File

@@ -51,13 +51,13 @@ struct HeaderTest {
const CONFIG_MILTER: &str = r#"
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/queue.db"
[session.rcpt]
@@ -77,13 +77,13 @@ stages = ["data"]
const CONFIG_JMILTER: &str = r#"
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/queue.db"
[session.rcpt]

View File

@@ -21,13 +21,13 @@ use crate::smtp::{
const CONFIG: &str = r#"
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/queue.db"
[directory."local"]

View File

@@ -42,11 +42,8 @@ max-connections = 10
min-connections = 0
idle-timeout = "5m"
[session.data.pipe."test"]
command = [ { if = "remote_ip = '10.0.0.123'", then = "'/bin/bash'" },
{ else = false } ]
arguments = "['{CFG_PATH}/pipe_me.sh', 'hello', 'world']"
timeout = "10s"
[spam-filter]
enable = false
[sieve.trusted]
from-name = "'Sieve Daemon'"
@@ -99,7 +96,6 @@ email = ["john@localdomain.org", "jdoe@localdomain.org", "john.doe@localdomain.o
email-list = ["info@localdomain.org"]
member-of = ["sales"]
"#;
#[tokio::test]
@@ -134,23 +130,8 @@ async fn sieve_scripts() {
}
// Prepare config
let tmp_dir = TempDir::new("smtp_sieve_test", true);
let mut config = Config::new(
tmp_dir.update_config(
config.replace(
"{CFG_PATH}",
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("resources")
.join("smtp")
.join("pipe")
.as_path()
.to_str()
.unwrap(),
),
),
)
.unwrap();
let mut config = Config::new(tmp_dir.update_config(config)).unwrap();
config.resolve_all_macros().await;
let stores = Stores::parse_all(&mut config).await;
let core = Core::parse(&mut config, stores, Default::default()).await;
@@ -392,24 +373,4 @@ async fn sieve_scripts() {
.assert_contains("Received: ")
.assert_contains("Authentication-Results: ");
qr.assert_no_events();
// Test pipes
session.data.remote_ip_str = "10.0.0.123".parse().unwrap();
session.data.remote_ip = session.data.remote_ip_str.parse().unwrap();
session
.send_message(
"test@example.net",
&["pipe@foobar.com"],
"test:no_dkim",
"250",
)
.await;
qr.expect_message()
.await
.read_lines(&qr)
.await
.assert_contains("X-My-Header: true")
.assert_contains("Authentication-Results");
qr.assert_no_events();
}

View File

@@ -75,13 +75,13 @@ set-body-length = false
const CONFIG: &str = r#"
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/queue.db"
[directory."local"]

View File

@@ -14,13 +14,13 @@ use utils::config::Config;
const CONFIG: &str = r#"
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/data.db"
[[session.throttle]]

View File

@@ -21,14 +21,14 @@ use crate::{
const CONFIG: &str = r#"
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
directory = "local"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/data.db"
[directory."local"]

View File

@@ -48,7 +48,6 @@ emails = "SELECT address FROM emails WHERE name = ? AND type != 'list' ORDER BY
verify = "SELECT address FROM emails WHERE address LIKE '%' || ? || '%' AND type = 'primary' ORDER BY address LIMIT 5"
expand = "SELECT p.address FROM emails AS p JOIN emails AS l ON p.name = l.name WHERE p.type = 'primary' AND l.address = ? AND l.type = 'list' ORDER BY p.address LIMIT 50"
domains = "SELECT 1 FROM emails WHERE address LIKE '%@' || ? LIMIT 1"
is_ip_allowed = "SELECT addr FROM allowed_ips WHERE addr = ? LIMIT 1"
[directory."sql"]
type = "sql"
@@ -73,7 +72,7 @@ relay = false
errors.wait = "5ms"
[session.extensions]
requiretls = [{if = "key_exists('sql/is_ip_allowed', remote_ip)", then = true},
requiretls = [{if = "sql_query('sql', 'SELECT addr FROM allowed_ips WHERE addr = ? LIMIT 1', remote_ip)", then = true},
{else = false}]
expn = true
vrfy = true
@@ -140,18 +139,6 @@ async fn lookup_sql() {
handle
.create_test_user_with_email("mike@foobar.net", "098765", "Mike")
.await;
/*handle
.link_test_address("jane@foobar.org", "sales@foobar.org", "list")
.await;
handle
.link_test_address("john@foobar.org", "sales@foobar.org", "list")
.await;
handle
.link_test_address("bill@foobar.org", "sales@foobar.org", "list")
.await;
handle
.link_test_address("mike@foobar.net", "support@foobar.org", "list")
.await;*/
for query in [
"CREATE TABLE domains (name TEXT PRIMARY KEY, description TEXT);",

View File

@@ -25,7 +25,7 @@ use crate::AssertConfig;
pub mod config;
pub mod inbound;
pub mod lookup;
pub mod management;
//pub mod management;
pub mod outbound;
pub mod queue;
pub mod reporting;
@@ -128,13 +128,13 @@ cert = '%{file:{CERT}}%'
private-key = '%{file:{PK}}%'
[storage]
data = "sqlite"
lookup = "sqlite"
blob = "sqlite"
fts = "sqlite"
data = "rocksdb"
lookup = "rocksdb"
blob = "rocksdb"
fts = "rocksdb"
[store."sqlite"]
type = "sqlite"
[store."rocksdb"]
type = "rocksdb"
path = "{TMP}/queue.db"
"#;