rustdesk/libs/hbb_common/src/socket_client.rs

202 lines
5.2 KiB
Rust
Raw Normal View History

use crate::{
2022-01-05 13:21:14 +08:00
config::{Config, NetworkType},
tcp::FramedStream,
udp::FramedSocket,
ResultType,
};
use anyhow::Context;
use std::net::SocketAddr;
2022-12-29 20:34:52 +08:00
use tokio::net::ToSocketAddrs;
2022-01-05 13:21:14 +08:00
use tokio_socks::{IntoTargetAddr, TargetAddr};
2022-01-05 13:21:14 +08:00
pub fn test_if_valid_server(host: &str) -> String {
let mut host = host.to_owned();
if !host.contains(":") {
host = format!("{}:{}", host, 0);
}
2022-12-29 20:34:52 +08:00
use std::net::ToSocketAddrs;
2022-01-05 13:21:14 +08:00
match Config::get_network_type() {
2022-12-29 20:34:52 +08:00
NetworkType::Direct => match host.to_socket_addrs() {
2022-01-05 13:21:14 +08:00
Err(err) => err.to_string(),
Ok(_) => "".to_owned(),
},
NetworkType::ProxySocks => match &host.into_target_addr() {
Err(err) => err.to_string(),
Ok(_) => "".to_owned(),
},
}
}
2022-12-29 20:34:52 +08:00
pub trait IsResolvedSocketAddr {
fn resolve(&self) -> Option<&SocketAddr>;
}
2022-12-29 20:34:52 +08:00
impl IsResolvedSocketAddr for SocketAddr {
fn resolve(&self) -> Option<&SocketAddr> {
Some(&self)
}
}
2022-12-29 20:34:52 +08:00
impl IsResolvedSocketAddr for String {
fn resolve(&self) -> Option<&SocketAddr> {
None
}
}
2022-12-29 20:34:52 +08:00
impl IsResolvedSocketAddr for &str {
fn resolve(&self) -> Option<&SocketAddr> {
None
}
}
2022-12-29 20:34:52 +08:00
#[inline]
pub async fn connect_tcp<
't,
T: IntoTargetAddr<'t> + ToSocketAddrs + IsResolvedSocketAddr + std::fmt::Display,
>(
target: T,
ms_timeout: u64,
) -> ResultType<FramedStream> {
2022-12-29 20:34:52 +08:00
connect_tcp_local(target, None, ms_timeout).await
2022-12-28 13:52:13 +08:00
}
2022-12-29 20:34:52 +08:00
pub async fn connect_tcp_local<
't,
T: IntoTargetAddr<'t> + ToSocketAddrs + IsResolvedSocketAddr + std::fmt::Display,
>(
2022-12-28 13:52:13 +08:00
target: T,
2022-12-29 20:34:52 +08:00
local: Option<SocketAddr>,
2022-12-28 13:52:13 +08:00
ms_timeout: u64,
) -> ResultType<FramedStream> {
if let Some(conf) = Config::get_socks() {
return FramedStream::connect(
conf.proxy.as_str(),
2022-12-29 20:34:52 +08:00
target,
local,
conf.username.as_str(),
conf.password.as_str(),
ms_timeout,
)
.await;
}
2022-12-29 20:34:52 +08:00
if let Some(target) = target.resolve() {
if let Some(local) = local {
if local.is_ipv6() && target.is_ipv4() {
let target = query_nip_io(&target).await?;
return Ok(FramedStream::new(target, Some(local), ms_timeout).await?);
}
}
2022-12-28 13:52:13 +08:00
}
2022-12-29 20:34:52 +08:00
Ok(FramedStream::new(target, local, ms_timeout).await?)
}
2022-12-28 13:52:13 +08:00
#[inline]
pub fn is_ipv4(target: &TargetAddr<'_>) -> bool {
match target {
TargetAddr::Ip(addr) => addr.is_ipv4(),
_ => true,
}
}
#[inline]
2022-12-29 20:34:52 +08:00
pub async fn query_nip_io(addr: &SocketAddr) -> ResultType<SocketAddr> {
tokio::net::lookup_host(format!("{}.nip.io:{}", addr.ip(), addr.port()))
.await?
.filter(|x| x.is_ipv6())
.next()
.context("Failed to get ipv6 from nip.io")
2022-12-28 13:52:13 +08:00
}
#[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
}
2022-12-29 20:34:52 +08:00
async fn test_is_ipv4(target: &str) -> bool {
if let Ok(Ok(s)) = super::timeout(1000, tokio::net::TcpStream::connect(target)).await {
return s.local_addr().map(|x| x.is_ipv4()).unwrap_or(true);
}
true
}
#[inline]
pub async fn new_udp_for(target: &str, ms_timeout: u64) -> ResultType<FramedSocket> {
new_udp(
Config::get_any_listen_addr(test_is_ipv4(target).await),
ms_timeout,
)
.await
2022-12-28 13:52:13 +08:00
}
async fn new_udp<T: ToSocketAddrs>(local: T, ms_timeout: u64) -> ResultType<FramedSocket> {
match Config::get_socks() {
2022-12-29 20:34:52 +08:00
None => Ok(FramedSocket::new(local).await?),
Some(conf) => {
2022-01-05 13:21:14 +08:00
let socket = FramedSocket::new_proxy(
conf.proxy.as_str(),
2022-12-29 20:34:52 +08:00
local,
conf.username.as_str(),
conf.password.as_str(),
ms_timeout,
)
.await?;
2022-01-05 13:21:14 +08:00
Ok(socket)
}
}
}
2022-12-29 20:34:52 +08:00
pub async fn rebind_udp_for(target: &str) -> ResultType<Option<FramedSocket>> {
match Config::get_network_type() {
2022-12-28 13:52:13 +08:00
NetworkType::Direct => Ok(Some(
2022-12-29 20:34:52 +08:00
FramedSocket::new(Config::get_any_listen_addr(test_is_ipv4(target).await)).await?,
2022-12-28 13:52:13 +08:00
)),
_ => Ok(None),
}
}
2022-12-28 13:52:13 +08:00
#[cfg(test)]
mod tests {
2022-12-29 20:34:52 +08:00
use std::net::ToSocketAddrs;
use super::*;
2022-12-28 13:52:13 +08:00
#[test]
fn test_nat64() {
2022-12-29 20:34:52 +08:00
test_nat64_async();
}
#[tokio::main(flavor = "current_thread")]
async fn test_nat64_async() {
2022-12-28 13:52:13 +08:00
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"
);
2022-12-29 20:34:52 +08:00
if ("rustdesk.com:80")
.to_socket_addrs()
.unwrap()
.next()
.unwrap()
.is_ipv6()
{
2022-12-28 13:52:13 +08:00
assert!(query_nip_io(&"1.1.1.1:80".parse().unwrap())
2022-12-29 20:34:52 +08:00
.await
2022-12-28 13:52:13 +08:00
.unwrap()
.is_ipv6());
return;
}
2022-12-29 20:34:52 +08:00
assert!(query_nip_io(&"1.1.1.1:80".parse().unwrap()).await.is_err());
2022-12-28 13:52:13 +08:00
}
}