vey-proxy: add nested connect escaper methods

Author: zh-jq

lib/vey-io-ext/src/lib.rs

@@ -3,6 +3,9 @@
  * SPDX-FileCopyrightText: 2023-2025 ByteDance and/or its affiliates.
  */
 
+#[derive(Default)]
+pub struct NilLimitedStats(());
+
 mod cache;
 mod limit;
 mod listen;

lib/vey-io-ext/src/stream/limited/mod.rs

@@ -4,12 +4,10 @@
  */
 
 mod read;
-pub use read::{
-    ArcLimitedReaderStats, LimitedReader, LimitedReaderStats, NilLimitedReaderStats, SizedReader,
-};
+pub use read::{ArcLimitedReaderStats, LimitedReader, LimitedReaderStats, SizedReader};
 
 mod stream;
 pub use stream::LimitedStream;
 
 mod write;
-pub use write::{ArcLimitedWriterStats, LimitedWriter, LimitedWriterStats, NilLimitedWriterStats};
+pub use write::{ArcLimitedWriterStats, LimitedWriter, LimitedWriterStats};

lib/vey-io-ext/src/stream/limited/read.rs

@@ -16,6 +16,7 @@ use pin_project_lite::pin_project;
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio::time::{Instant, Sleep};
 
+use crate::NilLimitedStats;
 use crate::limit::{GlobalLimitGroup, GlobalStreamLimit, StreamLimitAction, StreamLimiter};
 use crate::stream::AsyncStream;
 
@@ -24,10 +25,7 @@ pub trait LimitedReaderStats {
 }
 pub type ArcLimitedReaderStats = Arc<dyn LimitedReaderStats + Send + Sync>;
 
