fix: preserve provider auth metadata

2026-10-11 03:39:49 +08:00 · 2026-06-20 22:11:42 +08:00
parent 1123095bb7
commit f1c30439ff
51 changed files with 2098 additions and 420 deletions
@@ -1,6 +1,7 @@
 use super::super::helpers::admin_provider_oauth_key_name_from_auth_config;
 use super::super::token_import::{
-    build_provider_access_token_import_auth_config, provider_type_supports_access_token_import,
+    build_provider_access_token_import_auth_config, decode_access_token_expires_at,
+    provider_oauth_import_authorization_bearer_token, provider_type_supports_access_token_import,
 };
 use super::kiro_import::execute_admin_provider_oauth_kiro_batch_import;
 use super::parse::{
@@ -394,9 +395,22 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
        let AdminProviderOAuthResolvedBatchImport {
            access_token,
            mut auth_config,
-            expires_at,
+            mut expires_at,
        } = resolved_import;
        apply_admin_provider_oauth_batch_import_hints(provider_type, entry, &mut auth_config);
+        if let Some(header_access_token) =
+            provider_oauth_import_authorization_bearer_token(entry.request_headers.as_ref())
+        {
+            if let Some(header_expires_at) =
+                decode_access_token_expires_at(&header_access_token).or(entry.expires_at)
+            {
+                auth_config.insert("expires_at".to_string(), json!(header_expires_at));
+                expires_at = Some(header_expires_at);
+            } else {
+                auth_config.remove("expires_at");
+                expires_at = None;
+            }
+        }

        let duplicate =
            match find_duplicate_provider_oauth_key(state, provider_id, &auth_config, None).await {
@@ -1,4 +1,8 @@
-use super::super::token_import::{import_tokens_from_raw_token, normalize_provider_import_tokens};
+use super::super::token_import::{
+    import_tokens_from_raw_token, normalize_provider_import_tokens,
+    normalize_provider_oauth_import_headers_from_object,
+    provider_oauth_import_authorization_bearer_token,
+};
 use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
 use crate::handlers::admin::provider::oauth::state::{current_unix_secs, json_u64_value};
 use axum::{
@@ -9,6 +13,7 @@ use axum::{
 };
 use serde::Deserialize;
 use serde_json::json;
+use std::collections::BTreeMap;
 use std::time::{SystemTime, UNIX_EPOCH};

 #[derive(Debug, Clone, Deserialize)]
@@ -22,6 +27,7 @@ pub(super) struct AdminProviderOAuthBatchImportEntry {
    pub parse_error: Option<String>,
    pub refresh_token: Option<String>,
    pub access_token: Option<String>,
+    pub export_access_token: Option<String>,
    pub raw_credentials: Option<serde_json::Value>,
    pub expires_at: Option<u64>,
    pub account_id: Option<String>,
@@ -37,6 +43,7 @@ pub(super) struct AdminProviderOAuthBatchImportEntry {
    pub sso_rw_token: Option<String>,
    pub cf_cookies: Option<String>,
    pub cf_clearance: Option<String>,
+    pub request_headers: Option<BTreeMap<String, String>>,
    pub user_agent: Option<String>,
    pub browser_profile: Option<String>,
 }
@@ -183,6 +190,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
                    parse_error: None,
                    refresh_token,
                    access_token,
+                    export_access_token: None,
                    raw_credentials: None,
                    expires_at: None,
                    account_id: None,
@@ -198,6 +206,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
                    sso_rw_token: grok_cookie_value(raw_token, "sso-rw"),
                    cf_cookies: grok_cookie_profile(raw_token),
                    cf_clearance: grok_cookie_value(raw_token, "cf_clearance"),
+                    request_headers: None,
                    user_agent: None,
                    browser_profile: None,
                })
@@ -237,10 +246,24 @@ fn extract_admin_provider_oauth_batch_import_entry(
                    .as_deref()
                    .and_then(|cookie| grok_cookie_value(cookie, "sso"))
            });
+            let request_headers = normalize_provider_oauth_import_headers_from_object(object);
+            let header_bearer_token =
+                provider_oauth_import_authorization_bearer_token(request_headers.as_ref());
+            let export_access_token = if refresh_token.is_none() {
+                access_token
+                    .as_ref()
+                    .filter(|value| header_bearer_token.as_deref() != Some(value.as_str()))
+                    .cloned()
+            } else {
+                None
+            };
            let (refresh_token, access_token) = normalize_provider_import_tokens(
                provider_type,
                refresh_token.as_deref(),
-                access_token.as_deref().or(session_token.as_deref()),
+                access_token
+                    .as_deref()
+                    .or(session_token.as_deref())
+                    .or(header_bearer_token.as_deref()),
            );
            let windsurf_api_key = is_windsurf
                .then(|| {
@@ -386,6 +409,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
                parse_error: None,
                refresh_token,
                access_token,
+                export_access_token,
                raw_credentials,
                expires_at,
                account_id,
@@ -401,6 +425,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
                sso_rw_token,
                cf_cookies,
                cf_clearance,
+                request_headers,
                user_agent,
                browser_profile,
            })
@@ -479,6 +504,7 @@ fn parse_error_entry(error: String) -> AdminProviderOAuthBatchImportEntry {
        parse_error: Some(error),
        refresh_token: None,
        access_token: None,
+        export_access_token: None,
        raw_credentials: None,
        expires_at: None,
        account_id: None,
@@ -494,6 +520,7 @@ fn parse_error_entry(error: String) -> AdminProviderOAuthBatchImportEntry {
        sso_rw_token: None,
        cf_cookies: None,
        cf_clearance: None,
+        request_headers: None,
        user_agent: None,
        browser_profile: None,
    }
@@ -569,6 +596,11 @@ pub(super) fn apply_admin_provider_oauth_batch_import_hints(
            .entry("pool_tier".to_string())
            .or_insert_with(|| json!(pool_tier));
    }
+    if let Some(access_token) = entry.export_access_token.as_ref() {
+        auth_config
+            .entry("access_token".to_string())
+            .or_insert_with(|| json!(access_token));
+    }
    if let Some(user_id) = entry.user_id.as_ref() {
        auth_config
            .entry("user_id".to_string())
@@ -599,6 +631,11 @@ pub(super) fn apply_admin_provider_oauth_batch_import_hints(
            .entry("cf_clearance".to_string())
            .or_insert_with(|| json!(cf_clearance));
    }
+    if let Some(request_headers) = entry.request_headers.as_ref() {
+        auth_config
+            .entry("headers".to_string())
+            .or_insert_with(|| json!(request_headers));
+    }
    if let Some(user_agent) = entry.user_agent.as_ref() {
        auth_config
            .entry("user_agent".to_string())
@@ -721,11 +758,110 @@ mod tests {
        assert_eq!(entries.len(), 1);
        assert_eq!(entries[0].refresh_token, None);
        assert_eq!(entries[0].access_token.as_deref(), Some("at_1"));
+        assert_eq!(entries[0].export_access_token.as_deref(), Some("at_1"));
        assert_eq!(entries[0].expires_at, Some(2_100_000_000));
        assert_eq!(entries[0].account_id.as_deref(), Some("acc-1"));
        assert_eq!(entries[0].email.as_deref(), Some("[email protected]"));
    }

+    #[test]
+    fn parses_codex_import_header_overrides() {
+        let entries = parse_admin_provider_oauth_batch_import_entries(
+            "codex",
+            r#"[{"access_token":"jwt_access","headers":{"authorization":"Bearer session-1","chatgpt-account-id":"acc-1","host":"blocked.example"}}]"#,
+        );
+
+        assert_eq!(entries.len(), 1);
+        let headers = entries[0]
+            .request_headers
+            .as_ref()
+            .expect("headers should parse");
+        assert_eq!(
+            headers.get("authorization"),
+            Some(&"Bearer session-1".to_string())
+        );
+        assert_eq!(
+            headers.get("chatgpt-account-id"),
+            Some(&"acc-1".to_string())
+        );
+        assert!(!headers.contains_key("host"));
+    }
+
+    #[test]
+    fn treats_import_authorization_header_as_session_access_token() {
+        let entries = parse_admin_provider_oauth_batch_import_entries(
+            "codex",
+            r#"[{"access_token":"jwt_access","headers":{"authorization":"Bearer session-1"}}]"#,
+        );
+
+        assert_eq!(entries.len(), 1);
+        assert_eq!(entries[0].access_token.as_deref(), Some("jwt_access"));
+        assert_eq!(
+            entries[0].export_access_token.as_deref(),
+            Some("jwt_access")
+        );
+
+        let entries = parse_admin_provider_oauth_batch_import_entries(
+            "codex",
+            r#"[{"headers":{"authorization":"Bearer session-1"}}]"#,
+        );
+
+        assert_eq!(entries.len(), 1);
+        assert_eq!(entries[0].access_token.as_deref(), Some("session-1"));
+        assert!(entries[0].export_access_token.is_none());
+    }
+
+    #[test]
+    fn applies_codex_import_header_overrides_to_auth_config() {
+        let entries = parse_admin_provider_oauth_batch_import_entries(
+            "codex",
+            r#"{"accessToken":"at_1","headers":{"Authorization":"Bearer session-1"}}"#,
+        );
+        let mut auth_config = serde_json::Map::new();
+
+        apply_admin_provider_oauth_batch_import_hints("codex", &entries[0], &mut auth_config);
+
+        assert_eq!(
+            auth_config.get("headers"),
+            Some(&json!({"authorization":"Bearer session-1"}))
+        );
+        assert_eq!(auth_config.get("access_token"), Some(&json!("at_1")));
+    }
+
+    #[test]
+    fn does_not_export_access_token_when_it_is_the_header_bearer() {
+        let entries = parse_admin_provider_oauth_batch_import_entries(
+            "codex",
+            r#"{"accessToken":"session-1","headers":{"Authorization":"Bearer session-1"}}"#,
+        );
+        let mut auth_config = serde_json::Map::new();
+
+        apply_admin_provider_oauth_batch_import_hints("codex", &entries[0], &mut auth_config);
+
+        assert!(entries[0].export_access_token.is_none());
+        assert!(auth_config.get("access_token").is_none());
+        assert_eq!(
+            auth_config.get("headers"),
+            Some(&json!({"authorization":"Bearer session-1"}))
+        );
+    }
+
+    #[test]
+    fn does_not_export_access_token_from_refresh_token_payload() {
+        let entries = parse_admin_provider_oauth_batch_import_entries(
+            "codex",
+            r#"{"refreshToken":"refresh-1","accessToken":"refreshed-access"}"#,
+        );
+        let mut auth_config = serde_json::Map::new();
+
+        apply_admin_provider_oauth_batch_import_hints("codex", &entries[0], &mut auth_config);
+
+        assert_eq!(entries[0].refresh_token.as_deref(), Some("refresh-1"));
+        assert_eq!(entries[0].access_token.as_deref(), Some("refreshed-access"));
+        assert!(entries[0].export_access_token.is_none());
+        assert!(auth_config.get("access_token").is_none());
+    }
+
    #[test]
    fn parses_common_chatgpt_web_json_aliases() {
        let entries = parse_admin_provider_oauth_batch_import_entries(
@@ -16,7 +16,9 @@ use super::super::state::{
 };
 use super::helpers::admin_provider_oauth_key_name_from_auth_config;
 use super::token_import::{
-    build_provider_access_token_import_auth_config, normalize_provider_import_tokens,
+    build_provider_access_token_import_auth_config, decode_access_token_expires_at,
+    normalize_provider_import_tokens, normalize_provider_oauth_import_headers_from_object,
+    provider_oauth_import_authorization_bearer_token_from_object,
    provider_type_supports_access_token_import,
 };
 use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_import_provider_id;
@@ -151,6 +153,20 @@ fn import_payload_u64_any(
    })
 }

+fn import_payload_export_access_token(
+    payload: &serde_json::Map<String, serde_json::Value>,
+) -> Option<String> {
+    if import_payload_string(payload, "refresh_token", "refreshToken").is_some() {
+        return None;
+    }
+    let access_token = import_payload_string(payload, "access_token", "accessToken")?;
+    let header_bearer = provider_oauth_import_authorization_bearer_token_from_object(payload);
+    if header_bearer.as_deref() == Some(access_token.as_str()) {
+        return None;
+    }
+    Some(access_token)
+}
+
 fn apply_single_import_hints(
    provider_type: &str,
    payload: &serde_json::Map<String, serde_json::Value>,
@@ -267,6 +283,18 @@ fn apply_single_import_hints(
            }
        });
    }
+
+    if let Some(access_token) = import_payload_export_access_token(payload) {
+        auth_config
+            .entry("access_token".to_string())
+            .or_insert_with(|| json!(access_token));
+    }
+
+    if let Some(request_headers) = normalize_provider_oauth_import_headers_from_object(payload) {
+        auth_config
+            .entry("headers".to_string())
+            .or_insert_with(|| json!(request_headers));
+    }
 }

 async fn resolve_admin_provider_oauth_single_import_tokens(
@@ -497,7 +525,8 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
            "session_token",
            "sessionToken",
        ],
-    );
+    )
+    .or_else(|| provider_oauth_import_authorization_bearer_token_from_object(&raw_payload));
    let imported_expires_at =
        import_payload_u64_any(&raw_payload, &["expires_at", "expiresAt", "expired"]);
    let name = raw_payload
