Skip to content

Commit

Permalink
Added expression functions.
Browse files Browse the repository at this point in the history
  • Loading branch information
@mdecimus
mdecimus committed Feb 16, 2024
1 parent 7f5538a commit afe10e6
Show file tree
Hide file tree
Showing 2 changed files with 151 additions and 8 deletions.
81 changes: 79 additions & 2 deletions crates/smtp/src/core/eval.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,11 @@ pub const F_IS_LOCAL_DOMAIN: u32 = 0;
pub const F_IS_LOCAL_ADDRESS: u32 = 1;
pub const F_KEY_GET: u32 = 2;
pub const F_KEY_EXISTS: u32 = 3;
pub const F_SQL_QUERY: u32 = 4;
pub const F_DNS_QUERY: u32 = 5;
pub const F_KEY_SET: u32 = 4;
pub const F_COUNTER_INCR: u32 = 5;
pub const F_COUNTER_GET: u32 = 6;
pub const F_SQL_QUERY: u32 = 7;
pub const F_DNS_QUERY: u32 = 8;

pub const VARIABLES_MAP: &[(&str, u32)] = &[
("rcpt", V_RECIPIENT),
Expand All @@ -52,6 +55,9 @@ pub const FUNCTIONS_MAP: &[(&str, u32, u32)] = &[
("is_local_address", F_IS_LOCAL_ADDRESS, 2),
("key_get", F_KEY_GET, 2),
("key_exists", F_KEY_EXISTS, 2),
("key_set", F_KEY_SET, 3),
("counter_incr", F_COUNTER_INCR, 3),
("counter_get", F_COUNTER_GET, 2),
("dns_query", F_DNS_QUERY, 2),
("sql_query", F_SQL_QUERY, 3),
];
Expand Down Expand Up @@ -202,6 +208,73 @@ impl SMTP {
})
.into()
}
F_KEY_SET => {
let store = params.next_as_string();
let key = params.next_as_string();
let value = params.next_as_string();

self.get_lookup_store(store.as_ref())
.key_set(
key.into_owned().into_bytes(),
value.into_owned().into_bytes(),
None,
)
.await
.map(|_| true)
.unwrap_or_else(|err| {
tracing::warn!(
context = "eval_if",
event = "error",
property = property,
error = ?err,
"Failed to set key."
);

false
})
.into()
}
F_COUNTER_INCR => {
let store = params.next_as_string();
let key = params.next_as_string();
let value = params.next_as_integer();

self.get_lookup_store(store.as_ref())
.counter_incr(key.into_owned().into_bytes(), value, None)
.await
.map(Variable::Integer)
.unwrap_or_else(|err| {
tracing::warn!(
context = "eval_if",
event = "error",
property = property,
error = ?err,
"Failed to increment counter."
);

Variable::default()
})
}
F_COUNTER_GET => {
let store = params.next_as_string();
let key = params.next_as_string();

self.get_lookup_store(store.as_ref())
.counter_get(key.into_owned().into_bytes())
.await
.map(Variable::Integer)
.unwrap_or_else(|err| {
tracing::warn!(
context = "eval_if",
event = "error",
property = property,
error = ?err,
"Failed to increment counter."
);

Variable::default()
})
}
F_DNS_QUERY => self.dns_query(params).await,
F_SQL_QUERY => self.sql_query(params).await,
_ => Variable::default(),
Expand Down Expand Up @@ -450,6 +523,10 @@ impl<'x> FncParams<'x> {
self.params.next().unwrap().into_string()
}

pub fn next_as_integer(&mut self) -> i64 {
self.params.next().unwrap().to_integer().unwrap_or_default()
}

pub fn next(&mut self) -> Variable<'x> {
self.params.next().unwrap()
}
Expand Down
78 changes: 72 additions & 6 deletions tests/src/smtp/lookup/sql.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,16 @@
* for more details.
*/

use std::time::Duration;
use std::time::{Duration, Instant};

use directory::core::config::ConfigDirectory;
use mail_auth::MX;
use smtp_proto::{AUTH_LOGIN, AUTH_PLAIN};
use store::{config::ConfigStore, Store};
use utils::config::{if_block::IfBlock, Config};
use utils::{
config::{if_block::IfBlock, Config},
expr::Expression,
};

use crate::{
directory::DirectoryStore,
Expand All @@ -37,8 +41,9 @@ use crate::{
store::TempDir,
};
use smtp::{
config::{session::Mechanism, ConfigContext},
core::{Session, SMTP},
config::{map_expr_token, session::Mechanism, ConfigContext},
core::{eval::*, Session, SMTP},
queue::RecipientDomain,
};

const CONFIG: &str = r#"
Expand Down Expand Up @@ -76,9 +81,10 @@ type = "type"
#[tokio::test]
async fn lookup_sql() {
// Enable logging
/*tracing::subscriber::set_global_default(
/*let disable = true;
tracing::subscriber::set_global_default(
tracing_subscriber::FmtSubscriber::builder()
.with_max_level(tracing::Level::DEBUG)
.with_max_level(tracing::Level::TRACE)
.finish(),
)
.unwrap();*/
Expand All @@ -96,6 +102,14 @@ async fn lookup_sql() {
.await
.unwrap()
.directories;
core.resolvers.dns.mx_add(
"test.org",
vec![MX {
exchanges: vec!["mx.foobar.org".to_string()],
preference: 10,
}],
Instant::now() + Duration::from_secs(10),
);

// Obtain directory handle
let handle = DirectoryStore {
Expand Down Expand Up @@ -145,6 +159,58 @@ async fn lookup_sql() {
.unwrap();
}

// Test expression functions
for (expr, expected) in [
(
"sql_query('sql', 'SELECT description FROM domains WHERE name = ?', 'foobar.org')",
"Main domain",
),
("dns_query(rcpt_domain, 'mx')[0]", "mx.foobar.org"),
(
concat!(
"key_get('sql', 'hello') + '-' + key_exists('sql', 'hello') + '-' + ",
"key_set('sql', 'hello', 'world') + '-' + key_get('sql', 'hello') + ",
"'-' + key_exists('sql', 'hello')"
),
"0-0-1-world-1",
),
(
concat!(
"counter_get('sql', 'county') + '-' + counter_incr('sql', 'county', 1) + '-' ",
"+ counter_incr('sql', 'county', 1) + '-' + counter_get('sql', 'county')"
),
"0-0-0-2",
),
] {
let e = Expression::parse("test", expr, |name| {
map_expr_token::<Duration>(
name,
&[
V_RECIPIENT,
V_RECIPIENT_DOMAIN,
V_SENDER,
V_SENDER_DOMAIN,
V_MX,
V_HELO_DOMAIN,
V_AUTHENTICATED_AS,
V_LISTENER,
V_REMOTE_IP,
V_LOCAL_IP,
V_PRIORITY,
],
)
})
.unwrap();
assert_eq!(
core.eval_expr::<String, _>(&e, &RecipientDomain::new("test.org"), "text")
.await
.unwrap(),
expected,
"failed for '{}'",
expr
);
}

// Enable AUTH
let config = &mut core.session.config.auth;
config.directory = "\"'sql'\"".parse_if();
Expand Down

0 comments on commit afe10e6

Please sign in to comment.