Merge pull request #29 from 8shaws/middle

chore: adds middleware for verification checks

crates/api/src/auth/mod.rs

@@ -1,2 +1 @@
-pub mod middleware;
 pub mod utils;

crates/api/src/main.rs

@@ -6,7 +6,7 @@ use env_logger::Builder;
 use lazy_static::lazy_static;
 use num_cpus;
 use serde_json::json;
-use std::time::SystemTime;
+use std::time::{Duration, SystemTime};
 
 mod auth;
 mod db;
@@ -17,6 +17,7 @@ mod route;
 mod schema;
 mod types;
 
+use crate::middlewares::rate_limit::RateLimiter;
 use crate::route::user::user_config;
 
 lazy_static! {
@@ -47,18 +48,20 @@ async fn main() -> std::io::Result<()> {
     let dp_pool = initialize_db_pool();
     let redis_pool = redis::initialize_redis_pool();
 
-    let app_state = models::AppState {
+    let app_state = web::Data::new(models::AppState {
         db_pool: dp_pool,
         redis_pool: redis_pool,
-    };
+    });
+
+    let rate_limiter = web::Data::new(RateLimiter::new(10, Duration::from_secs(60)));
 
     println!("{:?}: Api Server is running on port: {}", *START_TIME, 8080);
 
     HttpServer::new(move || {
         App::new()
-            .app_data(web::Data::new(app_state.clone()))
+            .app_data(app_state.clone())
             .wrap(middleware::Logger::default())
-            .configure(user_config)
+            .configure(|cfg| user_config(cfg, app_state.clone(), rate_limiter.clone()))
             .route("/", web::get().to(root))
             .route("/_health", web::get().to(root))
     })

crates/api/src/auth/middleware.rs → crates/api/src/middlewares/extract_client_id.rs

@@ -1,4 +1,4 @@
-use super::utils::verify_token;
+use crate::auth::utils::verify_token;
 use actix_service::Service;
 use actix_web::{
     body::EitherBody,

crates/api/src/middlewares/mod.rs

@@ -1 +1,4 @@
-pub mod check_verify;
+pub mod extract_client_id;
+pub mod rate_limit;
+pub mod un_verify_user;
+pub mod verify_user;

crates/api/src/middlewares/rate_limit.rs

@@ -0,0 +1,53 @@
+use actix_service::Service;
+use actix_web::{web, App, Error, HttpRequest, HttpResponse};
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::time::Duration;
+use tokio::sync::Semaphore;
+
+pub struct RateLimiter {
+    semaphore: Arc<Semaphore>,
+    requests: Arc<Mutex<HashMap<String, usize>>>,
+    max_requests: usize,
+    window_duration: Duration,
+}
+
+impl RateLimiter {
+    pub fn new(max_requests: usize, window_duration: Duration) -> Self {
+        RateLimiter {
+            semaphore: Arc::new(Semaphore::new(max_requests)),
+            requests: Arc::new(Mutex::new(HashMap::new())),
+            max_requests,
+            window_duration,
+        }
+    }
+
+    pub async fn check_rate_limit(&self, key: &str) -> Result<(), HttpResponse> {
+        let mut requests = self.requests.lock().unwrap();
+        let count = requests.entry(key.to_string()).or_insert(0);
+        if *count >= self.max_requests {
+            Err(HttpResponse::TooManyRequests().finish())
+        } else {
+            *count += 1;
+            Ok(())
+        }
+    }
+
+    pub async fn handle_request(
+        &self,
+        key: &str,
+        request: HttpRequest,
+    ) -> Result<HttpResponse, Error> {
+        self.check_rate_limit(key).await;
+        Ok(HttpResponse::Ok().finish())
+    }
+}
+
+async fn rate_limited_handler(
+    req: HttpRequest,
+    rate_limiter: web::Data<RateLimiter>,
+) -> Result<HttpResponse, Error> {
+    let key = req.peer_addr().unwrap().to_string();
+    rate_limiter.handle_request(&key, req).await
+}

crates/api/src/middlewares/un_verify_user.rs

@@ -0,0 +1,110 @@
+use crate::db::user_db_fn::is_user_verified;
+use crate::{middlewares::extract_client_id::IdKey, models::AppState};
+use actix_service::Service;
+use actix_web::web::Data;
+use actix_web::{
+    body::EitherBody,
+    dev::{ServiceRequest, ServiceResponse, Transform},
+    Error, HttpMessage, HttpResponse,
+};
+use futures_util::future::{ok, LocalBoxFuture, Ready};
+use serde_json::json;
+use std::rc::Rc;
+use std::task::{Context, Poll};
+
+pub struct UnVerifyUser {
+    app_state: Data<AppState>,
+}
+
+impl UnVerifyUser {
+    pub fn new(app_state: Data<AppState>) -> Self {
+        Self { app_state }
+    }
+}
+
+impl<S, B: 'static> Transform<S, ServiceRequest> for UnVerifyUser
+where
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<EitherBody<B>>;
+    type Error = Error;
+    type InitError = ();
+    type Transform = UnVerifyUserMiddleware<S>;
+    type Future = Ready<Result<Self::Transform, Self::InitError>>;
+
+    fn new_transform(&self, service: S) -> Self::Future {
+        ok(UnVerifyUserMiddleware {
+            service: Rc::new(service),
+            app_state: self.app_state.clone(),
+        })
+    }
+}
+
+pub struct UnVerifyUserMiddleware<S> {
+    service: Rc<S>,
+    app_state: Data<AppState>,
+}
+
+impl<S, B: 'static> Service<ServiceRequest> for UnVerifyUserMiddleware<S>
+where
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<EitherBody<B>>;
+    type Error = Error;
+    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&self, ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.service.poll_ready(ctx)
+    }
+    fn call(&self, req: ServiceRequest) -> Self::Future {
+        let srv = Rc::clone(&self.service);
+        let id_key = req.extensions().get::<IdKey>().cloned();
+
+        let db_pool = self.app_state.db_pool.clone();
+
+        let fut = async move {
+            let mut conn = db_pool.get().unwrap();
+
+            match id_key {
+                Some(id) => {
+                    let verified = is_user_verified(&mut conn, id.0);
+                    match verified {
+                        Ok(verified) => {
+                            if !verified {
+                                let res = srv.call(req).await?;
+                                Ok(res.map_into_left_body())
+                            } else {
+                                let response = HttpResponse::Forbidden().json(json!({
+                                    "message": "User is already Verified!",
+                                    "status": "Error"
+                                }));
+
+                                Ok(req.into_response(response.map_into_right_body()))
+                            }
+                        }
+                        Err(_) => {
+                            let response = HttpResponse::InternalServerError().json(json!({
+                                "message": "User Verification Failed!",
+                                "status": "Error"
+                            }));
+                            Ok(req.into_response(response.map_into_right_body()))
+                        }
+                    }
+                }
+                None => {
+                    let response = HttpResponse::Forbidden().json(json!({
+                        "msg": "No authentication token found",
+                        "status": "Error"
+                    }));
+                    Ok(req.into_response(response.map_into_right_body()))
+                }
+            }
+        };
+
+        Box::pin(fut)
+    }
+}