add extra task_conf param to escaper tcp API

bytedance · Oct 25, 2024 · 6211595 · 6211595
1 parent 1ce19fb
commit 6211595
Show file tree

Hide file tree

Showing 98 changed files with 1,499 additions and 1,327 deletions.
diff --git a/g3proxy/src/escape/comply_audit/mod.rs b/g3proxy/src/escape/comply_audit/mod.rs
@@ -22,21 +22,23 @@ use async_trait::async_trait;
 
 use g3_daemon::stat::remote::ArcTcpConnectionTaskRemoteStats;
 use g3_types::metrics::MetricsName;
-use g3_types::net::{Host, OpensslClientConfig, UpstreamAddr};
+use g3_types::net::UpstreamAddr;
 
 use super::{ArcEscaper, Escaper, EscaperInternal, RouteEscaperStats};
 use crate::audit::{AuditContext, AuditHandle};
 use crate::config::escaper::comply_audit::ComplyAuditEscaperConfig;
 use crate::config::escaper::{AnyEscaperConfig, EscaperConfig};
 use crate::module::ftp_over_http::{
-    AnyFtpConnectContextParam, ArcFtpTaskRemoteControlStats, ArcFtpTaskRemoteTransferStats,
-    BoxFtpConnectContext, BoxFtpRemoteConnection,
+    ArcFtpTaskRemoteControlStats, ArcFtpTaskRemoteTransferStats, BoxFtpConnectContext,
+    BoxFtpRemoteConnection,
 };
 use crate::module::http_forward::{
     ArcHttpForwardTaskRemoteStats, BoxHttpForwardConnection, BoxHttpForwardContext,
     RouteHttpForwardContext,
 };
-use crate::module::tcp_connect::{TcpConnectError, TcpConnectResult, TcpConnectTaskNotes};
+use crate::module::tcp_connect::{
+    TcpConnectError, TcpConnectResult, TcpConnectTaskConf, TcpConnectTaskNotes, TlsConnectTaskConf,
+};
 use crate::module::udp_connect::{
     ArcUdpConnectTaskRemoteStats, UdpConnectResult, UdpConnectTaskNotes,
 };
@@ -109,6 +111,7 @@ impl Escaper for ComplyAuditEscaper {
 
     async fn tcp_setup_connection<'a>(
         &'a self,
+        task_conf: &TcpConnectTaskConf<'_>,
         tcp_notes: &'a mut TcpConnectTaskNotes,
         task_notes: &'a ServerTaskNotes,
         task_stats: ArcTcpConnectionTaskRemoteStats,
@@ -118,26 +121,23 @@ impl Escaper for ComplyAuditEscaper {
         self.stats.add_request_passed();
         self._update_audit_context(audit_ctx);
         self.next
-            .tcp_setup_connection(tcp_notes, task_notes, task_stats, audit_ctx)
+            .tcp_setup_connection(task_conf, tcp_notes, task_notes, task_stats, audit_ctx)
             .await
     }
 
     async fn tls_setup_connection<'a>(
         &'a self,
+        task_conf: &TlsConnectTaskConf<'_>,
         tcp_notes: &'a mut TcpConnectTaskNotes,
         task_notes: &'a ServerTaskNotes,
         task_stats: ArcTcpConnectionTaskRemoteStats,
         audit_ctx: &'a mut AuditContext,
-        tls_config: &'a OpensslClientConfig,
-        tls_name: &'a Host,
     ) -> TcpConnectResult {
         tcp_notes.escaper.clone_from(&self.config.name);
         self.stats.add_request_passed();
         self._update_audit_context(audit_ctx);
         self.next
-            .tls_setup_connection(
-                tcp_notes, task_notes, task_stats, audit_ctx, tls_config, tls_name,
-            )
+            .tls_setup_connection(task_conf, tcp_notes, task_notes, task_stats, audit_ctx)
             .await
     }
 
