Create TransformerContext before executing socket_transformer

Signed-off-by: Eval EXEC <execvy@gmail.com>
nervosnetwork · Dec 23, 2024 · 6b1f15f · 6b1f15f
1 parent 1eb0f49
commit 6b1f15f
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 9 deletions.
diff --git a/tentacle/src/builder.rs b/tentacle/src/builder.rs
@@ -1,6 +1,7 @@
 use std::{io, sync::Arc, time::Duration};
 
 use crate::service::config::ProxyConfig;
+use crate::service::config::TransformerContext;
 use nohash_hasher::IntMap;
 use tokio_util::codec::LengthDelimitedCodec;
 
@@ -218,7 +219,10 @@ where
     #[cfg(not(target_family = "wasm"))]
     pub fn tcp_config<F>(mut self, f: F) -> Self
     where
-        F: Fn(TcpSocket) -> Result<TcpSocket, std::io::Error> + Send + Sync + 'static,
+        F: Fn(TcpSocket, TransformerContext) -> Result<TcpSocket, std::io::Error>
+            + Send
+            + Sync
+            + 'static,
     {
         self.config.tcp_config.tcp.socket_transformer = Arc::new(f);
         self
@@ -236,7 +240,7 @@ where
     #[cfg_attr(docsrs, doc(cfg(feature = "ws")))]
     pub fn tcp_config_on_ws<F>(mut self, f: F) -> Self
     where
-        F: Fn(TcpSocket) -> Result<TcpSocket, std::io::Error> + Send + Sync + 'static,
+        F: Fn(TcpSocket, TransformerContext) -> Result<TcpSocket, std::io::Error> + Send + Sync + 'static,
     {
         self.config.tcp_config.ws.socket_transformer = Arc::new(f);
         self
@@ -260,7 +264,7 @@ where
     #[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
     pub fn tcp_config_on_tls<F>(mut self, f: F) -> Self
     where
-        F: Fn(TcpSocket) -> Result<TcpSocket, std::io::Error> + Send + Sync + 'static,
+        F: Fn(TcpSocket, TransformerContext) -> Result<TcpSocket, std::io::Error> + Send + Sync + 'static,
     {
         self.config.tcp_config.tls.socket_transformer = Arc::new(f);
         self

diff --git a/tentacle/src/runtime/async_runtime/os/mod.rs b/tentacle/src/runtime/async_runtime/os/mod.rs
@@ -3,7 +3,7 @@ mod time;
 use crate::{
     runtime::{proxy::socks5_config, CompatStream2},
     service::{
-        config::{TcpSocket, TcpSocketConfig, TcpSocketTransformer},
+        config::{TcpSocket, TcpSocketConfig, TcpSocketTransformer, TransformerContext},
         ProxyConfig,
     },
 };
@@ -138,7 +138,8 @@ pub(crate) fn listen(addr: SocketAddr, tcp_config: TcpSocketConfig) -> io::Resul
     let socket = Socket::new(domain, Type::STREAM, Some(SocketProtocol::TCP))?;
 
     let socket = {
-        let t = (tcp_config.socket_transformer)(TcpSocket { inner: socket })?;
+        let transformer_context = TransformerContext::new_listen(addr);
+        let t = (tcp_config.socket_transformer)(TcpSocket { inner: socket }, transformer_context)?;
         t.inner
     };
     // `bind` twice will return error
@@ -177,7 +178,8 @@ async fn connect_direct(
         // user can disable it on tcp_config
         #[cfg(not(windows))]
         socket.set_reuse_address(true)?;
-        let t = socket_transformer(TcpSocket { inner: socket })?;
+        let transformer_context = TransformerContext::new_dial(addr);
+        let t = socket_transformer(TcpSocket { inner: socket }, transformer_context)?;
         t.inner
     };
 

diff --git a/tentacle/src/runtime/tokio_runtime/mod.rs b/tentacle/src/runtime/tokio_runtime/mod.rs
@@ -7,7 +7,7 @@ pub use tokio::{
 };
 
 use crate::service::{
-    config::{TcpSocket, TcpSocketConfig, TcpSocketTransformer},
+    config::{TcpSocket, TcpSocketConfig, TcpSocketTransformer, TransformerContext},
     ProxyConfig,
 };
 use socket2::{Domain, Protocol as SocketProtocol, Socket, Type};
@@ -93,7 +93,8 @@ pub(crate) fn listen(addr: SocketAddr, tcp_config: TcpSocketConfig) -> io::Resul
         // user can disable it on tcp_config
         #[cfg(not(windows))]
         socket.set_reuse_address(true)?;
-        let t = (tcp_config.socket_transformer)(TcpSocket { inner: socket })?;
+        let transformer_context = TransformerContext::new_listen(addr);
+        let t = (tcp_config.socket_transformer)(TcpSocket { inner: socket }, transformer_context)?;
         t.inner.set_nonblocking(true)?;
         // safety: fd convert by socket2
         unsafe {
@@ -126,7 +127,8 @@ async fn connect_direct(
     let socket = Socket::new(domain, Type::STREAM, Some(SocketProtocol::TCP))?;
 
     let socket = {
-        let t = socket_transformer(TcpSocket { inner: socket })?;
+        let transformer_context = TransformerContext::new_dial(addr);
+        let t = socket_transformer(TcpSocket { inner: socket }, transformer_context)?;
         t.inner.set_nonblocking(true)?;
         // safety: fd convert by socket2
         unsafe {

diff --git a/tentacle/src/service.rs b/tentacle/src/service.rs
@@ -46,6 +46,7 @@ mod helper;
 pub use crate::service::{
     config::{
         HandshakeType, ProtocolHandle, ProtocolMeta, TargetProtocol, TargetSession, TcpSocket,
+        TransformerContext,
     },
     control::{ServiceAsyncControl, ServiceControl},
     event::{ServiceError, ServiceEvent},