Merge pull request #6 from HideyoshiNakazone/implements-request-limit

Implements Request Limit with Redis

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",
 ]

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"] }

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;
+    }
 }