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::().ok()) { return ip; } connect_ip } fn parse_forwarded_for(headers: &HeaderMap) -> Option { 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::().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, 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>, } 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, 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>, } 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 { type HmacSha256 = Hmac; 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::().unwrap() ); } }