From 4007c07dc57e138610efe6b182b10c0fd757ec24 Mon Sep 17 00:00:00 2001
From: Joshua Barretto <joshua@jsbarretto.com>
Date: Tue, 29 Apr 2025 20:57:25 +0100
Subject: [PATCH] Fetch client ip from X-Forwarded-For, if possible

---
 Cargo.lock  |  1 +
 Cargo.toml  |  1 +
 src/main.rs | 23 ++++++++++++++++-------
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 65482da..f0f1566 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -224,6 +224,7 @@ dependencies = [
  "clap",
  "flume",
  "hashbrown",
+ "http",
  "http-body",
  "itertools 0.14.0",
  "rand",
diff --git a/Cargo.toml b/Cargo.toml
index 3993f85..dd25afb 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,6 +11,7 @@ chrono = "0.4.40"
 clap = { version = "4.5.37", features = ["derive"] }
 flume = "0.11.1"
 hashbrown = "0.15.2"
+http = "1.3.1"
 http-body = "1.0.1"
 itertools = "0.14.0"
 rand = "0.9.1"
diff --git a/src/main.rs b/src/main.rs
index 4f499db..0898a60 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -10,6 +10,7 @@ use axum::{
 use axum_server::{bind_rustls, tls_rustls::RustlsConfig};
 use clap::Parser;
 use hashbrown::HashMap;
+use http::HeaderMap;
 use http_body::Frame;
 use itertools::Itertools;
 use rand::{prelude::*, Rng as _};
@@ -17,7 +18,7 @@ use rand_chacha::{rand_core::SeedableRng, ChaCha8Rng};
 use std::{
     borrow::Cow,
     convert::Infallible,
-    net::SocketAddr,
+    net::{IpAddr, SocketAddr},
     path::PathBuf,
     pin::Pin,
     sync::{
@@ -93,7 +94,7 @@ impl IntoResponse for SlowBody {
 }
 
 struct RequestStats {
-    sock: SocketAddr,
+    ip: IpAddr,
 }
 
 #[tokio::main]
@@ -124,11 +125,19 @@ async fn main() {
         Router::new().route(
             "/{id}",
             get(
-                |Path(id): Path<String>, ConnectInfo(sock): ConnectInfo<SocketAddr>| async move {
+                |Path(id): Path<String>,
+                 ConnectInfo(sock): ConnectInfo<SocketAddr>,
+                 headers: HeaderMap| async move {
                     // Create a RNG for this path (deterministic, to simulate static pages)
                     let mut rng = create_rng(id.bytes());
 
-                    stats_tx.send(RequestStats { sock }).unwrap();
+                    let ip = headers
+                        .get("X-Forwarded-For")
+                        .and_then(|h| h.to_str().ok())
+                        .and_then(|h| h.split(',').next())
+                        .and_then(|s| s.trim().parse().ok())
+                        .unwrap_or_else(|| sock.ip());
+                    stats_tx.send(RequestStats { ip }).unwrap();
 
                     // Count the request. Also doubles as the non-deterministic seed
                     let count = counter.fetch_add(1, Ordering::Relaxed);
@@ -197,7 +206,7 @@ async fn main() {
             interval.tick().await;
 
             while let Ok(stats) = stats_rx.try_recv() {
-                *worst_offenders.entry(stats.sock.ip()).or_default() += 1;
+                *worst_offenders.entry(stats.ip).or_default() += 1;
             }
 
             let count = counter.load(Ordering::Relaxed);
@@ -210,8 +219,8 @@ async fn main() {
                 let stats = worst_offenders
                     .iter()
                     .enumerate()
-                    .map(|(i, (sock, n))| format!("#{:>4} | {:>4} | {}", i + 1, n, sock))
-                    .join("\n");
+                    .map(|(i, (ip, n))| format!("#{:>4} | {:>4} | {}\n", i + 1, n, ip))
+                    .collect::<String>();
                 let _ = std::fs::write(STATS_FILE, &stats);
             }
         }