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

AI models

Browse Source
This commit is contained in:
mdecimus
2024-10-05 19:05:04 +02:00
parent d5c2dcb817
commit d0ce2b1a96
41 changed files with 1171 additions and 406 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tests/resources/jmap/sieve/test_mailbox.sieve
View File

@@ -1,4 +1,4 @@
require ["fileinto", "mailbox", "mailboxid", "special-use", "ihave", "imap4flags"];
require ["fileinto", "mailbox", "mailboxid", "special-use", "ihave", "imap4flags", "vnd.stalwart.expressions"];
# SpecialUse extension tests
if not specialuse_exists ["inbox", "trash"] {
@@ -61,3 +61,6 @@ if not mailboxexists "My" {
error "'My' not found.";
}
if eval "llm_prompt('echo-test', 'hello world') != 'hello world'" {
error "llm_prompt is unavailable.";
}

72
tests/resources/smtp/antispam/llm.test Normal file
View File

@@ -0,0 +1,72 @@
expect LLM_UNSOLICITED_HIGH
Subject: Unsolicited,High,Test
Test
<!-- NEXT TEST -->
expect LLM_COMMERCIAL_HIGH
Subject: Commercial,High,Test
Test
<!-- NEXT TEST -->
expect LLM_HARMFUL_HIGH
Subject: Harmful,High,Test
Test
<!-- NEXT TEST -->
expect LLM_LEGITIMATE_HIGH
Subject: Legitimate,High,Test
Test
<!-- NEXT TEST -->
expect LLM_UNSOLICITED_MEDIUM
Subject: Unsolicited,Medium,Test
Test
<!-- NEXT TEST -->
expect LLM_COMMERCIAL_MEDIUM
Subject: Commercial,Medium,Test
Test
<!-- NEXT TEST -->
expect LLM_HARMFUL_MEDIUM
Subject: Harmful,Medium,Test
Test
<!-- NEXT TEST -->
expect LLM_LEGITIMATE_MEDIUM
Subject: Legitimate,Medium,Test
Test
<!-- NEXT TEST -->
expect LLM_UNSOLICITED_LOW
Subject: Unsolicited,Low,Test
Test
<!-- NEXT TEST -->
expect LLM_COMMERCIAL_LOW
Subject: Commercial,Low,Test
Test
<!-- NEXT TEST -->
expect LLM_HARMFUL_LOW
Subject: Harmful,Low,Test
Test
<!-- NEXT TEST -->
expect LLM_LEGITIMATE_LOW
Subject: Legitimate,Low,Test
Test

10
tests/resources/smtp/lists/public-suffix.dat
View File

@@ -1,10 +0,0 @@
co.uk
org.uk
com
net
org
info
biz
*.wildcard
!test.wildcard
disposable.org

135
tests/src/directory/oidc.rs
View File

@@ -10,23 +10,18 @@
use std::sync::Arc;
use ahash::AHashMap;
use base64::{engine::general_purpose, Engine};
use common::{config::server::Listeners, listener::SessionData, Core, Data, Inner};
use directory::{backend::internal::PrincipalField, QueryBy};
use hyper::{body, server::conn::http1, service::service_fn, Method, StatusCode, Uri};
use hyper_util::rt::TokioIo;
use jmap::api::{
http::{fetch_body, ToHttpResponse},
HttpResponse, JsonResponse,
};
use hyper::{Method, StatusCode};
use jmap::api::{http::ToHttpResponse, JsonResponse};
use mail_send::Credentials;
use serde_json::json;
use tokio::sync::watch;
use trc::{AuthEvent, EventType};
use utils::config::Config;
use crate::{add_test_certs, directory::DirectoryTest, AssertConfig};
use crate::{
directory::DirectoryTest,
http_server::{spawn_mock_http_server, HttpMessage},
};
static TEST_TOKEN: &str = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ";
@@ -143,121 +138,3 @@ async fn oidc_directory() {
assert_eq!(principal.description(), Some("John Doe"));
}
}
const MOCK_HTTP_SERVER: &str = r#"
[server]
hostname = "'oidc.example.org'"
http.url = "'https://127.0.0.1:9090'"
[server.listener.jmap]
bind = ['127.0.0.1:9090']
protocol = 'http'
tls.implicit = true
[server.socket]
reuse-addr = true
[certificate.default]
cert = '%{file:{CERT}}%'
private-key = '%{file:{PK}}%'
default = true
"#;
#[derive(Clone)]
pub struct HttpSessionManager {
inner: HttpRequestHandler,
}
pub type HttpRequestHandler = Arc<dyn Fn(HttpMessage) -> HttpResponse + Sync + Send>;
#[derive(Debug)]
pub struct HttpMessage {
method: Method,
headers: AHashMap<String, String>,
uri: Uri,
body: Option<Vec<u8>>,
}
impl HttpMessage {
pub fn get_url_encoded(&self, key: &str) -> Option<String> {
form_urlencoded::parse(self.body.as_ref()?.as_slice())
.find(|(k, _)| k == key)
.map(|(_, v)| v.into_owned())
}
}
pub async fn spawn_mock_http_server(
handler: HttpRequestHandler,
) -> (watch::Sender<bool>, watch::Receiver<bool>) {
// Start mock push server
let mut settings = Config::new(add_test_certs(MOCK_HTTP_SERVER)).unwrap();
settings.resolve_all_macros().await;
let mock_inner = Arc::new(Inner {
shared_core: Core::parse(&mut settings, Default::default(), Default::default())
.await
.into_shared(),
data: Data::parse(&mut settings),
..Default::default()
});
settings.errors.clear();
settings.warnings.clear();
let mut servers = Listeners::parse(&mut settings);
servers.parse_tcp_acceptors(&mut settings, mock_inner.clone());
// Start JMAP server
servers.bind_and_drop_priv(&mut settings);
settings.assert_no_errors();
servers.spawn(|server, acceptor, shutdown_rx| {
server.spawn(
HttpSessionManager {
inner: handler.clone(),
},
mock_inner.clone(),
acceptor,
shutdown_rx,
);
})
}
impl common::listener::SessionManager for HttpSessionManager {
#[allow(clippy::manual_async_fn)]
fn handle<T: common::listener::SessionStream>(
self,
session: SessionData<T>,
) -> impl std::future::Future<Output = ()> + Send {
async move {
let sender = self.inner;
let _ = http1::Builder::new()
.keep_alive(false)
.serve_connection(
TokioIo::new(session.stream),
service_fn(|mut req: hyper::Request<body::Incoming>| {
let sender = sender.clone();
async move {
let response = sender(HttpMessage {
method: req.method().clone(),
uri: req.uri().clone(),
headers: req
.headers()
.iter()
.map(|(k, v)| {
(k.as_str().to_lowercase(), v.to_str().unwrap().to_string())
})
.collect(),
body: fetch_body(&mut req, 1024 * 1024, 0).await,
});
Ok::<_, hyper::Error>(response.build())
}
}),
)
.await;
}
}
#[allow(clippy::manual_async_fn)]
fn shutdown(&self) -> impl std::future::Future<Output = ()> + Send {
async {}
}
}

