AstroResearch/src/api/auth.rs

// src/api/auth.rs
use axum::{
    extract::State,
    http::{header, header::HeaderMap, header::HeaderValue, Request, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response},
    Json,
};
use serde::{Deserialize, Serialize};
use std::sync::Arc;

use super::error::{ApiResult, AppError};
use super::AppState;

#[derive(Debug, Deserialize)]
pub struct LoginRequest {
    pub password: Option<String>,
}

#[derive(Debug, Serialize)]
pub struct LoginResponse {
    pub status: String,
    pub bookmarklet_key: String,
}

// 辅助函数：从 Cookie 字符串中解析出 session_id 的值
fn parse_session_cookie(cookie_str: &str) -> Option<String> {
    for cookie in cookie_str.split(';') {
        let mut parts = cookie.splitn(2, '=');
        if let (Some(name), Some(value)) = (parts.next(), parts.next()) {
            if name.trim() == "session_id" {
                return Some(value.trim().to_string());
            }
        }
    }
    None
}

/// 从请求头中提取认证 Token。
/// 优先从 Authorization: Bearer 头读取，其次从 Cookie 中的 session_id 读取。
pub fn extract_auth_token_from_headers(headers: &HeaderMap) -> Option<String> {
    // 1. Authorization: Bearer <token>
    if let Some(auth_str) = headers
        .get(header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
    {
        if let Some(rest) = auth_str.strip_prefix("Bearer ") {
            return Some(rest.to_string());
        }
    }
    // 2. Cookie: session_id=<token>
    if let Some(cookie_header) = headers.get(header::COOKIE) {
        if let Ok(cookie_str) = cookie_header.to_str() {
            return parse_session_cookie(cookie_str);
        }
    }
    None
}

/// 恒定时间密码校验，防止时序侧信道攻击。
///
/// 普通 `==` 在字节不匹配时立即返回，攻击者可通过测量响应时间逐字节猜测密码。
/// `ct_eq` 无论是否匹配都遍历全部字节，耗时恒定，消除时序差异。
///
/// 长度不等时不能直接返回 false（会泄露长度信息），先用虚假比较平衡耗时。
fn verify_password(input: &str, expected: &str) -> bool {
    use subtle::ConstantTimeEq;

    let input_bytes = input.as_bytes();
    let expected_bytes = expected.as_bytes();

    if input_bytes.len() != expected_bytes.len() {
        // 长度不同：做一次虚假比较平衡耗时，再返回 false
        let _ = input_bytes.ct_eq(input_bytes);
        return false;
    }
    input_bytes.ct_eq(expected_bytes).into()
}

// 辅助函数：清除过期的会话记录（24小时未活跃即过期）
fn prune_expired_sessions(
    sessions: &mut std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>,
) {
    let now = chrono::Utc::now();
    let expiry_duration = chrono::Duration::hours(24);
    sessions.retain(|_, last_active| {
        if let Ok(dur) = now.signed_duration_since(*last_active).to_std() {
            dur < expiry_duration
                .to_std()
                .unwrap_or(std::time::Duration::from_secs(86400))
        } else {
            true // 尚未到达或时钟回拨，保留
        }
    });
}

/// 计算浏览器书签专用的永久 API 密钥（基于 admin_password + salt 的 SHA-256 哈希）
fn get_bookmarklet_api_key(password: &str) -> String {
    use sha2::{Digest, Sha256};
    let mut hasher = Sha256::new();
    hasher.update(password.as_bytes());
    hasher.update(b"_bookmarklet_salt_key");
    format!("{:x}", hasher.finalize())
}

/// 最大会话数上限，超过后拒绝新登录以防止内存耗尽
const MAX_SESSIONS: usize = 1000;

// 登录接口：验证密码成功后，写入 HttpOnly Cookie
pub async fn login(
    State(state): State<Arc<AppState>>,
    _headers: axum::http::header::HeaderMap,
    axum::extract::ConnectInfo(addr): axum::extract::ConnectInfo<std::net::SocketAddr>,
    Json(body): Json<LoginRequest>,
) -> Result<impl IntoResponse, AppError> {
    // 使用真实连接 IP 做限速（X-Forwarded-For 可伪造，不应用于安全判断）
    let client_ip = addr.ip().to_string();

    // 检查速率限制：5 次失败 / 5 分钟窗口
    const MAX_FAILURES: u32 = 5;
    const WINDOW_SECS: u64 = 300;
    const MAX_RATE_LIMITER_ENTRIES: usize = 10000;
    {
        let mut limiter = state.login_rate_limiter.lock().await;
        // 批量清理过期条目，防止内存无限增长
        limiter.retain(|_, (_, first_fail)| first_fail.elapsed().as_secs() < WINDOW_SECS);
        // 容量保护：超过上限时直接清空（比 O(n log n) 排序更高效，且不影响安全性）
        if limiter.len() > MAX_RATE_LIMITER_ENTRIES {
            limiter.clear();
        }

        if let Some((count, first_fail)) = limiter.get(&client_ip) {
            let elapsed = first_fail.elapsed().as_secs();
            if elapsed < WINDOW_SECS && *count >= MAX_FAILURES {
                let remaining = WINDOW_SECS - elapsed;
                return Err(AppError::unauthorized(format!(
                    "登录尝试过于频繁，请在 {} 秒后重试",
                    remaining
                )));
            }
        }
    }

    let password = body
        .password
        .ok_or_else(|| AppError::bad_request("缺少密码字段"))?;
    if password.is_empty() {
        return Err(AppError::bad_request("密码不能为空"));
    }
    if verify_password(&password, &state.config.admin_password) {
        // 成功：清除该 IP 的失败记录
        state.login_rate_limiter.lock().await.remove(&client_ip);

        let token = uuid::Uuid::new_v4().to_string();

        // 异常安全地记录活跃会话及其当前时间
        {
            let mut sessions = state.sessions.write().await;
            // 先清理过期会话
            prune_expired_sessions(&mut sessions);
            // 检查会话上限
            if sessions.len() >= MAX_SESSIONS {
                return Err(AppError::unauthorized(format!(
                    "活跃会话数已达上限 ({}), 请稍后重试",
                    MAX_SESSIONS
                )));
            }
            sessions.insert(token.clone(), chrono::Utc::now());
        }

        // 构建 Set-Cookie 报头
        // 全局统一使用 SameSite=Lax 且不带 Secure，以完美支持 HTTP、HTTPS 以及局域网 IP 访问
        let cookie_value = format!(
            "session_id={}; HttpOnly; SameSite=Lax; Path=/; Max-Age=86400",
            token
        );
        let mut headers = HeaderMap::new();
        if let Ok(val) = HeaderValue::from_str(&cookie_value) {
            headers.insert(header::SET_COOKIE, val);
        }

        let bookmarklet_key = get_bookmarklet_api_key(&state.config.admin_password);

        Ok((
            headers,
            Json(LoginResponse {
                status: "ok".to_string(),
                bookmarklet_key,
            }),
        ))
    } else {
        // 失败：记录该 IP 的失败次数
        let mut limiter = state.login_rate_limiter.lock().await;
        let entry = limiter
            .entry(client_ip)
            .or_insert((0, std::time::Instant::now()));
        entry.0 += 1;
        Err(AppError::unauthorized("密码错误"))
    }
}

// 登出接口：清除内存会话并使 Cookie 失效
pub async fn logout(
    State(state): State<Arc<AppState>>,
    req: Request<axum::body::Body>,
) -> ApiResult<impl IntoResponse> {
    let token_to_remove = extract_auth_token_from_headers(req.headers());

    if let Some(ref token) = token_to_remove {
        state.sessions.write().await.remove(token);
    }

    // 通过 Max-Age=0 清除浏览器端的 Cookie
    let mut headers = HeaderMap::new();
    let delete_cookie = "session_id=; HttpOnly; SameSite=Lax; Path=/; Max-Age=0";
    if let Ok(val) = HeaderValue::from_str(delete_cookie) {
        headers.insert(header::SET_COOKIE, val);
    }

    Ok((headers, Json(serde_json::json!({ "status": "ok" }))))
}

/// 鉴权状态检查接口。
///
/// 注：此端点由 `auth_middleware` 保护——只有携带有效会话 Token 的请求才能到达此处，
/// 因此始终返回 `authenticated: true`。如果中间件被移除，此端点将无条件返回已认证。
pub async fn check_auth() -> Result<impl IntoResponse, StatusCode> {
    Ok(Json(serde_json::json!({ "authenticated": true })))
}

// 鉴权验证中间件
pub async fn auth_middleware(
    State(state): State<Arc<AppState>>,
    req: axum::extract::Request,
    next: Next,
) -> Response {
    let token = extract_auth_token_from_headers(req.headers());

    let mut is_valid = false;
    if let Some(ref tok) = token {
        // 1. 优先校验内存中的动态 Session
        // 先用读锁快速检查会话是否存在，并获取其最后活跃时间
        let mut last_active_opt = None;
        {
            let sessions = state.sessions.read().await;
            if let Some(last_active) = sessions.get(tok) {
                is_valid = true;
                last_active_opt = Some(*last_active);
            }
        }

        // 会话存在时，如果最后活跃时间超过 60 秒，则获取写锁更新最后活跃时间 + 过期清理
        if is_valid {
            let needs_update = match last_active_opt {
                Some(last_active) => {
                    if let Ok(dur) = chrono::Utc::now()
                        .signed_duration_since(last_active)
                        .to_std()
                    {
                        dur.as_secs() > 60
                    } else {
                        true // 时钟回拨等异常情况，强制更新
                    }
                }
                None => true,
            };

            if needs_update {
                let mut sessions = state.sessions.write().await;
                // 顺便执行过期会话自动清理垃圾收集
                prune_expired_sessions(&mut sessions);
                if let Some(last_active) = sessions.get_mut(tok) {
                    *last_active = chrono::Utc::now();
                }
            }
        }

        // 2. 校验书签专属永久 API 密钥（严格隔离限制：仅允许用于特定的书签导入和同步路径）
        if !is_valid {
            let path = req.uri().path();
            if path == "/api/upload" || path == "/api/active_bibcode" {
                let expected_key = get_bookmarklet_api_key(&state.config.admin_password);
                if *tok == expected_key {
                    is_valid = true;
                }
            }
        }
    }

    if is_valid {
        next.run(req).await
    } else {
        StatusCode::UNAUTHORIZED.into_response()
    }
}

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

    #[test]
    fn test_verify_password_correct() {
        assert!(verify_password("secret123", "secret123"));
        assert!(verify_password("", ""));
    }

    #[test]
    fn test_verify_password_wrong() {
        assert!(!verify_password("secret123", "secret456"));
        assert!(!verify_password("secret", "secret123"));
        assert!(!verify_password("secret123", "secret"));
        assert!(!verify_password("wrong", ""));
    }
}