From 355689ed1926dde7e93e6489260ba62e87213f9b Mon Sep 17 00:00:00 2001
From: Geometrically <18202329+Geometrically@users.noreply.github.com>
Date: Thu, 9 Jun 2022 15:28:40 -0700
Subject: [PATCH] Replace ignore IP system with keys (#368)

---
 .env                        |  1 +
 src/main.rs                 | 11 +++--------
 src/ratelimit/middleware.rs | 30 ++++++++++++++++++------------
 3 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/.env b/.env
index b4aa007c..fefed064 100644
--- a/.env
+++ b/.env
@@ -4,6 +4,7 @@ RUST_LOG=info,sqlx::query=warn
 SITE_URL=https://modrinth.com
 CDN_URL=https://cdn.modrinth.com
 LABRINTH_ADMIN_KEY=feedbeef
+RATE_LIMIT_IGNORE_KEY=feedbeef
 
 MODERATION_DISCORD_WEBHOOK=
 CLOUDFLARE_INTEGRATION=false
diff --git a/src/main.rs b/src/main.rs
index de0314cf..b206b3ae 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -253,9 +253,8 @@ async fn main() -> std::io::Result<()> {
                     })
                     .with_interval(std::time::Duration::from_secs(60))
                     .with_max_requests(300)
-                    .with_ignore_ips(
-                        parse_strings_from_var("RATE_LIMIT_IGNORE_IPS")
-                            .unwrap_or_default(),
+                    .with_ignore_key(
+                        dotenv::var("RATE_LIMIT_IGNORE_KEY").ok(),
                     ),
             )
             .app_data(web::Data::new(pool.clone()))
@@ -292,11 +291,6 @@ fn check_env_vars() -> bool {
         check
     }
 
-    if parse_strings_from_var("RATE_LIMIT_IGNORE_IPS").is_none() {
-        warn!("Variable `RATE_LIMIT_IGNORE_IPS` missing in dotenv or not a json array of strings");
-        failed |= true;
-    }
-
     if parse_strings_from_var("WHITELISTED_MODPACK_DOMAINS").is_none() {
         warn!("Variable `WHITELISTED_MODPACK_DOMAINS` missing in dotenv or not a json array of strings");
         failed |= true;
@@ -305,6 +299,7 @@ fn check_env_vars() -> bool {
     failed |= check_var::<String>("SITE_URL");
     failed |= check_var::<String>("CDN_URL");
     failed |= check_var::<String>("LABRINTH_ADMIN_KEY");
+    failed |= check_var::<String>("RATE_LIMIT_IGNORE_KEY");
     failed |= check_var::<String>("DATABASE_URL");
     failed |= check_var::<String>("MEILISEARCH_ADDR");
     failed |= check_var::<String>("MEILISEARCH_KEY");
diff --git a/src/ratelimit/middleware.rs b/src/ratelimit/middleware.rs
index f70fab2c..1080f67f 100644
--- a/src/ratelimit/middleware.rs
+++ b/src/ratelimit/middleware.rs
@@ -30,7 +30,7 @@ where
     max_requests: usize,
     store: Addr<T>,
     identifier: RateLimiterIdentifier,
-    ignore_ips: Vec<String>,
+    ignore_key: Option<String>,
 }
 
 impl<T> RateLimiter<T>
@@ -51,7 +51,7 @@ where
             max_requests: 0,
             store,
             identifier: Rc::new(Box::new(identifier)),
-            ignore_ips: Vec::new(),
+            ignore_key: None
         }
     }
 
@@ -67,9 +67,9 @@ where
         self
     }
 
-    /// Sets IPs that should be ignored by the ratelimiter
-    pub fn with_ignore_ips(mut self, ignore_ips: Vec<String>) -> Self {
-        self.ignore_ips = ignore_ips;
+    /// Sets key which can be used to bypass rate-limiter
+    pub fn with_ignore_key(mut self, ignore_key: Option<String>) -> Self {
+        self.ignore_key = ignore_key;
         self
     }
 
@@ -107,7 +107,7 @@ where
             max_requests: self.max_requests,
             interval: self.interval.as_secs(),
             identifier: self.identifier.clone(),
-            ignore_ips: self.ignore_ips.clone(),
+            ignore_key: self.ignore_key.clone(),
         })
     }
 }
@@ -124,7 +124,7 @@ where
     max_requests: usize,
     interval: u64,
     identifier: RateLimiterIdentifier,
-    ignore_ips: Vec<String>,
+    ignore_key: Option<String>,
 }
 
 impl<T, S, B> Service<ServiceRequest> for RateLimitMiddleware<S, T>
@@ -154,14 +154,20 @@ where
         let max_requests = self.max_requests;
         let interval = Duration::from_secs(self.interval);
         let identifier = self.identifier.clone();
-        let ignore_ips = self.ignore_ips.clone();
+        let ignore_key = self.ignore_key.clone();
         Box::pin(async move {
             let identifier: String = (identifier)(&req)?;
-            if ignore_ips.contains(&identifier) {
-                let fut = srv.call(req);
-                let res = fut.await?;
-                return Ok(res);
+
+            if let Some(ignore_key) = ignore_key {
+                if let Some(key) = req.headers().get("x-ratelimit-key") {
+                    if key.to_str().ok().unwrap_or_default() == &*ignore_key {
+                        let fut = srv.call(req);
+                        let res = fut.await?;
+                        return Ok(res);
+                    }
+                }
             }
+
             let remaining: ActorResponse = store
                 .send(ActorMessage::Get(String::from(&identifier)))
                 .await