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

Improved tracing (part 2)

Browse Source
This commit is contained in:
mdecimus
2024-07-25 20:35:13 +02:00
parent ae7cadc27d
commit 52cb48353e
108 changed files with 3137 additions and 2307 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -13,6 +13,7 @@ use nlp::{
};
use sieve::{runtime::Variable, FunctionMap};
use store::{write::key::KeySerializer, LookupStore, U64_LEN};
use trc::AddContext;
use super::PluginContext;
@@ -32,42 +33,31 @@ pub fn register_is_balanced(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("bayes_is_balanced", plugin_id, 3);
}
pub async fn exec_train(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_train(ctx: PluginContext<'_>) -> trc::Result<Variable> {
train(ctx, true).await
}
pub async fn exec_untrain(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_untrain(ctx: PluginContext<'_>) -> trc::Result<Variable> {
train(ctx, false).await
}
async fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
async fn train(ctx: PluginContext<'_>, is_train: bool) -> trc::Result<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),
};
}
.ok_or_else(|| {
trc::SieveEvent::RuntimeError
.ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned())
.details("Unknown store")
})?;
let store = if let Some(store) = store {
store
} else {
tracing::warn!(
context = "sieve:bayes_train",
event = "failed",
reason = "Unknown store id",
lookup_store = ctx.arguments[0].to_string().as_ref(),
);
return false.into();
};
let text = ctx.arguments[1].to_string();
let is_spam = ctx.arguments[2].to_bool();
if text.is_empty() {
tracing::debug!(
context = "sieve:bayes_train",
event = "failed",
reason = "Empty message",
);
return false.into();
trc::bail!(trc::SpamEvent::TrainError
.into_err()
.reason("Empty message"));
}
// Train the model
@@ -80,28 +70,23 @@ async fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
is_spam,
);
if model.weights.is_empty() {
tracing::debug!(
context = "sieve:bayes_train",
event = "failed",
reason = "No weights found",
);
return false.into();
trc::bail!(trc::SpamEvent::TrainError
.into_err()
.reason("No weights found"));
}
tracing::debug!(
context = "sieve:bayes_train",
event = "train",
is_spam = is_spam,
num_tokens = model.weights.len(),
trc::event!(
Spam(trc::SpamEvent::Train),
SessionId = ctx.session_id,
Spam = is_spam,
Size = model.weights.len(),
);
// Update weight and invalidate cache
let bayes_cache = &ctx.cache.bayes_cache;
if is_train {
for (hash, weights) in model.weights {
if store
store
.counter_incr(
KeySerializer::new(U64_LEN)
.write(hash.h1)
@@ -112,10 +97,8 @@ async fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
false,
)
.await
.is_err()
{
return false.into();
}
.caused_by(trc::location!())?;
bayes_cache.invalidate(&hash);
}
@@ -125,7 +108,7 @@ async fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
} else {
Weights { spam: 0, ham: 1 }
};
if store
store
.counter_incr(
KeySerializer::new(U64_LEN)
.write(0u64)
@@ -136,41 +119,32 @@ async fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable {
false,
)
.await
.is_err()
{
return false.into();
}
.caused_by(trc::location!())?;
} else {
//TODO: Implement untrain
return false.into();
return Ok(false.into());
}
bayes_cache.invalidate(&TokenHash::default());
true.into()
Ok(true.into())
}
pub async fn exec_classify(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_classify(ctx: PluginContext<'_>) -> trc::Result<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),
};
let store = if let Some(store) = store {
store
} else {
tracing::warn!(
context = "sieve:bayes_classify",
event = "failed",
reason = "Unknown store id",
lookup_id = ctx.arguments[0].to_string().as_ref(),
);
return Variable::default();
};
}
.ok_or_else(|| {
trc::SieveEvent::RuntimeError
.ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned())
.details("Unknown store")
})?;
let text = ctx.arguments[1].to_string();
if text.is_empty() {
return Variable::default();
trc::bail!(trc::SpamEvent::ClassifyError
.into_err()
.reason("Empty message"));
}
// Create classifier from defaults
@@ -192,30 +166,21 @@ pub async fn exec_classify(ctx: PluginContext<'_>) -> Variable {
// Obtain training counts
let bayes_cache = &ctx.cache.bayes_cache;
let (spam_learns, ham_learns) =
if let Some(weights) = bayes_cache.get_or_update(TokenHash::default(), store).await {
(weights.spam, weights.ham)
} else {
tracing::warn!(
context = "sieve:classify",
event = "failed",
reason = "Failed to obtain training counts",
);
return Variable::default();
};
let (spam_learns, ham_learns) = bayes_cache
.get_or_update(TokenHash::default(), store)
.await
.map(|w| (w.spam, w.ham))?;
// Make sure we have enough training data
if spam_learns < classifier.min_learns || ham_learns < classifier.min_learns {
tracing::debug!(
context = "sieve:bayes_classify",
event = "skip-classify",
reason = "Not enough training data",
min_learns = classifier.min_learns,
spam_learns = %spam_learns,
ham_learns = %ham_learns);
return Variable::default();
trc::event!(
Spam(trc::SpamEvent::NotEnoughTrainingData),
SessionId = ctx.session_id,
MinLearns = classifier.min_learns,
SpamLearns = spam_learns,
HamLearns = ham_learns
);
return Ok(Variable::default());
}
// Classify the text
@@ -224,20 +189,27 @@ pub async fn exec_classify(ctx: PluginContext<'_>) -> Variable {
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,
});
}
let weights = bayes_cache.get_or_update(token.inner, store).await?;
tokens.push(OsbToken {
inner: weights,
idx: token.idx,
});
}
classifier
.classify(tokens.into_iter(), ham_learns, spam_learns)
.map(Variable::from)
.unwrap_or_default()
let result = classifier.classify(tokens.into_iter(), ham_learns, spam_learns);
trc::event!(
Spam(trc::SpamEvent::Classify),
SessionId = ctx.session_id,
MinLearns = classifier.min_learns,
SpamLearns = spam_learns,
HamLearns = ham_learns,
Result = result.unwrap_or_default()
);
Ok(result.map(Variable::from).unwrap_or_default())
}
pub async fn exec_is_balanced(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_is_balanced(ctx: PluginContext<'_>) -> trc::Result<Variable> {
let min_balance = match &ctx.arguments[2] {
Variable::Float(n) => *n,
Variable::Integer(n) => *n as f64,
@@ -245,42 +217,27 @@ pub async fn exec_is_balanced(ctx: PluginContext<'_>) -> Variable {
};
if min_balance == 0.0 {
return true.into();
return Ok(true.into());
}
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),
};
let store = if let Some(store) = store {
store
} else {
tracing::warn!(
context = "sieve:bayes_is_balanced",
event = "failed",
reason = "Unknown store id",
lookup_id = ctx.arguments[0].to_string().as_ref(),
);
return Variable::default();
};
}
.ok_or_else(|| {
trc::SieveEvent::RuntimeError
.ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned())
.details("Unknown store")
})?;
let learn_spam = ctx.arguments[1].to_bool();
// Obtain training counts
let bayes_cache = &ctx.cache.bayes_cache;
let (spam_learns, ham_learns) =
if let Some(weights) = bayes_cache.get_or_update(TokenHash::default(), store).await {
(weights.spam as f64, weights.ham as f64)
} else {
tracing::warn!(
context = "sieve:bayes_is_balanced",
event = "failed",
reason = "Failed to obtain training counts",
);
return Variable::default();
};
let (spam_learns, ham_learns) = bayes_cache
.get_or_update(TokenHash::default(), store)
.await
.map(|w| (w.spam as f64, w.ham as f64))?;
let result = if spam_learns > 0.0 || ham_learns > 0.0 {
if learn_spam {
@@ -292,37 +249,43 @@ pub async fn exec_is_balanced(ctx: PluginContext<'_>) -> Variable {
true
};
tracing::debug!(
context = "sieve:bayes_is_balanced",
event = "result",
is_balanced = %result,
learn_spam = %learn_spam,
min_balance = %min_balance,
spam_learns = %spam_learns,
ham_learns = %ham_learns);
trc::event!(
Spam(trc::SpamEvent::TrainBalance),
SessionId = ctx.session_id,
Spam = learn_spam,
MinBalance = min_balance,
SpamLearns = spam_learns,
HamLearns = ham_learns,
Result = result
);
result.into()
Ok(result.into())
}
trait LookupOrInsert {
async fn get_or_update(&self, hash: TokenHash, get_token: &LookupStore) -> Option<Weights>;
async fn get_or_update(&self, hash: TokenHash, get_token: &LookupStore)
-> trc::Result<Weights>;
}
impl LookupOrInsert for BayesTokenCache {
async fn get_or_update(&self, hash: TokenHash, get_token: &LookupStore) -> Option<Weights> {
async fn get_or_update(
&self,
hash: TokenHash,
get_token: &LookupStore,
) -> trc::Result<Weights> {
if let Some(weights) = self.get(&hash) {
weights.unwrap_or_default().into()
} else if let Ok(num) = get_token
.counter_get(
KeySerializer::new(U64_LEN)
.write(hash.h1)
.write(hash.h2)
.finalize(),
)
.await
{
if num != 0 {
Ok(weights.unwrap_or_default())
} else {
let num = get_token
.counter_get(
KeySerializer::new(U64_LEN)
.write(hash.h1)
.write(hash.h2)
.finalize(),
)
.await
.caused_by(trc::location!())?;
Ok(if num != 0 {
let weights = Weights::from(num);
self.insert_positive(hash, weights);
weights
@@ -330,10 +293,7 @@ impl LookupOrInsert for BayesTokenCache {
self.insert_negative(hash);
Weights::default()
}
.into()
} else {
// Something went wrong
None
.into())
}
}
}

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

@@ -19,11 +19,11 @@ pub fn register_exists(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("dns_exists", plugin_id, 2);
}
pub async fn exec(ctx: PluginContext<'_>) -> Variable {
pub async fn exec(ctx: PluginContext<'_>) -> trc::Result<Variable> {
let entry = ctx.arguments[0].to_string();
let record_type = ctx.arguments[1].to_string();
if record_type.eq_ignore_ascii_case("ip") {
Ok(if record_type.eq_ignore_ascii_case("ip") {
match ctx
.core
.smtp
@@ -56,7 +56,7 @@ pub async fn exec(ctx: PluginContext<'_>) -> Variable {
#[cfg(feature = "test_mode")]
{
if entry.contains("origin") {
return Variable::from("23028|US|arin|2002-01-04".to_string());
return Ok(Variable::from("23028|US|arin|2002-01-04".to_string()));
}
}
@@ -89,7 +89,7 @@ pub async fn exec(ctx: PluginContext<'_>) -> Variable {
{
if entry.contains(".168.192.") {
let parts = entry.split('.').collect::<Vec<_>>();
return vec![Variable::from(format!("127.0.{}.{}", parts[1], parts[0]))].into();
return Ok(vec![Variable::from(format!("127.0.{}.{}", parts[1], parts[0]))].into());
}
}
@@ -126,14 +126,14 @@ pub async fn exec(ctx: PluginContext<'_>) -> Variable {
}
} else {
Variable::default()
}
})
}
pub async fn exec_exists(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_exists(ctx: PluginContext<'_>) -> trc::Result<Variable> {
let entry = ctx.arguments[0].to_string();
let record_type = ctx.arguments[1].to_string();
if record_type.eq_ignore_ascii_case("ip") {
Ok(if record_type.eq_ignore_ascii_case("ip") {
match ctx
.core
.smtp
@@ -166,7 +166,7 @@ pub async fn exec_exists(ctx: PluginContext<'_>) -> Variable {
#[cfg(feature = "test_mode")]
{
if entry.starts_with("2.0.168.192.") {
return 1.into();
return Ok(1.into());
}
}
@@ -198,7 +198,7 @@ pub async fn exec_exists(ctx: PluginContext<'_>) -> Variable {
} else {
-1
}
.into()
.into())
}
trait ShortError {

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

@@ -14,36 +14,37 @@ pub fn register(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("exec", plugin_id, 2);
}
pub async fn exec(ctx: PluginContext<'_>) -> Variable {
pub async fn exec(ctx: PluginContext<'_>) -> trc::Result<Variable> {
let mut arguments = ctx.arguments.into_iter();
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()
let command = arguments
.next()
.map(|a| a.to_string().into_owned())
.unwrap_or_default();
match Command::new(&command)
.args(
arguments
.next()
.map(|a| a.into_string_array())
.unwrap_or_default(),
)
.output()
{
Ok(result) => result.status.success(),
Err(err) => {
tracing::warn!(
context = "sieve",
event = "execute-failed",
reason = %err,
);
false
}
Ok(result) => Ok(result.status.success()),
Err(err) => Err(trc::SieveEvent::RuntimeError
.ctx(trc::Key::Path, command)
.reason(err)
.details("Failed to execute command")),
}
})
.await
.unwrap_or_default()
.into()
.map_err(|err| {
trc::EventType::Server(trc::ServerEvent::ThreadError)
.reason(err)
.caused_by(trc::location!())
.details("Join Error")
})?
.map(Into::into)
}

6
crates/common/src/scripts/plugins/headers.rs
View File

@@ -14,8 +14,8 @@ pub fn register(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("add_header", plugin_id, 2);
}
pub fn exec(ctx: PluginContext<'_>) -> Variable {
if let (Variable::String(name), Variable::String(value)) =
pub fn exec(ctx: PluginContext<'_>) -> trc::Result<Variable> {
Ok(if let (Variable::String(name), Variable::String(value)) =
(&ctx.arguments[0], &ctx.arguments[1])
{
ctx.modifications.push(ScriptModification::AddHeader {
@@ -26,5 +26,5 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
} else {
false
}
.into()
.into())
}

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

@@ -15,7 +15,7 @@ pub fn register_header(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("http_header", plugin_id, 4);
}
pub async fn exec_header(ctx: PluginContext<'_>) -> Variable {
pub async fn exec_header(ctx: PluginContext<'_>) -> trc::Result<Variable> {
let url = ctx.arguments[0].to_string();
let header = ctx.arguments[1].to_string();
let agent = ctx.arguments[2].to_string();
@@ -23,30 +23,36 @@ pub async fn exec_header(ctx: PluginContext<'_>) -> Variable {
#[cfg(feature = "test_mode")]
if url.contains("redirect.") {
return Variable::from(url.split_once("/?").unwrap().1.to_string());
return Ok(Variable::from(url.split_once("/?").unwrap().1.to_string()));
}
if let Ok(client) = reqwest::Client::builder()
reqwest::Client::builder()
.user_agent(agent.as_ref())
.timeout(Duration::from_millis(timeout))
.redirect(Policy::none())
.danger_accept_invalid_certs(true)
.build()
{
client
.get(url.as_ref())
.send()
.await
.ok()
.and_then(|response| {
response
.headers()
.get(header.as_ref())
.and_then(|h| h.to_str().ok())
.map(|h| Variable::from(h.to_string()))
})
.unwrap_or_default()
} else {
false.into()
}
.map_err(|err| {
trc::SieveEvent::RuntimeError
.into_err()
.reason(err)
.details("Failed to build request")
})?
.get(url.as_ref())
.send()
.await
.map_err(|err| {
trc::SieveEvent::RuntimeError
.into_err()
.reason(err)
.details("Failed to send request")
})
.map(|response| {
response
.headers()
.get(header.as_ref())
.and_then(|h| h.to_str().ok())
.map(|h| Variable::from(h.to_string()))
.unwrap_or_default()
})
}

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

@@ -38,107 +38,113 @@ pub fn register_local_domain(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("is_local_domain", plugin_id, 2);
}
pub async fn exec(ctx: PluginContext<'_>) -> Variable {
pub async fn exec(ctx: PluginContext<'_>) -> trc::Result<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),
};
}
.ok_or_else(|| {
trc::SieveEvent::RuntimeError
.ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned())
.details("Unknown store")
})?;
if let Some(store) = store {
match &ctx.arguments[1] {
Variable::Array(items) => {
for item in items.iter() {
if !item.is_empty()
&& store
.key_exists(item.to_string().into_owned().into_bytes())
.await
.unwrap_or(false)
{
return true.into();
}
Ok(match &ctx.arguments[1] {
Variable::Array(items) => {
for item in items.iter() {
if !item.is_empty()
&& store
.key_exists(item.to_string().into_owned().into_bytes())
.await?
{
return Ok(true.into());
}
false
}
v if !v.is_empty() => store
.key_exists(v.to_string().into_owned().into_bytes())
.await
.unwrap_or(false),
_ => false,
false
}
} else {
tracing::debug!(
context = "sieve:lookup",
event = "failed",
reason = "Unknown lookup id",
lookup_id = ctx.arguments[0].to_string().as_ref(),
);
false
v if !v.is_empty() => {
store
.key_exists(v.to_string().into_owned().into_bytes())
.await?
}
_ => false,
}
.into()
.into())
}
pub async fn exec_get(ctx: PluginContext<'_>) -> Variable {
let store = match &ctx.arguments[0] {
pub async fn exec_get(ctx: PluginContext<'_>) -> trc::Result<Variable> {
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 {
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()
} else {
tracing::debug!(
context = "sieve:key_get",
event = "failed",
reason = "Unknown store or lookup id",
lookup_id = ctx.arguments[0].to_string().as_ref(),
);
Variable::default()
}
.ok_or_else(|| {
trc::SieveEvent::RuntimeError
.ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned())
.details("Unknown store")
})?
.key_get::<VariableWrapper>(ctx.arguments[1].to_string().into_owned().into_bytes())
.await
.map(|v| v.map(|v| v.into_inner()).unwrap_or_default())
}
pub async fn exec_set(ctx: PluginContext<'_>) -> Variable {
let store = match &ctx.arguments[0] {
pub async fn exec_set(ctx: PluginContext<'_>) -> trc::Result<Variable> {
let expires = match &ctx.arguments[3] {
Variable::Integer(v) => Some(*v as u64),
Variable::Float(v) => Some(*v as u64),
_ => None,
};
match &ctx.arguments[0] {
Variable::String(v) if !v.is_empty() => ctx.core.storage.lookups.get(v.as_ref()),
_ => Some(&ctx.core.storage.lookup),
};
}
.ok_or_else(|| {
trc::SieveEvent::RuntimeError
.ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned())
.details("Unknown 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()
} else {
vec![]
},
expires,
)
.await
.map(|_| true.into())
}
if let Some(store) = store {
let expires = match &ctx.arguments[3] {
Variable::Integer(v) => Some(*v as u64),
Variable::Float(v) => Some(*v as u64),
_ => None,
};
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);
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()
} else {
vec![]
},
expires,
)
.await
.is_ok()
.into()
} else {
tracing::warn!(
context = "sieve:key_set",
event = "failed",
reason = "Unknown store id",
store_id = ctx.arguments[0].to_string().as_ref(),
);
Variable::default()
let mut _lock = ctx.cache.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)
}
}
}
}
pub async fn exec_remote(ctx: PluginContext<'_>) -> Variable {
async fn exec_remote_(ctx: &PluginContext<'_>) -> trc::Result<Variable> {
let resource = ctx.arguments[0].to_string();
let item = ctx.arguments[1].to_string();
@@ -147,22 +153,21 @@ pub async fn exec_remote(ctx: PluginContext<'_>) -> Variable {
if (resource.contains("open") && item.contains("open"))
|| (resource.contains("tank") && item.contains("tank"))
{
return true.into();
return Ok(true.into());
}
}
if resource.is_empty() || item.is_empty() {
return false.into();
return Ok(false.into());
}
const TIMEOUT: Duration = Duration::from_secs(45);
const RETRY: Duration = Duration::from_secs(3600);
const MAX_ENTRY_SIZE: usize = 256;
const MAX_ENTRIES: usize = 100000;
match ctx.cache.remote_lists.read().get(resource.as_ref()) {
Some(remote_list) if remote_list.expires < Instant::now() => {
return remote_list.entries.contains(item.as_ref()).into()
return Ok(remote_list.entries.contains(item.as_ref()).into())
}
_ => {}
}
@@ -206,7 +211,7 @@ pub async fn exec_remote(ctx: PluginContext<'_>) -> Variable {
}
}
match reqwest::Client::builder()
let response = reqwest::Client::builder()
.timeout(TIMEOUT)
.user_agent(USER_AGENT)
.build()
@@ -214,181 +219,140 @@ pub async fn exec_remote(ctx: PluginContext<'_>) -> Variable {
.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[..])
};
.map_err(|err| {
trc::SieveEvent::RuntimeError
.into_err()
.reason(err)
.ctx(trc::Key::Url, resource.to_string())
.details("Failed to build request")
})?;
// Lock remote list for writing
let mut _lock = ctx.cache.remote_lists.write();
let list = _lock
.entry(resource.to_string())
.or_insert_with(|| RemoteList {
entries: HashSet::new(),
expires: Instant::now(),
});
if response.status().is_success() {
let bytes = response.bytes().await.map_err(|err| {
trc::SieveEvent::RuntimeError
.into_err()
.reason(err)
.ctx(trc::Key::Url, resource.to_string())
.details("Failed to fetch resource")
})?;
// Make sure that the list is still expired
if list.expires > Instant::now() {
return list.entries.contains(item.as_ref()).into();
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.cache.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 (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();
for ch in line_.chars() {
if ch != '"' {
if ch == *separator && !in_quote {
if col_num == *column {
break;
} else {
col_num += 1;
}
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);
}
}
_ => (),
} else if col_num == *column {
entry.push(ch);
if entry.len() > MAX_ENTRY_SIZE {
break;
}
}
Err(err) => {
tracing::warn!(
context = "sieve:key_exists_http",
event = "failed",
resource = resource.as_ref(),
reason = %err,
);
break;
}
}
if list.entries.len() == MAX_ENTRIES {
break;
} else {
in_quote = !in_quote;
}
}
tracing::debug!(
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();
if !entry.is_empty() {
list.entries.insert(entry);
}
}
Err(err) => {
tracing::warn!(
_ => (),
}
context = "sieve:key_exists_http",
event = "failed",
resource = resource.as_ref(),
reason = %err,
);
}
if list.entries.len() == MAX_ENTRIES {
break;
}
}
Ok(response) => {
tracing::warn!(
context = "sieve:key_exists_http",
event = "failed",
resource = resource.as_ref(),
status = %response.status(),
);
}
Err(err) => {
tracing::warn!(
trc::event!(
Spam(trc::SpamEvent::ListUpdated),
Url = resource.as_ref().to_string(),
Count = list.entries.len(),
);
context = "sieve:key_exists_http",
event = "failed",
resource = resource.as_ref(),
reason = %err,
);
}
}
// Something went wrong, try again in one hour
let mut _lock = ctx.cache.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()
// Update expiration
list.expires = Instant::now() + expires;
return Ok(list.entries.contains(item.as_ref()).into());
} else {
list.expires = Instant::now() + RETRY;
false.into()
trc::bail!(trc::SieveEvent::RuntimeError
.into_err()
.ctx(trc::Key::Status, 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<'_>) -> Variable {
pub async fn exec_local_domain(ctx: PluginContext<'_>) -> trc::Result<Variable> {
let domain = ctx.arguments[0].to_string();
if !domain.is_empty() {
let directory = match &ctx.arguments[0] {
return match &ctx.arguments[0] {
Variable::String(v) if !v.is_empty() => ctx.core.storage.directories.get(v.as_ref()),
_ => Some(&ctx.core.storage.directory),
};
if let Some(directory) = directory {
return directory
.is_local_domain(domain.as_ref())
.await
.unwrap_or_default()
.into();
} else {
tracing::warn!(
context = "sieve:is_local_domain",
event = "failed",
reason = "Unknown directory",
lookup_id = ctx.arguments[0].to_string().as_ref(),
);
}
.ok_or_else(|| {
trc::SieveEvent::RuntimeError
.ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned())
.details("Unknown directory")
})?
.is_local_domain(domain.as_ref())
.await
.map(Into::into);
}
Variable::default()
Ok(Variable::default())
}
#[derive(Debug, PartialEq, Eq)]

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

@@ -78,7 +78,8 @@ impl Core {
return test_print(ctx);
}
match id {
let session_id = ctx.session_id;
let result = match id {
0 => query::exec(ctx).await,
1 => exec::exec(ctx).await,
2 => lookup::exec(ctx).await,
@@ -98,8 +99,17 @@ impl Core {
16 => text::exec_tokenize(ctx),
17 => text::exec_domain_part(ctx),
_ => unreachable!(),
};
match result {
Ok(result) => result.into(),
Err(err) => {
trc::error!(err
.ctx(trc::Key::SessionId, session_id)
.details("Sieve runtime error"));
Input::FncResult(Variable::default())
}
}
.into()
}
}

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

@@ -35,7 +35,7 @@ pub fn register(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("pyzor_check", plugin_id, 2);
}
pub async fn exec(ctx: PluginContext<'_>) -> Variable {
pub async fn exec(ctx: PluginContext<'_>) -> trc::Result<Variable> {
// Make sure there is at least one text part
if !ctx
.message
@@ -43,7 +43,7 @@ pub async fn exec(ctx: PluginContext<'_>) -> Variable {
.iter()
.any(|p| matches!(p.body, PartType::Text(_) | PartType::Html(_)))
{
return Variable::default();
return Ok(Variable::default());
}
// Hash message
@@ -54,43 +54,43 @@ pub async fn exec(ctx: PluginContext<'_>) -> Variable {
#[cfg(feature = "test_mode")]
{
if request.contains("b5b476f0b5ba6e1c038361d3ded5818dd39c90a2") {
return PyzorResponse {
return Ok(PyzorResponse {
code: 200,
count: 1000,
wl_count: 0,
}
.into();
.into());
} else if request.contains("d67d4b8bfc3860449e3418bb6017e2612f3e2a99") {
return PyzorResponse {
return Ok(PyzorResponse {
code: 200,
count: 60,
wl_count: 10,
}
.into();
.into());
} else if request.contains("81763547012b75e57a20d18ce0b93014208cdfdb") {
return PyzorResponse {
return Ok(PyzorResponse {
code: 200,
count: 50,
wl_count: 20,
}
.into();
.into());
}
}
let address = ctx.arguments[0].to_string();
let timeout = Duration::from_secs((ctx.arguments[1].to_integer() as u64).clamp(5, 60));
// Send message to address
match pyzor_send_message(address.as_ref(), timeout, &request).await {
Ok(response) => response.into(),
Err(err) => {
tracing::debug!(
context = "sieve:pyzor_check",
event = "failed",
reason = %err,
);
Variable::default()
}
}
pyzor_send_message(address.as_ref(), timeout, &request)
.await
.map(Into::into)
.map_err(|err| {
trc::SpamEvent::PyzorError
.into_err()
.ctx(trc::Key::Url, address.to_string())
.reason(err)
.details("Pyzor failed")
})
}
impl From<PyzorResponse> for Variable {

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

@@ -16,34 +16,24 @@ pub fn register(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("query", plugin_id, 3);
}
pub async fn exec(ctx: PluginContext<'_>) -> Variable {
pub async fn exec(ctx: PluginContext<'_>) -> trc::Result<Variable> {
// Obtain store name
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),
};
let store = if let Some(store) = store {
store
} else {
tracing::warn!(
context = "sieve:query",
event = "failed",
reason = "Unknown store",
store = ctx.arguments[0].to_string().as_ref(),
);
return false.into();
};
}
.ok_or_else(|| {
trc::SieveEvent::RuntimeError
.ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned())
.details("Unknown store")
})?;
// Obtain query string
let query = ctx.arguments[1].to_string();
if query.is_empty() {
tracing::warn!(
context = "sieve:query",
event = "invalid",
reason = "Empty query string",
);
return false.into();
trc::bail!(trc::SieveEvent::RuntimeError
.ctx(trc::Key::Id, ctx.arguments[0].to_string().into_owned())
.details("Empty query string"));
}
// Obtain arguments
@@ -58,43 +48,40 @@ pub async fn exec(ctx: PluginContext<'_>) -> Variable {
.get(..6)
.map_or(false, |q| q.eq_ignore_ascii_case(b"SELECT"))
{
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;
match row.len().cmp(&1) {
Ordering::Equal if !matches!(row.first(), Some(Value::Null)) => {
row.pop().map(into_sieve_value).unwrap()
}
Ordering::Less => Variable::default(),
_ => Variable::Array(
row.into_iter()
.map(into_sieve_value)
.collect::<Vec<_>>()
.into(),
),
let mut rows = store.query::<Rows>(&query, arguments).await?;
Ok(match rows.rows.len().cmp(&1) {
Ordering::Equal => {
let mut row = rows.rows.pop().unwrap().values;
match row.len().cmp(&1) {
Ordering::Equal if !matches!(row.first(), Some(Value::Null)) => {
row.pop().map(into_sieve_value).unwrap()
}
Ordering::Less => Variable::default(),
_ => Variable::Array(
row.into_iter()
.map(into_sieve_value)
.collect::<Vec<_>>()
.into(),
),
}
Ordering::Less => Variable::default(),
Ordering::Greater => rows
.rows
.into_iter()
.map(|r| {
Variable::Array(
r.values
.into_iter()
.map(into_sieve_value)
.collect::<Vec<_>>()
.into(),
)
})
.collect::<Vec<_>>()
.into(),
}
} else {
false.into()
}
Ordering::Less => Variable::default(),
Ordering::Greater => rows
.rows
.into_iter()
.map(|r| {
Variable::Array(
r.values
.into_iter()
.map(into_sieve_value)
.collect::<Vec<_>>()
.into(),
)
})
.collect::<Vec<_>>()
.into(),
})
} else {
store.query::<usize>(&query, arguments).await.is_ok().into()
Ok(store.query::<usize>(&query, arguments).await.is_ok().into())
}
}

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

@@ -20,18 +20,18 @@ pub fn register_domain_part(plugin_id: u32, fnc_map: &mut FunctionMap) {
fnc_map.set_external_function("domain_part", plugin_id, 2);
}
pub fn exec_tokenize(ctx: PluginContext<'_>) -> Variable {
pub fn exec_tokenize(ctx: PluginContext<'_>) -> trc::Result<Variable> {
let mut v = ctx.arguments;
let (urls, urls_without_scheme, emails) = match v[1].to_string().as_ref() {
"html" => return html_to_tokens(v[0].to_string().as_ref()).into(),
"words" => return tokenize_words(&v[0]),
"html" => return Ok(html_to_tokens(v[0].to_string().as_ref()).into()),
"words" => return Ok(tokenize_words(&v[0])),
"uri" | "url" => (true, true, true),
"uri_strict" | "url_strict" => (true, false, false),
"email" => (false, false, true),
_ => return Variable::default(),
_ => return Ok(Variable::default()),
};
match v.remove(0) {
Ok(match v.remove(0) {
v @ (Variable::String(_) | Variable::Array(_)) => {
TypesTokenizer::new(v.to_string().as_ref(), &ctx.core.smtp.resolvers.psl)
.tokenize_numbers(false)
@@ -50,19 +50,19 @@ pub fn exec_tokenize(ctx: PluginContext<'_>) -> Variable {
.into()
}
v => v,
}
})
}
pub fn exec_domain_part(ctx: PluginContext<'_>) -> Variable {
pub fn exec_domain_part(ctx: PluginContext<'_>) -> trc::Result<Variable> {
let v = ctx.arguments;
let part = match v[1].to_string().as_ref() {
"sld" => DomainPart::Sld,
"tld" => DomainPart::Tld,
"host" => DomainPart::Host,
_ => return Variable::default(),
_ => return Ok(Variable::default()),
};
v[0].transform(|domain| {
Ok(v[0].transform(|domain| {
ctx.core
.smtp
.resolvers
@@ -70,5 +70,5 @@ pub fn exec_domain_part(ctx: PluginContext<'_>) -> Variable {
.domain_part(domain, part)
.map(Variable::from)
.unwrap_or_default()
})
}))
}