move forwarding logic to utils and reuse it in all services

xlmnxp · Aug 29, 2024 · 385e48f · 385e48f
1 parent df9741f
commit 385e48f
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 80 deletions.
diff --git a/src/services/http.rs b/src/services/http.rs
@@ -1,5 +1,5 @@
 use tokio::net::{TcpListener, TcpStream};
-use crate::utils::{get_bind_address, resolve_addr};
+use crate::utils::{forward, get_bind_address, resolve_addr};
 
 async fn handle_connection(client: TcpStream, port: u16) -> Option<()> {
     let src_addr = client.peer_addr().ok()?;
@@ -20,35 +20,16 @@ async fn handle_connection(client: TcpStream, port: u16) -> Option<()> {
             let resolved_address: Result<std::net::IpAddr, tokio::io::Error> =
                 resolve_addr(&host_string).await;
             if let Ok(ip) = resolved_address {
+                let distance = format!("[{}]:{}", ip, port);
+
                 log::info!(
                     "HTTP {} Choose AAAA record for {}: {}",
                     src_addr,
                     host_string,
-                    ip
+                    distance
                 );
 
-                let server: Result<TcpStream, tokio::io::Error> =
-                    TcpStream::connect(format!("[{}]:{}", ip, port)).await;
-                if server.is_err() {
-                    log::error!(
-                        "HTTP {} Failed to connect to upstream: {}",
-                        src_addr,
-                        format!("{}:{}", ip, port)
-                    );
-                    return None;
-                }
-
-                let server: TcpStream = server.ok()?;
-                let (mut eread, mut ewrite) = client.into_split();
-                let (mut oread, mut owrite) = server.into_split();
-                log::info!(
-                    "HTTP {} Connected to upstream: {}",
-                    src_addr,
-                    format!("[{}]:{}", ip, port)
-                );
-                tokio::spawn(async move { tokio::io::copy(&mut eread, &mut owrite).await });
-                tokio::spawn(async move { tokio::io::copy(&mut oread, &mut ewrite).await });
-                return Some(());
+                return forward(client, distance).await;
             } else {
                 if buf.len() > 4096 || last_buf_read_len == 0 {
                     log::error!(

diff --git a/src/services/https.rs b/src/services/https.rs
@@ -1,6 +1,6 @@
 use tokio::net::{TcpListener, TcpStream};
 use rustls::server::{Accepted, Acceptor};
-use crate::utils::{get_bind_address, resolve_addr};
+use crate::utils::{forward, get_bind_address, resolve_addr};
 
 pub async fn get_sni_from_packet(packet: Vec<u8>) -> Option<String> {
     let mut acceptor: Acceptor = Acceptor::default();
@@ -50,35 +50,16 @@ pub async fn handle_connection(client: TcpStream, port: u16) -> Option<()> {
             let resolved_address: Result<std::net::IpAddr, tokio::io::Error> =
                 resolve_addr(&sni_string).await;
             if let Ok(ip) = resolved_address {
+                let distance = format!("[{}]:{}", ip, port);
+
                 log::info!(
                     "HTTPS {} Choose AAAA record for {}: {}",
                     src_addr,
                     sni_string,
-                    ip
+                    distance
                 );
 
-                let server: Result<TcpStream, tokio::io::Error> =
-                    TcpStream::connect(format!("[{}]:{}", ip, port)).await;
-                if server.is_err() {
-                    log::error!(
-                        "HTTPS {} Failed to connect to upstream: {}",
-                        src_addr,
-                        format!("{}:{}", ip, port)
-                    );
-                    return None;
-                }
-
-                let server: TcpStream = server.ok()?;
-                let (mut eread, mut ewrite) = client.into_split();
-                let (mut oread, mut owrite) = server.into_split();
-                log::info!(
-                    "HTTPS {} Connected to upstream: {}",
-                    src_addr,
-                    format!("[{}]:{}", ip, port)
-                );
-                tokio::spawn(async move { tokio::io::copy(&mut eread, &mut owrite).await });
-                tokio::spawn(async move { tokio::io::copy(&mut oread, &mut ewrite).await });
-                return Some(());
+                return forward(client, distance).await
             } else {
                 log::error!(
                     "HTTPS {} Failed to resolve AAAA record for {}: {}",

diff --git a/src/services/minecraft.rs b/src/services/minecraft.rs
@@ -1,5 +1,5 @@
 use tokio::net::{TcpListener, TcpStream};
-use crate::utils::{get_bind_address, resolve_addr};
+use crate::utils::{forward, get_bind_address, resolve_addr};
 
 struct MinecraftServer {
     hostname: String,
@@ -88,36 +88,19 @@ async fn handle_connection(client: TcpStream, _port: u16) -> Option<()> {
 
     let server_info = MinecraftServer::read_server_info(&buf).expect("Failed to read server info");
 
-    log::info!(
-        "Minecraft {} Trying connect to: {}",
-        src_addr,
-        server_info.hostname
-    );
-
     let ip = resolve_addr(&server_info.hostname).await.expect("Failed to resolve hostname");
-    let upstream = format!("{}:{}", ip, server_info.port);
-    let server: Result<TcpStream, tokio::io::Error> =
-        TcpStream::connect(upstream.clone()).await;
-    if server.is_err() {
-        log::error!(
-            "Minecraft {} Failed to connect to upstream: {}",
-            src_addr,
-            upstream
-        );
-        return None;
-    }
+    let distance = format!("[{}]:{}", ip, server_info.port);
 
-    let server: TcpStream = server.ok()?;
-    let (mut eread, mut ewrite) = client.into_split();
-    let (mut oread, mut owrite) = server.into_split();
     log::info!(
-        "Minecraft {} Connected to upstream: {}",
+        "Minecraft {} Choose AAAA record for {}: {}",
         src_addr,
-        upstream
+        server_info.hostname,
+        distance
     );
-    tokio::spawn(async move { tokio::io::copy(&mut eread, &mut owrite).await });
-    tokio::spawn(async move { tokio::io::copy(&mut oread, &mut ewrite).await });
-    None
+
+    let distance = format!("[{}]:{}", ip, server_info.port);
+
+    forward(client, distance).await
 }
 
 pub async fn listener(port: u16) -> std::io::Result<()> {

diff --git a/src/utils.rs b/src/utils.rs
@@ -10,7 +10,9 @@ use tokio::net::TcpStream;
 
 pub async fn resolve_addr(addr: &str) -> std::io::Result<IpAddr> {
     let (stream, sender) = TcpClientStream::<AsyncIoTokioAsStd<TcpStream>>::new(
-        get_dns_server().parse().expect("Invalid DNS server address"),
+        get_dns_server()
+            .parse()
+            .expect("Invalid DNS server address"),
     );
     let dns_client = AsyncClient::new(stream, sender, None);
     let (mut dns_client, bg) = dns_client.await.expect("dns connection failed");
@@ -37,9 +39,9 @@ pub async fn resolve_addr(addr: &str) -> std::io::Result<IpAddr> {
     let bridge_ipv6 = get_bridge46_ipv6();
     if bridge_ipv4 != "" || bridge_ipv6 != "" {
         let response_ipv4: DnsResponse = dns_client
-        .query(Name::from_str(addr)?, DNSClass::IN, RecordType::A)
-        .await
-        .expect("Failed to query");
+            .query(Name::from_str(addr)?, DNSClass::IN, RecordType::A)
+            .await
+            .expect("Failed to query");
 
         let answers_ipv4: &[Record] = response_ipv4.answers();
 
@@ -81,7 +83,10 @@ pub async fn resolve_addr(addr: &str) -> std::io::Result<IpAddr> {
         if let Some(RData::AAAA(ref ip)) = answer.data() {
             let bridge_ipv6 = get_bridge46_ipv6();
             if bridge_ipv6 != "" && ip.to_string() == bridge_ipv6 {
-                log::info!("DNS Resolver: {} requested IPv6 is same as Bridge46 service IPv6", addr);
+                log::info!(
+                    "DNS Resolver: {} requested IPv6 is same as Bridge46 service IPv6",
+                    addr
+                );
                 continue;
             }
 
@@ -103,6 +108,28 @@ pub async fn resolve_addr(addr: &str) -> std::io::Result<IpAddr> {
     ))
 }
 
+pub async fn forward(source: TcpStream, distance: String) -> Option<()> {
+    let source_addr = source.peer_addr().ok()?;
+    let server: Result<TcpStream, tokio::io::Error> = TcpStream::connect(distance.clone()).await;
+    if server.is_err() {
+        log::error!(
+            "{} Failed to connect to upstream: {}",
+            source_addr,
+            distance
+        );
+        return None;
+    }
+
+    let server: TcpStream = server.ok()?;
+    let (mut eread, mut ewrite) = source.into_split();
+    let (mut oread, mut owrite) = server.into_split();
+    log::info!("{} Connected to upstream: {}", source_addr, distance);
+    tokio::spawn(async move { tokio::io::copy(&mut eread, &mut owrite).await });
+    tokio::spawn(async move { tokio::io::copy(&mut oread, &mut ewrite).await });
+
+    Some(())
+}
+
 pub fn get_dns_server() -> String {
     std::env::var("DNS_SERVER").unwrap_or_else(|_| "1.1.1.1:53".into())
 }
@@ -117,4 +144,4 @@ pub fn get_bridge46_ipv4() -> String {
 
 pub fn get_bridge46_ipv6() -> String {
     std::env::var("BRIDGE46_IPV6").unwrap_or_else(|_| "".into())
-}
+}