diff --git a/Cargo.toml b/Cargo.toml index a55c57c..2060a8c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ argon2 = "0.5.2" async-trait = "0.1.73" chrono = "0.4.31" tokio = { version = "1", features = ["full"] } -warp = { varsion = "0.3.3", features = ["tls"] } +warp = { version = "0.3.3", features = ["tls"] } validator = { version = "0.16", features = ["derive"] } tokio-stream = "0.1.14" futures-util = "0.3.28" diff --git a/modules/sneedstr.nix b/modules/sneedstr.nix index de4b222..e5c6858 100644 --- a/modules/sneedstr.nix +++ b/modules/sneedstr.nix @@ -98,6 +98,7 @@ in { locations."/" = { proxyPass = "http://${cfg.localAddress}:8080"; proxyWebsockets = true; # needed if you need to use WebSocket + recommendedProxySettings = true; }; }; }; diff --git a/src/relay/handler.rs b/src/relay/handler.rs index 26f8e81..6da5de8 100644 --- a/src/relay/handler.rs +++ b/src/relay/handler.rs @@ -6,7 +6,10 @@ use warp::{Rejection, Reply}; pub async fn ws_handler( ws: warp::ws::Ws, context: Context, - client_addr: Option, + client_ip: Option, + real_client_ip: Option, ) -> Result { - Ok(ws.on_upgrade(move |socket| ws::client_connection(socket, context, client_addr))) + Ok(ws.on_upgrade(move |socket| { + ws::client_connection(socket, context, client_ip, real_client_ip) + })) } diff --git a/src/relay/routes.rs b/src/relay/routes.rs index 0187217..bccf7bd 100644 --- a/src/relay/routes.rs +++ b/src/relay/routes.rs @@ -10,20 +10,32 @@ pub fn routes(context: Context) -> impl Filter impl Filter + Clone { - let client_addr = warp::addr::remote(); + let client_ip = warp::addr::remote(); + + let real_client_ip = warp::header::optional::("X-Real-IP") + .or(warp::header::optional::( + "X-Forwarded-For", + )) + .unify() + .map(|ip: Option| { + // Get the IP from either header, + // and unify into the inner type. + ip + }); warp::path::end() .and(warp::ws()) .and(with_context(context)) - .and(client_addr) + .and(client_ip) + .and(real_client_ip) .and_then(handler::ws_handler) } fn static_files() -> impl Filter + Clone { - let mut foo = std::env::current_exe().unwrap(); - foo.pop(); + let mut exe_dir = std::env::current_exe().unwrap(); + exe_dir.pop(); - let mut www = foo.clone(); + let mut www = exe_dir.clone(); www.pop(); www.push(std::path::Path::new("www/static")); diff --git a/src/relay/ws.rs b/src/relay/ws.rs index bf0d1dc..6b0aefd 100644 --- a/src/relay/ws.rs +++ b/src/relay/ws.rs @@ -13,14 +13,22 @@ use warp::ws::{Message, WebSocket}; use futures_util::SinkExt; -pub async fn client_connection(ws: WebSocket, context: Context, client_addr: Option) { +pub async fn client_connection( + ws: WebSocket, + context: Context, + client_ip: Option, + real_client_ip: Option, +) { let (mut ws_sender, mut ws_receiver) = ws.split(); let (client_sender, client_receiver) = mpsc::unbounded_channel(); let mut client_receiver = UnboundedReceiverStream::new(client_receiver); // Create and Add to the Context new Client and set its sender - let ip = client_addr.unwrap().ip().to_string(); + let ip = real_client_ip + .unwrap_or(client_ip.unwrap()) + .ip() + .to_string(); let mut client = Client::new(ip); client.client_connection = Some(client_sender);