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::ExternalApiAuthState,
};

#[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<ExternalApiAuthState>,
    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
        .authenticator()
        .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<ExternalApiAuthState>,
    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::*;
    use crate::state::external_api_auth::ExternalApiKeyAuthenticator;
    use axum::{
        Router,
        body::{Body, to_bytes},
        extract::Extension,
        middleware,
        routing::get,
    };
    use futures_util::future::BoxFuture;
    use spacetime_client::{ExternalApiKeyRecord, SpacetimeClientError};
    use std::sync::{
        Arc, Mutex,
        atomic::{AtomicUsize, Ordering},
    };
    use tower::ServiceExt;

    struct RecordingAuthenticator {
        requests: Mutex<Vec<ExternalApiKeyAuthenticateRecordInput>>,
        error: Option<&'static str>,
    }

    impl ExternalApiKeyAuthenticator for RecordingAuthenticator {
        fn authenticate_external_api_key(
            &self,
            input: ExternalApiKeyAuthenticateRecordInput,
        ) -> BoxFuture<'_, Result<ExternalApiKeyRecord, SpacetimeClientError>> {
            self.requests.lock().unwrap().push(input);
            Box::pin(async move {
                if let Some(error) = self.error {
                    return Err(SpacetimeClientError::Procedure(error.to_string()));
                }
                Ok(ExternalApiKeyRecord {
                    key_id: "key-from-store".to_string(),
                    owner_user_id: "owner-from-store".to_string(),
                    name: "测试密钥".to_string(),
                    key_prefix: "tnr_sk_fixture".to_string(),
                    scopes: vec!["editor:project".to_string()],
                    created_at: "0.000000Z".to_string(),
                    last_used_at: None,
                    revoked_at: None,
                    updated_at: "0.000000Z".to_string(),
                })
            })
        }
    }

    fn auth_test_router(
        dependency: Arc<RecordingAuthenticator>,
        entered: Arc<AtomicUsize>,
        mcp: bool,
    ) -> Router {
        let state = ExternalApiAuthState::new(dependency);
        let router = Router::new().route("/protected", get(move |Extension(principal): Extension<ExternalApiPrincipal>| {
            let entered = entered.clone();
            async move {
                entered.fetch_add(1, Ordering::Relaxed);
                axum::Json(json!({"owner": principal.owner_user_id(), "projectScope": principal.has_scope("editor:project")}))
            }
        }));
        if mcp {
            router.layer(middleware::from_fn_with_state(
                state,
                require_external_mcp_api_key,
            ))
        } else {
            router.layer(middleware::from_fn_with_state(
                state,
                require_external_api_key,
            ))
        }
    }

    #[tokio::test]
    async fn narrow_external_auth_forwards_store_identity_and_only_hashes_credentials() {
        let dependency = Arc::new(RecordingAuthenticator {
            requests: Mutex::default(),
            error: None,
        });
        let entered = Arc::new(AtomicUsize::new(0));
        let response = auth_test_router(dependency.clone(), entered.clone(), false)
            .oneshot(
                Request::builder()
                    .uri("/protected")
                    .header(AUTHORIZATION, "Bearer tnr_sk_fixture_secret")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        let principal = response.extensions().get::<ExternalApiPrincipal>().unwrap();
        assert_eq!(principal.owner_user_id(), "owner-from-store");
        assert_eq!(principal.key_id(), "key-from-store");
        let body: serde_json::Value =
            serde_json::from_slice(&to_bytes(response.into_body(), 1024).await.unwrap()).unwrap();
        assert_eq!(
            body,
            json!({"owner": "owner-from-store", "projectScope": true})
        );
        assert_eq!(entered.load(Ordering::Relaxed), 1);
        let requests = dependency.requests.lock().unwrap();
        assert_eq!(requests.len(), 1);
        assert_eq!(
            requests[0].key_hash,
            hash_external_api_key("tnr_sk_fixture_secret")
        );
        assert!(requests[0].used_at_micros > 0);
    }

    #[tokio::test]
    async fn narrow_external_auth_preserves_failure_mapping_and_mcp_guide() {
        for (message, expected_status) in [
            ("API Key 不存在", StatusCode::UNAUTHORIZED),
            ("无权使用此密钥", StatusCode::FORBIDDEN),
            ("校验请求失败", StatusCode::BAD_REQUEST),
        ] {
            for mcp in [false, true] {
                let dependency = Arc::new(RecordingAuthenticator {
                    requests: Mutex::default(),
                    error: Some(message),
                });
                let entered = Arc::new(AtomicUsize::new(0));
                let response = auth_test_router(dependency.clone(), entered.clone(), mcp)
                    .oneshot(
                        Request::builder()
                            .uri("/protected")
                            .header(AUTHORIZATION, "Bearer invalid-fixture-key")
                            .body(Body::empty())
                            .unwrap(),
                    )
                    .await
                    .unwrap();
                assert_eq!(response.status(), expected_status);
                assert!(
                    response
                        .extensions()
                        .get::<ExternalApiPrincipal>()
                        .is_none()
                );
                assert_eq!(entered.load(Ordering::Relaxed), 0);
                assert_eq!(dependency.requests.lock().unwrap().len(), 1);
                let body = String::from_utf8(
                    to_bytes(response.into_body(), 16_384)
                        .await
                        .unwrap()
                        .to_vec(),
                )
                .unwrap();
                if mcp && expected_status == StatusCode::UNAUTHORIZED {
                    assert!(body.contains("MCP_AUTHENTICATION_REQUIRED"));
                    assert!(!body.contains(message));
                } else {
                    assert!(body.contains(message));
                }
            }
        }
    }

    #[tokio::test]
    async fn narrow_external_auth_rejects_missing_credentials_before_store_access() {
        let dependency = Arc::new(RecordingAuthenticator {
            requests: Mutex::default(),
            error: None,
        });
        let entered = Arc::new(AtomicUsize::new(0));
        let response = auth_test_router(dependency.clone(), entered.clone(), false)
            .oneshot(
                Request::builder()
                    .uri("/protected")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
        assert!(dependency.requests.lock().unwrap().is_empty());
        assert_eq!(entered.load(Ordering::Relaxed), 0);
    }

    #[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("已失效"));
    }
}