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

Run Sieve scripts in async context

Browse Source
This commit is contained in:
mdecimus
2024-05-24 10:32:41 +02:00
parent ffdb7d766a
commit 4e7087d335
19 changed files with 363 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files

116
crates/common/src/scripts/plugins/bayes.rs
View File

@@ -30,7 +30,6 @@ use nlp::{
};
use sieve::{runtime::Variable, FunctionMap};
use store::{write::key::KeySerializer, LookupStore, U64_LEN};
use tokio::runtime::Handle;
use super::PluginContext;
@@ -50,15 +49,15 @@ pub fn register_is_balanced(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("bayes_is_balanced", plugin_id, 3);
}
pub fn exec_train(ctx: PluginContext<'_>) -> Variable {
train(ctx, true)
pub async fn exec_train(ctx: PluginContext<'_>) -> Variable {
train(ctx, true).await
}
pub fn exec_untrain(ctx: PluginContext<'_>) -> Variable {
train(ctx, false)
pub async fn exec_untrain(ctx: PluginContext<'_>) -> Variable {
train(ctx, false).await
}
fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
async fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
let span: &tracing::Span = ctx.span;
let store = match &ctx.arguments[0] {
Variable::String(v) if !v.is_empty() => ctx.core.storage.lookups.get(v.as_ref()),
@@ -82,7 +81,6 @@ fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
if text.is_empty() {
return false.into();
}
let handle = ctx.handle;
// Train the model
let mut model = BayesModel::default();
@@ -109,18 +107,17 @@ fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
let bayes_cache = &ctx.core.sieve.bayes_cache;
if is_train {
for (hash, weights) in model.weights {
if handle
.block_on(
store.counter_incr(
KeySerializer::new(U64_LEN)
.write(hash.h1)
.write(hash.h2)
.finalize(),
weights.into(),
None,
false,
),
if store
.counter_incr(
KeySerializer::new(U64_LEN)
.write(hash.h1)
.write(hash.h2)
.finalize(),
weights.into(),
None,
false,
)
.await
.is_err()
{
return false.into();
@@ -134,18 +131,17 @@ fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
} else {
Weights { spam: 0, ham: 1 }
};
if handle
.block_on(
store.counter_incr(
KeySerializer::new(U64_LEN)
.write(0u64)
.write(0u64)
.finalize(),
weights.into(),
None,
false,
),
if store
.counter_incr(
KeySerializer::new(U64_LEN)
.write(0u64)
.write(0u64)
.finalize(),
weights.into(),
None,
false,
)
.await
.is_err()
{
return false.into();
@@ -160,7 +156,7 @@ fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
true.into()
}
pub fn exec_classify(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_classify(ctx: PluginContext<'_>) -> Variable {
let span = ctx.span;
let store = match &ctx.arguments[0] {
Variable::String(v) if !v.is_empty() => ctx.core.storage.lookups.get(v.as_ref()),
@@ -200,12 +196,10 @@ pub fn exec_classify(ctx: PluginContext<'_>) -> Variable {
}
}
let handle = ctx.handle;
// Obtain training counts
let bayes_cache = &ctx.core.sieve.bayes_cache;
let (spam_learns, ham_learns) =
if let Some(weights) = bayes_cache.get_or_update(TokenHash::default(), handle, store) {
if let Some(weights) = bayes_cache.get_or_update(TokenHash::default(), store).await {
(weights.spam, weights.ham)
} else {
tracing::warn!(
@@ -231,27 +225,25 @@ pub fn exec_classify(ctx: PluginContext<'_>) -> Variable {
}
// Classify the text
let mut tokens = Vec::new();
for token in OsbTokenizer::<_, TokenHash>::new(
BayesTokenizer::new(text.as_ref(), &ctx.core.smtp.resolvers.psl),
5,
) {
if let Some(weights) = bayes_cache.get_or_update(token.inner, store).await {
tokens.push(OsbToken {
inner: weights,
idx: token.idx,
});
}
}
classifier
.classify(
OsbTokenizer::<_, TokenHash>::new(
BayesTokenizer::new(text.as_ref(), &ctx.core.smtp.resolvers.psl),
5,
)
.filter_map(|t| {
OsbToken {
inner: bayes_cache.get_or_update(t.inner, handle, store)?,
idx: t.idx,
}
.into()
}),
ham_learns,
spam_learns,
)
.classify(tokens.into_iter(), ham_learns, spam_learns)
.map(Variable::from)
.unwrap_or_default()
}
pub fn exec_is_balanced(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_is_balanced(ctx: PluginContext<'_>) -> Variable {
let min_balance = match &ctx.arguments[2] {
Variable::Float(n) => *n,
Variable::Integer(n) => *n as f64,
@@ -282,10 +274,9 @@ pub fn exec_is_balanced(ctx: PluginContext<'_>) -> Variable {
let learn_spam = ctx.arguments[1].to_bool();
// Obtain training counts
let handle = ctx.handle;
let bayes_cache = &ctx.core.sieve.bayes_cache;
let (spam_learns, ham_learns) =
if let Some(weights) = bayes_cache.get_or_update(TokenHash::default(), handle, store) {
if let Some(weights) = bayes_cache.get_or_update(TokenHash::default(), store).await {
(weights.spam as f64, weights.ham as f64)
} else {
tracing::warn!(
@@ -321,31 +312,22 @@ pub fn exec_is_balanced(ctx: PluginContext<'_>) -> Variable {
}
trait LookupOrInsert {
fn get_or_update(
&self,
hash: TokenHash,
handle: &Handle,
get_token: &LookupStore,
) -> Option<Weights>;
async fn get_or_update(&self, hash: TokenHash, get_token: &LookupStore) -> Option<Weights>;
}
impl LookupOrInsert for BayesTokenCache {
fn get_or_update(
&self,
hash: TokenHash,
handle: &Handle,
get_token: &LookupStore,
) -> Option<Weights> {
async fn get_or_update(&self, hash: TokenHash, get_token: &LookupStore) -> Option<Weights> {
if let Some(weights) = self.get(&hash) {
weights.unwrap_or_default().into()
} else if let Ok(num) = handle.block_on(
get_token.counter_get(
} else if let Ok(num) = get_token
.counter_get(
KeySerializer::new(U64_LEN)
.write(hash.h1)
.write(hash.h2)
.finalize(),
),
) {
)
.await
{
if num != 0 {
let weights = Weights::from(num);
self.insert_positive(hash, weights);

90
crates/common/src/scripts/plugins/dns.rs
View File

@@ -36,16 +36,19 @@ pub fn register_exists(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("dns_exists", plugin_id, 2);
}
pub fn exec(ctx: PluginContext<'_>) -> Variable {
pub async fn exec(ctx: PluginContext<'_>) -> Variable {
let entry = ctx.arguments[0].to_string();
let record_type = ctx.arguments[1].to_string();
if record_type.eq_ignore_ascii_case("ip") {
match ctx.handle.block_on(ctx.core.smtp.resolvers.dns.ip_lookup(
entry.as_ref(),
IpLookupStrategy::Ipv4thenIpv6,
10,
)) {
match ctx
.core
.smtp
.resolvers
.dns
.ip_lookup(entry.as_ref(), IpLookupStrategy::Ipv4thenIpv6, 10)
.await
{
Ok(result) => result
.iter()
.map(|ip| Variable::from(ip.to_string()))
@@ -54,10 +57,7 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
Err(err) => err.short_error().into(),
}
} else if record_type.eq_ignore_ascii_case("mx") {
match ctx
.handle
.block_on(ctx.core.smtp.resolvers.dns.mx_lookup(entry.as_ref()))
{
match ctx.core.smtp.resolvers.dns.mx_lookup(entry.as_ref()).await {
Ok(result) => result
.iter()
.flat_map(|mx| {
@@ -78,18 +78,19 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
}
match ctx
.handle
.block_on(ctx.core.smtp.resolvers.dns.txt_raw_lookup(entry.as_ref()))
.core
.smtp
.resolvers
.dns
.txt_raw_lookup(entry.as_ref())
.await
{
Ok(result) => Variable::from(String::from_utf8(result).unwrap_or_default()),
Err(err) => err.short_error().into(),
}
} else if record_type.eq_ignore_ascii_case("ptr") {
if let Ok(addr) = entry.parse::<IpAddr>() {
match ctx
.handle
.block_on(ctx.core.smtp.resolvers.dns.ptr_lookup(addr))
{
match ctx.core.smtp.resolvers.dns.ptr_lookup(addr).await {
Ok(result) => result
.iter()
.map(|host| Variable::from(host.to_string()))
@@ -110,8 +111,12 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
}
match ctx
.handle
.block_on(ctx.core.smtp.resolvers.dns.ipv4_lookup(entry.as_ref()))
.core
.smtp
.resolvers
.dns
.ipv4_lookup(entry.as_ref())
.await
{
Ok(result) => result
.iter()
@@ -122,8 +127,12 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
}
} else if record_type.eq_ignore_ascii_case("ipv6") {
match ctx
.handle
.block_on(ctx.core.smtp.resolvers.dns.ipv6_lookup(entry.as_ref()))
.core
.smtp
.resolvers
.dns
.ipv6_lookup(entry.as_ref())
.await
{
Ok(result) => result
.iter()
@@ -137,35 +146,32 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
}
}
pub fn exec_exists(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_exists(ctx: PluginContext<'_>) -> Variable {
let entry = ctx.arguments[0].to_string();
let record_type = ctx.arguments[1].to_string();
if record_type.eq_ignore_ascii_case("ip") {
match ctx.handle.block_on(ctx.core.smtp.resolvers.dns.ip_lookup(
entry.as_ref(),
IpLookupStrategy::Ipv4thenIpv6,
10,
)) {
match ctx
.core
.smtp
.resolvers
.dns
.ip_lookup(entry.as_ref(), IpLookupStrategy::Ipv4thenIpv6, 10)
.await
{
Ok(result) => i64::from(!result.is_empty()),
Err(Error::DnsRecordNotFound(_)) => 0,
Err(_) => -1,
}
} else if record_type.eq_ignore_ascii_case("mx") {
match ctx
.handle
.block_on(ctx.core.smtp.resolvers.dns.mx_lookup(entry.as_ref()))
{
match ctx.core.smtp.resolvers.dns.mx_lookup(entry.as_ref()).await {
Ok(result) => i64::from(result.iter().any(|mx| !mx.exchanges.is_empty())),
Err(Error::DnsRecordNotFound(_)) => 0,
Err(_) => -1,
}
} else if record_type.eq_ignore_ascii_case("ptr") {
if let Ok(addr) = entry.parse::<IpAddr>() {
match ctx
.handle
.block_on(ctx.core.smtp.resolvers.dns.ptr_lookup(addr))
{
match ctx.core.smtp.resolvers.dns.ptr_lookup(addr).await {
Ok(result) => i64::from(!result.is_empty()),
Err(Error::DnsRecordNotFound(_)) => 0,
Err(_) => -1,
@@ -182,8 +188,12 @@ pub fn exec_exists(ctx: PluginContext<'_>) -> Variable {
}
match ctx
.handle
.block_on(ctx.core.smtp.resolvers.dns.ipv4_lookup(entry.as_ref()))
.core
.smtp
.resolvers
.dns
.ipv4_lookup(entry.as_ref())
.await
{
Ok(result) => i64::from(!result.is_empty()),
Err(Error::DnsRecordNotFound(_)) => 0,
@@ -191,8 +201,12 @@ pub fn exec_exists(ctx: PluginContext<'_>) -> Variable {
}
} else if record_type.eq_ignore_ascii_case("ipv6") {
match ctx
.handle
.block_on(ctx.core.smtp.resolvers.dns.ipv6_lookup(entry.as_ref()))
.core
.smtp
.resolvers
.dns
.ipv6_lookup(entry.as_ref())
.await
{
Ok(result) => i64::from(!result.is_empty()),
Err(Error::DnsRecordNotFound(_)) => 0,

58
crates/common/src/scripts/plugins/exec.rs
View File

@@ -31,32 +31,38 @@ pub fn register(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("exec", plugin_id, 2);
}
pub fn exec(ctx: PluginContext<'_>) -> Variable {
let span = ctx.span;
pub async fn exec(ctx: PluginContext<'_>) -> Variable {
let span = ctx.span.clone();
let mut arguments = ctx.arguments.into_iter();
match Command::new(
arguments
.next()
.map(|a| a.to_string().into_owned())
.unwrap_or_default(),
)
.args(
arguments
.next()
.map(|a| a.into_string_array())
.unwrap_or_default(),
)
.output()
{
Ok(result) => result.status.success().into(),
Err(err) => {
tracing::warn!(
parent: span,
context = "sieve",
event = "execute-failed",
reason = %err,
);
false.into()
tokio::task::spawn_blocking(move || {
match Command::new(
arguments
.next()
.map(|a| a.to_string().into_owned())
.unwrap_or_default(),
)
.args(
arguments
.next()
.map(|a| a.into_string_array())
.unwrap_or_default(),
)
.output()
{
Ok(result) => result.status.success(),
Err(err) => {
tracing::warn!(
parent: span,
context = "sieve",
event = "execute-failed",
reason = %err,
);
false
}
}
}
})
.await
.unwrap_or_default()
.into()
}

9
crates/common/src/scripts/plugins/http.rs
View File

@@ -32,7 +32,7 @@ pub fn register_header(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("http_header", plugin_id, 4);
}
pub fn exec_header(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_header(ctx: PluginContext<'_>) -> Variable {
let url = ctx.arguments[0].to_string();
let header = ctx.arguments[1].to_string();
let agent = ctx.arguments[2].to_string();
@@ -50,9 +50,10 @@ pub fn exec_header(ctx: PluginContext<'_>) -> Variable {
.danger_accept_invalid_certs(true)
.build()
{
let _enter = ctx.handle.enter();
ctx.handle
.block_on(client.get(url.as_ref()).send())
client
.get(url.as_ref())
.send()
.await
.ok()
.and_then(|response| {
response

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

@@ -55,7 +55,7 @@ pub fn register_local_domain(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("is_local_domain", plugin_id, 2);
}
pub fn exec(ctx: PluginContext<'_>) -> Variable {
pub async fn exec(ctx: PluginContext<'_>) -> Variable {
let store = match &ctx.arguments[0] {
Variable::String(v) if !v.is_empty() => ctx.core.storage.lookups.get(v.as_ref()),
_ => Some(&ctx.core.storage.lookup),
@@ -66,9 +66,9 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
Variable::Array(items) => {
for item in items.iter() {
if !item.is_empty()
&& ctx
.handle
.block_on(store.key_exists(item.to_string().into_owned().into_bytes()))
&& store
.key_exists(item.to_string().into_owned().into_bytes())
.await
.unwrap_or(false)
{
return true.into();
@@ -76,9 +76,9 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
}
false
}
v if !v.is_empty() => ctx
.handle
.block_on(store.key_exists(v.to_string().into_owned().into_bytes()))
v if !v.is_empty() => store
.key_exists(v.to_string().into_owned().into_bytes())
.await
.unwrap_or(false),
_ => false,
}
@@ -95,19 +95,16 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
.into()
}
pub fn exec_get(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_get(ctx: PluginContext<'_>) -> Variable {
let store = match &ctx.arguments[0] {
Variable::String(v) if !v.is_empty() => ctx.core.storage.lookups.get(v.as_ref()),
_ => Some(&ctx.core.storage.lookup),
};
if let Some(store) = store {
ctx.handle
.block_on(
store.key_get::<VariableWrapper>(
ctx.arguments[1].to_string().into_owned().into_bytes(),
),
)
store
.key_get::<VariableWrapper>(ctx.arguments[1].to_string().into_owned().into_bytes())
.await
.unwrap_or_default()
.map(|v| v.into_inner())
.unwrap_or_default()
@@ -123,7 +120,7 @@ pub fn exec_get(ctx: PluginContext<'_>) -> Variable {
}
}
pub fn exec_set(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_set(ctx: PluginContext<'_>) -> Variable {
let store = match &ctx.arguments[0] {
Variable::String(v) if !v.is_empty() => ctx.core.storage.lookups.get(v.as_ref()),
_ => Some(&ctx.core.storage.lookup),
@@ -136,8 +133,8 @@ pub fn exec_set(ctx: PluginContext<'_>) -> Variable {
_ => None,
};
ctx.handle
.block_on(store.key_set(
store
.key_set(
ctx.arguments[1].to_string().into_owned().into_bytes(),
if !ctx.arguments[2].is_empty() {
bincode::serialize(&ctx.arguments[2]).unwrap_or_default()
@@ -145,7 +142,8 @@ pub fn exec_set(ctx: PluginContext<'_>) -> Variable {
vec![]
},
expires,
))
)
.await
.is_ok()
.into()
} else {
@@ -160,7 +158,7 @@ pub fn exec_set(ctx: PluginContext<'_>) -> Variable {
}
}
pub fn exec_remote(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_remote(ctx: PluginContext<'_>) -> Variable {
let resource = ctx.arguments[0].to_string();
let item = ctx.arguments[1].to_string();
@@ -228,126 +226,129 @@ pub fn exec_remote(ctx: PluginContext<'_>) -> Variable {
}
}
// Lock remote list for writing
let mut _lock = ctx.core.sieve.remote_lists.write();
let list = _lock
.entry(resource.to_string())
.or_insert_with(|| RemoteList {
entries: HashSet::new(),
expires: Instant::now(),
});
match reqwest::Client::builder()
.timeout(TIMEOUT)
.user_agent(USER_AGENT)
.build()
.unwrap_or_default()
.get(resource.as_ref())
.send()
.await
{
Ok(response) if response.status().is_success() => {
match response.bytes().await {
Ok(bytes) => {
let reader: Box<dyn std::io::Read> = if resource.ends_with(".gz") {
Box::new(flate2::read::GzDecoder::new(&bytes[..]))
} else {
Box::new(&bytes[..])
};
// Make sure that the list is still expired
if list.expires > Instant::now() {
return list.entries.contains(item.as_ref()).into();
}
// Lock remote list for writing
let mut _lock = ctx.core.sieve.remote_lists.write();
let list = _lock
.entry(resource.to_string())
.or_insert_with(|| RemoteList {
entries: HashSet::new(),
expires: Instant::now(),
});
let _enter = ctx.handle.enter();
match ctx
.handle
.block_on(
reqwest::Client::builder()
.timeout(TIMEOUT)
.user_agent(USER_AGENT)
.build()
.unwrap_or_default()
.get(resource.as_ref())
.send(),
)
.and_then(|r| {
if r.status().is_success() {
ctx.handle.block_on(r.bytes()).map(Ok)
} else {
Ok(Err(r))
}
}) {
Ok(Ok(bytes)) => {
let reader: Box<dyn std::io::Read> = if resource.ends_with(".gz") {
Box::new(flate2::read::GzDecoder::new(&bytes[..]))
} else {
Box::new(&bytes[..])
};
// Make sure that the list is still expired
if list.expires > Instant::now() {
return list.entries.contains(item.as_ref()).into();
}
for (pos, line) in BufReader::new(reader).lines().enumerate() {
match line {
Ok(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());
for (pos, line) in BufReader::new(reader).lines().enumerate() {
match line {
Ok(line_) => {
// Clear list once the first entry has been successfully fetched, decompressed and UTF8-decoded
if pos == 0 {
list.entries.clear();
}
}
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;
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 {
col_num += 1;
}
} else if col_num == *column {
entry.push(ch);
if entry.len() > MAX_ENTRY_SIZE {
break;
in_quote = !in_quote;
}
}
} else {
in_quote = !in_quote;
}
}
if !entry.is_empty() {
list.entries.insert(entry);
if !entry.is_empty() {
list.entries.insert(entry);
}
}
_ => (),
}
}
_ => (),
Err(err) => {
tracing::warn!(
parent: ctx.span,
context = "sieve:key_exists_http",
event = "failed",
resource = resource.as_ref(),
reason = %err,
);
break;
}
}
if list.entries.len() == MAX_ENTRIES {
break;
}
}
Err(err) => {
tracing::warn!(
parent: ctx.span,
context = "sieve:key_exists_http",
event = "failed",
resource = resource.as_ref(),
reason = %err,
);
break;
}
}
if list.entries.len() == MAX_ENTRIES {
break;
tracing::debug!(
parent: ctx.span,
context = "sieve:key_exists_http",
event = "fetch",
resource = resource.as_ref(),
num_entries = list.entries.len(),
);
// Update expiration
list.expires = Instant::now() + expires;
return list.entries.contains(item.as_ref()).into();
}
Err(err) => {
tracing::warn!(
parent: ctx.span,
context = "sieve:key_exists_http",
event = "failed",
resource = resource.as_ref(),
reason = %err,
);
}
}
tracing::debug!(
parent: ctx.span,
context = "sieve:key_exists_http",
event = "fetch",
resource = resource.as_ref(),
num_entries = list.entries.len(),
);
// Update expiration
list.expires = Instant::now() + expires;
return list.entries.contains(item.as_ref()).into();
}
Ok(Err(response)) => {
Ok(response) => {
tracing::warn!(
parent: ctx.span,
context = "sieve:key_exists_http",
@@ -368,11 +369,22 @@ pub fn exec_remote(ctx: PluginContext<'_>) -> Variable {
}
// Something went wrong, try again in one hour
list.expires = Instant::now() + RETRY;
false.into()
let mut _lock = ctx.core.sieve.remote_lists.write();
let list = _lock
.entry(resource.to_string())
.or_insert_with(|| RemoteList {
entries: HashSet::new(),
expires: Instant::now(),
});
if list.expires > Instant::now() {
list.entries.contains(item.as_ref()).into()
} else {
list.expires = Instant::now() + RETRY;
false.into()
}
}
pub fn exec_local_domain(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_local_domain(ctx: PluginContext<'_>) -> Variable {
let domain = ctx.arguments[0].to_string();
if !domain.is_empty() {
@@ -382,9 +394,9 @@ pub fn exec_local_domain(ctx: PluginContext<'_>) -> Variable {
};
if let Some(directory) = directory {
return ctx
.handle
.block_on(directory.is_local_domain(domain.as_ref()))
return directory
.is_local_domain(domain.as_ref())
.await
.unwrap_or_default()
.into();
} else {

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

@@ -33,44 +33,21 @@ pub mod text;
use mail_parser::Message;
use sieve::{runtime::Variable, FunctionMap, Input};
use tokio::runtime::Handle;
use crate::Core;
use super::ScriptModification;
type RegisterPluginFnc = fn(u32, &mut FunctionMap) -> ();
type ExecPluginFnc = fn(PluginContext<'_>) -> Variable;
pub struct PluginContext<'x> {
pub span: &'x tracing::Span,
pub handle: &'x Handle,
pub core: &'x Core,
pub message: &'x Message<'x>,
pub modifications: &'x mut Vec<ScriptModification>,
pub arguments: Vec<Variable>,
}
const PLUGINS_EXEC: [ExecPluginFnc; 18] = [
query::exec,
exec::exec,
lookup::exec,
lookup::exec_get,
lookup::exec_set,
lookup::exec_remote,
lookup::exec_local_domain,
dns::exec,
dns::exec_exists,
http::exec_header,
bayes::exec_train,
bayes::exec_untrain,
bayes::exec_classify,
bayes::exec_is_balanced,
pyzor::exec,
headers::exec,
text::exec_tokenize,
text::exec_domain_part,
];
const PLUGINS_REGISTER: [RegisterPluginFnc; 18] = [
query::register,
exec::register,
@@ -100,7 +77,7 @@ impl RegisterSievePlugins for FunctionMap {
fn register_plugins(mut self) -> Self {
#[cfg(feature = "test_mode")]
{
self.set_external_function("print", PLUGINS_EXEC.len() as u32, 1)
self.set_external_function("print", PLUGINS_REGISTER.len() as u32, 1)
}
for (i, fnc) in PLUGINS_REGISTER.iter().enumerate() {
@@ -111,17 +88,34 @@ impl RegisterSievePlugins for FunctionMap {
}
impl Core {
pub fn run_plugin_blocking(&self, id: u32, ctx: PluginContext<'_>) -> Input {
pub async fn run_plugin(&self, id: u32, ctx: PluginContext<'_>) -> Input {
#[cfg(feature = "test_mode")]
if id == PLUGINS_EXEC.len() as u32 {
if id == PLUGINS_REGISTER.len() as u32 {
return test_print(ctx);
}
PLUGINS_EXEC
.get(id as usize)
.map(|fnc| fnc(ctx))
.unwrap_or_default()
.into()
match id {
0 => query::exec(ctx).await,
1 => exec::exec(ctx).await,
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 => bayes::exec_train(ctx).await,
11 => bayes::exec_untrain(ctx).await,
12 => bayes::exec_classify(ctx).await,
13 => bayes::exec_is_balanced(ctx).await,
14 => pyzor::exec(ctx).await,
15 => headers::exec(ctx),
16 => text::exec_tokenize(ctx),
17 => text::exec_domain_part(ctx),
_ => unreachable!(),
}
.into()
}
}

7
crates/common/src/scripts/plugins/pyzor.rs
View File

@@ -52,7 +52,7 @@ pub fn register(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("pyzor_check", plugin_id, 2);
}
pub fn exec(ctx: PluginContext<'_>) -> Variable {
pub async fn exec(ctx: PluginContext<'_>) -> Variable {
// Make sure there is at least one text part
if !ctx
.message
@@ -101,10 +101,7 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
5,
));
// Send message to address
match ctx
.handle
.block_on(pyzor_send_message(address.as_ref(), timeout, &request))
{
match pyzor_send_message(address.as_ref(), timeout, &request).await {
Ok(response) => response.into(),
Err(err) => {
tracing::debug!(

9
crates/common/src/scripts/plugins/query.rs
View File

@@ -33,7 +33,7 @@ pub fn register(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("query", plugin_id, 3);
}
pub fn exec(ctx: PluginContext<'_>) -> Variable {
pub async fn exec(ctx: PluginContext<'_>) -> Variable {
let span = ctx.span;
// Obtain store name
@@ -79,7 +79,7 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
.get(..6)
.map_or(false, |q| q.eq_ignore_ascii_case(b"SELECT"))
{
if let Ok(mut rows) = ctx.handle.block_on(store.query::<Rows>(&query, arguments)) {
if let Ok(mut rows) = store.query::<Rows>(&query, arguments).await {
match rows.rows.len().cmp(&1) {
Ordering::Equal => {
let mut row = rows.rows.pop().unwrap().values;
@@ -116,9 +116,6 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
false.into()
}
} else {
ctx.handle
.block_on(store.query::<usize>(&query, arguments))
.is_ok()
.into()
store.query::<usize>(&query, arguments).await.is_ok().into()
}
}