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 + * + * 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 { + 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)