From ee6f2426672eec1eb9c81cb6aca1496b4af547fe Mon Sep 17 00:00:00 2001 From: Vitor Hideyoshi Nakazone Batista Date: Mon, 26 Feb 2024 01:03:26 -0300 Subject: [PATCH] Implements Request Limit with Redis --- Cargo.lock | 61 ++++++++++++++++++++++++++-- Cargo.toml | 1 + src/config/config_auth.rs | 2 + src/config/config_email.rs | 2 + src/config/config_limits.rs | 23 +++++++++++ src/config/config_redis.rs | 24 +++++++++++ src/config/config_server.rs | 2 + src/config/mod.rs | 2 + src/depends/depends_auth_service.rs | 15 +++++++ src/depends/depends_email_service.rs | 11 +++++ src/depends/mod.rs | 2 + src/handler/message.rs | 16 +++++++- src/main.rs | 2 +- src/route.rs | 8 ++-- src/service/auth_service.rs | 60 ++++++++++++++++++++++++++- 15 files changed, 220 insertions(+), 11 deletions(-) create mode 100644 src/config/config_limits.rs create mode 100644 src/config/config_redis.rs create mode 100644 src/depends/depends_auth_service.rs create mode 100644 src/depends/depends_email_service.rs create mode 100644 src/depends/mod.rs diff --git a/Cargo.lock b/Cargo.lock index f03f13a..83cd7ef 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -246,6 +246,20 @@ dependencies = [ "stacker", ] +[[package]] +name = "combine" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35ed6e9d84f0b51a7f52daf1c7d71dd136fd7a3f41a8462b8cdb8c78d920fad4" +dependencies = [ + "bytes", + "futures-core", + "memchr", + "pin-project-lite", + "tokio", + "tokio-util", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -495,6 +509,7 @@ checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" dependencies = [ "futures-core", "futures-io", + "futures-sink", "futures-task", "memchr", "pin-project-lite", @@ -715,7 +730,7 @@ dependencies = [ "httpdate", "itoa", "pin-project-lite", - "socket2", + "socket2 0.5.5", "tokio", "tower-service", "tracing", @@ -767,7 +782,7 @@ dependencies = [ "http-body 1.0.0", "hyper 1.2.0", "pin-project-lite", - "socket2", + "socket2 0.5.5", "tokio", ] @@ -891,7 +906,7 @@ dependencies = [ "quoted_printable", "rustls", "rustls-pemfile 2.1.0", - "socket2", + "socket2 0.5.5", "tokio", "tokio-rustls", "url", @@ -956,6 +971,7 @@ dependencies = [ "http 1.0.0", "lettre", "log", + "redis", "reqwest", "serde", "serde_json", @@ -1214,6 +1230,27 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "79ec282e887b434b68c18fe5c121d38e72a5cf35119b59e54ec5b992ea9c8eb0" +[[package]] +name = "redis" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c580d9cbbe1d1b479e8d67cf9daf6a62c957e6846048408b80b43ac3f6af84cd" +dependencies = [ + "async-trait", + "bytes", + "combine", + "futures-util", + "itoa", + "percent-encoding", + "pin-project-lite", + "ryu", + "sha1_smol", + "socket2 0.4.10", + "tokio", + "tokio-util", + "url", +] + [[package]] name = "redox_syscall" version = "0.4.1" @@ -1491,6 +1528,12 @@ dependencies = [ "digest", ] +[[package]] +name = "sha1_smol" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae1a47186c03a32177042e55dbc5fd5aee900b8e0069a8d70fba96a9375cd012" + [[package]] name = "signal-hook-registry" version = "1.4.1" @@ -1515,6 +1558,16 @@ version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6ecd384b10a64542d77071bd64bd7b231f4ed5940fba55e98c3de13824cf3d7" +[[package]] +name = "socket2" +version = "0.4.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7916fc008ca5542385b89a3d3ce689953c143e9304a9bf8beec1de48994c0d" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "socket2" version = "0.5.5" @@ -1697,7 +1750,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2", + "socket2 0.5.5", "tokio-macros", "windows-sys 0.48.0", ] diff --git a/Cargo.toml b/Cargo.toml index 767e02e..35fc75d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,3 +20,4 @@ cached = "0.49.2" dotenv = "0.15.0" log = "0.4.20" lettre = { version = "0.11.4", default-features = false, features = ["smtp-transport", "tokio1-rustls-tls", "hostname", "builder"] } +redis = { version = "0.24.0", features = ["aio", "tokio-comp"] } diff --git a/src/config/config_auth.rs b/src/config/config_auth.rs index 2e3bbcf..cfb8666 100644 --- a/src/config/config_auth.rs +++ b/src/config/config_auth.rs @@ -8,6 +8,8 @@ pub struct ConfigAuth { #[cached] pub fn get_config_auth() -> ConfigAuth { + dotenv::dotenv().ok(); + let url = env::var("AUTH_URL").expect("AUTH_URL must be set"); ConfigAuth { auth_url: url } diff --git a/src/config/config_email.rs b/src/config/config_email.rs index c967998..d5c87d9 100644 --- a/src/config/config_email.rs +++ b/src/config/config_email.rs @@ -13,6 +13,8 @@ pub struct ConfigEmail { #[cached] pub fn get_config_email() -> ConfigEmail { + dotenv::dotenv().ok(); + let server = env::var("SMTP_SERVER").expect("SMTP_SERVER must be set"); let port = env::var("SMTP_PORT").expect("SMTP_PORT must be set"); let username = env::var("SMTP_USERNAME").expect("SMTP_USERNAME must be set"); diff --git a/src/config/config_limits.rs b/src/config/config_limits.rs new file mode 100644 index 0000000..9bc2cfc --- /dev/null +++ b/src/config/config_limits.rs @@ -0,0 +1,23 @@ +use cached::proc_macro::cached; +use std::env; + +#[derive(Clone)] +pub struct ConfigLimits { + pub max_requests: u32, + pub expiration_time: usize, +} + +#[cached] +pub fn get_config_limits() -> ConfigLimits { + dotenv::dotenv().ok(); + + let max_requests = env::var("MAX_REQUESTS").unwrap_or("10".to_string()) + .parse::().unwrap(); + let expiration_time = env::var("EXPIRATION_TIME").unwrap_or("604800".to_string()) + .parse::().unwrap(); + + ConfigLimits { + max_requests, + expiration_time, + } +} diff --git a/src/config/config_redis.rs b/src/config/config_redis.rs new file mode 100644 index 0000000..4a183fe --- /dev/null +++ b/src/config/config_redis.rs @@ -0,0 +1,24 @@ +use cached::proc_macro::cached; +use std::env; + +#[derive(Clone)] +pub struct ConfigRedis { + pub redis_url: String, + pub redis_port: u16, + pub redis_password: Option, +} + +#[cached] +pub fn get_config_redis() -> ConfigRedis { + dotenv::dotenv().ok(); + + let url = env::var("REDIS_URL").unwrap_or("localhost".to_string()); + let port = env::var("REDIS_PORT").unwrap_or("6379".to_string()); + let password = env::var("REDIS_PASSWORD").ok(); + + ConfigRedis { + redis_url: url, + redis_port: port.parse::().unwrap(), + redis_password: password + } +} diff --git a/src/config/config_server.rs b/src/config/config_server.rs index 02c7bf1..e0798c8 100644 --- a/src/config/config_server.rs +++ b/src/config/config_server.rs @@ -9,6 +9,8 @@ pub struct ConfigServer { #[cached] pub fn get_config_server() -> ConfigServer { + dotenv::dotenv().ok(); + let h = option_env!("HOST").unwrap_or("localhost").to_string(); let p = option_env!("PORT") diff --git a/src/config/mod.rs b/src/config/mod.rs index 54585ba..50630ae 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1,3 +1,5 @@ pub mod config_auth; pub mod config_email; pub mod config_server; +pub mod config_redis; +pub mod config_limits; diff --git a/src/depends/depends_auth_service.rs b/src/depends/depends_auth_service.rs new file mode 100644 index 0000000..772d411 --- /dev/null +++ b/src/depends/depends_auth_service.rs @@ -0,0 +1,15 @@ +use cached::proc_macro::cached; +use crate::config::config_auth::get_config_auth; +use crate::config::config_limits::get_config_limits; +use crate::config::config_redis::get_config_redis; +use crate::service::auth_service::AuthService; + + +#[cached] +pub fn get_depends_auth_service() -> AuthService { + AuthService::new( + get_config_auth(), + get_config_redis(), + get_config_limits(), + ) +} diff --git a/src/depends/depends_email_service.rs b/src/depends/depends_email_service.rs new file mode 100644 index 0000000..86a412d --- /dev/null +++ b/src/depends/depends_email_service.rs @@ -0,0 +1,11 @@ +use cached::proc_macro::cached; +use crate::config::config_email::get_config_email; +use crate::service::email_service::EmailService; + + +#[cached] +pub fn get_depends_email_service() -> EmailService { + EmailService::new( + get_config_email() + ) +} diff --git a/src/depends/mod.rs b/src/depends/mod.rs new file mode 100644 index 0000000..1bba9bd --- /dev/null +++ b/src/depends/mod.rs @@ -0,0 +1,2 @@ +pub mod depends_auth_service; +pub mod depends_email_service; \ No newline at end of file diff --git a/src/handler/message.rs b/src/handler/message.rs index 3f51c15..7ebc1bd 100644 --- a/src/handler/message.rs +++ b/src/handler/message.rs @@ -2,15 +2,27 @@ use crate::model::generic_response::GenericResponse; use crate::model::send_message::{MessageAuthor, SendMessage}; use crate::service::email_service::EmailService; use axum::{http::StatusCode, response::IntoResponse, Extension, Json}; +use crate::service::auth_service::AuthService; pub async fn send_message( + Extension(auth_service): Extension, Extension(email_service): Extension, Extension(author): Extension, Json(payload): Json, ) -> impl IntoResponse { let mut package = payload.clone(); - package.author = Some(author).clone(); + package.author = Some(author.clone()).clone(); + + if auth_service.has_user_reached_limit(&author).await { + return ( + StatusCode::TOO_MANY_REQUESTS, + Json(GenericResponse { + status: StatusCode::TOO_MANY_REQUESTS.to_string(), + message: "User has reached the limit of messages".to_string(), + }), + ); + } match email_service.send_email_smtp(package).await { Ok(_) => {}, @@ -25,6 +37,8 @@ pub async fn send_message( }, }; + auth_service.increase_user_request(&author).await; + return ( StatusCode::OK, Json(GenericResponse { diff --git a/src/main.rs b/src/main.rs index 1a050c9..7585b24 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,12 +5,12 @@ mod model; mod route; mod service; mod utils; +mod depends; use crate::config::config_server; #[tokio::main] async fn main() { - dotenv::dotenv().ok(); let server_config = config_server::get_config_server(); let app = route::create_route(); diff --git a/src/route.rs b/src/route.rs index 1f87b41..69a3c5a 100644 --- a/src/route.rs +++ b/src/route.rs @@ -10,15 +10,15 @@ use axum::{ routing::{get, post}, Extension, Router, }; +use crate::depends::depends_auth_service::get_depends_auth_service; +use crate::depends::depends_email_service::get_depends_email_service; fn configure_message_endpoint(router: Router) -> Router { router .route("/message", post(send_message)) .layer(middleware::from_fn(auth_middleware)) - .layer(Extension(AuthService::new(config_auth::get_config_auth()))) - .layer(Extension(EmailService::new( - config_email::get_config_email(), - ))) + .layer(Extension(get_depends_auth_service())) + .layer(Extension(get_depends_email_service())) } fn configure_health_endpoint(router: Router) -> Router { diff --git a/src/service/auth_service.rs b/src/service/auth_service.rs index e8aa4d0..114b279 100644 --- a/src/service/auth_service.rs +++ b/src/service/auth_service.rs @@ -1,16 +1,30 @@ +use std::collections::BTreeMap; +use redis::{AsyncCommands, ExistenceCheck, SetExpiry, SetOptions}; use crate::config::config_auth::ConfigAuth; use crate::model::send_message::MessageAuthor; use reqwest::header::AUTHORIZATION; +use crate::config::config_limits::ConfigLimits; +use crate::config::config_redis::ConfigRedis; #[derive(Clone)] pub struct AuthService { auth_url: String, + redis: redis::Client, + max_requests: u32, + expiration_time: usize, } impl AuthService { - pub fn new(config_auth: ConfigAuth) -> Self { + pub fn new(config_auth: ConfigAuth, config_redis: ConfigRedis, limits: ConfigLimits) -> Self { + let client = redis::Client::open( + format!("redis://{}:{}", config_redis.redis_url, config_redis.redis_port).as_str() + ).unwrap(); + AuthService { auth_url: config_auth.auth_url, + redis: client, + max_requests: limits.max_requests, + expiration_time: limits.expiration_time, } } @@ -34,4 +48,48 @@ impl AuthService { None } + + pub async fn has_user_reached_limit(&self, user: &MessageAuthor) -> bool { + let user_requests = self.count_user_requests(user).await; + return user_requests >= self.max_requests; + } + + pub async fn increase_user_request(&self, user: &MessageAuthor) -> bool { + let mut con = self.redis.get_async_connection().await.unwrap(); + let current_request_key= format!( + "user-message:{}:requests:{}", + user.email, + chrono::Utc::now().timestamp() + ); + + let set_options = SetOptions::default() + .with_expiration(SetExpiry::EX(self.expiration_time)) + .conditional_set(ExistenceCheck::NX) + .get(false); + + return con.set_options( + ¤t_request_key, + 1, + set_options + ).await.expect("Error setting key"); + + } + + async fn count_user_requests(&self, user: &MessageAuthor) -> u32 { + let mut con = self.redis.get_async_connection().await.unwrap(); + let query_user_requests = format!("user-message:{}:requests:*", user.email); + + let results: Vec; + match con.keys(query_user_requests).await { + Ok(r) => { + results = r; + }, + Err(e) => { + return 0; + } + + }; + + return results.len() as u32; + } }