g3proxy: allow SMTP EHLO at later stage

bytedance · Jun 4, 2024 · 006016e · 006016e
1 parent baf708c
commit 006016e
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 10 deletions.
diff --git a/g3proxy/src/inspect/smtp/forward.rs b/g3proxy/src/inspect/smtp/forward.rs
@@ -23,13 +23,17 @@ use g3_io_ext::{LimitedWriteExt, LineRecvBuf};
 use g3_smtp_proto::command::{Command, MailParam};
 use g3_smtp_proto::response::{ReplyCode, ResponseEncoder, ResponseParser};
 
-use super::{CommandLineRecvExt, ResponseLineRecvExt, ResponseParseExt, SmtpRelayBuf};
+use super::{
+    CommandLineRecvExt, InitializedExtensions, Initiation, ResponseLineRecvExt, ResponseParseExt,
+    SmtpRelayBuf,
+};
 use crate::serve::{ServerTaskError, ServerTaskResult};
 
 pub(super) enum ForwardNextAction {
     Quit,
     StartTls,
     ReverseConnection,
+    SetExtensions(InitializedExtensions),
     MailTransport(MailParam),
 }
 
@@ -83,7 +87,6 @@ impl<'a> Forward<'a> {
                     |cmd| {
                         match &cmd {
                             Command::Hello(_)
-                            | Command::ExtendHello(_)
                             | Command::Recipient(_)
                             | Command::Data
                             | Command::BinaryData(_)
@@ -144,6 +147,17 @@ impl<'a> Forward<'a> {
                         return Ok(ForwardNextAction::ReverseConnection);
                     }
                 }
+                Command::ExtendHello(_host) => {
+                    let mut initialization = Initiation::new(self.config, self.local_ip, true);
+                    if initialization
+                        .recv_relay_check_rsp(&mut buf.rsp_recv_buf, ups_r, clt_w)
+                        .await?
+                        .is_some()
+                    {
+                        let (_, extensions) = initialization.into_parts();
+                        return Ok(ForwardNextAction::SetExtensions(extensions));
+                    }
+                }
                 Command::Mail(param) => {
                     let rsp = self.recv_relay_rsp(buf, ups_r, clt_w).await?;
                     if rsp == ReplyCode::OK {

diff --git a/g3proxy/src/inspect/smtp/initiation.rs b/g3proxy/src/inspect/smtp/initiation.rs
@@ -132,7 +132,7 @@ impl<'a> Initiation<'a> {
         }
     }
 
-    async fn recv_relay_check_rsp<CW, UR>(
+    pub(super) async fn recv_relay_check_rsp<CW, UR>(
         &mut self,
         rsp_recv_buf: &mut LineRecvBuf<{ ResponseParser::MAX_LINE_SIZE }>,
         ups_r: &mut UR,

diff --git a/g3proxy/src/inspect/smtp/mod.rs b/g3proxy/src/inspect/smtp/mod.rs
@@ -39,7 +39,7 @@ mod ending;
 use ending::{EndQuitServer, EndWaitClient};
 
 mod initiation;
-use initiation::Initiation;
+use initiation::{InitializedExtensions, Initiation};
 
 mod forward;
 use forward::{Forward, ForwardNextAction};
@@ -242,17 +242,14 @@ where
         initiation
             .relay(&mut clt_r, &mut clt_w, &mut ups_r, &mut ups_w)
             .await?;
-        let (client_host, server_ext) = initiation.into_parts();
+        let (client_host, mut server_ext) = initiation.into_parts();
         self.client_host = Some(client_host);
 
-        let allow_odmr = server_ext.allow_odmr(interception_config);
-        let allow_starttls = server_ext.allow_starttls(self.from_starttls);
-        let allow_chunking = server_ext.allow_chunking();
-        let allow_burl = server_ext.allow_burl(interception_config);
-
         let mut relay_buf = SmtpRelayBuf::default();
 
         loop {
+            let allow_odmr = server_ext.allow_odmr(interception_config);
+            let allow_starttls = server_ext.allow_starttls(self.from_starttls);
             let mut forward =
                 Forward::new(interception_config, local_ip, allow_odmr, allow_starttls);
             let next_action = forward
@@ -304,7 +301,11 @@ where
                     .await
                     .map(|_| None);
                 }
+                ForwardNextAction::SetExtensions(ext) => server_ext = ext,
                 ForwardNextAction::MailTransport(param) => {
+                    let allow_chunking = server_ext.allow_chunking();
+                    let allow_burl = server_ext.allow_burl(interception_config);
+
                     let mut transaction =
                         Transaction::new(&self.ctx, local_ip, allow_chunking, allow_burl, param);
                     transaction