From 9261937098fdc293e5e9438551dc5ded789186d1 Mon Sep 17 00:00:00 2001
From: Condorra <kvox990h@anonaddy.me>
Date: Sun, 3 Nov 2024 22:23:31 +1100
Subject: [PATCH] Add option to receive real IP first from select hosts.

---
 .ci/deploy-baseimage/build-push |  4 ++--
 blastmud_listener/rustfmt.toml  |  1 +
 blastmud_listener/src/main.rs   | 18 +++++++++++++++++-
 3 files changed, 20 insertions(+), 3 deletions(-)
 create mode 100644 blastmud_listener/rustfmt.toml

diff --git a/.ci/deploy-baseimage/build-push b/.ci/deploy-baseimage/build-push
index df41dd9..c9c5a33 100755
--- a/.ci/deploy-baseimage/build-push
+++ b/.ci/deploy-baseimage/build-push
@@ -1,2 +1,2 @@
-docker build -t blasthavers/deploy-base:latest .
-docker push blasthavers/deploy-base:latest
+buildah build -t registry:8080/blasthavers/deploy-base:latest .
+podman push registry:8080/blasthavers/deploy-base:latest
diff --git a/blastmud_listener/rustfmt.toml b/blastmud_listener/rustfmt.toml
new file mode 100644
index 0000000..3a26366
--- /dev/null
+++ b/blastmud_listener/rustfmt.toml
@@ -0,0 +1 @@
+edition = "2021"
diff --git a/blastmud_listener/src/main.rs b/blastmud_listener/src/main.rs
index 23f2c9a..b1caf60 100644
--- a/blastmud_listener/src/main.rs
+++ b/blastmud_listener/src/main.rs
@@ -31,6 +31,7 @@ use warp::{self, filters::ws, Filter, Reply};
 #[derive(Deserialize, Debug)]
 struct Config {
     listeners: Vec<String>,
+    ips_trusted_to_send_real_ip: Vec<String>,
     ws_listener: String,
     gameserver: String,
     pidfile: String,
@@ -349,6 +350,7 @@ async fn handle_client_socket(
     active_sessions: SessionMap,
     mut stream: TcpStream,
     addr: SocketAddr,
+    trusted_ips_for_realip: Vec<String>,
 ) {
     let (rstream, mut wstream) = stream.split();
     let mut rbuf = codec::FramedRead::new(
@@ -364,7 +366,18 @@ async fn handle_client_socket(
     let (discon_sender, mut discon_receiver) = mpsc::unbounded_channel();
 
     let mut sess_idx_lock = active_sessions.lock().await;
-    let addr_str = addr.ip().to_string();
+    let mut addr_str = addr.ip().to_string();
+
+    if trusted_ips_for_realip.contains(&addr_str) {
+        if let Some(Ok(real_ip)) = rbuf.next().await {
+            info!(
+                "Real IP sent by upstream for session {} is {}",
+                session, real_ip
+            );
+            addr_str = real_ip;
+        }
+    }
+
     if *sess_idx_lock.count_by_source.get(&addr_str).unwrap_or(&0) >= MAX_CONNS_PER_IP {
         drop(sess_idx_lock);
         info!(
@@ -788,6 +801,7 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
         for listener in config.listeners.clone() {
             let server_sender_for_listener = server_sender.clone();
             let active_sessions_for_listener = active_sessions.clone();
+            let trusted_ips_for_listener = config.ips_trusted_to_send_real_ip.clone();
             listen_handles.push(task::spawn(async move {
                 match TcpListener::bind(&listener).await {
                     Err(e) => {
@@ -802,12 +816,14 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
                                 let server_sender_for_client = server_sender_for_listener.clone();
                                 let active_sessions_for_client =
                                     active_sessions_for_listener.clone();
+                                let trusted_ips_for_client = trusted_ips_for_listener.clone();
                                 task::spawn(async move {
                                     handle_client_socket(
                                         server_sender_for_client,
                                         active_sessions_for_client,
                                         stream,
                                         addr,
+                                        trusted_ips_for_client,
                                     )
                                     .await;
                                 });