From 8e5b921771bbc02a1f61f042251f4d0ec7426ea9 Mon Sep 17 00:00:00 2001
From: svatantrya <svatantrya@bitcoin-after.life>
Date: Wed, 19 Aug 2026 02:54:46 -0400
Subject: [PATCH] fix: rate limit per real client IP behind reverse proxy

Replace PeerIpKeyExtractor with custom RealIpKeyExtractor that extracts
the real client IP from X-Real-IP/X-Forwarded-For headers when the
request comes from a trusted proxy (default 127.0.0.1).

This fixes rate limiting being applied to the proxy IP instead of
individual clients when deployed behind Nginx.
---
 src/bin/bal-server.rs | 59 +++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 57 insertions(+), 2 deletions(-)

diff --git a/src/bin/bal-server.rs b/src/bin/bal-server.rs
index 43cf646..c8bac58 100644
--- a/src/bin/bal-server.rs
+++ b/src/bin/bal-server.rs
@@ -1,4 +1,5 @@
-use actix_governor::{Governor, GovernorConfigBuilder};
+use actix_governor::{Governor, GovernorConfigBuilder, KeyExtractor, SimpleKeyExtractionError};
+use actix_web::dev::ServiceRequest;
 use actix_web::middleware;
 use actix_web::web::Bytes;
 use actix_web::{App, HttpResponse, HttpServer, Responder, web};
@@ -10,6 +11,8 @@ use serde::{Deserialize, Serialize};
 use std::collections::{HashMap, HashSet};
 use std::env;
 use std::fs;
+use std::net::IpAddr;
+use std::str::FromStr;
 
 use bal_server::db::{
     DatabasePool, InsertInpData, InsertOutData, InsertTxData, check_duplicate_txids,
@@ -130,6 +133,7 @@ struct ActixConfig {
     rate_limit_default: (u64, u32),
     workers: usize,
     max_connections: usize,
+    trusted_proxy: IpAddr,
 }
 
 fn parse_actix_config() -> ActixConfig {
@@ -190,6 +194,10 @@ fn parse_actix_config() -> ActixConfig {
             .unwrap_or("100".to_string())
             .parse::<usize>()
             .unwrap_or(100),
+        trusted_proxy: env::var("BAL_SERVER_TRUSTED_PROXY")
+            .unwrap_or_else(|_| "127.0.0.1".to_string())
+            .parse::<IpAddr>()
+            .unwrap_or(IpAddr::from_str("127.0.0.1").unwrap()),
     }
 }
 
@@ -223,6 +231,51 @@ fn is_valid_ip(ip: &str) -> bool {
     ip.parse::<std::net::IpAddr>().is_ok()
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+struct RealIpKeyExtractor;
+
+impl KeyExtractor for RealIpKeyExtractor {
+    type Key = IpAddr;
+    type KeyExtractionError = SimpleKeyExtractionError<&'static str>;
+
+    fn extract(&self, req: &ServiceRequest) -> Result<Self::Key, Self::KeyExtractionError> {
+        let proxy_ip = req
+            .app_data::<web::Data<IpAddr>>()
+            .map(|ip| *ip.get_ref())
+            .unwrap_or_else(|| IpAddr::from_str("0.0.0.0").unwrap());
+
+        let peer_ip = req.peer_addr().map(|socket| socket.ip());
+        let connection_info = req.connection_info();
+
+        match peer_ip {
+            Some(peer) if peer == proxy_ip => connection_info
+                .realip_remote_addr()
+                .ok_or_else(|| {
+                    SimpleKeyExtractionError::new("Could not extract real IP address from request")
+                })
+                .and_then(|str| {
+                    str.parse::<IpAddr>().map_err(|_| {
+                        SimpleKeyExtractionError::new(
+                            "Could not extract real IP address from request",
+                        )
+                    })
+                }),
+            _ => connection_info
+                .peer_addr()
+                .ok_or_else(|| {
+                    SimpleKeyExtractionError::new("Could not extract peer IP address from request")
+                })
+                .and_then(|str| {
+                    str.parse::<IpAddr>().map_err(|_| {
+                        SimpleKeyExtractionError::new(
+                            "Could not extract peer IP address from request",
+                        )
+                    })
+                }),
+        }
+    }
+}
+
 fn extract_client_ip(req: &actix_web::HttpRequest) -> String {
     if let Some(val) = req.headers().get("X-Real-IP")
         && let Ok(s) = val.to_str()
@@ -776,9 +829,10 @@ async fn main() -> std::io::Result<()> {
     let bind_address = data.cfg.bind_address.clone();
     let bind_port = data.cfg.bind_port;
 
-    let governor_conf = GovernorConfigBuilder::const_default()
+    let governor_conf = GovernorConfigBuilder::default()
         .seconds_per_request(actix_cfg.rate_limit_pushtxs.0)
         .burst_size(actix_cfg.rate_limit_pushtxs.1)
+        .key_extractor(RealIpKeyExtractor)
         .finish()
         .unwrap();
 
@@ -788,6 +842,7 @@ async fn main() -> std::io::Result<()> {
         App::new()
             .app_data(web::PayloadConfig::default().limit(actix_cfg.max_body_size))
             .app_data(data.clone())
+            .app_data(web::Data::new(actix_cfg.trusted_proxy))
             .wrap(middleware::Logger::default())
             .wrap(middleware::Compress::default())
             .wrap(Governor::new(&governor_conf))