eric / rustwarden

/// Update Last Login Middleware
/// Feature 009: Admin & Settings
///
/// Middleware to update last_login_at and last_login_ip after successful authentication.
/// Extracts IP from request headers (handles X-Forwarded-For, X-Real-IP for reverse proxy).
use axum::{
    extract::{Request, State},
    http::HeaderMap,
    middleware::Next,
    response::Response,
};
use chrono::Utc;
use sea_orm::{ActiveModelTrait, DatabaseConnection, EntityTrait, Set};
use tower_sessions::Session;

use crate::middleware::auth_required::extract_user_from_jwt;
use crate::models::user;

/// Extract client IP from request headers
///
/// Checks headers in order:
/// 1. X-Forwarded-For (first IP if multiple)
/// 2. X-Real-IP
/// 3. Direct connection IP (from request) - currently not available, returns None
///
/// # Arguments
///
/// * `headers` - HTTP request headers
///
/// # Returns
///
/// IP address as String, or None if cannot be determined
pub fn extract_client_ip(headers: &HeaderMap) -> Option<String> {
    // Check X-Forwarded-For header (first IP if multiple)
    if let Some(forwarded_for) = headers.get("X-Forwarded-For") {
        if let Ok(forwarded_str) = forwarded_for.to_str() {
            // X-Forwarded-For can contain multiple IPs: "client, proxy1, proxy2"
            // Take the first one (original client)
            let first_ip = forwarded_str.split(',').next().unwrap_or("").trim();
            if !first_ip.is_empty() {
                return Some(first_ip.to_string());
            }
        }
    }

    // Check X-Real-IP header
    if let Some(real_ip) = headers.get("X-Real-IP") {
        if let Ok(real_ip_str) = real_ip.to_str() {
            if !real_ip_str.trim().is_empty() {
                return Some(real_ip_str.trim().to_string());
            }
        }
    }

    // For direct connections, we would need access to the connection info
    // This is not easily available in Axum middleware, so we return None
    // In production behind a reverse proxy, X-Forwarded-For or X-Real-IP should be set
    None
}

/// Middleware that updates last_login_at and last_login_ip after authentication
///
/// This middleware should be applied after auth_required middleware to ensure
/// the user is authenticated. It updates the user's last login information
/// without blocking the request if the update fails.
///
/// # Arguments
///
/// * `State(db)` - Database connection
/// * `session` - Tower session
/// * `request` - HTTP request
/// * `next` - Next middleware or handler
///
/// # Returns
///
/// Always passes the request through (non-blocking)
pub async fn update_last_login(
    State(db): State<DatabaseConnection>,
    _session: Session,
    request: Request,
    next: Next,
) -> Response {
    // Extract user from JWT (reuse auth_required logic)
    if let Some(claims) = extract_user_from_jwt(&db, &request).await {
        let user_id = claims.sub;
        let headers = request.headers();
        let client_ip = extract_client_ip(headers);

        // Update user's last login information
        if let Ok(Some(user_model)) = user::Entity::find_by_id(user_id).one(&db).await {
            let mut user_active: user::ActiveModel = user_model.into();
            user_active.last_login_at = Set(Some(Utc::now()));
            if let Some(ip) = client_ip {
                user_active.last_login_ip = Set(Some(ip));
            }

            // Update without blocking request if it fails
            if let Err(e) = user_active.update(&db).await {
                tracing::warn!("Failed to update last_login for user {}: {}", user_id, e);
            } else {
                tracing::debug!("Updated last_login for user {}", user_id);
            }
        }
    }

    // Always proceed to handler (non-blocking)
    next.run(request).await
}

/// Internal implementation that takes db and session as parameters
pub async fn update_last_login_impl(
    db: &DatabaseConnection,
    session: Session,
    request: Request,
    next: Next,
) -> Response {
    update_last_login(State(db.clone()), session, request, next).await
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_extract_client_ip_from_forwarded_for() {
        let mut headers = HeaderMap::new();
        headers.insert("X-Forwarded-For", "192.168.1.1, 10.0.0.1".parse().unwrap());

        let ip = extract_client_ip(&headers);

        assert_eq!(ip, Some("192.168.1.1".to_string()));
    }

    #[test]
    fn test_extract_client_ip_from_real_ip() {
        let mut headers = HeaderMap::new();
        headers.insert("X-Real-IP", "192.168.1.2".parse().unwrap());

        let ip = extract_client_ip(&headers);

        assert_eq!(ip, Some("192.168.1.2".to_string()));
    }
}