diff --git a/crates/smtp/src/core/eval.rs b/crates/smtp/src/core/eval.rs index c1022817..c0479052 100644 --- a/crates/smtp/src/core/eval.rs +++ b/crates/smtp/src/core/eval.rs @@ -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), @@ -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), ]; @@ -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(), @@ -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() } diff --git a/tests/src/smtp/lookup/sql.rs b/tests/src/smtp/lookup/sql.rs index 7b327836..39ccdb37 100644 --- a/tests/src/smtp/lookup/sql.rs +++ b/tests/src/smtp/lookup/sql.rs @@ -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, @@ -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#" @@ -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();*/ @@ -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 { @@ -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::( + 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::(&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();