Include queueId in MTA Hooks (fixes #708)

stalwartlabs · Aug 23, 2024 · 18a24f7 · 18a24f7
1 parent 0186c21
commit 18a24f7
Show file tree

Hide file tree

Showing 8 changed files with 18 additions and 10 deletions.
diff --git a/crates/smtp/src/inbound/data.rs b/crates/smtp/src/inbound/data.rs
@@ -398,7 +398,7 @@ impl<T: SessionStream> Session<T> {
 
         // Run MTA Hooks
         match self
-            .run_mta_hooks(Stage::Data, (&auth_message).into())
+            .run_mta_hooks(Stage::Data, (&auth_message).into(), message_id.into())
             .await
         {
             Ok(modifications_) => {

diff --git a/crates/smtp/src/inbound/ehlo.rs b/crates/smtp/src/inbound/ehlo.rs
@@ -115,7 +115,7 @@ impl<T: SessionStream> Session<T> {
             }
 
             // MTAHook filtering
-            if let Err(message) = self.run_mta_hooks(Stage::Ehlo, None).await {
+            if let Err(message) = self.run_mta_hooks(Stage::Ehlo, None, None).await {
                 self.data.mail_from = None;
                 self.data.helo_domain = prev_helo_domain;
                 self.data.spf_ehlo = None;

diff --git a/crates/smtp/src/inbound/hooks/message.rs b/crates/smtp/src/inbound/hooks/message.rs
@@ -24,15 +24,17 @@ use crate::{
         milter::Modification,
         FilterResponse,
     },
+    queue::QueueId,
 };
 
-use super::{client::send_mta_hook_request, Action, Response};
+use super::{client::send_mta_hook_request, Action, Queue, Response};
 
 impl<T: SessionStream> Session<T> {
     pub async fn run_mta_hooks(
         &self,
         stage: Stage,
         message: Option<&AuthenticatedMessage<'_>>,
+        queue_id: Option<QueueId>,
     ) -> Result<Vec<Modification>, FilterResponse> {
         let mta_hooks = &self.core.core.smtp.session.hooks;
         if mta_hooks.is_empty() {
@@ -53,7 +55,7 @@ impl<T: SessionStream> Session<T> {
             }
 
             let time = Instant::now();
-            match self.run_mta_hook(stage, mta_hook, message).await {
+            match self.run_mta_hook(stage, mta_hook, message, queue_id).await {
                 Ok(response) => {
                     trc::event!(
                         MtaHook(match response.action {
@@ -174,6 +176,7 @@ impl<T: SessionStream> Session<T> {
         stage: Stage,
         mta_hook: &MTAHook,
         message: Option<&AuthenticatedMessage<'_>>,
+        queue_id: Option<QueueId>,
     ) -> Result<Response, String> {
         // Build request
         let (tls_version, tls_cipher) = self.stream.tls_version_and_cipher();
@@ -210,7 +213,9 @@ impl<T: SessionStream> Session<T> {
                     port: self.data.local_port,
                     ip: self.data.local_ip.to_string().into(),
                 },
-                queue: None,
+                queue: queue_id.map(|id| Queue {
+                    id: format!("{:x}", id),
+                }),
                 protocol: Protocol { version: 1 },
             },
             envelope: self.data.mail_from.as_ref().map(|from| Envelope {

diff --git a/crates/smtp/src/inbound/mail.rs b/crates/smtp/src/inbound/mail.rs
@@ -169,7 +169,7 @@ impl<T: SessionStream> Session<T> {
         }
 
         // MTAHook filtering
-        if let Err(message) = self.run_mta_hooks(Stage::Mail, None).await {
+        if let Err(message) = self.run_mta_hooks(Stage::Mail, None, None).await {
             self.data.mail_from = None;
             return self.write(message.message.as_bytes()).await;
         }

diff --git a/crates/smtp/src/inbound/rcpt.rs b/crates/smtp/src/inbound/rcpt.rs
@@ -139,7 +139,7 @@ impl<T: SessionStream> Session<T> {
             }
 
             // MTAHook filtering
-            if let Err(message) = self.run_mta_hooks(Stage::Rcpt, None).await {
+            if let Err(message) = self.run_mta_hooks(Stage::Rcpt, None, None).await {
                 self.data.rcpt_to.pop();
                 return self.write(message.message.as_bytes()).await;
             }

diff --git a/crates/smtp/src/inbound/spawn.rs b/crates/smtp/src/inbound/spawn.rs
@@ -116,7 +116,7 @@ impl<T: SessionStream> Session<T> {
         }
 
         // MTAHook filtering
-        if let Err(message) = self.run_mta_hooks(Stage::Connect, None).await {
+        if let Err(message) = self.run_mta_hooks(Stage::Connect, None, None).await {
             let _ = self.write(message.message.as_bytes()).await;
             return false;
         }

diff --git a/tests/src/smtp/inbound/antispam.rs b/tests/src/smtp/inbound/antispam.rs
@@ -409,7 +409,10 @@ async fn antispam() {
             // Run script
             let core_ = core.clone();
             let script = script.clone();
-            match core_.run_script(script, params, 0).await {
+            match core_
+                .run_script("test".to_string(), script, params, 0)
+                .await
+            {
                 ScriptResult::Accept { modifications } => {
                     if modifications.len() != expected_headers.len() {
                         panic!(

diff --git a/tests/src/smtp/inbound/scripts.rs b/tests/src/smtp/inbound/scripts.rs
@@ -177,7 +177,7 @@ async fn sieve_scripts() {
             .with_envelope(&core.core, &session, 0)
             .await;
         let core_ = core.clone();
-        match core_.run_script(script, params, 0).await {
+        match core_.run_script(name.to_string(), script, params, 0).await {
             ScriptResult::Accept { .. } => (),
             ScriptResult::Reject(message) => panic!("{}", message),
             err => {