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

Port Spam filter to Rust - part 3

Browse Source
This commit is contained in:
mdecimus
2024-12-09 17:49:11 +01:00
parent 4453dc8f3d
commit f0d84c8e68
34 changed files with 1791 additions and 653 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
crates/common/src/scripts/functions/mod.rs
View File

@@ -4,15 +4,15 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
mod array;
pub mod array;
mod email;
mod header;
pub mod html;
mod image;
mod misc;
pub mod image;
pub mod misc;
pub mod text;
mod unicode;
mod url;
pub mod unicode;
pub mod url;
use sieve::{runtime::Variable, FunctionMap};

2
crates/common/src/scripts/functions/unicode.rs
View File

@@ -43,7 +43,7 @@ pub fn fn_has_obscured<'x>(_: &'x Context<'x>, v: Vec<Variable>) -> Variable {
.into()
}
trait CharUtils {
pub trait CharUtils {
fn is_zwsp(&self) -> bool;
fn is_obscured(&self) -> bool;
}

251
crates/common/src/scripts/plugins/lookup.rs
View File

@@ -4,19 +4,10 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use std::{
collections::HashSet,
io::{BufRead, BufReader},
time::{Duration, Instant},
};
use mail_auth::flate2;
use sieve::{runtime::Variable, FunctionMap};
use store::{Deserialize, Value};
use crate::{
config::scripts::RemoteList, scripts::into_sieve_value, HttpLimitResponse, USER_AGENT,
};
use crate::scripts::into_sieve_value;
use super::PluginContext;
@@ -32,10 +23,6 @@ pub fn register_set(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("key_set", plugin_id, 4);
}
pub fn register_remote(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("key_exists_http", plugin_id, 3);
}
pub fn register_local_domain(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("is_local_domain", plugin_id, 2);
}
@@ -118,242 +105,6 @@ pub async fn exec_set(ctx: PluginContext<'_>) -> trc::Result<Variable> {
.map(|_| true.into())
}
pub async fn exec_remote(ctx: PluginContext<'_>) -> trc::Result<Variable> {
match exec_remote_(&ctx).await {
Ok(result) => Ok(result),
Err(err) => {
// Something went wrong, try again in one hour
const RETRY: Duration = Duration::from_secs(3600);
let mut _lock = ctx.server.inner.data.remote_lists.write();
let list = _lock
.entry(ctx.arguments[0].to_string().to_string())
.or_insert_with(|| RemoteList {
entries: HashSet::new(),
expires: Instant::now(),
});
if list.expires > Instant::now() {
Ok(list
.entries
.contains(ctx.arguments[1].to_string().as_ref())
.into())
} else {
list.expires = Instant::now() + RETRY;
Err(err)
}
}
}
}
const MAX_RESOURCE_SIZE: usize = 10 * 1024 * 1024;
async fn exec_remote_(ctx: &PluginContext<'_>) -> trc::Result<Variable> {
let resource = ctx.arguments[0].to_string();
let item = ctx.arguments[1].to_string();
#[cfg(feature = "test_mode")]
{
if (resource.contains("open") && item.contains("open"))
|| (resource.contains("tank") && item.contains("tank"))
{
return Ok(true.into());
}
}
if resource.is_empty() || item.is_empty() {
return Ok(false.into());
}
const TIMEOUT: Duration = Duration::from_secs(45);
const MAX_ENTRY_SIZE: usize = 256;
const MAX_ENTRIES: usize = 100000;
match ctx
.server
.inner
.data
.remote_lists
.read()
.get(resource.as_ref())
{
Some(remote_list) if remote_list.expires < Instant::now() => {
return Ok(remote_list.entries.contains(item.as_ref()).into())
}
_ => {}
}
enum Format {
List,
Csv {
column: u32,
separator: char,
skip_first: bool,
},
}
// Obtain parameters
let mut format = Format::List;
let mut expires = Duration::from_secs(12 * 3600);
if let Some(arr) = ctx.arguments[2].as_array() {
// Obtain expiration
match arr.first() {
Some(Variable::Integer(v)) if *v > 0 => {
expires = Duration::from_secs(*v as u64);
}
Some(Variable::Float(v)) if *v > 0.0 => {
expires = Duration::from_secs(*v as u64);
}
_ => (),
}
// Obtain list type
if matches!(arr.get(1), Some(Variable::String(list_type)) if list_type.eq_ignore_ascii_case("csv"))
{
format = Format::Csv {
column: arr.get(2).map(|v| v.to_integer()).unwrap_or_default() as u32,
separator: arr
.get(3)
.and_then(|v| v.to_string().chars().next())
.unwrap_or(','),
skip_first: arr.get(4).map_or(false, |v| v.to_bool()),
};
}
}
let response = reqwest::Client::builder()
.timeout(TIMEOUT)
.user_agent(USER_AGENT)
.build()
.unwrap_or_default()
.get(resource.as_ref())
.send()
.await
.map_err(|err| {
trc::SieveEvent::RuntimeError
.into_err()
.reason(err)
.ctx(trc::Key::Url, resource.to_string())
.details("Failed to build request")
})?;
if response.status().is_success() {
let bytes = response
.bytes_with_limit(MAX_RESOURCE_SIZE)
.await
.map_err(|err| {
trc::SieveEvent::RuntimeError
.into_err()
.reason(err)
.ctx(trc::Key::Url, resource.to_string())
.details("Failed to fetch resource")
})?
.ok_or_else(|| {
trc::SieveEvent::RuntimeError
.into_err()
.ctx(trc::Key::Url, resource.to_string())
.details("Resource is too large")
})?;
let reader: Box<dyn std::io::Read> = if resource.ends_with(".gz") {
Box::new(flate2::read::GzDecoder::new(&bytes[..]))
} else {
Box::new(&bytes[..])
};
// Lock remote list for writing
let mut _lock = ctx.server.inner.data.remote_lists.write();
let list = _lock
.entry(resource.to_string())
.or_insert_with(|| RemoteList {
entries: HashSet::new(),
expires: Instant::now(),
});
// Make sure that the list is still expired
if list.expires > Instant::now() {
return Ok(list.entries.contains(item.as_ref()).into());
}
for (pos, line) in BufReader::new(reader).lines().enumerate() {
let line_ = line.map_err(|err| {
trc::SieveEvent::RuntimeError
.into_err()
.reason(err)
.ctx(trc::Key::Url, resource.to_string())
.details("Failed to read line")
})?;
// Clear list once the first entry has been successfully fetched, decompressed and UTF8-decoded
if pos == 0 {
list.entries.clear();
}
match &format {
Format::List => {
let line = line_.trim();
if !line.is_empty() {
list.entries.insert(line.to_string());
}
}
Format::Csv {
column,
separator,
skip_first,
} if pos > 0 || !*skip_first => {
let mut in_quote = false;
let mut col_num = 0;
let mut entry = String::new();
for ch in line_.chars() {
if ch != '"' {
if ch == *separator && !in_quote {
if col_num == *column {
break;
} else {
col_num += 1;
}
} else if col_num == *column {
entry.push(ch);
if entry.len() > MAX_ENTRY_SIZE {
break;
}
}
} else {
in_quote = !in_quote;
}
}
if !entry.is_empty() {
list.entries.insert(entry);
}
}
_ => (),
}
if list.entries.len() == MAX_ENTRIES {
break;
}
}
trc::event!(
Spam(trc::SpamEvent::ListUpdated),
Url = resource.as_ref().to_string(),
Total = list.entries.len(),
);
// Update expiration
list.expires = Instant::now() + expires;
Ok(list.entries.contains(item.as_ref()).into())
} else {
trc::bail!(trc::SieveEvent::RuntimeError
.into_err()
.ctx(trc::Key::Code, response.status().as_u16())
.ctx(trc::Key::Url, resource.to_string())
.details("Failed to fetch remote list"));
}
}
pub async fn exec_local_domain(ctx: PluginContext<'_>) -> trc::Result<Variable> {
let domain = ctx.arguments[1].to_string();

20
crates/common/src/scripts/plugins/mod.rs
View File

@@ -31,13 +31,12 @@ pub struct PluginContext<'x> {
pub arguments: Vec<Variable>,
}
const PLUGINS_REGISTER: [RegisterPluginFnc; 14] = [
const PLUGINS_REGISTER: [RegisterPluginFnc; 13] = [
query::register,
exec::register,
lookup::register,
lookup::register_get,
lookup::register_set,
lookup::register_remote,
lookup::register_local_domain,
dns::register,
dns::register_exists,
@@ -86,15 +85,14 @@ impl Core {
2 => lookup::exec(ctx).await,
3 => lookup::exec_get(ctx).await,
4 => lookup::exec_set(ctx).await,
5 => lookup::exec_remote(ctx).await,
6 => lookup::exec_local_domain(ctx).await,
7 => dns::exec(ctx).await,
8 => dns::exec_exists(ctx).await,
9 => http::exec_header(ctx).await,
10 => headers::exec(ctx),
11 => text::exec_tokenize(ctx),
12 => text::exec_domain_part(ctx),
13 => llm_prompt::exec(ctx).await,
5 => lookup::exec_local_domain(ctx).await,
6 => dns::exec(ctx).await,
7 => dns::exec_exists(ctx).await,
8 => http::exec_header(ctx).await,
9 => headers::exec(ctx),
10 => text::exec_tokenize(ctx),
11 => text::exec_domain_part(ctx),
12 => llm_prompt::exec(ctx).await,
_ => unreachable!(),
};