hostbasket/actix_mvc_app/src/middleware/mod.rs

use actix_web::{
    dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform},
    Error, HttpMessage,
};
use futures::future::{ready, LocalBoxFuture, Ready};
use std::{
    future::Future,
    pin::Pin,
    time::Instant,
};
use actix_session::SessionExt;

/// Middleware for logging request duration
pub struct RequestTimer;

impl<S, B> Transform<S, ServiceRequest> for RequestTimer
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Transform = RequestTimerMiddleware<S>;
    type InitError = ();
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(RequestTimerMiddleware { service }))
    }
}

pub struct RequestTimerMiddleware<S> {
    service: S,
}

impl<S, B> Service<ServiceRequest> for RequestTimerMiddleware<S>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;

    forward_ready!(service);

    fn call(&self, req: ServiceRequest) -> Self::Future {
        let start = Instant::now();
        let path = req.path().to_owned();
        let method = req.method().to_string();

        let fut = self.service.call(req);

        Box::pin(async move {
            let res = fut.await?;
            let duration = start.elapsed();
            log::info!(
                "{} {} - {} - {:?}",
                method,
                path,
                res.status().as_u16(),
                duration
            );
            Ok(res)
        })
    }
}

/// Middleware for adding security headers
pub struct SecurityHeaders;

impl<S, B> Transform<S, ServiceRequest> for SecurityHeaders
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Transform = SecurityHeadersMiddleware<S>;
    type InitError = ();
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(SecurityHeadersMiddleware { service }))
    }
}

pub struct SecurityHeadersMiddleware<S> {
    service: S,
}

impl<S, B> Service<ServiceRequest> for SecurityHeadersMiddleware<S>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Future = Pin<Box<dyn Future<Output = Result<ServiceResponse<B>, Error>>>>;

    forward_ready!(service);

    fn call(&self, req: ServiceRequest) -> Self::Future {
        let fut = self.service.call(req);

        Box::pin(async move {
            let mut res = fut.await?;
            
            // Add security headers
            res.headers_mut().insert(
                actix_web::http::header::X_CONTENT_TYPE_OPTIONS,
                actix_web::http::header::HeaderValue::from_static("nosniff"),
            );
            res.headers_mut().insert(
                actix_web::http::header::X_FRAME_OPTIONS,
                actix_web::http::header::HeaderValue::from_static("DENY"),
            );
            res.headers_mut().insert(
                actix_web::http::header::HeaderName::from_static("x-xss-protection"),
                actix_web::http::header::HeaderValue::from_static("1; mode=block"),
            );
            
            Ok(res)
        })
    }
}

/// Middleware for JWT authentication
pub struct JwtAuth;

impl<S, B> Transform<S, ServiceRequest> for JwtAuth
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Transform = JwtAuthMiddleware<S>;
    type InitError = ();
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(JwtAuthMiddleware { service }))
    }
}

pub struct JwtAuthMiddleware<S> {
    service: S,
}

impl<S, B> Service<ServiceRequest> for JwtAuthMiddleware<S>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;

    forward_ready!(service);

    fn call(&self, req: ServiceRequest) -> Self::Future {
        // Define public routes that don't require authentication
        let path = req.path().to_string();
        let public_routes = vec![
            "/login", 
            "/register", 
            "/static", 
            "/favicon.ico",
            "/",
            "/about",
            "/contact"
        ];
        
        // Check if the current path is a public route
        let is_public_route = public_routes.iter().any(|route| path.starts_with(route));
        
        if is_public_route {
            // For public routes, just pass through without authentication check
            let fut = self.service.call(req);
            return Box::pin(async move {
                fut.await
            });
        }
        
        // First try to get token from cookie
        let cookie_token = crate::controllers::auth::AuthController::extract_token_from_cookie(req.request());
        
        // If no cookie token, try to get from session
        let session = req.get_session();
        let session_token = crate::controllers::auth::AuthController::extract_token_from_session(&session);
        
        // Use cookie token if available, otherwise use session token
        let token = cookie_token.or(session_token);
        
        // Process based on token availability
        if let Some(token_str) = token {
            // Validate the token
            let validation_result = crate::controllers::auth::AuthController::validate_token(&token_str);
            
            match validation_result {
                Ok(claims) => {
                    // Token is valid, store claims in request extensions
                    req.extensions_mut().insert(claims.clone());
                    
                    // Create a user from claims and store in session
                    let mut user = crate::models::User::new(
                        claims.sub.clone(),
                        claims.sub.clone()
                    );
                    
                    // Set the user ID and role
                    user.id = Some(1);
                    user.role = if claims.role == "admin" {
                        crate::models::user::UserRole::Admin
                    } else {
                        crate::models::user::UserRole::User
                    };
                    
                    // Store user data in session
                    if let Ok(user_json) = serde_json::to_string(&user) {
                        let _ = session.insert("user", &user_json);
                        let _ = session.insert("auth_token", &token_str);
                    }
                    
                    let fut = self.service.call(req);
                    Box::pin(async move {
                        fut.await
                    })
                }
                Err(_) => {
                    // Token is invalid, redirect to login
                    Box::pin(async move {
                        // Return an error that will be handled by the error handlers
                        Err(actix_web::error::InternalError::from_response(
                            "JWT validation failed",
                            actix_web::HttpResponse::Found()
                                .append_header((actix_web::http::header::LOCATION, "/login"))
                                .finish()
                        ).into())
                    })
                }
            }
        } else {
            // No token found, redirect to login
            Box::pin(async move {
                // Return an error that will be handled by the error handlers
                Err(actix_web::error::InternalError::from_response(
                    "No JWT token found",
                    actix_web::HttpResponse::Found()
                        .append_header((actix_web::http::header::LOCATION, "/login"))
                        .finish()
                ).into())
            })
        }
    }
}