From 82554327ae6442ef3bc0fe74fcc256e1d7f9cab6 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Wed, 12 Jun 2024 18:50:01 +0800
Subject: [PATCH] Fix hostname issues in CDN

---
 src/client.rs   |  3 +--
 src/lib.rs      |  8 --------
 src/weirduri.rs | 15 +++++++++++++--
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 2ba572c..e1d9e88 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -261,8 +261,7 @@ pub(crate) async fn create_ws_stream<S: AsyncRead + AsyncWrite + Unpin>(
 
     let b64_dst = dst_addr.as_ref().map(|dst_addr| addess_to_b64str(dst_addr, false));
 
-    let server_ip = client.server_ip_addr.ok_or("server ip addr")?.to_string();
-    let uri = format!("ws://{}/{}/", server_ip, tunnel_path);
+    let uri = format!("ws://{}/{}/", client.server_host, tunnel_path);
 
     let uri = WeirdUri::new(&uri, b64_dst, udp_tunnel, client.client_id.clone());
 
diff --git a/src/lib.rs b/src/lib.rs
index 81c9d71..ccd6e93 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -63,11 +63,3 @@ pub(crate) fn b64str_to_address(s: &str, url_safe: bool) -> Result<Address> {
     };
     Address::try_from(&buf[..]).map_err(|e| e.into())
 }
-
-pub(crate) fn combine_addr_and_port(addr: &str, port: u16) -> String {
-    if addr.contains(':') {
-        format!("[{}]:{}", addr, port)
-    } else {
-        format!("{}:{}", addr, port)
-    }
-}
diff --git a/src/weirduri.rs b/src/weirduri.rs
index bbc475b..f5b48c1 100644
--- a/src/weirduri.rs
+++ b/src/weirduri.rs
@@ -33,12 +33,23 @@ impl WeirdUri {
     }
 }
 
+fn combine_addr_and_port(addr: &url::Host<&str>, port: Option<u16>) -> String {
+    match port {
+        None => addr.to_string(),
+        Some(port) => match addr {
+            url::Host::Domain(domain) => format!("{}:{}", domain, port),
+            url::Host::Ipv4(ip) => format!("{}:{}", ip, port),
+            url::Host::Ipv6(ip) => format!("[{}]:{}", ip, port),
+        },
+    }
+}
+
 impl IntoClientRequest for WeirdUri {
     fn into_client_request(self) -> Result<Request> {
         let uri = url::Url::parse(&self.uri).map_err(|_| Error::Url(UrlError::NoPathOrQuery))?;
 
-        let host = uri.host_str().ok_or(Error::Url(UrlError::EmptyHostName))?;
-        let host = crate::combine_addr_and_port(host, uri.port().unwrap_or(80));
+        let host = uri.host().ok_or(Error::from(UrlError::EmptyHostName))?;
+        let host = combine_addr_and_port(&host, uri.port());
 
         let mut builder = Request::builder()
             .method("GET")