From 247c059f29cad387068146976cb1cb2c6995d92a Mon Sep 17 00:00:00 2001
From: Paul Makles <paulmakles@gmail.com>
Date: Thu, 28 Apr 2022 15:17:04 +0100
Subject: [PATCH] fix(ratelimiter): read CF-Connecting-IP for identifying user

---
 src/util/ratelimiter.rs | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/util/ratelimiter.rs b/src/util/ratelimiter.rs
index 9db47408..4b9852d9 100644
--- a/src/util/ratelimiter.rs
+++ b/src/util/ratelimiter.rs
@@ -5,7 +5,6 @@
 
 use std::collections::hash_map::DefaultHasher;
 use std::hash::Hasher;
-use std::net::SocketAddr;
 use std::ops::Add;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
@@ -144,6 +143,27 @@ fn resolve_bucket_limit(bucket: &str) -> u8 {
     }
 }
 
+/// Find the remote IP of the client
+fn to_ip(request: &'_ rocket::Request<'_>) -> String {
+    request
+        .remote()
+        .map(|x| x.ip().to_string())
+        .unwrap_or_default()
+}
+
+/// Find the actual IP of the client
+fn to_real_ip(request: &'_ rocket::Request<'_>) -> String {
+    if let Ok(true) = std::env::var("TRUST_CLOUDFLARE").map(|x| x == "1") {
+        request
+            .headers()
+            .get_one("CF-Connecting-IP")
+            .map(|x| x.to_string())
+            .unwrap_or_else(|| to_ip(request))
+    } else {
+        to_ip(request)
+    }
+}
+
 impl Ratelimiter {
     /// Generate guard from identifier and target bucket
     pub fn from(
@@ -191,10 +211,8 @@ impl<'r> FromRequest<'r> for Ratelimiter {
                     request.guard::<rauth::entities::Session>().await
                 {
                     session.id.expect("`id` on User")
-                } else if let Outcome::Success(addr) = request.guard::<SocketAddr>().await {
-                    addr.ip().to_string()
                 } else {
-                    panic!("No identifier!");
+                    to_real_ip(request)
                 };
 
                 Ratelimiter::from(&identifier, resolve_bucket(request))
@@ -236,7 +254,7 @@ impl Fairing for RatelimitFairing {
             info!(
                 "User rate-limited on route {}! (IP = {:?})",
                 request.uri(),
-                request.remote().map(|x| x.ip())
+                to_real_ip(request)
             );
 
             request.set_method(Method::Get);