From 6f7a5e877ae7fd41f23374660f316ba5db3d6830 Mon Sep 17 00:00:00 2001
From: Frederic Henrichs <frederic@tinkerforge.com>
Date: Mon, 16 Dec 2024 13:29:38 +0100
Subject: [PATCH] backend: add rate limitation using actix-governor. fix #16

---
 backend/src/key_extractor.rs | 50 ++++++++++++++++++++++++++++++++++++
 backend/src/main.rs          | 12 +++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 backend/src/key_extractor.rs

diff --git a/backend/src/key_extractor.rs b/backend/src/key_extractor.rs
new file mode 100644
index 0000000..a2ce5fd
--- /dev/null
+++ b/backend/src/key_extractor.rs
@@ -0,0 +1,50 @@
+/* esp32-remote-access
+ * Copyright (C) 2024 Frederic Henrichs <frederic@tinkerforge.com>
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2 of the License, or (at your option) any later version.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this library; if not, write to the
+ * Free Software Foundation, Inc., 59 Temple Place - Suite 330,
+ * Boston, MA 02111-1307, USA.
+ */
+
+use actix_governor::{KeyExtractor, SimpleKeyExtractionError};
+
+/**
+ * The struct used to extract ip for ratelimiting
+ */
+#[derive(Clone)]
+pub struct Extractor;
+
+impl KeyExtractor for Extractor {
+    type Key = String;
+    type KeyExtractionError = SimpleKeyExtractionError<&'static str>;
+
+    fn extract(&self, req: &actix_web::dev::ServiceRequest) -> Result<Self::Key, Self::KeyExtractionError> {
+        let info = req.connection_info();
+        if let Some(ip) = info.realip_remote_addr() {
+            Ok(ip.to_string())
+        } else {
+            Err(SimpleKeyExtractionError::new("Invalid real IP"))
+        }
+    }
+
+    fn name(&self) -> &'static str {
+        "KeyExtractor"
+    }
+}
+
+impl Extractor {
+    pub fn new() -> Self {
+        Self
+    }
+}
diff --git a/backend/src/main.rs b/backend/src/main.rs
index 2dba306..bc226db 100644
--- a/backend/src/main.rs
+++ b/backend/src/main.rs
@@ -18,6 +18,7 @@
  */
 
 mod monitoring;
+mod key_extractor;
 
 use std::{
     collections::HashMap,
@@ -26,12 +27,14 @@ use std::{
     time::Duration,
 };
 
+use actix_governor::{Governor, GovernorConfigBuilder};
 use backend::utils::get_connection;
 pub use backend::*;
 
 use actix_web::{middleware::Logger, web, App, HttpServer};
 use db_connector::{get_connection_pool, run_migrations, Pool};
 use diesel::prelude::*;
+use key_extractor::Extractor;
 use lettre::{transport::smtp::authentication::Credentials, SmtpTransport};
 use simplelog::{ColorChoice, CombinedLogger, Config, LevelFilter, TermLogger, TerminalMode};
 use udp_server::packet::{
@@ -167,11 +170,20 @@ async fn main() -> std::io::Result<()> {
 
     udp_server::start_server(bridge_state.clone()).unwrap();
 
+    // Config for rate limitation
+    let governor_config = GovernorConfigBuilder::default()
+        .key_extractor(Extractor::new())
+        .requests_per_second(2)
+        .burst_size(20)
+        .finish()
+        .unwrap();
+
     HttpServer::new(move || {
         let cors = actix_cors::Cors::permissive();
         App::new()
             .wrap(cors)
             .wrap(Logger::default())
+            .wrap(Governor::new(&governor_config))
             .app_data(state.clone())
             .app_data(bridge_state.clone())
             .configure(routes::configure)