g3proxy: change default idle check interval and max count value

bytedance · Feb 14, 2025 · a76bda6 · a76bda6
1 parent a393f2f
commit a76bda6
Show file tree

Hide file tree

Showing 32 changed files with 180 additions and 144 deletions.
diff --git a/g3proxy/src/audit/detour/mod.rs b/g3proxy/src/audit/detour/mod.rs
@@ -26,6 +26,7 @@ use g3_io_ext::IdleWheel;
 use g3_types::net::UpstreamAddr;
 
 use crate::config::audit::AuditStreamDetourConfig;
+use crate::config::server::ServerConfig;
 use crate::inspect::StreamInspectTaskNotes;
 
 mod connect;
@@ -64,6 +65,7 @@ pub(crate) struct StreamDetourContext<'a, SC> {
     protocol: Protocol,
     payload: Vec<u8>,
     request_timeout: Duration,
+    max_idle_count: usize,
 }
 
 impl<SC> StreamDetourContext<'_, SC> {
@@ -99,7 +101,14 @@ impl StreamDetourClient {
         task_notes: &'a StreamInspectTaskNotes,
         upstream: &'a UpstreamAddr,
         protocol: Protocol,
-    ) -> StreamDetourContext<'a, SC> {
+    ) -> StreamDetourContext<'a, SC>
+    where
+        SC: ServerConfig,
+    {
+        let max_idle_count = task_notes
+            .user()
+            .and_then(|u| u.task_max_idle_count())
+            .unwrap_or(server_config.task_max_idle_count());
         StreamDetourContext {
             server_config,
             server_quit_policy,
@@ -109,6 +118,7 @@ impl StreamDetourClient {
             protocol,
             payload: Vec::new(),
             request_timeout: self.config.request_timeout,
+            max_idle_count,
         }
     }
 

diff --git a/g3proxy/src/audit/detour/stream.rs b/g3proxy/src/audit/detour/stream.rs
@@ -175,16 +175,13 @@ where
                     if clt_to_d.is_idle() && d_to_clt.is_idle() && ups_to_d.is_idle() && d_to_ups.is_idle() {
                         idle_count += n;
 
-                        let quit = if let Some(user) = self.task_notes.user() {
+                        if let Some(user) = self.task_notes.user() {
                             if user.is_blocked() {
                                 return Err(ServerTaskError::CanceledAsUserBlocked);
                             }
-                            idle_count >= user.task_max_idle_count()
-                        } else {
-                            idle_count >= self.server_config.task_max_idle_count()
-                        };
+                        }
 
-                        if quit {
+                        if idle_count >= self.max_idle_count {
                             return Err(ServerTaskError::Idle(idle_interval.period(), idle_count));
                         }
                     } else {

diff --git a/g3proxy/src/auth/user.rs b/g3proxy/src/auth/user.rs
@@ -67,7 +67,7 @@ pub(crate) struct User {
 
 impl User {
     #[inline]
-    pub(crate) fn task_max_idle_count(&self) -> usize {
+    pub(crate) fn task_max_idle_count(&self) -> Option<usize> {
         self.config.task_idle_max_count
     }
 

diff --git a/g3proxy/src/config/auth/user/json.rs b/g3proxy/src/config/auth/user/json.rs
@@ -217,8 +217,9 @@ impl UserConfig {
                 Ok(())
             }
             "task_idle_max_count" => {
-                self.task_idle_max_count = g3_json::value::as_usize(v)
+                let count = g3_json::value::as_usize(v)
                     .context(format!("invalid usize value for key {k}"))?;
+                self.task_idle_max_count = Some(count);
                 Ok(())
             }
             "socks_use_udp_associate" => {

diff --git a/g3proxy/src/config/auth/user/mod.rs b/g3proxy/src/config/auth/user/mod.rs
@@ -74,7 +74,7 @@ pub(crate) struct UserConfig {
     pub(crate) http_user_agent_filter: Option<AclUserAgentRule>,
     pub(crate) resolve_strategy: Option<ResolveStrategy>,
     pub(crate) resolve_redirection: Option<ResolveRedirectionBuilder>,
-    pub(crate) task_idle_max_count: usize,
+    pub(crate) task_idle_max_count: Option<usize>,
     pub(crate) socks_use_udp_associate: bool,
     pub(crate) egress_path_selection: Option<EgressPathSelection>,
     pub(crate) explicit_sites: BTreeMap<NodeName, Arc<UserSiteConfig>>,
@@ -114,7 +114,7 @@ impl Default for UserConfig {
             http_user_agent_filter: None,
             resolve_strategy: None,
             resolve_redirection: None,
-            task_idle_max_count: 1,
+            task_idle_max_count: None,
             socks_use_udp_associate: false,
             egress_path_selection: None,
             explicit_sites: BTreeMap::new(),

diff --git a/g3proxy/src/config/auth/user/yaml.rs b/g3proxy/src/config/auth/user/yaml.rs
@@ -223,8 +223,9 @@ impl UserConfig {
                 Ok(())
             }
             "task_idle_max_count" => {
-                self.task_idle_max_count = g3_yaml::value::as_usize(v)
+                let count = g3_yaml::value::as_usize(v)
                     .context(format!("invalid usize value for key {k}"))?;
+                self.task_idle_max_count = Some(count);
                 Ok(())
             }
             "socks_use_udp_associate" => {

diff --git a/g3proxy/src/config/server/http_proxy.rs b/g3proxy/src/config/server/http_proxy.rs
@@ -38,7 +38,7 @@ use g3_yaml::YamlDocPosition;
 
 use super::{
     AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_DEFAULT_DURATION,
-    IDLE_CHECK_MAXIMUM_DURATION,
+    IDLE_CHECK_DEFAULT_MAX_COUNT, IDLE_CHECK_MAXIMUM_DURATION,
 };
 
 const SERVER_CONFIG_TYPE: &str = "HttpProxy";
@@ -129,7 +129,7 @@ impl HttpProxyServerConfig {
             tcp_sock_speed_limit: TcpSockSpeedLimitConfig::default(),
             timeout: HttpProxyServerTimeoutConfig::default(),
             task_idle_check_duration: IDLE_CHECK_DEFAULT_DURATION,
-            task_idle_max_count: 1,
+            task_idle_max_count: IDLE_CHECK_DEFAULT_MAX_COUNT,
             flush_task_log_on_created: false,
             flush_task_log_on_connected: false,
             task_log_flush_interval: None,

diff --git a/g3proxy/src/config/server/http_rproxy/mod.rs b/g3proxy/src/config/server/http_rproxy/mod.rs
@@ -35,7 +35,7 @@ use g3_yaml::YamlDocPosition;
 
 use super::{
     AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_DEFAULT_DURATION,
-    IDLE_CHECK_MAXIMUM_DURATION,
+    IDLE_CHECK_DEFAULT_MAX_COUNT, IDLE_CHECK_MAXIMUM_DURATION,
 };
 
 mod host;
@@ -116,7 +116,7 @@ impl HttpRProxyServerConfig {
             tcp_sock_speed_limit: TcpSockSpeedLimitConfig::default(),
             timeout: HttpRProxyServerTimeoutConfig::default(),
             task_idle_check_duration: IDLE_CHECK_DEFAULT_DURATION,
-            task_idle_max_count: 1,
+            task_idle_max_count: IDLE_CHECK_DEFAULT_MAX_COUNT,
             flush_task_log_on_created: false,
             flush_task_log_on_connected: false,
             task_log_flush_interval: None,

diff --git a/g3proxy/src/config/server/mod.rs b/g3proxy/src/config/server/mod.rs
@@ -60,7 +60,8 @@ const CONFIG_KEY_SERVER_TYPE: &str = "type";
 const CONFIG_KEY_SERVER_NAME: &str = "name";
 
 const IDLE_CHECK_MAXIMUM_DURATION: Duration = Duration::from_secs(1800);
-const IDLE_CHECK_DEFAULT_DURATION: Duration = Duration::from_secs(300);
+const IDLE_CHECK_DEFAULT_DURATION: Duration = Duration::from_secs(60);
+const IDLE_CHECK_DEFAULT_MAX_COUNT: usize = 5;
 
 pub(crate) enum ServerConfigDiffAction {
     NoAction,

diff --git a/g3proxy/src/config/server/sni_proxy/mod.rs b/g3proxy/src/config/server/sni_proxy/mod.rs
@@ -29,7 +29,10 @@ use g3_types::net::{TcpListenConfig, TcpMiscSockOpts, TcpSockSpeedLimitConfig};
 use g3_types::route::HostMatch;
 use g3_yaml::YamlDocPosition;
 
-use super::{AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_MAXIMUM_DURATION};
+use super::{
+    AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_DEFAULT_DURATION,
+    IDLE_CHECK_DEFAULT_MAX_COUNT, IDLE_CHECK_MAXIMUM_DURATION,
+};
 
 mod host;
 pub(crate) use host::SniHostConfig;
@@ -76,8 +79,8 @@ impl SniProxyServerConfig {
             listen_in_worker: false,
             ingress_net_filter: None,
             tcp_sock_speed_limit: TcpSockSpeedLimitConfig::default(),
-            task_idle_check_duration: Duration::from_secs(300),
-            task_idle_max_count: 1,
+            task_idle_check_duration: IDLE_CHECK_DEFAULT_DURATION,
+            task_idle_max_count: IDLE_CHECK_DEFAULT_MAX_COUNT,
             flush_task_log_on_created: false,
             flush_task_log_on_connected: false,
             task_log_flush_interval: None,

diff --git a/g3proxy/src/config/server/socks_proxy.rs b/g3proxy/src/config/server/socks_proxy.rs
@@ -35,7 +35,7 @@ use g3_yaml::YamlDocPosition;
 
 use super::{
     AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_DEFAULT_DURATION,
-    IDLE_CHECK_MAXIMUM_DURATION,
+    IDLE_CHECK_DEFAULT_MAX_COUNT, IDLE_CHECK_MAXIMUM_DURATION,
 };
 
 const SERVER_CONFIG_TYPE: &str = "SocksProxy";
@@ -116,7 +116,7 @@ impl SocksProxyServerConfig {
             udp_sock_speed_limit: UdpSockSpeedLimitConfig::default(),
             timeout: SocksProxyServerTimeoutConfig::default(),
             task_idle_check_duration: IDLE_CHECK_DEFAULT_DURATION,
-            task_idle_max_count: 1,
+            task_idle_max_count: IDLE_CHECK_DEFAULT_MAX_COUNT,
             flush_task_log_on_created: false,
             flush_task_log_on_connected: false,
             task_log_flush_interval: None,

diff --git a/g3proxy/src/config/server/tcp_stream.rs b/g3proxy/src/config/server/tcp_stream.rs
@@ -31,7 +31,10 @@ use g3_types::net::{
 };
 use g3_yaml::YamlDocPosition;
 
-use super::{AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_MAXIMUM_DURATION};
+use super::{
+    AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_DEFAULT_DURATION,
+    IDLE_CHECK_DEFAULT_MAX_COUNT, IDLE_CHECK_MAXIMUM_DURATION,
+};
 
 const SERVER_CONFIG_TYPE: &str = "TcpStream";
 
@@ -76,8 +79,8 @@ impl TcpStreamServerConfig {
             upstream_pick_policy: SelectivePickPolicy::Random,
             upstream_tls_name: None,
             tcp_sock_speed_limit: TcpSockSpeedLimitConfig::default(),
-            task_idle_check_duration: Duration::from_secs(300),
-            task_idle_max_count: 1,
+            task_idle_check_duration: IDLE_CHECK_DEFAULT_DURATION,
+            task_idle_max_count: IDLE_CHECK_DEFAULT_MAX_COUNT,
             flush_task_log_on_created: false,
             flush_task_log_on_connected: false,
             task_log_flush_interval: None,

diff --git a/g3proxy/src/config/server/tcp_tproxy.rs b/g3proxy/src/config/server/tcp_tproxy.rs
@@ -27,7 +27,10 @@ use g3_types::metrics::{NodeName, StaticMetricsTags};
 use g3_types::net::{TcpListenConfig, TcpMiscSockOpts, TcpSockSpeedLimitConfig};
 use g3_yaml::YamlDocPosition;
 
-use super::{AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_MAXIMUM_DURATION};
+use super::{
+    AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_DEFAULT_DURATION,
+    IDLE_CHECK_DEFAULT_MAX_COUNT, IDLE_CHECK_MAXIMUM_DURATION,
+};
 
 const SERVER_CONFIG_TYPE: &str = "TcpTProxy";
 
@@ -64,8 +67,8 @@ impl TcpTProxyServerConfig {
             listen_in_worker: false,
             ingress_net_filter: None,
             tcp_sock_speed_limit: TcpSockSpeedLimitConfig::default(),
-            task_idle_check_duration: Duration::from_secs(300),
-            task_idle_max_count: 1,
+            task_idle_check_duration: IDLE_CHECK_DEFAULT_DURATION,
+            task_idle_max_count: IDLE_CHECK_DEFAULT_MAX_COUNT,
             flush_task_log_on_created: false,
             flush_task_log_on_connected: false,
             task_log_flush_interval: None,

diff --git a/g3proxy/src/config/server/tls_stream.rs b/g3proxy/src/config/server/tls_stream.rs
@@ -32,7 +32,10 @@ use g3_types::net::{
 };
 use g3_yaml::YamlDocPosition;
 
-use super::{AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_MAXIMUM_DURATION};
+use super::{
+    AnyServerConfig, ServerConfig, ServerConfigDiffAction, IDLE_CHECK_DEFAULT_DURATION,
+    IDLE_CHECK_DEFAULT_MAX_COUNT, IDLE_CHECK_MAXIMUM_DURATION,
+};
 
 const SERVER_CONFIG_TYPE: &str = "TlsStream";
 
@@ -81,8 +84,8 @@ impl TlsStreamServerConfig {
             upstream_pick_policy: SelectivePickPolicy::Random,
             upstream_tls_name: None,
             tcp_sock_speed_limit: TcpSockSpeedLimitConfig::default(),
-            task_idle_check_duration: Duration::from_secs(300),
-            task_idle_max_count: 1,
+            task_idle_check_duration: IDLE_CHECK_DEFAULT_DURATION,
+            task_idle_max_count: IDLE_CHECK_DEFAULT_MAX_COUNT,
             flush_task_log_on_created: false,
             flush_task_log_on_connected: false,
             task_log_flush_interval: None,

diff --git a/g3proxy/src/inspect/http/v1/connect/mod.rs b/g3proxy/src/inspect/http/v1/connect/mod.rs
@@ -411,7 +411,6 @@ where
 
         let mut idle_interval = self.ctx.idle_wheel.get();
         let mut idle_count = 0;
-        let max_idle_count = self.ctx.task_max_idle_count();
 
         loop {
             tokio::select! {
@@ -433,7 +432,7 @@ where
                 n = idle_interval.tick() => {
                     if ups_to_clt.is_idle() {
                         idle_count += n;
-                        if idle_count >= max_idle_count {
+                        if idle_count >= self.ctx.max_idle_count {
                             return if ups_to_clt.no_cached_data() {
                                 Err(ServerTaskError::UpstreamAppTimeout("idle while reading response body"))
                             } else {

diff --git a/g3proxy/src/inspect/http/v1/forward/mod.rs b/g3proxy/src/inspect/http/v1/forward/mod.rs
@@ -512,7 +512,6 @@ impl<'a, SC: ServerConfig> H1ForwardTask<'a, SC> {
 
         let mut idle_interval = self.ctx.idle_wheel.get();
         let mut idle_count = 0;
-        let max_idle_count = self.ctx.task_max_idle_count();
 
         loop {
             tokio::select! {
@@ -549,7 +548,7 @@ impl<'a, SC: ServerConfig> H1ForwardTask<'a, SC> {
                 n = idle_interval.tick() => {
                     if clt_to_ups.is_idle() {
                         idle_count += n;
-                        if idle_count >= max_idle_count {
+                        if idle_count >= self.ctx.max_idle_count {
                             return if clt_to_ups.no_cached_data() {
                                 Err(ServerTaskError::ClientAppTimeout("idle while reading request body"))
                             } else {
@@ -825,7 +824,6 @@ impl<'a, SC: ServerConfig> H1ForwardTask<'a, SC> {
 
         let mut idle_interval = self.ctx.idle_wheel.get();
         let mut idle_count = 0;
-        let max_idle_count = self.ctx.task_max_idle_count();
 
         loop {
             tokio::select! {
@@ -848,7 +846,7 @@ impl<'a, SC: ServerConfig> H1ForwardTask<'a, SC> {
                 n = idle_interval.tick() => {
                     if ups_to_clt.is_idle() {
                         idle_count += n;
-                        if idle_count >= max_idle_count {
+                        if idle_count >= self.ctx.max_idle_count {
                             return if ups_to_clt.no_cached_data() {
                                 Err(ServerTaskError::UpstreamAppTimeout("idle while reading response body"))
                             } else {

diff --git a/g3proxy/src/inspect/http/v1/upgrade/mod.rs b/g3proxy/src/inspect/http/v1/upgrade/mod.rs
@@ -511,7 +511,6 @@ where
 
         let mut idle_interval = self.ctx.idle_wheel.get();
         let mut idle_count = 0;
-        let max_idle_count = self.ctx.task_max_idle_count();
 
         loop {
             tokio::select! {
@@ -533,7 +532,7 @@ where
                 n = idle_interval.tick() => {
                     if ups_to_clt.is_idle() {
                         idle_count += n;
-                        if idle_count >= max_idle_count {
+                        if idle_count >= self.ctx.max_idle_count {
                             return if ups_to_clt.no_cached_data() {
                                 Err(ServerTaskError::UpstreamAppTimeout("idle while reading response body"))
                             } else {

diff --git a/g3proxy/src/inspect/http/v2/connect/mod.rs b/g3proxy/src/inspect/http/v2/connect/mod.rs
@@ -413,7 +413,6 @@ impl<'a, SC: ServerConfig> ExchangeHead<'a, SC> {
 
             let mut idle_interval = self.ctx.idle_wheel.get();
             let mut idle_count = 0;
-            let max_idle_count = self.ctx.task_max_idle_count();
 
             loop {
                 tokio::select! {
@@ -429,7 +428,7 @@ impl<'a, SC: ServerConfig> ExchangeHead<'a, SC> {
                         if rsp_body_transfer.is_idle() {
                             idle_count += n;
 
-                            if idle_count > max_idle_count {
+                            if idle_count > self.ctx.max_idle_count {
                                 return Err(H2StreamTransferError::Idle(idle_interval.period(), idle_count));
                             }
                         } else {

diff --git a/g3proxy/src/inspect/http/v2/forward/mod.rs b/g3proxy/src/inspect/http/v2/forward/mod.rs
@@ -462,7 +462,6 @@ where
 
         let mut idle_interval = self.ctx.idle_wheel.get();
         let mut idle_count = 0;
-        let max_idle_count = self.ctx.task_max_idle_count();
 
         let mut ups_rsp: Option<Response<RecvStream>> = None;
 
@@ -497,7 +496,7 @@ where
                     if req_body_transfer.is_idle() {
                         idle_count += n;
 
-                        if idle_count > max_idle_count {
+                        if idle_count > self.ctx.max_idle_count {
                             return Err(H2StreamTransferError::Idle(idle_interval.period(), idle_count));
                         }
                     } else {
@@ -652,7 +651,6 @@ where
 
             let mut idle_interval = self.ctx.idle_wheel.get();
             let mut idle_count = 0;
-            let max_idle_count = self.ctx.task_max_idle_count();
 
             loop {
                 tokio::select! {
@@ -671,7 +669,7 @@ where
                         if rsp_body_transfer.is_idle() {
                             idle_count += n;
 
-                            if idle_count > max_idle_count {
+                            if idle_count > self.ctx.max_idle_count {
                                 return Err(H2StreamTransferError::Idle(idle_interval.period(), idle_count));
                             }
                         } else {