From 4b9698034cf46b18089042c365f4001944273921 Mon Sep 17 00:00:00 2001
From: iven <iven_h@qq.com>
Date: Tue, 31 Mar 2026 16:24:02 +0800
Subject: [PATCH] fix(saas): support X-Forwarded-For from trusted reverse
 proxies

---
 crates/zclaw-saas/src/middleware.rs | 82 +++++++++++++++++++++++++++--
 1 file changed, 79 insertions(+), 3 deletions(-)

diff --git a/crates/zclaw-saas/src/middleware.rs b/crates/zclaw-saas/src/middleware.rs
index 30b758c..85790be 100644
--- a/crates/zclaw-saas/src/middleware.rs
+++ b/crates/zclaw-saas/src/middleware.rs
@@ -130,13 +130,34 @@ pub async fn public_rate_limit_middleware(
     };
 
     // 从连接信息提取客户端 IP
-    // 安全策略: 仅使用 TCP 连接层 IP，不信任 X-Forwarded-For / X-Real-IP 头
-    // 反向代理场景下应使用 ConnectInfo<SocketAddr> 或在代理层做限流
-    let client_ip = req.extensions()
+    // 安全策略: 仅对配置的 trusted_proxies 解析 X-Forwarded-For 头
+    // 反向代理场景下，ConnectInfo 返回代理 IP，需从 XFF 获取真实客户端 IP
+    let connect_ip = req.extensions()
         .get::<axum::extract::ConnectInfo<std::net::SocketAddr>>()
         .map(|ci| ci.0.ip().to_string())
         .unwrap_or_else(|| "unknown".to_string());
 
+    let client_ip = {
+        let config = state.config.read().await;
+        let xff = req.headers()
+            .get("x-forwarded-for")
+            .and_then(|v| v.to_str().ok());
+
+        if let Some(xff_value) = xff {
+            if config.server.trusted_proxies.iter().any(|p| p == &connect_ip) {
+                xff_value.split(',')
+                    .next()
+                    .map(|s| s.trim().to_string())
+                    .filter(|s| !s.is_empty())
+                    .unwrap_or(connect_ip)
+            } else {
+                connect_ip
+            }
+        } else {
+            connect_ip
+        }
+    };
+
     let key = format!("{}:{}", key_prefix, client_ip);
     let now = Instant::now();
     let window_start = now - std::time::Duration::from_secs(window_secs);
@@ -160,3 +181,58 @@ pub async fn public_rate_limit_middleware(
 
     next.run(req).await
 }
+
+#[cfg(test)]
+mod tests {
+    // Imports kept for potential future use in integration tests
+    #[allow(unused_imports)]
+    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+
+    fn extract_client_ip(
+        connect_ip: &str,
+        xff_header: Option<&str>,
+        trusted_proxies: &[&str],
+    ) -> String {
+        if let Some(xff) = xff_header {
+            if trusted_proxies.iter().any(|p| *p == connect_ip) {
+                if let Some(client_ip) = xff.split(',').next() {
+                    let trimmed = client_ip.trim();
+                    if !trimmed.is_empty() {
+                        return trimmed.to_string();
+                    }
+                }
+            }
+        }
+        connect_ip.to_string()
+    }
+
+    #[test]
+    fn trusted_proxy_with_xff_uses_header_ip() {
+        let ip = extract_client_ip("127.0.0.1", Some("203.0.113.50"), &["127.0.0.1"]);
+        assert_eq!(ip, "203.0.113.50");
+    }
+
+    #[test]
+    fn trusted_proxy_without_xff_uses_connect_ip() {
+        let ip = extract_client_ip("127.0.0.1", None, &["127.0.0.1"]);
+        assert_eq!(ip, "127.0.0.1");
+    }
+
+    #[test]
+    fn untrusted_source_ignores_xff() {
+        let ip = extract_client_ip("198.51.100.1", Some("10.0.0.1"), &["127.0.0.1"]);
+        assert_eq!(ip, "198.51.100.1");
+    }
+
+    #[test]
+    fn empty_trusted_proxies_uses_connect_ip() {
+        let ip = extract_client_ip("127.0.0.1", Some("203.0.113.50"), &[]);
+        assert_eq!(ip, "127.0.0.1");
+    }
+
+    #[test]
+    fn xff_multiple_proxies_takes_first() {
+        let ip = extract_client_ip("127.0.0.1", Some("203.0.113.50, 10.0.0.1"), &["127.0.0.1"]);
+        assert_eq!(ip, "203.0.113.50");
+    }
+}