@@ -175,12 +175,12 @@ impl Escaper for ComplyAuditEscaper {
     async fn new_ftp_connect_context<'a>(
         &'a self,
         escaper: ArcEscaper,
+        task_conf: &TcpConnectTaskConf<'_>,
         task_notes: &'a ServerTaskNotes,
-        upstream: &'a UpstreamAddr,
     ) -> BoxFtpConnectContext {
         self.stats.add_request_passed();
         self.next
-            .new_ftp_connect_context(Arc::clone(&escaper), task_notes, upstream)
+            .new_ftp_connect_context(Arc::clone(&escaper), task_conf, task_notes)
             .await
     }
 }
@@ -233,6 +233,7 @@ impl EscaperInternal for ComplyAuditEscaper {
 
     async fn _new_http_forward_connection<'a>(
         &'a self,
+        _task_conf: &TcpConnectTaskConf<'_>,
         tcp_notes: &'a mut TcpConnectTaskNotes,
         _task_notes: &'a ServerTaskNotes,
         _task_stats: ArcHttpForwardTaskRemoteStats,
@@ -243,18 +244,18 @@ impl EscaperInternal for ComplyAuditEscaper {
 
     async fn _new_https_forward_connection<'a>(
         &'a self,
+        _task_conf: &TlsConnectTaskConf<'_>,
         tcp_notes: &'a mut TcpConnectTaskNotes,
         _task_notes: &'a ServerTaskNotes,
         _task_stats: ArcHttpForwardTaskRemoteStats,
-        _tls_config: &'a OpensslClientConfig,
-        _tls_name: &'a Host,
     ) -> Result<BoxHttpForwardConnection, TcpConnectError> {
         tcp_notes.escaper.clone_from(&self.config.name);
         Err(TcpConnectError::MethodUnavailable)
     }
 
     async fn _new_ftp_control_connection<'a>(
         &'a self,
+        _task_conf: &TcpConnectTaskConf<'_>,
         tcp_notes: &'a mut TcpConnectTaskNotes,
         _task_notes: &'a ServerTaskNotes,
         _task_stats: ArcFtpTaskRemoteControlStats,
@@ -265,11 +266,12 @@ impl EscaperInternal for ComplyAuditEscaper {
 
     async fn _new_ftp_transfer_connection<'a>(
         &'a self,
+        _task_conf: &TcpConnectTaskConf<'_>,
         transfer_tcp_notes: &'a mut TcpConnectTaskNotes,
         _control_tcp_notes: &'a TcpConnectTaskNotes,
         _task_notes: &'a ServerTaskNotes,
         _task_stats: ArcFtpTaskRemoteTransferStats,
-        _context: AnyFtpConnectContextParam,
+        _ftp_server: &UpstreamAddr,
     ) -> Result<BoxFtpRemoteConnection, TcpConnectError> {
         transfer_tcp_notes.escaper.clone_from(&self.config.name);
         Err(TcpConnectError::MethodUnavailable)

diff --git a/g3proxy/src/escape/direct_fixed/ftp_connect/mod.rs b/g3proxy/src/escape/direct_fixed/ftp_connect/mod.rs
@@ -17,23 +17,27 @@
 use std::sync::Arc;
 
 use g3_io_ext::LimitedStream;
+use g3_types::net::UpstreamAddr;
 
 use super::DirectFixedEscaper;
 use crate::module::ftp_over_http::{
     ArcFtpTaskRemoteControlStats, ArcFtpTaskRemoteTransferStats, BoxFtpRemoteConnection,
     FtpControlRemoteWrapperStats, FtpTransferRemoteWrapperStats,
 };
-use crate::module::tcp_connect::{TcpConnectError, TcpConnectTaskNotes};
+use crate::module::tcp_connect::{TcpConnectError, TcpConnectTaskConf, TcpConnectTaskNotes};
 use crate::serve::ServerTaskNotes;
 
 impl DirectFixedEscaper {
     pub(super) async fn new_ftp_control_connection<'a>(
         &'a self,
+        task_conf: &TcpConnectTaskConf<'_>,
         tcp_notes: &'a mut TcpConnectTaskNotes,
         task_notes: &'a ServerTaskNotes,
         task_stats: ArcFtpTaskRemoteControlStats,
     ) -> Result<BoxFtpRemoteConnection, TcpConnectError> {
-        let stream = self.tcp_connect_to(tcp_notes, task_notes).await?;
+        let stream = self
+            .tcp_connect_to(task_conf, tcp_notes, task_notes)
+            .await?;
 
         let mut wrapper_stats = FtpControlRemoteWrapperStats::new(&self.stats, task_stats);
         wrapper_stats.push_user_io_stats(self.fetch_user_upstream_io_stats(task_notes));
@@ -53,13 +57,21 @@ impl DirectFixedEscaper {
 
     pub(super) async fn new_ftp_transfer_connection<'a>(
         &'a self,
+        task_conf: &TcpConnectTaskConf<'_>,
         transfer_tcp_notes: &'a mut TcpConnectTaskNotes,
         control_tcp_notes: &'a TcpConnectTaskNotes,
         task_notes: &'a ServerTaskNotes,
         task_stats: ArcFtpTaskRemoteTransferStats,
+        ftp_server: &UpstreamAddr,
     ) -> Result<BoxFtpRemoteConnection, TcpConnectError> {
         let stream = self
-            .tcp_connect_to_again(transfer_tcp_notes, control_tcp_notes, task_notes)
+            .tcp_connect_to_again(
+                task_conf,
+                ftp_server,
+                transfer_tcp_notes,
+                control_tcp_notes,
+                task_notes,
+            )
             .await?;
 
         let mut wrapper_stats = FtpTransferRemoteWrapperStats::new(&self.stats, task_stats);

diff --git a/g3proxy/src/escape/direct_fixed/http_forward/mod.rs b/g3proxy/src/escape/direct_fixed/http_forward/mod.rs
@@ -17,15 +17,16 @@
 use std::sync::Arc;
 
 use g3_io_ext::{AsyncStream, LimitedBufReader, LimitedWriter, NilLimitedReaderStats};
-use g3_types::net::{Host, OpensslClientConfig};
 
 use super::{DirectFixedEscaper, DirectFixedEscaperStats};
 use crate::log::escape::tls_handshake::TlsApplication;
 use crate::module::http_forward::{
     ArcHttpForwardTaskRemoteStats, BoxHttpForwardConnection, HttpForwardRemoteWrapperStats,
     HttpForwardTaskRemoteWrapperStats,
 };
-use crate::module::tcp_connect::{TcpConnectError, TcpConnectTaskNotes};
+use crate::module::tcp_connect::{
+    TcpConnectError, TcpConnectTaskConf, TcpConnectTaskNotes, TlsConnectTaskConf,
+};
 use crate::serve::ServerTaskNotes;
 
 mod reader;
@@ -37,11 +38,14 @@ pub(crate) use writer::DirectHttpForwardWriter;
 impl DirectFixedEscaper {
     pub(super) async fn http_forward_new_connection<'a>(
         &'a self,
+        task_conf: &TcpConnectTaskConf<'_>,
         tcp_notes: &'a mut TcpConnectTaskNotes,
         task_notes: &'a ServerTaskNotes,
         task_stats: ArcHttpForwardTaskRemoteStats,
     ) -> Result<BoxHttpForwardConnection, TcpConnectError> {
-        let stream = self.tcp_connect_to(tcp_notes, task_notes).await?;
+        let stream = self
+            .tcp_connect_to(task_conf, tcp_notes, task_notes)
+            .await?;
 
         let (ups_r, ups_w) = stream.into_split();
 
@@ -73,18 +77,16 @@ impl DirectFixedEscaper {
 
     pub(super) async fn https_forward_new_connection<'a>(
         &'a self,
+        task_conf: &TlsConnectTaskConf<'_>,
         tcp_notes: &'a mut TcpConnectTaskNotes,
         task_notes: &'a ServerTaskNotes,
         task_stats: ArcHttpForwardTaskRemoteStats,
-        tls_config: &'a OpensslClientConfig,
-        tls_name: &'a Host,
     ) -> Result<BoxHttpForwardConnection, TcpConnectError> {
         let tls_stream = self
             .tls_connect_to(
+                task_conf,
                 tcp_notes,
                 task_notes,
-                tls_config,
-                tls_name,
                 TlsApplication::HttpForward,
             )
             .await?;