Implement compression quota refunds and admin manual subscription

src/api/context.rs

@@ -0,0 +1,229 @@
+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 {
+    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 (jar, session_id) = ensure_session_cookie(jar);
+    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, "账号已被禁用"));
+    }
+
+    Ok(Some(Principal::User {
+        user_id: user.id,
+        role: user.role,
+        email_verified: user.email_verified_at.is_some(),
+    }))
+}
+
+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;
+
+    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(),
+    }))
+}
+
+pub fn ensure_session_cookie(jar: CookieJar) -> (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)
+        .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))
+}
+