Skip to content

Commit

Permalink
backend: Add rate limiter for routes that need charger credentials.
Browse files Browse the repository at this point in the history
  • Loading branch information
ffreddow committed Dec 23, 2024
1 parent 8744dfd commit ecf06b1
Show file tree
Hide file tree
Showing 6 changed files with 139 additions and 26 deletions.
4 changes: 3 additions & 1 deletion backend/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()> {
Expand Down Expand Up @@ -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);
Expand Down
4 changes: 3 additions & 1 deletion backend/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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();
Expand All @@ -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)
})
Expand Down
118 changes: 101 additions & 17 deletions backend/src/rate_limit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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 {
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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());
}
}
13 changes: 12 additions & 1 deletion backend/src/routes/charger/allow_user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand All @@ -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,
Expand Down Expand Up @@ -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?;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down
11 changes: 7 additions & 4 deletions backend/src/routes/management.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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>,
Expand Down Expand Up @@ -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;

Expand All @@ -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());
Expand All @@ -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)? {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
Loading

0 comments on commit ecf06b1

Please sign in to comment.