sz-rust-middleware-facade 0.6.3

//! SSO 中间件 — 本地验签 + 远程校验
//!
//! 对齐 spec.md FR-6 ~ FR-7，design.md §3.2。
//!
//! ## 本地验签（默认）
//!
//! 业务系统与 SSO 认证中心共享 JWT secret，本地 `SsoJwtCodec::decode` 验签，零网络开销。
//!
//! ## 远程校验（feature = "remote-validate"）
//!
//! 密钥不共享场景，通过 HTTP 调用 SSO 认证中心 `/sso/validate` 端点校验。

use axum::extract::State;
use axum::http::{Request, StatusCode};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use std::sync::Arc;

use sz_rust_auth_facade::refresh::{
    MemoryRefreshTokenStore, MemoryTokenBlacklist, RefreshTokenStore, RefreshTokenVerifier,
    SsoJwtCodec, TokenBlacklist,
};

// ── AuthenticatedUser ──

/// 认证后的用户信息（注入 request extensions）
#[derive(Debug, Clone)]
pub struct AuthenticatedUser {
    /// 用户 ID
    pub user_id: i64,
    /// 用户名
    pub username: String,
}

// ── SsoMiddlewareConfig ──

/// SSO 中间件配置
pub struct SsoMiddlewareConfig {
    /// JWT 编解码器（本地验签）
    codec: SsoJwtCodec,
    /// Token 黑名单
    blacklist: Arc<dyn TokenBlacklist>,
    /// Token 版本存储
    store: Arc<dyn RefreshTokenStore>,
    /// JWT 签发人
    issuer: String,
    /// 白名单路由（支持 `*` 通配符）
    allow_all_action: Vec<String>,
}

impl SsoMiddlewareConfig {
    /// 创建本地验签配置
    pub fn local(
        secret: impl Into<String>,
        issuer: impl Into<String>,
        blacklist: Arc<dyn TokenBlacklist>,
        store: Arc<dyn RefreshTokenStore>,
        allow_all_action: Vec<String>,
    ) -> Self {
        Self {
            codec: SsoJwtCodec::new(secret),
            blacklist,
            store,
            issuer: issuer.into(),
            allow_all_action,
        }
    }

    /// 创建本地验签配置（内存黑名单 + 内存存储，测试用）
    pub fn local_memory(
        secret: impl Into<String>,
        issuer: impl Into<String>,
        allow_all_action: Vec<String>,
    ) -> Self {
        Self::local(
            secret,
            issuer,
            Arc::new(MemoryTokenBlacklist::new()),
            Arc::new(MemoryRefreshTokenStore::new()),
            allow_all_action,
        )
    }

    /// 检查路由是否在白名单中
    fn is_allowed(&self, path: &str) -> bool {
        for pattern in &self.allow_all_action {
            if pattern == "*" || pattern == path {
                return true;
            }
            if pattern.ends_with('*') && path.starts_with(&pattern[..pattern.len() - 1]) {
                return true;
            }
        }
        false
    }
}

impl std::fmt::Debug for SsoMiddlewareConfig {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SsoMiddlewareConfig")
            .field("codec", &self.codec)
            .field("issuer", &self.issuer)
            .field("allow_all_action", &self.allow_all_action)
            .finish_non_exhaustive()
    }
}

// ── sso_middleware ──

/// SSO 中间件
///
/// 从 `Authorization: Bearer <token>` 提取 accessToken，
/// 执行本地验签 + 黑名单查询 + 版本校验，通过后注入 `AuthenticatedUser`。
pub async fn sso_middleware(
    State(config): State<Arc<SsoMiddlewareConfig>>,
    req: Request<axum::body::Body>,
    next: Next,
) -> Response {
    let path = req.uri().path().to_string();

    if config.is_allowed(&path) {
        return next.run(req).await;
    }

    let auth_header = match req.headers().get("authorization") {
        Some(v) => v.to_str().unwrap_or(""),
        None => return unauthorized("missing authorization header"),
    };

    let token = match auth_header.strip_prefix("Bearer ") {
        Some(t) => t,
        None => return unauthorized("invalid authorization scheme"),
    };

    let verifier = RefreshTokenVerifier::new(
        config.codec.clone(),
        config.blacklist.clone(),
        config.store.clone(),
        config.issuer.clone(),
    );

    match verifier.verify_access(token).await {
        Ok(claims) => {
            let user_id = claims.user_id.unwrap_or(0);
            let username = claims.sub.clone();
            let mut req = req;
            req.extensions_mut().insert(AuthenticatedUser { user_id, username });
            next.run(req).await
        }
        Err(e) => {
            tracing::warn!(error = %e, "SSO token validation failed");
            unauthorized(&e.to_string())
        }
    }
}

fn unauthorized(msg: &str) -> Response {
    (
        StatusCode::UNAUTHORIZED,
        [("Cache-Control", "no-store"), ("Pragma", "no-cache")],
        format!("{{\"code\":-1,\"msg\":\"{msg}\"}}"),
    )
        .into_response()
}

// ── 远程校验（feature = "remote-validate"） ──

/// 远程校验配置
#[cfg(feature = "remote-validate")]
pub struct RemoteValidateConfig {
    /// SSO 认证中心校验端点
    pub endpoint: String,
    /// 超时时间
    pub timeout: std::time::Duration,
    /// HTTP 客户端（单例复用连接池）
    client: reqwest::Client,
    /// 白名单路由
    pub allow_all_action: Vec<String>,
}

