ystp/src/api/context.rs

use crate::auth;
use crate::error::{AppError, ErrorCode};
use crate::state::AppState;

use axum::http::HeaderMap;
use axum_extra::extract::cookie::{Cookie, CookieJar, SameSite};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use chrono::{DateTime, Utc};
use hmac::{Hmac, Mac};
use rand::RngCore;
use serde::Serialize;
use sha2::Sha256;
use sqlx::FromRow;
use std::net::IpAddr;
use time::Duration as TimeDuration;
use uuid::Uuid;

#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Principal {
    Anonymous {
        session_id: String,
    },
    User {
        user_id: Uuid,
        role: String,
        email_verified: bool,
    },
    ApiKey {
        user_id: Uuid,
        api_key_id: Uuid,
        role: String,
        email_verified: bool,
    },
}

pub fn client_ip(headers: &HeaderMap, connect_ip: IpAddr) -> IpAddr {
    resolve_client_ip(headers, connect_ip, crate::config::trust_proxy_headers())
}

fn resolve_client_ip(headers: &HeaderMap, connect_ip: IpAddr, trust_proxy: bool) -> IpAddr {
    if !trust_proxy {
        return connect_ip;
    }

    if let Some(ip) = parse_forwarded_for(headers) {
        return ip;
    }
    if let Some(ip) = headers
        .get("x-real-ip")
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.parse::<IpAddr>().ok())
    {
        return ip;
    }
    connect_ip
}

fn parse_forwarded_for(headers: &HeaderMap) -> Option<IpAddr> {
    let value = headers.get("x-forwarded-for")?.to_str().ok()?;
    value
        .split(',')
        .next()
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .and_then(|s| s.parse::<IpAddr>().ok())
}

pub async fn authenticate(
    state: &AppState,
    jar: CookieJar,
    headers: &HeaderMap,
    ip: IpAddr,
) -> Result<(CookieJar, Principal), AppError> {
    if let Some(principal) = try_jwt(state, headers).await? {
        return Ok((jar, principal));
    }
    if let Some(principal) = try_api_key(state, headers, ip).await? {
        return Ok((jar, principal));
    }

    if !state.config.allow_anonymous_upload {
        return Err(AppError::new(ErrorCode::Unauthorized, "未登录"));
    }

    let cookie_secure = state
        .config
        .public_base_url
        .trim()
        .to_ascii_lowercase()
        .starts_with("https://");
    let (jar, session_id) = ensure_session_cookie(jar, cookie_secure);
    Ok((jar, Principal::Anonymous { session_id }))
}