135
tests/src/http_server.rs Normal file
View File

@@ -0,0 +1,135 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use std::sync::Arc;
use ahash::AHashMap;
use common::{config::server::Listeners, listener::SessionData, Core, Data, Inner};
use hyper::{body, server::conn::http1, service::service_fn, Method, Uri};
use hyper_util::rt::TokioIo;
use jmap::api::{http::fetch_body, HttpResponse};
use tokio::sync::watch;
use utils::config::Config;
use crate::{add_test_certs, AssertConfig};
const MOCK_HTTP_SERVER: &str = r#"
[server]
hostname = "'oidc.example.org'"
http.url = "'https://127.0.0.1:9090'"
[server.listener.jmap]
bind = ['127.0.0.1:9090']
protocol = 'http'
tls.implicit = true
[server.socket]
reuse-addr = true
[certificate.default]
cert = '%{file:{CERT}}%'
private-key = '%{file:{PK}}%'
default = true
"#;
#[derive(Clone)]
pub struct HttpSessionManager {
inner: HttpRequestHandler,
}
pub type HttpRequestHandler = Arc<dyn Fn(HttpMessage) -> HttpResponse + Sync + Send>;
#[derive(Debug)]
pub struct HttpMessage {
pub method: Method,
pub headers: AHashMap<String, String>,
pub uri: Uri,
pub body: Option<Vec<u8>>,
}
impl HttpMessage {
pub fn get_url_encoded(&self, key: &str) -> Option<String> {
form_urlencoded::parse(self.body.as_ref()?.as_slice())
.find(|(k, _)| k == key)
.map(|(_, v)| v.into_owned())
}
}
pub async fn spawn_mock_http_server(
handler: HttpRequestHandler,
) -> (watch::Sender<bool>, watch::Receiver<bool>) {
// Start mock push server
let mut settings = Config::new(add_test_certs(MOCK_HTTP_SERVER)).unwrap();
settings.resolve_all_macros().await;
let mock_inner = Arc::new(Inner {
shared_core: Core::parse(&mut settings, Default::default(), Default::default())
.await
.into_shared(),
data: Data::parse(&mut settings),
..Default::default()
});
settings.errors.clear();
settings.warnings.clear();
let mut servers = Listeners::parse(&mut settings);
servers.parse_tcp_acceptors(&mut settings, mock_inner.clone());
// Start JMAP server
servers.bind_and_drop_priv(&mut settings);
settings.assert_no_errors();
servers.spawn(|server, acceptor, shutdown_rx| {
server.spawn(
HttpSessionManager {
inner: handler.clone(),
},
mock_inner.clone(),
acceptor,
shutdown_rx,
);
})
}
impl common::listener::SessionManager for HttpSessionManager {
#[allow(clippy::manual_async_fn)]
fn handle<T: common::listener::SessionStream>(
self,
session: SessionData<T>,
) -> impl std::future::Future<Output = ()> + Send {
async move {
let sender = self.inner;
let _ = http1::Builder::new()
.keep_alive(false)
.serve_connection(
TokioIo::new(session.stream),
service_fn(|mut req: hyper::Request<body::Incoming>| {
let sender = sender.clone();
async move {
let response = sender(HttpMessage {
method: req.method().clone(),
uri: req.uri().clone(),
headers: req
.headers()
.iter()
.map(|(k, v)| {
(k.as_str().to_lowercase(), v.to_str().unwrap().to_string())
})
.collect(),
body: fetch_body(&mut req, 1024 * 1024, 0).await,
});
Ok::<_, hyper::Error>(response.build())
}
}),
)
.await;
}
}
#[allow(clippy::manual_async_fn)]
fn shutdown(&self) -> impl std::future::Future<Output = ()> + Send {
async {}
}
}

