eric / rustwarden

/// Root Handler
/// Feature 004: UI Post-Auth
///
/// Handles the root path `/` with smart redirection based on authentication status.
use std::sync::Arc;

use axum::{
    extract::Request,
    response::{IntoResponse, Redirect},
};
use rustwarden_core::config::auth::AuthConfig;

/// Handler for GET `/` - Root path with smart redirection
///
/// Redirects based on authentication status:
/// - If authenticated → redirects to `/home`
/// - If not authenticated → redirects to the configured login path
///   ([`AuthConfig::login_path`], default `/login`)
///
/// This provides a seamless entry point for users.
///
/// The [`AuthConfig`] is read from the request extensions, where
/// `RustwardenBuilder::build()` injects the configuration resolved at
/// startup (including builder overrides such as `login_path`, RFC 0042).
/// `AuthConfig::from_env()` is only a defence for stacks that do not pose
/// the extension.
///
/// # Arguments
///
/// * `request` - HTTP request to check for JWT token
///
/// # Returns
///
/// * `303 See Other` redirect to `/home` or to the login path
///
/// # Example
///
/// ```bash
/// # Without authentication
/// curl -I http://localhost:3000/
/// # → 303 Redirect to /login
///
/// # With valid JWT token
/// curl -I http://localhost:3000/ -H "Cookie: jwt_token=..."
/// # → 303 Redirect to /home
/// ```
pub async fn root_redirect(request: Request) -> impl IntoResponse {
    use rustwarden_core::auth::jwt::JwtService;

    // Prefer the config resolved at build time; keep `from_env()` as a
    // defence for stacks that do not pose the extension.
    let auth_config = match request.extensions().get::<Arc<AuthConfig>>().cloned() {
        Some(config) => config,
        None => match AuthConfig::from_env() {
            Ok(config) => Arc::new(config),
            Err(_) => {
                // If config can't be loaded, redirect to login as fallback
                tracing::warn!("root_redirect: failed to load auth config, redirecting to /login");
                return Redirect::to("/login");
            }
        },
    };
    let login_path = auth_config.login_path.as_str();

    // Try to extract JWT token from request
    match rustwarden_core::auth::jwt_extractor::extract_jwt_from_request(&request) {
        Ok(jwt_token) => {
            // Validate the token without checking the blacklist: it is
            // re-validated on protected routes anyway.
            match JwtService::validate_token(&auth_config, &jwt_token) {
                Ok(_claims) => {
                    tracing::debug!("root_redirect: valid JWT token found, redirecting to /home");
                    Redirect::to("/home")
                }
                Err(_) => {
                    tracing::debug!(
                        "root_redirect: invalid JWT token, redirecting to {login_path}"
                    );
                    Redirect::to(login_path)
                }
            }
        }
        Err(_) => {
            tracing::debug!("root_redirect: no JWT token found, redirecting to {login_path}");
            Redirect::to(login_path)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::root_redirect;
    use axum::body::Body;
    use axum::http::{Request as HttpRequest, StatusCode, header::LOCATION};
    use axum::{Extension, Router, routing::get};
    use rustwarden_core::auth::jwt::JwtService;
    use rustwarden_core::config::auth::{AuthConfig, AuthFailureMode, SessionBackend};
    use std::sync::Arc;
    use tower::ServiceExt; // for `oneshot`

    /// A `login_path` that cannot come from the environment of the test
    /// process, so a redirect to it proves the extension was used.
    const INJECTED_LOGIN_PATH: &str = "/auth/signin";

    fn injected_config() -> AuthConfig {
        AuthConfig {
            jwt_secret: "0123456789abcdef0123456789abcdef".to_string(),
            jwt_expiration_minutes: 15,
            jwt_issuer: "test".to_string(),
            refresh_token_expiration_days: 7,
            remember_me_expiration_days: 30,
            session_timeout_minutes: 30,
            session_backend: SessionBackend::Memory,
            redis_url: None,
            rate_limit_login_per_minute: 5,
            rate_limit_refresh_per_minute: 10,
            rate_limit_general_per_minute: 100,
            cookie_secure: false,
            cookie_same_site: "Strict".to_string(),
            login_path: INJECTED_LOGIN_PATH.to_string(),
            auth_failure_mode: AuthFailureMode::Redirect,
        }
    }

    fn router_with_injected_config() -> Router {
        Router::new()
            .route("/", get(root_redirect))
            .layer(Extension(Arc::new(injected_config())))
    }

    async fn location_of(router: Router, bearer: Option<&str>) -> (StatusCode, String) {
        let mut builder = HttpRequest::builder().uri("/");
        if let Some(token) = bearer {
            builder = builder.header("authorization", format!("Bearer {token}"));
        }
        let req = builder.body(Body::empty()).expect("request");
        let resp = router.oneshot(req).await.expect("oneshot");
        let location = resp
            .headers()
            .get(LOCATION)
            .and_then(|v| v.to_str().ok())
            .unwrap_or_default()
            .to_string();
        (resp.status(), location)
    }

    #[tokio::test]
    async fn anonymous_visitor_is_sent_to_the_injected_login_path() {
        let (status, location) = location_of(router_with_injected_config(), None).await;
        assert_eq!(status, StatusCode::SEE_OTHER);
        assert_eq!(location, INJECTED_LOGIN_PATH);
    }

    #[tokio::test]
    async fn invalid_token_is_sent_to_the_injected_login_path() {
        let (status, location) =
            location_of(router_with_injected_config(), Some("not-a-jwt")).await;
        assert_eq!(status, StatusCode::SEE_OTHER);
        assert_eq!(location, INJECTED_LOGIN_PATH);
    }

    #[tokio::test]
    async fn valid_token_signed_with_the_injected_secret_goes_home() {
        let token = JwtService::generate_token(
            &injected_config(),
            uuid::Uuid::new_v4(),
            "alice",
            vec!["user".to_string()],
            Vec::new(),
        )
        .expect("token");
        let (status, location) = location_of(router_with_injected_config(), Some(&token)).await;
        assert_eq!(status, StatusCode::SEE_OTHER);
        assert_eq!(location, "/home");
    }
}