async fn try_jwt(state: &AppState, headers: &HeaderMap) -> Result<Option<Principal>, AppError> {
    let auth_header = headers
        .get(axum::http::header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .unwrap_or("");

    if !auth_header.starts_with("Bearer ") {
        return Ok(None);
    }

    let claims = auth::require_jwt(&state.config.jwt_secret, headers)?;

    #[derive(Debug, FromRow)]
    struct UserAuthRow {
        id: Uuid,
        role: String,
        is_active: bool,
        email_verified_at: Option<DateTime<Utc>>,
    }

    let user = sqlx::query_as::<_, UserAuthRow>(
        r#"
        SELECT id, role::text AS role, is_active, email_verified_at
        FROM users
        WHERE id = $1
        "#,
    )
    .bind(claims.sub)
    .fetch_optional(&state.db)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?
    .ok_or_else(|| AppError::new(ErrorCode::Unauthorized, "用户不存在或未登录"))?;

    if !user.is_active {
        return Err(AppError::new(ErrorCode::Forbidden, "账号已被禁用"));
    }

    let verification_required =
        crate::services::settings::email_verification_required(state).await?;

    Ok(Some(Principal::User {
        user_id: user.id,
        role: user.role,
        email_verified: user.email_verified_at.is_some() || !verification_required,
    }))
}

async fn try_api_key(
    state: &AppState,
    headers: &HeaderMap,
    ip: IpAddr,
) -> Result<Option<Principal>, AppError> {
    let key = headers
        .get("x-api-key")
        .or_else(|| headers.get("X-API-Key"))
        .and_then(|v| v.to_str().ok())
        .map(str::trim)
        .filter(|v| !v.is_empty());

    let Some(full_key) = key else {
        return Ok(None);
    };

    let key_prefix = full_key
        .get(0..16)
        .ok_or_else(|| AppError::new(ErrorCode::Unauthorized, "API Key 格式错误"))?;

    #[derive(Debug, FromRow)]
    struct ApiKeyAuthRow {
        id: Uuid,
        user_id: Uuid,
        key_hash: String,
        is_active: bool,
        user_role: String,
        user_is_active: bool,
        email_verified_at: Option<DateTime<Utc>>,
    }

    let row = sqlx::query_as::<_, ApiKeyAuthRow>(
        r#"
        SELECT
          k.id,
          k.user_id,
          k.key_hash,
          k.is_active,
          u.role::text AS user_role,
          u.is_active AS user_is_active,
          u.email_verified_at
        FROM api_keys k
        JOIN users u ON u.id = k.user_id
        WHERE k.key_prefix = $1
        "#,
    )
    .bind(key_prefix)
    .fetch_optional(&state.db)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "查询 API Key 失败").with_source(err))?
    .ok_or_else(|| AppError::new(ErrorCode::Unauthorized, "API Key 无效"))?;

    if !row.user_is_active || !row.is_active {
        return Err(AppError::new(ErrorCode::Forbidden, "API Key 已禁用"));
    }

    let expected = api_key_hash(full_key, &state.config.api_key_pepper)?;
    if expected != row.key_hash {
        return Err(AppError::new(ErrorCode::Unauthorized, "API Key 无效"));
    }

    let _ =
        sqlx::query("UPDATE api_keys SET last_used_at = NOW(), last_used_ip = $2 WHERE id = $1")
            .bind(row.id)
            .bind(ip.to_string())
            .execute(&state.db)
            .await;

    let verification_required =
        crate::services::settings::email_verification_required(state).await?;

    Ok(Some(Principal::ApiKey {
        user_id: row.user_id,
        api_key_id: row.id,
        role: row.user_role,
        email_verified: row.email_verified_at.is_some() || !verification_required,
    }))
}

pub fn ensure_session_cookie(jar: CookieJar, secure: bool) -> (CookieJar, String) {
    if let Some(cookie) = jar.get("if_session") {
        let session_id = cookie.value().trim().to_string();
        if !session_id.is_empty() {
            return (jar, session_id);
        }
    }

    let session_id = generate_session_id();
    let cookie = Cookie::build(("if_session", session_id.clone()))
        .path("/")
        .http_only(true)
        .secure(secure)
        .same_site(SameSite::Lax)
        .max_age(TimeDuration::days(7))
        .build();

    (jar.add(cookie), session_id)
}

fn generate_session_id() -> String {
    let mut bytes = [0u8; 32];
    rand::rngs::OsRng.fill_bytes(&mut bytes);
    URL_SAFE_NO_PAD.encode(bytes)
}

pub fn api_key_hash(full_key: &str, pepper: &str) -> Result<String, AppError> {
    type HmacSha256 = Hmac<Sha256>;

    let mut mac = HmacSha256::new_from_slice(pepper.as_bytes()).map_err(|err| {
        AppError::new(ErrorCode::Internal, "API Key pepper 错误").with_source(err)
    })?;
    mac.update(full_key.as_bytes());
    let result = mac.finalize().into_bytes();
    Ok(hex::encode(result))
}

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

    #[test]
    fn untrusted_proxy_headers_cannot_override_peer_ip() {
        let mut headers = HeaderMap::new();
        headers.insert("x-forwarded-for", "203.0.113.9".parse().unwrap());
        let peer = "192.0.2.10".parse().unwrap();
        assert_eq!(resolve_client_ip(&headers, peer, false), peer);
    }

    #[test]
    fn trusted_proxy_headers_use_forwarded_client_ip() {
        let mut headers = HeaderMap::new();
        headers.insert(
            "x-forwarded-for",
            "203.0.113.9, 192.0.2.20".parse().unwrap(),
        );
        let peer = "192.0.2.10".parse().unwrap();
        assert_eq!(
            resolve_client_ip(&headers, peer, true),
            "203.0.113.9".parse::<IpAddr>().unwrap()
        );
    }
}