diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 4bed4eebb0..482f346557 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -4488,6 +4488,63 @@ fn hosted_actor_credentials_are_only_available_to_host_owned_mcp_servers() { ); } +#[test] +fn gov_codex_apps_retain_local_chatgpt_auth_provider() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let provider = codex_model_provider::auth_provider_from_auth(&auth); + let mut config = crate::mcp::tests::test_mcp_config(codex_home.path().to_path_buf()); + config.chatgpt_base_url = "https://gov.chatgpt.com/backend-api/".to_string(); + config.apps_enabled = true; + + let local_config = crate::codex_apps_mcp_server_config( + &config.chatgpt_base_url, + /*apps_mcp_product_sku*/ None, + /*originator*/ None, + ); + let mut remote_config = local_config.clone(); + remote_config.environment_id = "customer-executor".to_string(); + let mut custom_oauth = local_config.clone(); + custom_oauth.auth = McpServerAuth::OAuth; + + let mut catalog = crate::ResolvedMcpCatalog::builder(); + catalog.register(crate::McpServerRegistration::from_hosted_apps( + "test-host", + /*contribution_order*/ 0, + local_config, + )); + catalog.register(crate::McpServerRegistration::from_config( + "custom-oauth".to_string(), + custom_oauth, + )); + catalog.register(crate::McpServerRegistration::from_config( + "remote-gov".to_string(), + remote_config, + )); + config.mcp_server_catalog = catalog.build(); + + let servers = crate::effective_mcp_servers(&config, Some(&auth)); + let local = &servers[CODEX_APPS_MCP_SERVER_NAME]; + assert_eq!(local.config().auth, McpServerAuth::ChatGpt); + let retained = chatgpt_auth_provider_for_server(local, Some(Arc::clone(&provider))) + .expect("local Gov Apps must retain ChatGPT authentication"); + assert!( + Arc::ptr_eq(&provider, &retained), + "server selection must preserve the supplied provider" + ); + + assert_eq!(servers["custom-oauth"].config().auth, McpServerAuth::OAuth); + assert_eq!(servers["remote-gov"].config().auth, McpServerAuth::ChatGpt); + for name in ["custom-oauth", "remote-gov"] { + assert!( + chatgpt_auth_provider_for_server(&servers[name], Some(Arc::clone(&provider))).is_none(), + "{name} must not receive the host ChatGPT provider" + ); + } + + Ok(()) +} + #[tokio::test] async fn executor_owned_chatgpt_mcp_accepts_only_safe_explicit_authorization() -> anyhow::Result<()> { diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index c41c66e062..137f1bff04 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -374,9 +374,10 @@ fn is_trusted_chatgpt_mcp_server( return false; } - if url::Url::parse(CHATGPT_CODEX_BASE_URL) - .ok() - .is_some_and(|chatgpt_url| server_url.origin() == chatgpt_url.origin()) + if [CHATGPT_CODEX_BASE_URL, "https://gov.chatgpt.com"] + .into_iter() + .filter_map(|base_url| url::Url::parse(base_url).ok()) + .any(|trusted_url| server_url.origin() == trusted_url.origin()) { return true; } diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index 248e0956e5..8332b18048 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -437,6 +437,10 @@ fn codex_apps_mcp_url_for_base_url_uses_plugin_service_paths() { codex_apps_mcp_url_for_base_url("https://chat.openai.com"), "https://chat.openai.com/backend-api/ps/mcp" ); + assert_eq!( + codex_apps_mcp_url_for_base_url("https://gov.chatgpt.com/backend-api/"), + "https://gov.chatgpt.com/backend-api/ps/mcp" + ); assert_eq!( codex_apps_mcp_url_for_base_url("http://localhost:8080/api/codex"), "http://localhost:8080/api/codex/ps/mcp" @@ -595,10 +599,52 @@ fn effective_mcp_servers_require_catalog_registration() { } #[test] -fn effective_mcp_servers_preserve_chatgpt_auth_for_staging() { +fn effective_mcp_servers_preserve_chatgpt_auth_for_trusted_origins() { + for (base_url, server_url) in [ + ("https://chatgpt-staging.com", "https://chatgpt-staging.com"), + ( + "https://preview.chatgpt-staging.com", + "https://preview.chatgpt-staging.com", + ), + ("https://gov.chatgpt.com", "https://chatgpt.com"), + ( + "https://chatgpt.com", + "https://gov.chatgpt.com/backend-api/", + ), + ("https://chatgpt.com", "https://GOV.CHATGPT.COM:443"), + ] { + let mut config = test_mcp_config(PathBuf::new()); + config.chatgpt_base_url = base_url.to_string(); + let server = codex_apps_mcp_server_config( + server_url, /*apps_mcp_product_sku*/ None, /*originator*/ None, + ); + let mut catalog = ResolvedMcpCatalog::builder(); + catalog.register(McpServerRegistration::from_config( + "trusted".to_string(), + server, + )); + config.mcp_server_catalog = catalog.build(); + + let effective = effective_mcp_servers(&config, /*auth*/ None); + + assert_eq!( + effective["trusted"].config().auth, + McpServerAuth::ChatGpt, + "backend {base_url}, server {server_url}", + ); + } +} + +#[test] +fn effective_mcp_servers_reject_untrusted_origins_for_chatgpt_auth() { for url in [ - "https://chatgpt-staging.com", - "https://preview.chatgpt-staging.com", + "http://gov.chatgpt.com", + "https://gov.chatgpt.com:8443", + "https://gov.chatgpt.com.evil.example", + "https://sub.gov.chatgpt.com", + "https://gov.chatgpt.com.", + "https://gov.chatgpt.com@evil.example", + "https://example.com", ] { let mut config = test_mcp_config(PathBuf::new()); config.chatgpt_base_url = url.to_string(); @@ -607,15 +653,18 @@ fn effective_mcp_servers_preserve_chatgpt_auth_for_staging() { ); let mut catalog = ResolvedMcpCatalog::builder(); catalog.register(McpServerRegistration::from_config( - "staging".to_string(), - server.clone(), + "untrusted".to_string(), + server, )); config.mcp_server_catalog = catalog.build(); - let configured = HashMap::from([("staging".to_string(), server)]); - let effective = - effective_mcp_servers_from_configured(configured, &config, /*auth*/ None); - assert_eq!(effective["staging"].config().auth, McpServerAuth::ChatGpt); + let effective = effective_mcp_servers(&config, /*auth*/ None); + + assert_eq!( + effective["untrusted"].config().auth, + McpServerAuth::OAuth, + "server {url}", + ); } }