@@ -614,9 +643,22 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
    let AdminProviderOAuthSingleImportTokens {
        access_token,
        mut auth_config,
-        expires_at,
+        mut expires_at,
    } = resolved_import;
    apply_single_import_hints(&provider_type, &raw_payload, &mut auth_config);
+    if let Some(header_access_token) =
+        provider_oauth_import_authorization_bearer_token_from_object(&raw_payload)
+    {
+        if let Some(header_expires_at) =
+            decode_access_token_expires_at(&header_access_token).or(imported_expires_at)
+        {
+            auth_config.insert("expires_at".to_string(), json!(header_expires_at));
+            expires_at = Some(header_expires_at);
+        } else {
+            auth_config.remove("expires_at");
+            expires_at = None;
+        }
+    }
    let has_refresh_token = auth_config
        .get("refresh_token")
        .and_then(serde_json::Value::as_str)
@@ -776,6 +818,69 @@ mod tests {
        assert_eq!(auth_config.get("user_agent"), Some(&json!("antigravity")));
    }

+    #[test]
+    fn single_import_preserves_distinct_top_level_access_token_for_export() {
+        let payload = json!({
+            "access_token": "jwt-access-token",
+            "headers": {
+                "authorization": "Bearer imported-session-token"
+            },
+        })
+        .as_object()
+        .cloned()
+        .expect("payload should be an object");
+        let mut auth_config = serde_json::Map::new();
+
+        apply_single_import_hints("codex", &payload, &mut auth_config);
+
+        assert_eq!(
+            auth_config.get("access_token"),
+            Some(&json!("jwt-access-token"))
+        );
+        assert_eq!(
+            auth_config.get("headers"),
+            Some(&json!({"authorization":"Bearer imported-session-token"}))
+        );
+    }
+
+    #[test]
+    fn single_import_does_not_preserve_header_bearer_as_export_access_token() {
+        let payload = json!({
+            "access_token": "imported-session-token",
+            "headers": {
+                "authorization": "Bearer imported-session-token"
+            },
+        })
+        .as_object()
+        .cloned()
+        .expect("payload should be an object");
+        let mut auth_config = serde_json::Map::new();
+
+        apply_single_import_hints("codex", &payload, &mut auth_config);
+
+        assert!(auth_config.get("access_token").is_none());
+        assert_eq!(
+            auth_config.get("headers"),
+            Some(&json!({"authorization":"Bearer imported-session-token"}))
+        );
+    }
+
+    #[test]
+    fn single_import_does_not_preserve_access_token_from_refresh_token_payload_for_export() {
+        let payload = json!({
+            "refresh_token": "refresh-token",
+            "access_token": "refreshed-access-token",
+        })
+        .as_object()
+        .cloned()
+        .expect("payload should be an object");
+        let mut auth_config = serde_json::Map::new();
+
+        apply_single_import_hints("codex", &payload, &mut auth_config);
+
+        assert!(auth_config.get("access_token").is_none());
+    }
+
    #[test]
    fn windsurf_import_error_redacts_http_body() {
        let error = OAuthError::HttpStatus {
@@ -5,6 +5,7 @@ use base64::{
    Engine as _,
 };
 use serde_json::{json, Map, Value};
+use std::collections::BTreeMap;

 fn decode_base64_url_part(value: &str) -> Option<Vec<u8>> {
    URL_SAFE_NO_PAD
@@ -117,6 +118,95 @@ pub(super) fn decode_access_token_expires_at(access_token: &str) -> Option<u64>
    json_u64_value(claims.get("exp"))
 }

+fn normalize_import_header_name(raw: &str) -> Option<String> {
+    let value = raw.trim().to_ascii_lowercase();
+    if value.is_empty()
+        || matches!(
+            value.as_str(),
+            "host"
+                | "content-length"
+                | "connection"
+                | "keep-alive"
+                | "proxy-authenticate"
+                | "proxy-authorization"
+                | "proxy-connection"
+                | "te"
+                | "trailer"
+                | "transfer-encoding"
+                | "upgrade"
+        )
+    {
+        return None;
+    }
+    http::header::HeaderName::from_bytes(value.as_bytes())
+        .ok()
+        .map(|name| name.as_str().to_string())
+}
+
+fn normalize_import_header_value(value: &Value) -> Option<String> {
+    let value = match value {
+        Value::String(raw) => raw.trim().to_string(),
+        Value::Number(raw) => raw.to_string(),
+        Value::Bool(raw) => raw.to_string(),
+        _ => return None,
+    };
+    if value.is_empty() || http::header::HeaderValue::from_str(&value).is_err() {
+        return None;
+    }
+    Some(value)
+}
+
+pub(super) fn normalize_provider_oauth_import_headers(
+    value: Option<&Value>,
+) -> Option<BTreeMap<String, String>> {
+    let object = value?.as_object()?;
+    let mut headers = BTreeMap::new();
+    for (raw_key, raw_value) in object {
+        let Some(key) = normalize_import_header_name(raw_key) else {
+            continue;
+        };
+        let Some(value) = normalize_import_header_value(raw_value) else {
+            continue;
+        };
+        headers.insert(key, value);
+    }
+    (!headers.is_empty()).then_some(headers)
+}
+
+pub(super) fn normalize_provider_oauth_import_headers_from_object(
+    object: &Map<String, Value>,
+) -> Option<BTreeMap<String, String>> {
+    normalize_provider_oauth_import_headers(
+        object
+            .get("headers")
+            .or_else(|| object.get("request_headers"))
+            .or_else(|| object.get("requestHeaders"))
+            .or_else(|| object.get("header_overrides"))
+            .or_else(|| object.get("headerOverrides"))
+            .or_else(|| object.get("extra_headers"))
+            .or_else(|| object.get("extraHeaders")),
+    )
+}
+
+pub(super) fn provider_oauth_import_authorization_bearer_token(
+    headers: Option<&BTreeMap<String, String>>,
+) -> Option<String> {
+    let value = headers?.get("authorization")?.trim();
+    let (scheme, token) = value.split_once(' ')?;
+    if !scheme.eq_ignore_ascii_case("bearer") {
+        return None;
+    }
+    let token = token.trim();
+    (!token.is_empty()).then(|| token.to_string())
+}
+
+pub(super) fn provider_oauth_import_authorization_bearer_token_from_object(
+    object: &Map<String, Value>,
+) -> Option<String> {
+    let headers = normalize_provider_oauth_import_headers_from_object(object)?;
+    provider_oauth_import_authorization_bearer_token(Some(&headers))
+}
+
 pub(super) fn provider_type_supports_access_token_import(provider_type: &str) -> bool {
    matches!(
        provider_type.trim().to_ascii_lowercase().as_str(),
@@ -176,7 +266,9 @@ pub(super) fn build_provider_access_token_import_auth_config(
 mod tests {
    use super::{
        build_provider_access_token_import_auth_config, decode_access_token_expires_at,
-        looks_like_access_token, normalize_provider_import_tokens, normalize_single_import_tokens,
+        looks_like_access_token, normalize_provider_import_tokens,
+        normalize_provider_oauth_import_headers, normalize_single_import_tokens,
+        provider_oauth_import_authorization_bearer_token,
    };
    use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
    use serde_json::json;
@@ -247,6 +339,38 @@ mod tests {
        );
    }

+    #[test]
+    fn normalizes_import_header_overrides_for_auth_headers() {
+        let headers = normalize_provider_oauth_import_headers(Some(&json!({
+            "Authorization": " Bearer session-token ",
+            "X-Feature": true,
+            "Host": "evil.example",
+            "X-Bad": "line\nbreak",
+        })))
+        .expect("headers should normalize");
+
+        assert_eq!(
+            headers.get("authorization"),
+            Some(&"Bearer session-token".to_string())
+        );
+        assert_eq!(headers.get("x-feature"), Some(&"true".to_string()));
+        assert!(!headers.contains_key("host"));
+        assert!(!headers.contains_key("x-bad"));
+    }
+
+    #[test]
+    fn extracts_import_authorization_header_bearer_token() {
+        let headers = normalize_provider_oauth_import_headers(Some(&json!({
+            "Authorization": "Bearer at-session-token",
+        })))
+        .expect("headers should normalize");
+
+        assert_eq!(
+            provider_oauth_import_authorization_bearer_token(Some(&headers)).as_deref(),
+            Some("at-session-token")
+        );
+    }
+
    #[test]
    fn builds_chatgpt_web_temporary_auth_config_from_access_token() {
        let token = unsigned_jwt(json!({