2
tests/src/jmap/enterprise.rs
View File

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

8
tests/src/jmap/mod.rs
View File

@@ -384,12 +384,12 @@ pub async fn jmap_tests() {
mailbox::test(&mut params).await;
delivery::test(&mut params).await;
auth_acl::test(&mut params).await;
auth_limits::test(&mut params).await;*/
auth_limits::test(&mut params).await;
auth_oauth::test(&mut params).await;
/*event_source::test(&mut params).await;
push_subscription::test(&mut params).await;
event_source::test(&mut params).await;
push_subscription::test(&mut params).await;*/
sieve_script::test(&mut params).await;
vacation_response::test(&mut params).await;
/*vacation_response::test(&mut params).await;
email_submission::test(&mut params).await;
websocket::test(&mut params).await;
quota::test(&mut params).await;

2
tests/src/lib.rs
View File

@@ -18,6 +18,8 @@ static GLOBAL: Jemalloc = Jemalloc;
#[cfg(test)]
pub mod directory;
#[cfg(test)]
pub mod http_server;
#[cfg(test)]
pub mod imap;
#[cfg(test)]
pub mod jmap;

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

@@ -10,12 +10,17 @@ use std::{
use ahash::AHashMap;
use common::{
auth::AccessToken,
enterprise::llm::{
AiApiConfig, ChatCompletionChoice, ChatCompletionRequest, ChatCompletionResponse, Message,
},
scripts::{
functions::html::{get_attribute, html_attr_tokens, html_img_area, html_to_tokens},
ScriptModification,
},
Core,
};
use hyper::Method;
use jmap::api::{http::ToHttpResponse, JsonResponse};
use mail_auth::{dmarc::Policy, DkimResult, DmarcResult, IprevResult, SpfResult, MX};
use sieve::runtime::Variable;
use smtp::{
@@ -26,7 +31,11 @@ use smtp::{
use store::Stores;
use utils::config::Config;
use crate::smtp::{session::TestSession, TempDir, TestSMTP};
use crate::{
http_server::{spawn_mock_http_server, HttpMessage},
jmap::enterprise::EnterpriseCore,
smtp::{session::TestSession, TempDir, TestSMTP},
};
const CONFIG: &str = r#"
[spam.header]
@@ -46,6 +55,10 @@ threshold-discard = 0
threshold-reject = 0
directory = ""
lookup = ""
llm-model = "dummy"
llm-prompt = "You are an AI assistant specialized in analyzing email content to detect unsolicited, commercial, or harmful messages. Format your response as follows, separated by commas: Category,Confidence,Explanation
Here's the email to analyze, please provide your analysis based on the above instructions, ensuring your response is in the specified comma-separated format:"
add-llm-result = false
[session.rcpt]
relay = true
@@ -70,6 +83,11 @@ data = "spamdb"
lookup = "spamdb"
blob = "spamdb"
fts = "spamdb"
directory = "spamdb"
[directory."spamdb"]
type = "internal"
store = "spamdb"
[store."spamdb"]
type = "sqlite"
@@ -79,6 +97,12 @@ path = "{PATH}/test_antispam.db"
#type = "redis"
#url = "redis://127.0.0.1"
[enterprise.ai.dummy]
endpoint = "https://127.0.0.1:9090/v1/chat/completions"
type = "chat"
model = "gpt-dummy"
allow-invalid-certs = true
[lookup]
"spam-free" = {"gmail.com", "googlemail.com", "yahoomail.com", "*freemail.org"}
"spam-disposable" = {"guerrillamail.com", "*disposable.org"}
@@ -94,9 +118,6 @@ path = "{PATH}/test_antispam.db"
"spam-trap" = {"spamtrap@*"}
"spam-allow" = {"stalw.art"}
[resolver]
public-suffix = "file://{LIST_PATH}/public-suffix.dat"
[sieve.trusted.scripts]
"#;
@@ -129,6 +150,7 @@ async fn antispam() {
"bayes_classify",
"reputation",
"pyzor",
"llm",
];
let tmp_dir = TempDir::new("smtp_antispam_test", true);
let base_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
@@ -192,8 +214,14 @@ async fn antispam() {
let mut config = Config::new(&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;
//config.assert_no_errors();
let mut core = Core::parse(&mut config, stores, Default::default())
.await
.enable_enterprise();
core.enterprise.as_mut().unwrap().ai_apis.insert(
"dummy".to_string(),
AiApiConfig::parse(&mut config, "dummy").unwrap(),
);
crate::AssertConfig::assert_no_errors(config);
// Add mock DNS entries
for (domain, ip) in [
@@ -252,6 +280,34 @@ async fn antispam() {
let server = TestSMTP::from_core(core).server;
// Spawn mock OpenAI server
let _tx = spawn_mock_http_server(Arc::new(|req: HttpMessage| {
assert_eq!(req.uri.path(), "/v1/chat/completions");
assert_eq!(req.method, Method::POST);
let req =
serde_json::from_slice::<ChatCompletionRequest>(req.body.as_ref().unwrap()).unwrap();
assert_eq!(req.model, "gpt-dummy");
let message = &req.messages[0].content;
assert!(message.contains("You are an AI assistant specialized in analyzing email"));
JsonResponse::new(&ChatCompletionResponse {
created: 0,
object: String::new(),
id: String::new(),
model: req.model,
choices: vec![ChatCompletionChoice {
index: 0,
finish_reason: "stop".to_string(),
message: Message {
role: "assistant".to_string(),
content: message.split_once("Subject: ").unwrap().1.to_string(),
},
}],
})
.into_http_response()
}))
.await;
// Run tests
let base_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("resources")
@@ -420,10 +476,7 @@ async fn antispam() {
// Run script
let server_ = server.clone();
let script = script.clone();
match server_
.run_script("test".to_string(), script, params, 0)
.await
{
match server_.run_script("test".to_string(), script, params).await {
ScriptResult::Accept { modifications } => {
if modifications.len() != expected_headers.len() {
panic!(

2
tests/src/smtp/inbound/scripts.rs
View File

@@ -177,7 +177,7 @@ async fn sieve_scripts() {
.await;
match test
.server
.run_script(name.to_string(), script, params, 0)
.run_script(name.to_string(), script, params)
.await
{
ScriptResult::Accept { .. } => (),