PicoBot/src/gateway/auth.rs

//! 网关认证与访问控制。
//!
//! 设计目标（第一性原理）：
//! - 本地单机部署（host 为 loopback）：免认证，仅靠 CORS 防御 DNS rebinding / CSRF。
//! - 远程访问（host 非 loopback）：必须配置 `auth_token`，所有 `/api/*` 与 `/ws` 强制校验。
//! - token 通过 `Authorization: Bearer <token>`（HTTP）或 `?token=<token>`（WS）传递。
//! - 校验使用常量时间比较，避免计时侧信道。

use axum::extract::Request;
use axum::http::{HeaderMap, StatusCode};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use axum::Json;
use serde_json::json;
use subtle::ConstantTimeEq;

/// 判定给定 host 是否为 loopback 地址。
/// 通过 `std::net::IpAddr` 解析，覆盖 IPv4/IPv6 的所有等价表示。
/// 也接受 `localhost` 主机名。
pub fn is_loopback_host(host: &str) -> bool {
    let h = host.trim();
    if h.eq_ignore_ascii_case("localhost") {
        return true;
    }
    // 去除 IPv6 方括号（如 `[::1]`）
    let bare = h
        .strip_prefix('[')
        .and_then(|s| s.strip_suffix(']'))
        .unwrap_or(h);
    // 尝试解析为 IpAddr
    match bare.parse::<std::net::IpAddr>() {
        Ok(std::net::IpAddr::V4(v4)) => v4.is_loopback(),
        Ok(std::net::IpAddr::V6(v6)) => v6.is_loopback(),
        Err(_) => {
            // 解析失败（如域名），保守判定为非 loopback，强制认证
            false
        }
    }
}

/// 判定当前配置是否需要强制认证。
/// - host 非 loopback：必须认证（且 auth_token 必须存在，否则启动会报错）
/// - host 为 loopback 但显式配置了 auth_token：也启用认证（用户主动加固）
pub fn requires_auth(host: &str, auth_token: &Option<String>) -> bool {
    !is_loopback_host(host) || auth_token.is_some()
}

/// 校验请求携带的 token 是否与配置的 auth_token 匹配（常量时间）。
/// 返回 true 表示通过（含「未配置 token 则放行」的兜底，调用方应先用 requires_auth 判定）。
pub fn token_matches(provided: Option<&str>, expected: &Option<String>) -> bool {
    match expected {
        // 未配置 token：放行（调用方已通过 requires_auth 保证只在 loopback 下到达此处）
        None => true,
        Some(expected_str) => match provided {
            None => false,
            Some(provided_str) => {
                // subtle::ConstantTimeEq 要求两端等长；长度不等时 ct_eq 返回 0（false），
                // 但为避免长度差异导致的提前退出计时泄露，我们确保比较路径不因长度而分支提前返回。
                let p = provided_str.as_bytes();
                let e = expected_str.as_bytes();
                // ct_eq 内部在长度不等时仍会遍历较短的一侧，返回 0，无提前退出
                p.ct_eq(e).unwrap_u8() == 1
            }
        },
    }
}

/// 从 Authorization 头提取 Bearer token（RFC 6750，scheme 不区分大小写）。
pub fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
    headers
        .get(axum::http::header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .and_then(|s| {
            // RFC 6750: scheme 名不区分大小写（Bearer / bearer / BEARER 均合法）
            // 找到第一个空格分隔 scheme 与 token，比较 scheme 部分忽略大小写
            let trimmed = s.trim_start();
            let split = trimmed.find(char::is_whitespace)?;
            let (scheme, rest) = trimmed.split_at(split);
            if scheme.eq_ignore_ascii_case("Bearer") {
                Some(rest.trim())
            } else {
                None
            }
        })
}

/// axum 中间件：对 `/api/*` 路由强制 Bearer token 校验。
/// 仅在 `requires_auth` 为 true 时挂载。
/// `/health`、`/ws`、静态资源放行；`/ws` 的 token 校验在 ws_handler 内完成。
pub async fn require_bearer_auth(
    headers: HeaderMap,
    request: Request,
    next: Next,
) -> Response {
    let path = request.uri().path();

    // 仅对 /api/ 前缀的请求强制认证
    if !path.starts_with("/api/") {
        return next.run(request).await;
    }

    // expected_token 通过 extension 注入（见 mod.rs 装配处）
    let expected = request
        .extensions()
        .get::<AuthConfig>()
        .map(|c| c.token.clone())
        .unwrap_or(None);

    let provided = extract_bearer_token(&headers);
    if token_matches(provided, &expected) {
        next.run(request).await
    } else {
        (
            StatusCode::UNAUTHORIZED,
            Json(json!({ "error": "unauthorized", "message": "missing or invalid token" })),
        )
            .into_response()
    }
}

