g3proxy: do more clean shutdown in interception

bytedance · Jan 19, 2025 · a8245f2 · a8245f2
1 parent 1be1807
commit a8245f2
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 8 deletions.
diff --git a/g3proxy/src/inspect/http/v1/mod.rs b/g3proxy/src/inspect/http/v1/mod.rs
@@ -19,7 +19,7 @@ use std::sync::Arc;
 use async_recursion::async_recursion;
 use http::{Method, Version};
 use slog::slog_info;
-use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 
 use g3_dpi::Protocol;
 use g3_io_ext::{FlexBufReader, LimitedBufReadExt};
@@ -150,7 +150,11 @@ where
 
                 r = req_acceptor.accept() => match r {
                     Some(r) => r,
-                    None => return Ok(None),
+                    None => {
+                        let _ = rsp_io.ups_w.shutdown().await;
+                        let _ = rsp_io.clt_w.shutdown().await;
+                        return Ok(None);
+                    }
                 },
                 r = rsp_io.ups_r.fill_wait_data() => {
                     req_acceptor.close();
@@ -164,7 +168,10 @@ where
 
             self.req_id += 1;
             match r {
-                HttpRecvRequest::ClientConnectionClosed => return Ok(None),
+                HttpRecvRequest::ClientConnectionClosed => {
+                    let _ = rsp_io.ups_w.shutdown().await;
+                    return Ok(None);
+                }
                 HttpRecvRequest::ClientConnectionError(e) => return Err(e),
                 HttpRecvRequest::ClientRequestError(e) => {
                     if let Some(rsp) =

diff --git a/g3proxy/src/inspect/imap/mod.rs b/g3proxy/src/inspect/imap/mod.rs
@@ -393,10 +393,14 @@ where
             CloseReason::Client => {
                 self.handle_client_logout(&mut clt_w, &mut ups_r, &mut relay_buf.rsp_recv_buf)
                     .await?;
+                let _ = ups_w.shutdown().await;
+                let _ = clt_w.shutdown().await;
                 Ok(())
             }
             CloseReason::Server => {
                 self.mark_close_by_server();
+                let _ = ups_w.shutdown().await;
+                let _ = clt_w.shutdown().await;
                 Ok(())
             }
             CloseReason::Local(e) => {

diff --git a/g3proxy/src/inspect/smtp/mod.rs b/g3proxy/src/inspect/smtp/mod.rs
@@ -342,7 +342,11 @@ where
                 )
                 .await?;
             match next_action {
-                ForwardNextAction::Quit => return Ok(None),
+                ForwardNextAction::Quit => {
+                    let _ = ups_w.shutdown().await;
+                    let _ = clt_w.shutdown().await;
+                    return Ok(None);
+                }
                 ForwardNextAction::StartTls => {
                     return if let Some(tls_interception) = self.ctx.tls_interception() {
                         let mut start_tls_obj =
@@ -393,6 +397,8 @@ where
                         )
                         .await?;
                     if transaction.quit() {
+                        let _ = ups_w.shutdown().await;
+                        let _ = clt_w.shutdown().await;
                         return Ok(None);
                     }
                 }

diff --git a/g3proxy/src/module/http_forward/connection/eof_poller.rs b/g3proxy/src/module/http_forward/connection/eof_poller.rs
@@ -40,16 +40,12 @@ impl HttpConnectionEofCheck {
             _ = conn.1.fill_wait_data() => {
                 // close early when EOF or unexpected data, to avoid waiting at other side
                 wait_channel.close();
-                // make sure we correctly shutdown tls connection
-                // FIXME use async drop at escaper side when supported
                 let _ = conn.0.shutdown().await;
             }
             v = &mut wait_channel => {
                 if matches!(v, Ok(true)) {
                     let _ = send_channel.send(conn);
                 } else {
-                    // make sure we correctly shutdown tls connection
-                    // FIXME use async drop at escaper side when supported
                     let _ = conn.0.shutdown().await;
                 }
             }