chore(proxy): merge functionalities into construct_outbound

XOR-op · Oct 24, 2024 · 041eab8 · 041eab8
1 parent 1441267
commit 041eab8
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 61 deletions.
diff --git a/boltconn/src/proxy/dispatcher.rs b/boltconn/src/proxy/dispatcher.rs
@@ -226,6 +226,39 @@ impl Dispatcher {
         Ok(ChainOutbound::new(chain_name, res))
     }
 
+    pub async fn construct_outbound(
+        &self,
+        src_addr: SocketAddr,
+        dst_addr: &NetworkAddr,
+        proxy_config: &ProxyImpl,
+        proxy_name: &str,
+        iface_name: &str,
+        resolved_dst: Option<&SocketAddr>,
+    ) -> Result<(Box<dyn Outbound>, OutboundType), DispatchError> {
+        Ok(match proxy_config {
+            ProxyImpl::Chain(vec) => (
+                Box::new(
+                    self.create_chain(proxy_name, vec, src_addr, dst_addr, iface_name)
+                        .map_err(|_| DispatchError::BadChain)?,
+                ),
+                OutboundType::Chain,
+            ),
+            ProxyImpl::BlackHole => {
+                return Err(DispatchError::BlackHole);
+            }
+            _ => self
+                .build_normal_outbound(
+                    proxy_name,
+                    iface_name,
+                    proxy_config,
+                    src_addr,
+                    dst_addr,
+                    resolved_dst,
+                )
+                .map_err(|_| DispatchError::Reject)?,
+        })
+    }
+
     pub async fn submit_tcp<S: StreamOutboundTrait>(
         &self,
         inbound: InboundInfo,
@@ -251,33 +284,28 @@ impl Dispatcher {
         let iface_name = iface
             .as_ref()
             .map_or(self.iface_name.as_str(), |s| s.as_str());
-        let (outbounding, proxy_type): (Box<dyn Outbound>, OutboundType) =
-            match proxy_config.as_ref() {
-                ProxyImpl::Chain(vec) => (
-                    Box::new(
-                        self.create_chain(&proxy_name, vec, src_addr, &dst_addr, iface_name)
-                            .map_err(|_| DispatchError::BadChain)?,
-                    ),
-                    OutboundType::Chain,
-                ),
-                ProxyImpl::BlackHole => {
-                    tokio::spawn(async move {
-                        tokio::time::sleep(Duration::from_secs(30)).await;
-                        drop(stream)
-                    });
-                    return Err(DispatchError::BlackHole);
-                }
-                _ => self
-                    .build_normal_outbound(
-                        &proxy_name,
-                        iface_name,
-                        proxy_config.as_ref(),
-                        src_addr,
-                        &dst_addr,
-                        conn_info.resolved_dst.as_ref(),
-                    )
-                    .map_err(|_| DispatchError::Reject)?,
-            };
+        let (outbounding, proxy_type): (Box<dyn Outbound>, OutboundType) = match self
+            .construct_outbound(
+                src_addr,
+                &dst_addr,
+                &proxy_config,
+                &proxy_name,
+                iface_name,
+                conn_info.resolved_dst.as_ref(),
+            )
+            .await
+        {
+            Ok(r) => r,
+            Err(DispatchError::Reject) => return Err(DispatchError::Reject),
+            Err(DispatchError::BlackHole) => {
+                tokio::spawn(async move {
+                    tokio::time::sleep(Duration::from_secs(30)).await;
+                    drop(stream)
+                });
+                return Err(DispatchError::BlackHole);
+            }
+            Err(e) => return Err(e),
+        };
 
         // conn info
         let abort_handle = ConnAbortHandle::new();

diff --git a/boltconn/src/proxy/mod.rs b/boltconn/src/proxy/mod.rs
@@ -9,10 +9,10 @@ mod socks5_inbound;
 mod tun_inbound;
 mod tun_udp_inbound;
 
-use crate::adapter::{Connector, Outbound};
+use crate::adapter::Connector;
 use crate::common::create_tls_connector;
 use crate::common::duplex_chan::DuplexChan;
-use crate::dispatch::{Latency, Proxy, ProxyImpl};
+use crate::dispatch::{Latency, Proxy};
 use crate::proxy::error::RuntimeError;
 use bytes::Bytes;
 pub use context::*;
@@ -100,9 +100,17 @@ pub async fn latency_test(
     );
     let src_addr = get_random_local_addr(&dst_addr, rng.gen_range(32768..65535));
     // create outbound
-    let creator = construct_outbound(dispatcher, &proxy, src_addr, &dst_addr, &iface)
+    let (creator, _) = dispatcher
+        .construct_outbound(
+            src_addr,
+            &dst_addr,
+            &proxy.get_impl(),
+            &proxy.get_name(),
+            &iface,
+            None,
+        )
         .await
-        .ok_or_else(|| {
+        .map_err(|_| {
             proxy.set_latency(Latency::Failed);
             RuntimeError::LatencyTest("Create outbound failed")
         })?;
@@ -145,33 +153,3 @@ pub async fn latency_test(
     });
     Ok(timeout_future)
 }
-
-async fn construct_outbound(
-    dispatcher: &Dispatcher,
-    proxy: &Proxy,
-    src_addr: SocketAddr,
-    dst_addr: &NetworkAddr,
-    iface: &str,
-) -> Option<Box<dyn Outbound>> {
-    match proxy.get_impl().as_ref() {
-        ProxyImpl::Chain(vec) => {
-            match dispatcher.create_chain(&proxy.get_name(), vec, src_addr, dst_addr, iface) {
-                Ok(o) => Some(Box::new(o)),
-                Err(_) => None,
-            }
-        }
-        proxy_config => {
-            match dispatcher.build_normal_outbound(
-                &proxy.get_name(),
-                iface,
-                proxy_config,
-                src_addr,
-                dst_addr,
-                None,
-            ) {
-                Ok((o, _)) => Some(o),
-                Err(_) => None,
-            }
-        }
-    }
-}