eric / rustwarden

/// Admin Service Middleware
///
/// Tower Service wrapper that applies admin_required middleware to /admin/* routes.
/// This works around Axum 0.8 Sync limitations with from_fn_with_state.
use axum::{
    extract::Request,
    response::{IntoResponse, Response},
};
use sea_orm::DatabaseConnection;
use std::{
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};
use tower::Service;

use super::auth_service;
use crate::config::auth::AuthConfig;

/// Service wrapper that applies admin_required middleware to /admin/* routes
///
/// The [`AuthConfig`] used to verify the JWT is resolved by
/// `auth_service::resolve_auth_config`: the config captured at construction
/// first, then `Extension(Arc<AuthConfig>)`, then
/// [`AuthConfig::from_env`]. The last two steps are the legacy path of the
/// deprecated [`AdminService::new`] constructor and ignore builder overrides
/// whenever this middleware is layered outside the `Extension` layer.
#[derive(Clone)]
pub struct AdminService<S> {
    inner: S,
    db: DatabaseConnection,
    /// Config captured at construction. `None` only on the deprecated
    /// [`AdminService::new`] path.
    auth_config: Option<Arc<AuthConfig>>,
}

impl<S> AdminService<S> {
    /// Create a new `AdminService` middleware wrapping the given inner service.
    #[deprecated(
        since = "1.2.0",
        note = "use `AdminService::with_auth_config(inner, db, auth_config)`. Removal no earlier than 1.5.0, effective in the next major."
    )]
    pub fn new(inner: S, db: DatabaseConnection) -> Self {
        Self {
            inner,
            db,
            auth_config: None,
        }
    }

    /// Create an `AdminService` that carries its own resolved
    /// [`AuthConfig`] — the supported constructor since RFC 0042.
    ///
    /// The config is captured **once at construction**, the same pattern
    /// `security_headers_middleware` uses in
    /// [`crate::app::RustwardenBuilder::build`]. Consequences:
    ///
    /// - Builder overrides (`login_path`, `auth_failure_mode`, `headless`)
    ///   are honoured regardless of tower layer order. A middleware layered
    ///   *outside* `Extension(Arc<AuthConfig>)` can never observe that
    ///   extension, which is what made [`AdminService::new`] fall back to
    ///   [`AuthConfig::from_env`] and discard those overrides.
    /// - The captured config takes precedence over any
    ///   `Extension(Arc<AuthConfig>)` present in the request.
    /// - On every request the middleware inserts the captured `Arc` into the
    ///   request extensions, so downstream extractors (`AuthUser`,
    ///   `auth_required`, `admin_required`) observe the same instance even in
    ///   hand-composed tower stacks. In the canonical wiring the inner
    ///   `Extension` layer re-inserts the very same `Arc`, a no-op.
    ///
    /// ```no_run
    /// # use std::sync::Arc;
    /// # use rustwarden_core::config::auth::AuthConfig;
    /// # use rustwarden_core::middleware::admin_service::AdminService;
    /// # fn wire(db: sea_orm::DatabaseConnection, router: axum::Router) -> Result<(), String> {
    /// let auth_config = Arc::new(AuthConfig::from_env()?);
    /// let app = router.layer(tower::layer::layer_fn(move |inner| {
    ///     AdminService::with_auth_config(inner, db.clone(), auth_config.clone())
    /// }));
    /// # let _ = app; Ok(()) }
    /// ```
    pub fn with_auth_config(
        inner: S,
        db: DatabaseConnection,
        auth_config: Arc<AuthConfig>,
    ) -> Self {
        Self {
            inner,
            db,
            auth_config: Some(auth_config),
        }
    }
}

