backend: Add rate limiter for routes that need charger credentials.

ffreddow · ffreddow · commit ecf06b1bb5c1 · 2024-12-23T10:00:16.000+01:00
diff --git a/backend/src/lib.rs b/backend/src/lib.rs
@@ -206,7 +206,7 @@ pub(crate) mod tests {
     use lru::LruCache;
     use rand::RngCore;
     use rand_core::OsRng;
-    use rate_limit::LoginRateLimiter;
+    use rate_limit::{ChargerRateLimiter, LoginRateLimiter};
     use routes::user::tests::{get_test_uuid, TestUser};
 
     pub struct ScopeCall<F: FnMut()> {
@@ -279,7 +279,9 @@ pub(crate) mod tests {
         let state = web::Data::new(state);
         let bridge_state = web::Data::new(bridge_state);
         let login_rate_limiter = web::Data::new(LoginRateLimiter::new());
+        let charger_rate_limiter = web::Data::new(ChargerRateLimiter::new());
         cfg.app_data(login_rate_limiter);
+        cfg.app_data(charger_rate_limiter);
         cfg.app_data(state);
         cfg.app_data(bridge_state);
         cfg.app_data(cache);
diff --git a/backend/src/main.rs b/backend/src/main.rs
@@ -35,7 +35,7 @@ use db_connector::{get_connection_pool, run_migrations, Pool};
 use diesel::prelude::*;
 use lettre::{transport::smtp::authentication::Credentials, SmtpTransport};
 use lru::LruCache;
-use rate_limit::LoginRateLimiter;
+use rate_limit::{ChargerRateLimiter, LoginRateLimiter};
 use simplelog::{ColorChoice, CombinedLogger, Config, LevelFilter, TermLogger, TerminalMode};
 use udp_server::packet::{
     ManagementCommand, ManagementCommandId, ManagementCommandPacket, ManagementPacket,
@@ -175,6 +175,7 @@ async fn main() -> std::io::Result<()> {
         web::Data::new(Mutex::new(LruCache::new(NonZeroUsize::new(10000).unwrap())));
 
     let login_ratelimiter = web::Data::new(LoginRateLimiter::new());
+    let charger_ratelimiter = web::Data::new(ChargerRateLimiter::new());
 
     HttpServer::new(move || {
         let cors = actix_cors::Cors::permissive();
@@ -184,6 +185,7 @@ async fn main() -> std::io::Result<()> {
             .app_data(cache.clone())
             .app_data(state.clone())
             .app_data(login_ratelimiter.clone())
+            .app_data(charger_ratelimiter.clone())
             .app_data(bridge_state.clone())
             .configure(routes::configure)
     })
diff --git a/backend/src/rate_limit.rs b/backend/src/rate_limit.rs
@@ -60,6 +60,17 @@ impl IPExtractor {
     }
 }
 
+fn ip_from_req(req: &HttpRequest) -> actix_web::Result<String> {
+    let ip = if let Some(ip) = req.connection_info().realip_remote_addr() {
+        ip.to_string()
+    } else {
+        println!("No ip found for route {}", req.path());
+        return Err(crate::error::Error::InternalError.into());
+    };
+
+    Ok(ip)
+}
+
 #[derive(Clone, PartialEq, Eq, Hash, Debug)]
 pub struct LoginRateLimitKey {
     user: String,
@@ -79,37 +90,69 @@ const REQUESTS_PER_SECOND: u32 = 5;
 const REQUESTS_BURST: u32 = 25;
 
 // RateLimiter for the login route
-pub struct LoginRateLimiter {
-    rate_limiter: RateLimiter<
+pub struct LoginRateLimiter(
+    RateLimiter<
         LoginRateLimitKey,
         dashmap::DashMap<LoginRateLimitKey, InMemoryState>,
         QuantaClock,
         governor::middleware::NoOpMiddleware<governor::clock::QuantaInstant>,
     >,
-}
+);
 
 impl LoginRateLimiter {
     pub fn new() -> Self {
-        Self {
-            rate_limiter: RateLimiter::keyed(
-                Quota::per_second(NonZeroU32::new(REQUESTS_PER_SECOND).unwrap())
-                    .allow_burst(NonZeroU32::new(REQUESTS_BURST).unwrap()),
-            ),
-        }
+        Self(RateLimiter::keyed(
+            Quota::per_second(NonZeroU32::new(REQUESTS_PER_SECOND).unwrap())
+                .allow_burst(NonZeroU32::new(REQUESTS_BURST).unwrap()),
+        ))
     }
 
     pub fn check(&self, email: String, req: &HttpRequest) -> actix_web::Result<()> {
-        let ip = if let Some(ip) = req.connection_info().realip_remote_addr() {
-            ip.to_string()
-        } else {
-            println!("No ip found for route {}", req.path());
-            return Err(crate::error::Error::InternalError.into());
-        };
+        let ip = ip_from_req(req)?;
 
         let key = LoginRateLimitKey { user: email, ip };
-        if let Err(err) = self.rate_limiter.check_key(&key) {
+        if let Err(err) = self.0.check_key(&key) {
             log::warn!("RateLimiter triggered for {:?}", key);
-            let now = self.rate_limiter.clock().now();
+            let now = self.0.clock().now();
+
+            Err(RateLimitError::new(err, now).into())
+        } else {
+            Ok(())
+        }
+    }
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct ChargerRateLimitKey {
+    charger_id: String,
+    ip: String,
+}
+
+// Rate limiter for all routes that get called by chargers
+pub struct ChargerRateLimiter(
+    RateLimiter<
+        ChargerRateLimitKey,
+        dashmap::DashMap<ChargerRateLimitKey, InMemoryState>,
+        QuantaClock,
+        governor::middleware::NoOpMiddleware<governor::clock::QuantaInstant>,
+    >,
+);
+
+impl ChargerRateLimiter {
+    pub fn new() -> Self {
+        Self(RateLimiter::keyed(
+            Quota::per_second(NonZeroU32::new(REQUESTS_PER_SECOND).unwrap())
+                .allow_burst(NonZeroU32::new(REQUESTS_BURST).unwrap()),
+        ))
+    }
+
+    pub fn check(&self, charger_id: String, req: &HttpRequest) -> actix_web::Result<()> {
+        let ip = ip_from_req(req)?;
+
+        let key = ChargerRateLimitKey { charger_id, ip };
+        if let Err(err) = self.0.check_key(&key) {
+            log::warn!("RateLimiter triggered for {:?}", key);
+            let now = self.0.clock().now();
 
             Err(RateLimitError::new(err, now).into())
         } else {
@@ -155,6 +198,8 @@ impl ResponseError for RateLimitError {
 mod tests {
     use actix_web::test;
 
+    use crate::rate_limit::ChargerRateLimiter;
+
     use super::LoginRateLimiter;
 
     #[actix_web::test]
@@ -195,4 +240,43 @@ mod tests {
         let ret = limiter.check(email.clone(), &req);
         assert!(ret.is_ok());
     }
+
+    #[actix_web::test]
+    async fn test_charger_rate_limiter() {
+        let limiter = ChargerRateLimiter::new();
+        let req = test::TestRequest::get()
+            .uri("/login")
+            .insert_header(("X-Forwarded-For", "123.123.123.2"))
+            .to_http_request();
+        let email = uuid::Uuid::new_v4().to_string();
+
+        let ret = limiter.check(email.clone(), &req);
+        assert!(ret.is_ok());
+
+        let ret = limiter.check(email.clone(), &req);
+        assert!(ret.is_ok());
+
+        let ret = limiter.check(email.clone(), &req);
+        assert!(ret.is_ok());
+
+        let ret = limiter.check(email.clone(), &req);
+        assert!(ret.is_ok());
+
+        let ret = limiter.check(email.clone(), &req);
+        assert!(ret.is_ok());
+
+        let ret = limiter.check(email.clone(), &req);
+        assert!(ret.is_err());
+
+        let email2 = uuid::Uuid::new_v4().to_string();
+        let ret = limiter.check(email2.clone(), &req);
+        assert!(ret.is_ok());
+
+        let req = test::TestRequest::get()
+            .uri("/login")
+            .insert_header(("X-Forwarded-For", "123.123.123.3"))
+            .to_http_request();
+        let ret = limiter.check(email.clone(), &req);
+        assert!(ret.is_ok());
+    }
 }
diff --git a/backend/src/routes/charger/allow_user.rs b/backend/src/routes/charger/allow_user.rs
@@ -17,7 +17,7 @@
  * Boston, MA 02111-1307, USA.
  */
 
-use actix_web::{error::ErrorBadRequest, put, web, HttpResponse, Responder};
+use actix_web::{error::ErrorBadRequest, put, web, HttpRequest, HttpResponse, Responder};
 use base64::Engine;
 use db_connector::models::{allowed_users::AllowedUser, wg_keys::WgKey};
 use diesel::prelude::*;
@@ -26,6 +26,7 @@ use utoipa::ToSchema;
 
 use crate::{
     error::Error,
+    rate_limit::ChargerRateLimiter,
     routes::{
         auth::login::{validate_password, FindBy},
         charger::add::get_charger_from_db,
@@ -124,7 +125,11 @@ async fn authenticate_user(
 pub async fn allow_user(
     state: web::Data<AppState>,
     allow_user: web::Json<AllowUserSchema>,
+    rate_limiter: web::Data<ChargerRateLimiter>,
+    req: HttpRequest,
 ) -> Result<impl Responder, actix_web::Error> {
+    rate_limiter.check(allow_user.charger_id.clone(), &req)?;
+
     let cid = parse_uuid(&allow_user.charger_id)?;
 
     let charger = get_charger_from_db(cid, &state).await?;
@@ -226,6 +231,7 @@ pub mod tests {
         };
         let req = test::TestRequest::put()
             .uri("/allow_user")
+            .append_header(("X-Forwarded-For", "123.123.123.3"))
             .set_json(body)
             .to_request();
         let resp = test::call_service(&app, req).await;
@@ -257,6 +263,7 @@ pub mod tests {
         };
         let req = test::TestRequest::put()
             .uri("/allow_user")
+            .append_header(("X-Forwarded-For", "123.123.123.3"))
             .set_json(allow)
             .to_request();
         let resp = test::call_service(&app, req).await;
@@ -287,6 +294,7 @@ pub mod tests {
         };
         let req = test::TestRequest::put()
             .uri("/allow_user")
+            .append_header(("X-Forwarded-For", "123.123.123.3"))
             .set_json(allow)
             .to_request();
         let resp = test::call_service(&app, req).await;
@@ -316,6 +324,7 @@ pub mod tests {
         };
         let req = test::TestRequest::put()
             .uri("/allow_user")
+            .append_header(("X-Forwarded-For", "123.123.123.3"))
             .set_json(allow)
             .to_request();
         let resp = test::call_service(&app, req).await;
@@ -344,6 +353,7 @@ pub mod tests {
         };
         let req = test::TestRequest::put()
             .uri("/allow_user")
+            .append_header(("X-Forwarded-For", "123.123.123.3"))
             .set_json(allow)
             .to_request();
         let resp = test::call_service(&app, req).await;
@@ -361,6 +371,7 @@ pub mod tests {
         };
         let req = test::TestRequest::put()
             .uri("/allow_user")
+            .append_header(("X-Forwarded-For", "123.123.123.3"))
             .set_json(allow)
             .to_request();
         let resp = test::call_service(&app, req).await;
diff --git a/backend/src/routes/management.rs b/backend/src/routes/management.rs
@@ -31,6 +31,7 @@ use utoipa::ToSchema;
 
 use crate::{
     error::Error,
+    rate_limit::ChargerRateLimiter,
     routes::{auth::login::FindBy, charger::add::get_charger_from_db, user::get_user_id},
     utils::{get_charger_by_uid, get_connection, parse_uuid, web_block_unpacked},
     AppState, BridgeState,
@@ -52,7 +53,7 @@ pub enum ManagementDataVersion {
     V2(ManagementDataVersion2),
 }
 
-#[derive(Serialize, Deserialize, ToSchema)]
+#[derive(Serialize, Deserialize, ToSchema, Debug)]
 pub struct ConfiguredUser {
     pub email: String,
     pub name: Option<String>,
@@ -233,6 +234,7 @@ pub async fn management(
     state: web::Data<AppState>,
     data: web::Json<ManagementSchema>,
     bridge_state: web::Data<BridgeState>,
+    rate_limiter: web::Data<ChargerRateLimiter>,
 ) -> actix_web::Result<impl Responder> {
     use db_connector::schema::chargers::dsl as chargers;
 
@@ -248,6 +250,8 @@ pub async fn management(
     let charger_id;
     let mut output_uuid = None;
     let charger = if let Some(charger_uid) = data.id {
+        rate_limiter.check(charger_uid.to_string(), &req)?;
+
         let charger = get_charger_by_uid(charger_uid, data.password.clone(), &state).await?;
         charger_id = charger.id;
         output_uuid = Some(charger_id.to_string());
@@ -256,6 +260,8 @@ pub async fn management(
         match &data.data {
             ManagementDataVersion::V1(_) => return Err(Error::ChargerCredentialsWrong.into()),
             ManagementDataVersion::V2(data) => {
+                rate_limiter.check(data.id.clone(), &req)?;
+
                 charger_id = parse_uuid(&data.id)?;
                 let charger = get_charger_from_db(charger_id, &state).await?;
                 if !password_matches(&data.password, &charger.password)? {
@@ -415,7 +421,6 @@ mod tests {
         let req = test::TestRequest::put()
             .uri("/management")
             .append_header(("X-Forwarded-For", "123.123.123.3"))
-            .cookie(Cookie::new("X-Forwarded-For", "123.123.123.3"))
             .set_json(body)
             .to_request();
         let resp: ManagementResponseSchema = test::call_and_read_body_json(&app, req).await;
@@ -447,7 +452,6 @@ mod tests {
         let req = test::TestRequest::put()
             .uri("/management")
             .append_header(("X-Forwarded-For", "123.123.123.3"))
-            .cookie(Cookie::new("X-Forwarded-For", "123.123.123.3"))
             .set_json(body)
             .to_request();
         let resp: ManagementResponseSchema = test::call_and_read_body_json(&app, req).await;
@@ -482,7 +486,6 @@ mod tests {
         let req = test::TestRequest::put()
             .uri("/management")
             .append_header(("X-Forwarded-For", "123.123.123.3"))
-            .cookie(Cookie::new("X-Forwarded-For", "123.123.123.3"))
             .set_json(body)
             .to_request();
         let resp: ManagementResponseSchema = test::call_and_read_body_json(&app, req).await;
diff --git a/backend/src/routes/selfdestruct.rs b/backend/src/routes/selfdestruct.rs
