more ratelimits

Author: 0x7d8

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "web"
-version = "3.1.19"
+version = "3.1.20"
 edition = "2024"
 
 [dependencies]
@@ -14,7 +14,7 @@ axum = "0.8.1"
 colored = "3.0.0"
 sqlx = { version = "0.8", features = ["runtime-tokio", "tls-rustls-ring-webpki", "postgres", "chrono", "ipnetwork", "uuid"] }
 dotenvy = "0.15.7"
-rustis = "0.16.1"
+rustis = "0.19.1"
 serde = { version = "1.0.218", features = ["derive"] }
 serde_json = { version = "1.0.140", features = ["preserve_order"] }
 tokio = { version = "1.43.0", features = ["full"] }

apps/backend/Cargo.toml

@@ -1,8 +1,9 @@
-use crate::env::RedisMode;
+use crate::{env::RedisMode, response::ApiResponse};
+use axum::http::StatusCode;
 use colored::Colorize;
 use rustis::{
     client::Client,
-    commands::{GenericCommands, SetCondition, SetExpiration, StringCommands},
+    commands::{GenericCommands, SetExpiration, StringCommands},
     resp::{BulkString, cmd},
 };
 use serde::{Serialize, de::DeserializeOwned};
@@ -59,6 +60,43 @@ impl Cache {
         instance
     }
 
+    pub async fn ratelimit(
+        &self,
+        limit_identifier: impl AsRef<str>,
+        limit: u64,
+        limit_window: u64,
+        client: impl AsRef<str>,
+    ) -> Result<(), ApiResponse> {
+        let key = format!(
+            "ratelimit::{}::{}",
+            limit_identifier.as_ref(),
+            client.as_ref()
+        );
+
+        let now = chrono::Utc::now().timestamp();
+        let expiry = self.client.expiretime(&key).await.unwrap_or_default();
+        let expire_unix: u64 = if expiry > now + 2 {
+            expiry as u64
+        } else {
+            now as u64 + limit_window
+        };
+
+        let limit_used = self.client.get::<u64>(&key).await.unwrap_or_default() + 1;
+        self.client
+            .set_with_options(key, limit_used, None, SetExpiration::Exat(expire_unix))
+            .await?;
+
+        if limit_used >= limit {
+            return Err(ApiResponse::error(&format!(
+                "you are ratelimited, retry in {}s",
+                expiry - now
+            ))
+            .with_status(StatusCode::TOO_MANY_REQUESTS));
+        }
+
+        Ok(())
+    }
+
     #[tracing::instrument(skip(self, fn_compute))]
     pub async fn cached<T, F, Fut>(
         &self,
@@ -83,13 +121,7 @@ impl Cache {
 
                 let serialized = rmp_serde::to_vec(&result)?;
                 self.client
-                    .set_with_options(
-                        key,
-                        serialized,
-                        SetCondition::None,
-                        SetExpiration::Ex(ttl),
-                        false,
-                    )
+                    .set_with_options(key, serialized, None, SetExpiration::Ex(ttl))
                     .await?;
 
                 Ok(result)

apps/backend/src/cache.rs

@@ -65,7 +65,7 @@ pub struct Env {
     pub bind: String,
     pub port: u16,
 
-    pub telemetry_ratelimit_per_day: i64,
+    pub telemetry_ratelimit_per_day: u64,
 
     pub update_prices: bool,
     pub sxc_token: Option<String>,

apps/backend/src/env.rs

@@ -53,6 +53,11 @@ mod post {
                 .ok();
         }
 
+        state
+            .cache
+            .ratelimit("login", 5, 60, ip.to_string())
+            .await?;
+
         if let Err(error) = state.captcha.verify(ip, data.captcha).await {
             return ApiResponse::error(&error)
                 .with_status(StatusCode::BAD_REQUEST)

apps/backend/src/routes/auth/login/mod.rs

@@ -39,6 +39,15 @@ mod post {
                 .ok();
         }
 
+        state
+            .cache
+            .ratelimit("password_forgot", 2, 5 * 60, ip.to_string())
+            .await?;
+        state
+            .cache
+            .ratelimit("password_forgot", 2, 5 * 60, &data.email)
+            .await?;
+
         if let Err(error) = state.captcha.verify(ip, data.captcha).await {
             return ApiResponse::error(&error)
                 .with_status(StatusCode::BAD_REQUEST)

apps/backend/src/routes/auth/password/forgot.rs

@@ -57,6 +57,11 @@ mod post {
                 .ok();
         }
 
+        state
+            .cache
+            .ratelimit("register", 2, 5 * 60, ip.to_string())
+            .await?;
+
         if let Err(error) = state.captcha.verify(ip, data.captcha).await {
             return ApiResponse::error(&error)
                 .with_status(StatusCode::BAD_REQUEST)

apps/backend/src/routes/auth/register.rs

@@ -22,12 +22,9 @@ pub fn router(state: &State) -> OpenApiRouter<State> {
                         }
                     };
 
-                    let telemetry = state.telemetry.log(ip, data).await;
-                    if telemetry.is_none() {
-                        return ApiResponse::error("too many requests")
-                            .with_status(StatusCode::TOO_MANY_REQUESTS)
-                            .ok();
-                    }
+                    state.cache.ratelimit("telemetry", state.env.telemetry_ratelimit_per_day, 24 * 60 * 60, ip.to_string()).await?;
+
+                    state.telemetry.log(ip, data).await;
 
                     ApiResponse::new_serialized(json!({})).ok()
                 },

apps/backend/src/routes/telemetry.rs

@@ -41,6 +41,15 @@ mod patch {
                 .ok();
         }
 
+        state
+            .cache
+            .ratelimit("update_email", 3, 5 * 60, ip.to_string())
+            .await?;
+        state
+            .cache
+            .ratelimit("update_email", 2, 5 * 60, user.id.to_string())
+            .await?;
+
         if let Err(error) = state.captcha.verify(ip, data.captcha).await {
             return ApiResponse::error(&error)
                 .with_status(StatusCode::BAD_REQUEST)

apps/backend/src/routes/user/email/mod.rs

@@ -1,6 +1,5 @@
 use chrono::NaiveDateTime;
 use colored::Colorize;
-use rustis::commands::{ExpireOption, GenericCommands, StringCommands};
 use serde::{Deserialize, Serialize};
 use sha2::Digest;
 use sqlx::types::Uuid;
@@ -101,24 +100,7 @@ impl TelemetryLogger {
     }
 
     #[inline]
-    pub async fn log(&self, ip: IpAddr, telemetry: TelemetryData) -> Option<()> {
-        let mut processing = self.processing.lock().await;
-
-        let ratelimit_key = format!("blueprint_api::ratelimit::{ip}");
-
-        let count = self.cache.client.incr(&ratelimit_key).await.ok()?;
-        if count == 1 {
-            self.cache
-                .client
-                .expire(&ratelimit_key, 86400, ExpireOption::None)
-                .await
-                .ok()?;
-        }
-
-        if count > self.env.telemetry_ratelimit_per_day {
-            return None;
-        }
-
+    pub async fn log(&self, ip: IpAddr, telemetry: TelemetryData) {
         let data = Telemetry {
             panel_id: telemetry.id,
             telemetry_version: telemetry.telemetry_version as i16,
@@ -129,9 +111,7 @@ impl TelemetryLogger {
             created: chrono::Utc::now().naive_utc(),
         };
 
-        processing.push(data);
-
-        Some(())
+        self.processing.lock().await.push(data);
     }
 
     #[inline]