From 385e48f3abd9d2d95f4847c5871ce3e4899571df Mon Sep 17 00:00:00 2001 From: Salem Yaslem Date: Fri, 30 Aug 2024 02:05:59 +0300 Subject: [PATCH] move forwarding logic to utils and reuse it in all services --- src/services/http.rs | 29 +++++------------------------ src/services/https.rs | 29 +++++------------------------ src/services/minecraft.rs | 35 +++++++++-------------------------- src/utils.rs | 39 +++++++++++++++++++++++++++++++++------ 4 files changed, 52 insertions(+), 80 deletions(-) diff --git a/src/services/http.rs b/src/services/http.rs index a925e9f..61b54b7 100644 --- 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 = 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::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 index d2920df..52868db 100644 --- 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) -> Option { let mut acceptor: Acceptor = Acceptor::default(); @@ -50,35 +50,16 @@ pub async fn handle_connection(client: TcpStream, port: u16) -> Option<()> { let resolved_address: Result = 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::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 index 089a29e..69f6b6a 100644 --- 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::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 index 15d5a93..f8026ec 100644 --- 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 { let (stream, sender) = TcpClientStream::>::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 { 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 { 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 { )) } +pub async fn forward(source: TcpStream, distance: String) -> Option<()> { + let source_addr = source.peer_addr().ok()?; + let server: Result = 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()) -} \ No newline at end of file +}