v1.1.121.1.12

2 files changed, 37 insertions, 10 deletions

diff --git a/src/relay_server.rs b/src/relay_server.rs
index b55d544..abe6fc2 100644
--- a/src/relay_server.rs
+++ b/src/relay_server.rs
@@ -392,19 +392,30 @@ async fn handle_connection(
 
 async fn make_pair(
     stream: TcpStream,
-    addr: SocketAddr,
+    mut addr: SocketAddr,
     key: &str,
     limiter: Limiter,
     ws: bool,
 ) -> ResultType<()> {
     if ws {
-        make_pair_(
-            tokio_tungstenite::accept_async(stream).await?,
-            addr,
-            key,
-            limiter,
-        )
-        .await;
+        use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
+        let callback = |req: &Request, response: Response| {
+            let headers = req.headers();
+            let real_ip = headers
+                .get("X-Real-IP")
+                .or_else(|| headers.get("X-Forwarded-For"))
+                .and_then(|header_value| header_value.to_str().ok());
+            if let Some(ip) = real_ip {
+                if ip.contains('.') {
+                    addr = format!("{ip}:0").parse().unwrap_or(addr);
+                } else {
+                    addr = format!("[{ip}]:0").parse().unwrap_or(addr);
+                }
+            }
+            Ok(response)
+        };
+        let ws_stream = tokio_tungstenite::accept_hdr_async(stream, callback).await?;
+        make_pair_(ws_stream, addr, key, limiter).await;
     } else {
         make_pair_(FramedStream::from(stream, addr), addr, key, limiter).await;
     }
diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs
index 3449e75..9394731 100644
--- a/src/rendezvous_server.rs
+++ b/src/rendezvous_server.rs
@@ -1106,13 +1106,29 @@ impl RendezvousServer {
     async fn handle_listener_inner(
         &mut self,
         stream: TcpStream,
-        addr: SocketAddr,
+        mut addr: SocketAddr,
         key: &str,
         ws: bool,
     ) -> ResultType<()> {
         let mut sink;
         if ws {
-            let ws_stream = tokio_tungstenite::accept_async(stream).await?;
+            use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
+            let callback = |req: &Request, response: Response| {
+                let headers = req.headers();
+                let real_ip = headers
+                    .get("X-Real-IP")
+                    .or_else(|| headers.get("X-Forwarded-For"))
+                    .and_then(|header_value| header_value.to_str().ok());
+                if let Some(ip) = real_ip {
+                    if ip.contains('.') {
+                        addr = format!("{ip}:0").parse().unwrap_or(addr);
+                    } else {
+                        addr = format!("[{ip}]:0").parse().unwrap_or(addr);
+                    }
+                }
+                Ok(response)
+            };
+            let ws_stream = tokio_tungstenite::accept_hdr_async(stream, callback).await?;
             let (a, mut b) = ws_stream.split();
             sink = Some(Sink::Ws(a));
             while let Ok(Some(Ok(msg))) = timeout(30_000, b.next()).await {