Implements Request Limit with Redis

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 }

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");

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::<u32>().unwrap();
+    let expiration_time = env::var("EXPIRATION_TIME").unwrap_or("604800".to_string())
+        .parse::<usize>().unwrap();
+
+    ConfigLimits {
+        max_requests,
+        expiration_time,
+    }
+}

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<String>,
+}
+
+#[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::<u16>().unwrap(),
+        redis_password: password
+    }
+}

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")

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;

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(),
+    )
+}

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()
+    )
+}

src/depends/mod.rs

@@ -0,0 +1,2 @@
+pub mod depends_auth_service;
+pub mod depends_email_service;

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<AuthService>,
     Extension(email_service): Extension<EmailService>,
     Extension(author): Extension<MessageAuthor>,
     Json(payload): Json<SendMessage>,
 ) -> 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 {

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();

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 {

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(
+            &current_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<String>;
+        match con.keys(query_user_requests).await {
+            Ok(r) => {
+                results = r;
+            },
+            Err(e) => {
+                return 0;
+            }
+
+        };
+
+        return results.len() as u32;
+    }
 }