diff --git a/src/proxy.rs b/src/proxy.rs index 2e67e0328..a31ea6f10 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -1,7 +1,4 @@ -use std::{ - io::Error as IoError, - net::{SocketAddr, ToSocketAddrs}, -}; +use std::{io::Error as IoError, net::SocketAddr}; use anyhow::bail; use async_recursion::async_recursion; @@ -234,12 +231,10 @@ impl ProxyScheme { } async fn resolve_host(&self, host: &str, default_port: u16) -> Result { - let (host_str, port) = match host.split_once(':') { - Some((h, p)) => (h, p.parse::().ok()), - None => (host, None), - }; - let addr = (host_str, port.unwrap_or(default_port)) - .to_socket_addrs()? + let addr = self + .http_url(host)? + .socket_addrs(|| Some(default_port))? + .into_iter() .next() .ok_or_else(|| ProxyError::AddressResolutionFailed(host.to_string()))?; Ok(addr) @@ -248,11 +243,11 @@ impl ProxyScheme { pub fn get_domain(&self) -> Result { match self { ProxyScheme::Http { host, .. } | ProxyScheme::Https { host, .. } => { - let domain = host - .split(':') - .next() - .ok_or_else(|| ProxyError::AddressResolutionFailed(host.clone()))?; - Ok(domain.to_string()) + match self.http_url(host)?.host() { + Some(url::Host::Ipv6(ip)) => Ok(ip.to_string()), + Some(host) => Ok(host.to_string()), + None => Err(ProxyError::AddressResolutionFailed(host.clone())), + } } ProxyScheme::Socks5 { addr, .. } => match addr { SocketAddr::V4(addr_v4) => Ok(addr_v4.ip().to_string()), @@ -262,17 +257,30 @@ impl ProxyScheme { } pub fn get_host_and_port(&self) -> Result { match self { - ProxyScheme::Http { host, .. } => Ok(self.append_default_port(host, 80)), - ProxyScheme::Https { host, .. } => Ok(self.append_default_port(host, 443)), + ProxyScheme::Http { host, .. } => self.append_default_port(host, 80), + ProxyScheme::Https { host, .. } => self.append_default_port(host, 443), ProxyScheme::Socks5 { addr, .. } => Ok(format!("{}", addr)), } } - fn append_default_port(&self, host: &str, default_port: u16) -> String { - if host.contains(':') { - host.to_string() + fn append_default_port(&self, host: &str, default_port: u16) -> Result { + let url = self.http_url(host)?; + let host = url + .host_str() + .ok_or_else(|| ProxyError::AddressResolutionFailed(host.to_string()))?; + Ok(format!( + "{}:{}", + host, + url.port_or_known_default().unwrap_or(default_port) + )) + } + + fn http_url(&self, host: &str) -> Result { + let scheme = if matches!(self, Self::Https { .. }) { + "https" } else { - format!("{}:{}", host, default_port) - } + "http" + }; + Ok(Url::parse(&format!("{scheme}://{host}"))?) } } @@ -717,3 +725,102 @@ where None => Err(ProxyError::NoHttpCode), }; } + +#[cfg(test)] +mod tests { + use super::ProxyScheme; + + #[test] + fn http_url_parses_hosts_and_ports() { + for host in ["localhost", "127.0.0.1", "[::1]"] { + for (proxy, scheme, default_port) in [ + (ProxyScheme::http(host).unwrap(), "http", 80), + (ProxyScheme::https(host).unwrap(), "https", 443), + ] { + for port in [None, Some(80), Some(443), Some(8443)] { + let authority = match port { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + let url = proxy.http_url(&authority).unwrap(); + assert_eq!(url.scheme(), scheme); + assert_eq!(url.host_str(), Some(host)); + assert_eq!( + url.port_or_known_default(), + Some(port.unwrap_or(default_port)) + ); + } + } + } + } + + #[test] + fn append_default_port_formats_hosts_and_ports() { + for host in ["localhost", "127.0.0.1", "[::1]"] { + for (proxy, default_port) in [ + (ProxyScheme::http(host).unwrap(), 80), + (ProxyScheme::https(host).unwrap(), 443), + ] { + for port in [None, Some(80), Some(443), Some(8443)] { + let authority = match port { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + assert_eq!( + proxy.append_default_port(&authority, default_port).unwrap(), + format!("{host}:{}", port.unwrap_or(default_port)) + ); + } + } + } + } + + #[tokio::test] + async fn resolve_host_resolves_hosts_and_ports() { + for (host, expected_ip) in [ + ("localhost", None), + ("127.0.0.1", Some("127.0.0.1")), + ("[::1]", Some("::1")), + ] { + for (proxy, default_port) in [ + (ProxyScheme::http(host).unwrap(), 80), + (ProxyScheme::https(host).unwrap(), 443), + ] { + for port in [None, Some(80), Some(443), Some(8443)] { + let authority = match port { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + let addr = proxy.resolve_host(&authority, default_port).await.unwrap(); + assert_eq!(addr.port(), port.unwrap_or(default_port)); + assert!(addr.ip().is_loopback()); + if let Some(ip) = expected_ip { + assert_eq!(addr.ip().to_string(), ip); + } + } + } + } + } + + #[test] + fn get_domain_strips_ports_and_ipv6_brackets() { + for (host, domain) in [ + ("localhost", "localhost"), + ("127.0.0.1", "127.0.0.1"), + ("[::1]", "::1"), + ] { + for port in [None, Some(80), Some(443), Some(8443)] { + let authority = match port { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + for proxy in [ + ProxyScheme::http(&authority).unwrap(), + ProxyScheme::https(&authority).unwrap(), + ] { + assert_eq!(proxy.get_domain().unwrap(), domain); + } + } + } + } +} diff --git a/tests/proxy_transport.rs b/tests/proxy_transport.rs index 67e9f979d..2faedf023 100644 --- a/tests/proxy_transport.rs +++ b/tests/proxy_transport.rs @@ -1,6 +1,7 @@ use hbb_common::{ config::{APP_NAME, OVERWRITE_SETTINGS}, proxy::{Proxy, ProxyError}, + tcp::DynTcpStream, tls::{upsert_tls_cache, TlsType}, }; use rustls_pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer}; @@ -28,24 +29,32 @@ fn tls_acceptor() -> TlsAcceptor { TlsAcceptor::from(Arc::new(config)) } -async fn https_connection(tls_type: TlsType, fixed_local_port: bool) { +async fn proxy_connection(bind: &str, tls_type: TlsType, fixed_local_port: bool) { *APP_NAME.write().unwrap() = format!("RustDesk-proxy-review-{}", std::process::id()); OVERWRITE_SETTINGS .write() .unwrap() .insert("allow-insecure-tls-fallback".into(), "Y".into()); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let url = format!("https://{}", listener.local_addr().unwrap()); + let listener = TcpListener::bind(bind).await.unwrap(); + let scheme = if matches!(tls_type, TlsType::Plain) { + "http" + } else { + "https" + }; + let url = format!("{scheme}://{}", listener.local_addr().unwrap()); upsert_tls_cache(&url, tls_type, true); let local = if fixed_local_port { - let reserved = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let reserved = TcpListener::bind(bind).await.unwrap(); Some(reserved.local_addr().unwrap()) } else { None }; let server = tokio::spawn(async move { let (tcp, peer) = listener.accept().await.unwrap(); - let mut stream = tls_acceptor().accept(tcp).await.unwrap(); + let mut stream = DynTcpStream(Box::new(tcp)); + if !matches!(tls_type, TlsType::Plain) { + stream = DynTcpStream(Box::new(tls_acceptor().accept(stream).await.unwrap())); + } let mut headers = Vec::new(); while !headers.ends_with(b"\r\n\r\n") { assert!(headers.len() < 4096); @@ -93,22 +102,22 @@ async fn https_connection(tls_type: TlsType, fixed_local_port: bool) { #[tokio::test] async fn https_rustls_reuses_initial_connection() { - https_connection(TlsType::Rustls, false).await; + proxy_connection("127.0.0.1:0", TlsType::Rustls, false).await; } #[tokio::test] async fn https_nativetls_reuses_initial_connection() { - https_connection(TlsType::NativeTls, false).await; + proxy_connection("127.0.0.1:0", TlsType::NativeTls, false).await; } #[tokio::test] async fn https_rustls_with_fixed_local_port() { - https_connection(TlsType::Rustls, true).await; + proxy_connection("127.0.0.1:0", TlsType::Rustls, true).await; } #[tokio::test] async fn https_nativetls_with_fixed_local_port() { - https_connection(TlsType::NativeTls, true).await; + proxy_connection("127.0.0.1:0", TlsType::NativeTls, true).await; } #[tokio::test] @@ -149,3 +158,45 @@ async fn https_reconnects_for_tls_fallback() { .unwrap() .unwrap(); } + +#[tokio::test] +async fn proxy_addresses_and_tls_names() { + for (scheme, default_port) in [("http", 80), ("https", 443)] { + for (host, tls_name) in [ + ("[::1]", "::1"), + ("127.0.0.1", "127.0.0.1"), + ("localhost", "localhost"), + ] { + for port in [None, Some(80), Some(443), Some(8443)] { + let authority = match port { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + let proxy = Proxy::new(format!("{scheme}://{authority}"), TIMEOUT).unwrap(); + let port = port.unwrap_or(default_port); + assert_eq!( + proxy.intercept.get_domain().unwrap(), + tls_name, + "{scheme}://{authority}" + ); + assert_eq!( + proxy.intercept.get_host_and_port().unwrap(), + format!("{host}:{port}") + ); + let resolved = proxy.proxy_addrs().await.unwrap(); + assert_eq!(resolved.port(), port); + assert!(resolved.ip().is_loopback()); + if host == "[::1]" { + assert!(resolved.is_ipv6()); + } + } + } + } +} + +#[tokio::test] +async fn ipv6_proxy_connect_and_data_exchange() { + for tls_type in [TlsType::Plain, TlsType::Rustls, TlsType::NativeTls] { + proxy_connection("[::1]:0", tls_type, false).await; + } +}