Genarrative/server-rs/crates/api-server/src/external_api_auth.rs

use axum::{
    extract::{Request, State},
    http::{
        HeaderMap, HeaderValue, StatusCode,
        header::{AUTHORIZATION, WWW_AUTHENTICATE},
    },
    middleware::Next,
    response::Response,
};
use serde_json::json;
use spacetime_client::ExternalApiKeyAuthenticateRecordInput;
use tracing::warn;

use crate::{
    editor_project::current_utc_micros,
    external_api_keys::{hash_external_api_key, map_external_api_key_error},
    http_error::AppError,
    request_context::RequestContext,
    state::AppState,
};

#[derive(Clone, Debug)]
pub struct ExternalApiPrincipal {
    owner_user_id: String,
    key_id: String,
    scopes: Vec<String>,
}

impl ExternalApiPrincipal {
    #[cfg(test)]
    pub(crate) fn for_test(owner_user_id: &str, scopes: &[&str]) -> Self {
        Self {
            owner_user_id: owner_user_id.to_string(),
            key_id: "external-api-key-test".to_string(),
            scopes: scopes.iter().map(|scope| (*scope).to_string()).collect(),
        }
    }

    pub fn owner_user_id(&self) -> &str {
        self.owner_user_id.as_str()
    }

    pub fn key_id(&self) -> &str {
        self.key_id.as_str()
    }

    pub fn has_scope(&self, scope: &str) -> bool {
        self.scopes.iter().any(|item| item == scope)
    }
}

pub async fn require_external_api_key(
    State(state): State<AppState>,
    mut request: Request,
    next: Next,
) -> Result<Response, AppError> {
    let request_id = request
        .extensions()
        .get::<RequestContext>()
        .map(|context| context.request_id().to_string())
        .unwrap_or_else(|| "unknown".to_string());
    let raw_key = extract_external_api_bearer(request.headers())?;
    let key = state
        .spacetime_client()
        .authenticate_external_api_key(ExternalApiKeyAuthenticateRecordInput {
            key_hash: hash_external_api_key(raw_key.as_str()),
            used_at_micros: current_utc_micros(),
        })
        .await
        .map_err(|error| {
            warn!(
                %request_id,
                error = %error,
                "外部 API Key 校验失败"
            );
            map_external_api_key_error(error)
        })?;
    let principal = ExternalApiPrincipal {
        owner_user_id: key.owner_user_id,
        key_id: key.key_id,
        scopes: key.scopes,
    };
    request.extensions_mut().insert(principal.clone());

    let mut response = next.run(request).await;
    response.extensions_mut().insert(principal);
    Ok(response)
}

pub async fn require_external_mcp_api_key(
    State(state): State<AppState>,
    request: Request,
    next: Next,
) -> Result<Response, AppError> {
    let request_context = request.extensions().get::<RequestContext>().cloned();
    match require_external_api_key(State(state), request, next).await {
        Ok(response) => Ok(response),
        Err(error) if error.status_code() == StatusCode::UNAUTHORIZED => {
            Ok(map_external_mcp_authentication_error(error)
                .into_response_with_context(request_context.as_ref()))
        }
        Err(error) => Err(error),
    }
}

fn map_external_mcp_authentication_error(error: AppError) -> AppError {
    debug_assert_eq!(error.status_code(), StatusCode::UNAUTHORIZED);
    external_mcp_authentication_guide_error()
}

fn external_mcp_authentication_guide_error() -> AppError {
    AppError::from_status(StatusCode::UNAUTHORIZED)
        .with_message("连接陶泥儿托管 MCP 需要开发者 API Key")
        .with_details(json!({
            "guide": {
                "reason": "MCP_AUTHENTICATION_REQUIRED",
                "action": "CONFIGURE_BEARER_API_KEY",
                "authentication": {
                    "scheme": "Bearer",
                    "header": "Authorization",
                    "valueFormat": "Bearer <tnr_sk_...>"
                },
                "keyManagement": {
                    "navigationLabel": "开发者 API Key",
                    "rawKeyShownOnce": true
                },
                "retry": {
                    "method": "POST",
                    "path": "/api/external/v1/mcp",
                    "rpcMethod": "initialize"
                },
                "steps": [
                    "登录陶泥儿，在「开发者 API Key」中创建密钥；原始密钥只显示一次",
                    "把密钥配置为 MCP 连接的 Bearer token；不要粘贴到聊天或写入仓库",
                    "使用相同 MCP URL 重新发送 initialize"
                ],
                "credentialSafety": {
                    "rawKeyShownOnce": true,
                    "neverPasteIntoChat": true,
                    "neverStoreInRepository": true
                },
                "publicDiscovery": {
                    "manifest": "/api/external/v1/agent-integration.json",
                    "skill": "/api/external/v1/skill/SKILL.md",
                    "openapi": "/api/external/v1/openapi.json"
                }
            }
        }))
        .with_header(
            WWW_AUTHENTICATE.as_str(),
            HeaderValue::from_static("Bearer realm=\"genarrative-external-editor\""),
        )
}

fn extract_external_api_bearer(headers: &HeaderMap) -> Result<String, AppError> {
    let authorization = headers
        .get(AUTHORIZATION)
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .ok_or_else(|| AppError::from_status(StatusCode::UNAUTHORIZED))?;

    authorization
        .strip_prefix("Bearer ")
        .or_else(|| authorization.strip_prefix("bearer "))
        .map(str::trim)
        .filter(|token| !token.is_empty())
        .map(ToOwned::to_owned)
        .ok_or_else(|| AppError::from_status(StatusCode::UNAUTHORIZED))
}

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

    #[test]
    fn mcp_authentication_guide_replaces_sensitive_key_diagnostics() {
        let error = AppError::from_status(StatusCode::UNAUTHORIZED).with_details(json!({
            "provider": "external-api-key",
            "message": "SENSITIVE_KEY_LURE 不存在或已失效"
        }));

        let mapped = map_external_mcp_authentication_error(error);
        let serialized = serde_json::to_string(
            mapped
                .details()
                .expect("mapped authentication error should contain guide details"),
        )
        .expect("guide should serialize");
        assert_eq!(mapped.status_code(), StatusCode::UNAUTHORIZED);
        assert_eq!(mapped.message(), "连接陶泥儿托管 MCP 需要开发者 API Key");
        assert!(serialized.contains("MCP_AUTHENTICATION_REQUIRED"));
        assert!(serialized.contains("CONFIGURE_BEARER_API_KEY"));
        assert!(!serialized.contains("SENSITIVE_KEY_LURE"));
        assert!(!serialized.contains("provider"));
        assert!(!serialized.contains("不存在"));
        assert!(!serialized.contains("已失效"));
    }
}