use crate::{ config::{Config, NetworkType}, tcp::FramedStream, udp::FramedSocket, ResultType, }; use anyhow::Context; use std::net::SocketAddr; use std::net::ToSocketAddrs; use tokio_socks::{IntoTargetAddr, TargetAddr}; fn to_socket_addr(host: T) -> ResultType { let mut addr_ipv4 = None; let mut addr_ipv6 = None; for addr in host.to_socket_addrs()? { if addr.is_ipv4() && addr_ipv4.is_none() { addr_ipv4 = Some(addr); } if addr.is_ipv6() && addr_ipv6.is_none() { addr_ipv6 = Some(addr); } } if let Some(addr) = addr_ipv4 { Ok(addr) } else { addr_ipv6.context("Failed to solve") } } pub fn get_target_addr(host: &str) -> ResultType> { let addr = match Config::get_network_type() { NetworkType::Direct => to_socket_addr(&host)?.into_target_addr()?, NetworkType::ProxySocks => host.into_target_addr()?, } .to_owned(); Ok(addr) } pub fn test_if_valid_server(host: &str) -> String { let mut host = host.to_owned(); if !host.contains(":") { host = format!("{}:{}", host, 0); } match Config::get_network_type() { NetworkType::Direct => match to_socket_addr(&host) { Err(err) => err.to_string(), Ok(_) => "".to_owned(), }, NetworkType::ProxySocks => match &host.into_target_addr() { Err(err) => err.to_string(), Ok(_) => "".to_owned(), }, } } pub trait IntoTargetAddr2<'a> { /// Converts the value of self to a `TargetAddr`. fn into_target_addr2(&self) -> ResultType>; } impl<'a> IntoTargetAddr2<'a> for SocketAddr { fn into_target_addr2(&self) -> ResultType> { Ok(TargetAddr::Ip(*self)) } } impl<'a> IntoTargetAddr2<'a> for TargetAddr<'a> { fn into_target_addr2(&self) -> ResultType> { Ok(self.clone()) } } impl<'a> IntoTargetAddr2<'a> for String { fn into_target_addr2(&self) -> ResultType> { Ok(to_socket_addr(self)?.into_target_addr()?) } } impl<'a> IntoTargetAddr2<'a> for &str { fn into_target_addr2(&self) -> ResultType> { Ok(to_socket_addr(self)?.into_target_addr()?) } } pub async fn connect_tcp<'t, T: IntoTargetAddr2<'t> + std::fmt::Debug>( target: T, ms_timeout: u64, ) -> ResultType { let target_addr = target.into_target_addr2()?; let local = Config::get_any_listen_addr(is_ipv4(&target_addr)); connect_tcp_local(target_addr, local, ms_timeout) .await .context(format!("Invalid target addr: {:?}", target)) } pub async fn connect_tcp_local<'t, T: IntoTargetAddr<'t> + std::fmt::Debug>( target: T, local: SocketAddr, ms_timeout: u64, ) -> ResultType { let target_addr = target.into_target_addr()?; if let Some(conf) = Config::get_socks() { return FramedStream::connect( conf.proxy.as_str(), target_addr, local, conf.username.as_str(), conf.password.as_str(), ms_timeout, ) .await; } let mut addr = ToSocketAddrs::to_socket_addrs(&target_addr)? .next() .context(format!("Invalid target addr: {:?}", target_addr))?; if local.is_ipv6() && addr.is_ipv4() { addr = query_nip_io(&addr)?; } Ok(FramedStream::new(addr, local, ms_timeout).await?) } #[inline] pub fn is_ipv4(target: &TargetAddr<'_>) -> bool { match target { TargetAddr::Ip(addr) => addr.is_ipv4(), _ => true, } } #[inline] pub fn query_nip_io(addr: &SocketAddr) -> ResultType { to_socket_addr(format!("{}.nip.io:{}", addr.ip(), addr.port())) } #[inline] pub fn ipv4_to_ipv6(addr: String, ipv4: bool) -> String { if !ipv4 && crate::is_ipv4_str(&addr) { if let Some(ip) = addr.split(":").next() { return addr.replace(ip, &format!("{}.nip.io", ip)); } } addr } pub async fn new_udp_for(target: &TargetAddr<'_>, ms_timeout: u64) -> ResultType { new_udp(Config::get_any_listen_addr(is_ipv4(target)), ms_timeout).await } async fn new_udp(local: T, ms_timeout: u64) -> ResultType { match Config::get_socks() { None => Ok(FramedSocket::new(to_socket_addr(&local)?).await?), Some(conf) => { let socket = FramedSocket::new_proxy( conf.proxy.as_str(), to_socket_addr(local)?, conf.username.as_str(), conf.password.as_str(), ms_timeout, ) .await?; Ok(socket) } } } pub async fn rebind_udp_for(target: &TargetAddr<'_>) -> ResultType> { match Config::get_network_type() { NetworkType::Direct => Ok(Some( FramedSocket::new(Config::get_any_listen_addr(is_ipv4(target))).await?, )), _ => Ok(None), } } #[cfg(test)] mod tests { use super::*; #[test] fn test_to_socket_addr() { assert_eq!( to_socket_addr("127.0.0.1:8080").unwrap(), "127.0.0.1:8080".parse().unwrap() ); assert!(to_socket_addr("[ff::]:0").unwrap().is_ipv6()); assert!(to_socket_addr("xx").is_err()); } #[test] fn test_nat64() { assert_eq!(ipv4_to_ipv6("1.1.1.1".to_owned(), true), "1.1.1.1"); assert_eq!(ipv4_to_ipv6("1.1.1.1".to_owned(), false), "1.1.1.1.nip.io"); assert_eq!( ipv4_to_ipv6("1.1.1.1:8080".to_owned(), false), "1.1.1.1.nip.io:8080" ); assert_eq!( ipv4_to_ipv6("rustdesk.com".to_owned(), false), "rustdesk.com" ); if to_socket_addr("rustdesk.com:80").unwrap().is_ipv6() { assert!(query_nip_io(&"1.1.1.1:80".parse().unwrap()) .unwrap() .is_ipv6()); return; } assert_eq!( query_nip_io(&"1.1.1.1:80".parse().unwrap()).unwrap(), "1.1.1.1:80".parse().unwrap() ); } }