-#[derive(Default)]
-pub struct NilLimitedReaderStats(());
-
-impl LimitedReaderStats for NilLimitedReaderStats {
+impl LimitedReaderStats for NilLimitedStats {
     fn add_read_bytes(&self, _size: usize) {}
 }
 

lib/vey-io-ext/src/stream/limited/write.rs

@@ -16,17 +16,15 @@ use pin_project_lite::pin_project;
 use tokio::io::AsyncWrite;
 use tokio::time::{Instant, Sleep};
 
+use crate::NilLimitedStats;
 use crate::limit::{GlobalLimitGroup, GlobalStreamLimit, StreamLimitAction, StreamLimiter};
 
 pub trait LimitedWriterStats {
     fn add_write_bytes(&self, size: usize);
 }
 pub type ArcLimitedWriterStats = Arc<dyn LimitedWriterStats + Send + Sync>;
 
-#[derive(Default)]
-pub struct NilLimitedWriterStats(());
-
-impl LimitedWriterStats for NilLimitedWriterStats {
+impl LimitedWriterStats for NilLimitedStats {
     fn add_write_bytes(&self, _size: usize) {}
 }
 

lib/vey-io-ext/src/udp/stats.rs

@@ -5,6 +5,8 @@
 
 use std::sync::Arc;
 
+use crate::NilLimitedStats;
+
 pub trait LimitedRecvStats {
     fn add_recv_bytes(&self, size: usize);
     fn add_recv_packet(&self) {
@@ -14,6 +16,14 @@ pub trait LimitedRecvStats {
 }
 pub type ArcLimitedRecvStats = Arc<dyn LimitedRecvStats + Send + Sync>;
 
+impl LimitedRecvStats for NilLimitedStats {
+    fn add_recv_bytes(&self, _size: usize) {}
+
+    fn add_recv_packet(&self) {}
+
+    fn add_recv_packets(&self, _n: usize) {}
+}
+
 pub trait LimitedSendStats {
     fn add_send_bytes(&self, size: usize);
     fn add_send_packet(&self) {
@@ -23,6 +33,14 @@ pub trait LimitedSendStats {
 }
 pub type ArcLimitedSendStats = Arc<dyn LimitedSendStats + Send + Sync>;
 
+impl LimitedSendStats for NilLimitedStats {
+    fn add_send_bytes(&self, _size: usize) {}
+
+    fn add_send_packet(&self) {}
+
+    fn add_send_packets(&self, _n: usize) {}
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::atomic::{AtomicUsize, Ordering};

vey-proxy/src/escape/comply_audit/mod.rs

@@ -215,6 +215,34 @@ impl EscaperInternal for ComplyAuditEscaper {
         audit_ctx.set_handle(self.audit_handle.clone());
     }
 
+    async fn _nested_tcp_connect(
+        &self,
+        task_conf: &TcpConnectTaskConf<'_>,
+        egress_notes: &mut EgressNotes,
+        task_notes: &ServerTaskNotes,
+        audit_ctx: &mut AuditContext,
+    ) -> TcpConnectResult {
+        egress_notes.escaper.clone_from(&self.config.name);
+        self.stats.add_request_passed();
+        self._update_audit_context(audit_ctx);
+        self.next
+            ._nested_tcp_connect(task_conf, egress_notes, task_notes, audit_ctx)
+            .await
+    }
+
+    async fn _nested_udp_connect(
+        &self,
+        task_conf: &UdpConnectTaskConf<'_>,
+        egress_notes: &mut EgressNotes,
+        task_notes: &ServerTaskNotes,
+    ) -> UdpConnectResult {
+        egress_notes.escaper.clone_from(&self.config.name);
+        self.stats.add_request_passed();
+        self.next
+            ._nested_udp_connect(task_conf, egress_notes, task_notes)
+            .await
+    }
+
     async fn _new_http_forward_connection(
         &self,
         _task_conf: &TcpConnectTaskConf<'_>,

vey-proxy/src/escape/comply_context/mod.rs

@@ -240,6 +240,35 @@ impl EscaperInternal for ComplyContextEscaper {
         }
     }
 
+    async fn _nested_tcp_connect(
+        &self,
+        task_conf: &TcpConnectTaskConf<'_>,
+        egress_notes: &mut EgressNotes,
+        task_notes: &ServerTaskNotes,
+        audit_ctx: &mut AuditContext,
+    ) -> TcpConnectResult {
+        egress_notes.escaper.clone_from(&self.config.name);
+        self._update_egress_path(task_notes);
+        self.stats.add_request_passed();
+        self.next
+            ._nested_tcp_connect(task_conf, egress_notes, task_notes, audit_ctx)
+            .await
+    }
+
+    async fn _nested_udp_connect(
+        &self,
+        task_conf: &UdpConnectTaskConf<'_>,
+        egress_notes: &mut EgressNotes,
+        task_notes: &ServerTaskNotes,
+    ) -> UdpConnectResult {
+        egress_notes.escaper.clone_from(&self.config.name);
+        self._update_egress_path(task_notes);
+        self.stats.add_request_passed();
+        self.next
+            ._nested_udp_connect(task_conf, egress_notes, task_notes)
+            .await
+    }
+
     async fn _new_http_forward_connection(
         &self,
         _task_conf: &TcpConnectTaskConf<'_>,

vey-proxy/src/escape/direct_fixed/http_forward/mod.rs

@@ -6,7 +6,7 @@
 
 use std::sync::Arc;
 
-use vey_io_ext::{AsyncStream, LimitedBufReader, LimitedWriter, NilLimitedReaderStats};
+use vey_io_ext::{AsyncStream, LimitedBufReader, LimitedWriter, NilLimitedStats};
 
 use super::{DirectFixedEscaper, DirectFixedEscaperStats};
 use crate::escape::EgressNotes;
@@ -94,7 +94,7 @@ impl DirectFixedEscaper {
 
         let ups_r = LimitedBufReader::new_unlimited(
             ups_r,
-            Arc::new(NilLimitedReaderStats::default()),
+            Arc::new(NilLimitedStats::default()),
             wrapper_stats.clone(),
         );
         let ups_w = LimitedWriter::new(ups_w, wrapper_stats);

vey-proxy/src/escape/direct_fixed/mod.rs

@@ -363,7 +363,7 @@ impl Escaper for DirectFixedEscaper {
     ) -> UdpConnectResult {
         self.stats.interface.add_udp_connect_attempted();
         egress_notes.escaper.clone_from(&self.config.name);
-        self.udp_connect_to(task_conf, egress_notes, task_notes, task_stats)
+        self.new_udp_connection(task_conf, egress_notes, task_notes, task_stats)
             .await
     }
 
@@ -428,6 +428,31 @@ impl EscaperInternal for DirectFixedEscaper {
         capability
     }
 
+    async fn _nested_tcp_connect(
+        &self,
+        task_conf: &TcpConnectTaskConf<'_>,
+        egress_notes: &mut EgressNotes,
+        task_notes: &ServerTaskNotes,
+        _audit_ctx: &mut AuditContext,
+    ) -> TcpConnectResult {
+        self.stats.interface.add_tcp_connect_attempted();
+        egress_notes.escaper.clone_from(&self.config.name);
+        self.nested_tcp_connect(task_conf, egress_notes, task_notes)
+            .await
+    }
+
+    async fn _nested_udp_connect(
+        &self,
+        task_conf: &UdpConnectTaskConf<'_>,
+        egress_notes: &mut EgressNotes,
+        task_notes: &ServerTaskNotes,
+    ) -> UdpConnectResult {
+        self.stats.interface.add_udp_connect_attempted();
+        egress_notes.escaper.clone_from(&self.config.name);
+        self.nested_udp_connect(task_conf, egress_notes, task_notes)
+            .await
+    }
+
     async fn _new_http_forward_connection(
         &self,
         task_conf: &TcpConnectTaskConf<'_>,

vey-proxy/src/escape/direct_fixed/stats.rs

@@ -9,7 +9,7 @@ use std::sync::Arc;
 use arc_swap::ArcSwapOption;
 
 use vey_daemon::stat::remote::{TcpConnectionTaskRemoteStats, UdpConnectTaskRemoteStats};
-use vey_io_ext::{LimitedReaderStats, LimitedWriterStats};
+use vey_io_ext::{LimitedReaderStats, LimitedRecvStats, LimitedSendStats, LimitedWriterStats};
 use vey_types::metrics::{MetricTagMap, NodeName};
 use vey_types::stats::{StatId, TcpIoSnapshot, UdpIoSnapshot};
 
@@ -124,6 +124,26 @@ impl LimitedWriterStats for DirectFixedEscaperStats {
     }
 }
 
+impl LimitedRecvStats for DirectFixedEscaperStats {
+    fn add_recv_bytes(&self, size: usize) {
+        self.udp.io.add_in_bytes(size as u64);
+    }
+
+    fn add_recv_packets(&self, n: usize) {
+        self.udp.io.add_in_packets(n);
+    }
+}
+
+impl LimitedSendStats for DirectFixedEscaperStats {
+    fn add_send_bytes(&self, size: usize) {
+        self.udp.io.add_out_bytes(size as u64);
+    }
+
+    fn add_send_packets(&self, n: usize) {
+        self.udp.io.add_out_packets(n);
+    }
+}
+
 impl TcpConnectionTaskRemoteStats for DirectFixedEscaperStats {
     fn add_read_bytes(&self, size: u64) {
         self.tcp.io.add_in_bytes(size);