From 5836ecc35ed301026fcacc12eeda7603fdccf5af Mon Sep 17 00:00:00 2001
From: Tony Klink <antonyklink@gmail.com>
Date: Mon, 15 Jan 2024 16:12:22 -0600
Subject: [PATCH] Use only forwarded client IP

---
 src/relay/handler.rs |  5 +----
 src/relay/routes.rs  | 14 +-------------
 src/relay/ws.rs      | 11 ++++++-----
 3 files changed, 8 insertions(+), 22 deletions(-)

diff --git a/src/relay/handler.rs b/src/relay/handler.rs
index 6da5de8..5189103 100644
--- a/src/relay/handler.rs
+++ b/src/relay/handler.rs
@@ -6,10 +6,7 @@ use warp::{Rejection, Reply};
 pub async fn ws_handler(
     ws: warp::ws::Ws,
     context: Context,
-    client_ip: Option<SocketAddr>,
     real_client_ip: Option<SocketAddr>,
 ) -> Result<impl Reply, Rejection> {
-    Ok(ws.on_upgrade(move |socket| {
-        ws::client_connection(socket, context, client_ip, real_client_ip)
-    }))
+    Ok(ws.on_upgrade(move |socket| ws::client_connection(socket, context, real_client_ip)))
 }
diff --git a/src/relay/routes.rs b/src/relay/routes.rs
index bccf7bd..538b84c 100644
--- a/src/relay/routes.rs
+++ b/src/relay/routes.rs
@@ -10,23 +10,11 @@ pub fn routes(context: Context) -> impl Filter<Extract = impl Reply, Error = Rej
 }
 
 fn index(context: Context) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
-    let client_ip = warp::addr::remote();
-
-    let real_client_ip = warp::header::optional::<std::net::SocketAddr>("X-Real-IP")
-        .or(warp::header::optional::<std::net::SocketAddr>(
-            "X-Forwarded-For",
-        ))
-        .unify()
-        .map(|ip: Option<std::net::SocketAddr>| {
-            // Get the IP from either header,
-            // and unify into the inner type.
-            ip
-        });
+    let real_client_ip = warp::header::optional::<std::net::SocketAddr>("X-Real-IP");
 
     warp::path::end()
         .and(warp::ws())
         .and(with_context(context))
-        .and(client_ip)
         .and(real_client_ip)
         .and_then(handler::ws_handler)
 }
diff --git a/src/relay/ws.rs b/src/relay/ws.rs
index 6b0aefd..41dc941 100644
--- a/src/relay/ws.rs
+++ b/src/relay/ws.rs
@@ -16,7 +16,6 @@ use futures_util::SinkExt;
 pub async fn client_connection(
     ws: WebSocket,
     context: Context,
-    client_ip: Option<SocketAddr>,
     real_client_ip: Option<SocketAddr>,
 ) {
     let (mut ws_sender, mut ws_receiver) = ws.split();
@@ -25,10 +24,12 @@ pub async fn client_connection(
     let mut client_receiver = UnboundedReceiverStream::new(client_receiver);
 
     // Create and Add to the Context new Client and set its sender
-    let ip = real_client_ip
-        .unwrap_or(client_ip.unwrap())
-        .ip()
-        .to_string();
+    let ip = if real_client_ip.is_some() {
+        real_client_ip.unwrap().to_string()
+    } else {
+        "".to_string()
+    };
+
     let mut client = Client::new(ip);
 
     client.client_connection = Some(client_sender);