/// 通过 extension 注入到 Router 的认证配置。
#[derive(Clone, Debug)]
pub struct AuthConfig {
    pub token: Option<String>,
}

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

    #[test]
    fn loopback_detection() {
        // IPv4 loopback 整个 127.0.0.0/8 段
        assert!(is_loopback_host("127.0.0.1"));
        assert!(is_loopback_host("127.0.0.2"));
        assert!(is_loopback_host("127.1.2.3"));
        assert!(is_loopback_host("127.255.255.255"));
        // IPv6 loopback 的各种等价表示
        assert!(is_loopback_host("::1"));
        assert!(is_loopback_host("[::1]"));
        assert!(is_loopback_host("0:0:0:0:0:0:0:1"));
        // localhost 主机名（大小写不敏感）
        assert!(is_loopback_host("localhost"));
        assert!(is_loopback_host("LOCALHOST"));
        assert!(is_loopback_host("Localhost"));
        // 空白容错
        assert!(is_loopback_host("  127.0.0.1  "));
        assert!(is_loopback_host("  localhost  "));
        // 非 loopback
        assert!(!is_loopback_host("0.0.0.0"));
        assert!(!is_loopback_host("192.168.1.1"));
        assert!(!is_loopback_host("10.0.0.1"));
        assert!(!is_loopback_host("172.16.0.1"));
        assert!(!is_loopback_host("example.com"));
        assert!(!is_loopback_host("picobot.local"));
        assert!(!is_loopback_host("::"));
        assert!(!is_loopback_host(""));
    }

    #[test]
    fn requires_auth_logic() {
        // 非 loopback 无论 token 是否配置都要认证
        assert!(requires_auth("0.0.0.0", &None));
        assert!(requires_auth("192.168.1.1", &None));
        assert!(requires_auth("0.0.0.0", &Some("t".into())));
        // loopback 无 token：免认证
        assert!(!requires_auth("127.0.0.1", &None));
        assert!(!requires_auth("::1", &None));
        assert!(!requires_auth("localhost", &None));
        // loopback 有 token：认证（用户主动加固）
        assert!(requires_auth("127.0.0.1", &Some("t".into())));
        // 域名始终需认证
        assert!(requires_auth("myserver.com", &None));
    }

    #[test]
    fn token_matching() {
        // 未配置 token：始终放行
        assert!(token_matches(None, &None));
        assert!(token_matches(Some("anything"), &None));
        // 配置了 token
        let expected = Some("s3cr3t".to_string());
        assert!(token_matches(Some("s3cr3t"), &expected));
        assert!(!token_matches(Some("wrong"), &expected));
        assert!(!token_matches(None, &expected));
        // 长度不同
        assert!(!token_matches(Some("s3cr3t-extra"), &expected));
        assert!(!token_matches(Some("s3"), &expected));
        assert!(!token_matches(Some(""), &expected));
        // 空字符串 token（配置了但为空 — 等同于未配置的放行语义由 requires_auth 控制）
        let empty_expected = Some(String::new());
        assert!(token_matches(Some(""), &empty_expected));
        assert!(!token_matches(Some("x"), &empty_expected));
    }

    #[test]
    fn extract_bearer_token_rfc6750() {
        fn make_auth_header(value: &str) -> HeaderMap {
            let mut h = HeaderMap::new();
            h.insert(
                axum::http::header::AUTHORIZATION,
                axum::http::HeaderValue::from_str(value).unwrap(),
            );
            h
        }

        // 标准格式
        assert_eq!(
            extract_bearer_token(&make_auth_header("Bearer abc123")),
            Some("abc123")
        );
        // scheme 大小写不敏感
        assert_eq!(
            extract_bearer_token(&make_auth_header("bearer abc123")),
            Some("abc123")
        );
        assert_eq!(
            extract_bearer_token(&make_auth_header("BEARER abc123")),
            Some("abc123")
        );
        // 多空格容忍
        assert_eq!(
            extract_bearer_token(&make_auth_header("Bearer  abc123")),
            Some("abc123")
        );
        assert_eq!(
            extract_bearer_token(&make_auth_header("Bearer\tabc123")),
            Some("abc123")
        );
        // 非 Bearer scheme
        assert_eq!(
            extract_bearer_token(&make_auth_header("Basic dXNlcjpwYXNz")),
            None
        );
        // 无 Authorization 头
        let empty = HeaderMap::new();
        assert_eq!(extract_bearer_token(&empty), None);
        // 缺少 token 部分
        assert_eq!(extract_bearer_token(&make_auth_header("Bearer")), None);
        assert_eq!(extract_bearer_token(&make_auth_header("Bearer ")), Some(""));
    }
}