impl<S, B> Service<Request<B>> for AdminService<S>
where
    S: Service<Request<B>, Response = Response> + Clone + Send + 'static,
    S::Future: Send + 'static,
    B: Send + 'static,
{
    type Response = Response;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut req: Request<B>) -> Self::Future {
        let uri = req.uri().clone();

        // Seed the captured config into the request extensions so downstream
        // extractors (`AuthUser`, `auth_required`, `admin_required`) resolve
        // the exact same instance. In the canonical wiring the later
        // `Extension(Arc<AuthConfig>)` layer re-inserts the same `Arc`, which
        // is a no-op; seeding only matters for hand-composed tower stacks.
        if let Some(config) = &self.auth_config {
            req.extensions_mut().insert(config.clone());
        }
        let captured = self.auth_config.clone();

        let path = uri.path();

        // Only apply admin_required middleware to /admin/* and /api/v1/admin/* routes
        if path.starts_with("/admin") || path.starts_with("/api/v1/admin") {
            let db = self.db.clone();
            let mut inner = self.inner.clone();

            tracing::debug!("AdminService: intercepting admin route: {}", path);

            Box::pin(async move {
                // Extract JWT and validate it
                // We use parts to extract JWT (no need for body)
                let (parts, body) = req.into_parts();

                use crate::auth::jwt::JwtService;
                use crate::services::jwt_blacklist_service::JwtBlacklistService;

                // Resolve auth config up-front so every auth-failure branch
                // honours `login_path` / `auth_failure_mode`.
                let auth_config = match crate::middleware::auth_service::resolve_auth_config(
                    captured.as_ref(),
                    &parts.extensions,
                ) {
                    Ok(config) => config,
                    Err(err) => {
                        tracing::error!(
                            "AdminService: AuthConfig unavailable (no captured config, no extension and from_env failed): {err}"
                        );
                        return Ok((
                            axum::http::StatusCode::INTERNAL_SERVER_ERROR,
                            "Internal server error",
                        )
                            .into_response());
                    }
                };

                // Extract JWT token from parts
                let jwt_token =
                    match crate::auth::jwt_extractor::extract_jwt_from_request_parts(&parts) {
                        Ok(token) => {
                            tracing::debug!("AdminService: JWT token found");
                            token
                        }
                        Err(_) => {
                            tracing::warn!("AdminService: No JWT token found");
                            return Ok(auth_service::auth_failure_response(&auth_config, &uri));
                        }
                    };

                if JwtBlacklistService::is_blacklisted(&db, &jwt_token)
                    .await
                    .is_ok_and(|b| b)
                {
                    return Ok(auth_service::auth_failure_response(&auth_config, &uri));
                }

                let claims = match JwtService::validate_token(&auth_config, &jwt_token) {
                    Ok(claims) => claims,
                    Err(_) => {
                        return Ok(auth_service::auth_failure_response(&auth_config, &uri));
                    }
                };

                // Check if user has admin role
                if !claims.roles.contains(&"admin".to_string()) {
                    tracing::warn!(
                        "Non-admin user {} attempted to access admin route: {}",
                        claims.sub,
                        uri
                    );
                    return Ok((
                        axum::http::StatusCode::FORBIDDEN,
                        "Access denied: Admin role required",
                    )
                        .into_response());
                }

                tracing::debug!("Admin user {} accessing: {}", claims.sub, uri);

                // Store claims in request extensions for handlers to use
                let mut parts_with_claims = parts;
                parts_with_claims.extensions.insert(claims.clone());
                tracing::debug!(
                    "AdminService: Claims inserted into extensions for user {}",
                    claims.sub
                );
                let req_with_claims = Request::from_parts(parts_with_claims, body);

                // Call inner service
                inner.call(req_with_claims).await
            })
        } else {
            // Not an admin route, proceed normally
            let mut inner = self.inner.clone();
            Box::pin(async move { inner.call(req).await })
        }
    }
}

#[cfg(test)]
mod tests {
    use super::AdminService;
    use crate::config::auth::{AuthConfig, AuthFailureMode, SessionBackend};
    use axum::body::Body;
    use axum::http::{Request as HttpRequest, StatusCode, header::LOCATION};
    use axum::response::Response;
    use sea_orm::{DatabaseBackend, DatabaseConnection, MockDatabase};
    use std::sync::Arc;
    use tower::ServiceExt; // for `oneshot`

    fn test_config(mode: AuthFailureMode, login_path: &str) -> 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: login_path.to_string(),
            auth_failure_mode: mode,
        }
    }

    fn mock_db() -> DatabaseConnection {
        MockDatabase::new(DatabaseBackend::Postgres).into_connection()
    }

    /// Router wrapped with `AdminService::with_auth_config` and **no**
    /// `Extension(Arc<AuthConfig>)` layer: the captured config is the only
    /// source of truth, exactly as RFC 0042 requires.
    fn admin_router(cfg: AuthConfig) -> axum::Router {
        let db = mock_db();
        let cfg = Arc::new(cfg);
        axum::Router::new()
            .route("/admin", axum::routing::get(|| async { "admin" }))
            .route(
                "/api/v1/admin/x",
                axum::routing::get(|| async { "admin api" }),
            )
            .layer(tower::layer::layer_fn(move |inner| {
                AdminService::with_auth_config(inner, db.clone(), cfg.clone())
            }))
    }

    async fn get(router: axum::Router, uri: &str) -> Response {
        let req = HttpRequest::builder()
            .uri(uri)
            .body(Body::empty())
            .expect("request");
        router.oneshot(req).await.expect("oneshot")
    }

    async fn body_string(resp: Response) -> String {
        let bytes = axum::body::to_bytes(resp.into_body(), 64 * 1024)
            .await
            .expect("body");
        String::from_utf8(bytes.to_vec()).expect("utf-8 body")
    }

    /// T7 — both admin prefixes honour the captured `Json401` mode without
    /// any environment variable or `Extension` layer.
    #[tokio::test]
    async fn admin_routes_without_jwt_return_json_401_from_the_captured_config() {
        for uri in ["/admin", "/api/v1/admin/x"] {
            let resp = get(
                admin_router(test_config(AuthFailureMode::Json401, "/login")),
                uri,
            )
            .await;
            assert_eq!(resp.status(), StatusCode::UNAUTHORIZED, "uri = {uri}");
            assert!(resp.headers().get(LOCATION).is_none(), "uri = {uri}");
            let body = body_string(resp).await;
            assert!(body.contains("unauthorized"), "uri = {uri}, body = {body}");
        }
    }

    /// T7 — `Redirect` mode honours the captured `login_path` override that
    /// the deprecated `AdminService::new` path used to discard.
    #[tokio::test]
    async fn admin_route_redirect_mode_uses_the_captured_login_path() {
        let resp = get(
            admin_router(test_config(AuthFailureMode::Redirect, "/auth/signin")),
            "/admin",
        )
        .await;

        assert_eq!(resp.status(), StatusCode::SEE_OTHER);
        assert_eq!(
            resp.headers().get(LOCATION).and_then(|v| v.to_str().ok()),
            Some("/auth/signin?redirect=%2Fadmin")
        );
    }
}