#[cfg(feature = "remote-validate")]
impl RemoteValidateConfig {
    /// 创建远程校验配置
    pub fn new(
        endpoint: impl Into<String>,
        timeout: std::time::Duration,
        allow_all_action: Vec<String>,
    ) -> Self {
        let client = reqwest::Client::builder()
            .timeout(timeout)
            .pool_max_idle_per_host(32)
            .build()
            .expect("failed to build reqwest client");
        Self {
            endpoint: endpoint.into(),
            timeout,
            client,
            allow_all_action,
        }
    }

    /// 检查路由是否在白名单中
    fn is_allowed(&self, path: &str) -> bool {
        for pattern in &self.allow_all_action {
            if pattern == "*" || pattern == path {
                return true;
            }
            if pattern.ends_with('*') && path.starts_with(&pattern[..pattern.len() - 1]) {
                return true;
            }
        }
        false
    }
}

/// 远程校验中间件
#[cfg(feature = "remote-validate")]
pub async fn sso_middleware_remote(
    State(config): State<Arc<RemoteValidateConfig>>,
    req: Request<axum::body::Body>,
    next: Next,
) -> Response {
    let path = req.uri().path().to_string();

    if config.is_allowed(&path) {
        return next.run(req).await;
    }

    let auth_header = match req.headers().get("authorization") {
        Some(v) => v.to_str().unwrap_or(""),
        None => return unauthorized("missing authorization header"),
    };

    let token = match auth_header.strip_prefix("Bearer ") {
        Some(t) => t,
        None => return unauthorized("invalid authorization scheme"),
    };

    let validate_url = format!("{}?token={}", config.endpoint, token);
    match config.client.get(&validate_url).send().await {
        Ok(resp) if resp.status().is_success() => {
            match resp.json::<serde_json::Value>().await {
                Ok(json) => {
                    let data = json.get("data").cloned().unwrap_or_default();
                    let user_id = data.get("user_id").and_then(|v| v.as_i64()).unwrap_or(0);
                    let mut req = req;
                    req.extensions_mut().insert(AuthenticatedUser {
                        user_id,
                        username: String::new(),
                    });
                    next.run(req).await
                }
                Err(_) => service_unavailable(),
            }
        }
        Ok(resp) if resp.status() == StatusCode::UNAUTHORIZED => {
            unauthorized("token invalid or expired")
        }
        _ => service_unavailable(),
    }
}

#[cfg(feature = "remote-validate")]
fn service_unavailable() -> Response {
    (
        StatusCode::SERVICE_UNAVAILABLE,
        [("Cache-Control", "no-store"), ("Pragma", "no-cache")],
        r#"{"code":-1,"msg":"认证服务暂时不可用"}"#,
    )
        .into_response()
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::Request;
    use axum::routing::get;
    use axum::Router;
    use sz_rust_auth_facade::refresh::{
        RefreshTokenConfig, RefreshTokenIssuer,
    };
    use tower::ServiceExt;

    async fn handler(req: Request<Body>) -> String {
        let user = req.extensions().get::<AuthenticatedUser>();
        match user {
            Some(u) => format!("user_id={}, username={}", u.user_id, u.username),
            None => "no user".to_string(),
        }
    }

    fn make_app() -> (Router, RefreshTokenIssuer) {
        let codec = SsoJwtCodec::new("test-secret");
        let blacklist: Arc<dyn TokenBlacklist> = Arc::new(MemoryTokenBlacklist::new());
        let store: Arc<dyn RefreshTokenStore> = Arc::new(MemoryRefreshTokenStore::new());
        let config = RefreshTokenConfig::default();
        let issuer = RefreshTokenIssuer::new(codec.clone(), blacklist.clone(), store.clone(), config.clone());
        let mw_config = Arc::new(SsoMiddlewareConfig::local(
            "test-secret",
            config.issuer.clone(),
            blacklist,
            store,
            vec!["/public/*".to_string()],
        ));
        let app = Router::new()
            .route("/protected", get(handler))
            .route("/public/health", get(handler))
            .layer(axum::middleware::from_fn_with_state(mw_config, sso_middleware));
        (app, issuer)
    }

    #[tokio::test]
    async fn test_middleware_allows_whitelist() {
        let (app, _) = make_app();
        let resp = app
            .oneshot(Request::builder().uri("/public/health").body(Body::empty()).unwrap())
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_middleware_rejects_missing_token() {
        let (app, _) = make_app();
        let resp = app
            .oneshot(Request::builder().uri("/protected").body(Body::empty()).unwrap())
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
    }

    #[tokio::test]
    async fn test_middleware_accepts_valid_token() {
        let (app, issuer) = make_app();
        let pair = issuer.issue(42, "alice").await.unwrap();
        let resp = app
            .oneshot(
                Request::builder()
                    .uri("/protected")
                    .header("authorization", format!("Bearer {}", pair.access_token))
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX).await.unwrap();
        let text = String::from_utf8(body.to_vec()).unwrap();
        assert!(text.contains("user_id=42"));
        assert!(text.contains("username=alice"));
    }

    #[tokio::test]
    async fn test_middleware_rejects_refresh_token_as_access() {
        let (app, issuer) = make_app();
        let pair = issuer.issue(42, "alice").await.unwrap();
        let resp = app
            .oneshot(
                Request::builder()
                    .uri("/protected")
                    .header("authorization", format!("Bearer {}", pair.refresh_token))
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
    }

    #[tokio::test]
    async fn test_middleware_rejects_invalid_token() {
        let (app, _) = make_app();
        let resp = app
            .oneshot(
                Request::builder()
                    .uri("/protected")
                    .header("authorization", "Bearer invalid.token.here")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
    }
}