openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-11 20:36:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add staged enterprise OIDC login and coordinated logout (#43844)

Browse Source 
## Why

Enterprise sign-in needs to keep browser completion separate from credential storage so callers can recheck the active account and configuration before saving a grant. Logout must also prevent an earlier sign-in from restoring credentials, including from another process sharing `CODEX_HOME`.

## What changed

- Add enterprise login APIs in `rmcp-client` that return an authorization URL and stage validated credentials for an explicit `commit_if` call. Store grants only in the keyring after rechecking caller authority under the credential lock.
- Require a registered client ID, published metadata matching the configured issuer, HTTP loopback callbacks, a refresh token, and a valid OIDC identity assertion. Request `openid` and `offline_access` with `prompt=consent`, and omit MCP resource indicators from authorization and code exchange.
- Persist a login generation under the credential lock so logout invalidates pending and staged sign-ins across processes, even when no grant is stored.
- Keep credentials and account identifiers out of enterprise error chains and logs, and avoid logging callback payloads when the receiver has closed.

## Testing

Add coverage for discovery validation, loopback callbacks, PKCE, staged keyring storage, cancellation, stale attempts, cross-process logout, and error/log privacy. Preserve ordinary MCP OAuth login without a refresh token, and adjust the terminal polling test deadline to include the minimum empty-poll wait.

GitOrigin-RevId: 2a27b9a26505a2f6fdecce8877f6e2c21e148f72
This commit is contained in:
Nick Steele
2026-09-08 15:33:06 +00:00
committed by copyberry
parent cbfa321ecd
commit b090e901f8
13 changed files with 1435 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
codex-rs/Cargo.lock generated
View File

@@ -4327,6 +4327,7 @@ dependencies = [
"tiny_http",
"tokio",
"tracing",
"tracing-test",
"url",
"urlencoding",
"webbrowser",

3
codex-rs/core/src/unified_exec/mod_tests.rs
View File

@@ -951,8 +951,9 @@ async fn stdin_approval_preserves_the_reviewed_terminal() -> anyhow::Result<()>
assert!(original.interaction_lock().try_lock_owned().is_err());
}
// Empty polling must complete without an approval response.
// The test deadline must allow the minimum empty-poll wait.
tokio::time::timeout(
Duration::from_secs(/*secs*/ 5),
Duration::from_millis(MIN_EMPTY_YIELD_TIME_MS) + Duration::from_secs(/*secs*/ 5),
write_stdin(&session, &turn, process_id, "", /*yield_time_ms*/ 250),
)
.await??;

1
codex-rs/rmcp-client/Cargo.toml
View File

@@ -72,6 +72,7 @@ codex-utils-cargo-bin = { workspace = true }
pretty_assertions = { workspace = true }
serial_test = { workspace = true }
tempfile = { workspace = true }
tracing-test = { workspace = true, features = ["no-env-filter"] }
wiremock = { workspace = true }
[target.'cfg(unix)'.dependencies]

60
codex-rs/rmcp-client/src/ema_identity.rs
View File

@@ -1,4 +1,5 @@
//! An enterprise IdP session is independent of Codex account authentication.
//! Shared enterprise IdP discovery and stored refresh-token resolution.
//! Login and token exchange require published metadata bound to the configured issuer.
use std::sync::Arc;
@@ -11,6 +12,8 @@ use codex_keyring_store::DefaultKeyringStore;
use codex_keyring_store::KeyringStore;
use oauth2::TokenResponse;
use rmcp::transport::AuthorizationManager;
use rmcp::transport::auth::AuthorizationMetadata;
use rmcp::transport::auth::OAuthHttpClient;
use crate::ema_auth_policy::advertised_capability;
use crate::ema_auth_policy::ema_reauthentication_required;
@@ -55,6 +58,36 @@ pub fn stored_ema_identity_is_usable(
&& stored_oidc_identity(tokens).is_ok()
}
pub(crate) async fn resolve_ema_idp_authorization_manager(
issuer: &str,
http_client: Arc<dyn OAuthHttpClient>,
) -> Result<(AuthorizationManager, AuthorizationMetadata)> {
validate_ema_oauth_endpoint(issuer, "enterprise IdP issuer")?;
let mut manager = AuthorizationManager::new_with_oauth_http_client(issuer, http_client)
.await
.context("failed to create enterprise IdP metadata discovery client")?;
manager.set_allow_missing_issuer(false);
let resolution = manager
.resolve_metadata()
.await
.context("failed to discover enterprise IdP authorization metadata")?;
if !resolution.source.is_discovered() {
bail!("enterprise IdP must publish authorization metadata");
}
let metadata = resolution.metadata;
if metadata.issuer.as_deref() != Some(issuer) {
bail!("enterprise IdP authorization metadata issuer does not match configuration");
}
validate_ema_oauth_endpoint(&metadata.token_endpoint, "enterprise IdP token endpoint")?;
validate_ema_public_client_auth(
metadata
.additional_fields
.get("token_endpoint_auth_methods_supported"),
"enterprise IdP",
)?;
Ok((manager, metadata))
}
/// Resolve a stored refresh-token subject against the configured enterprise IdP metadata.
pub async fn resolve_ema_idp_identity(
request: EmaIdpIdentityRequest<'_>,
@@ -69,7 +102,6 @@ async fn resolve_ema_idp_identity_in<K: KeyringStore + Clone + 'static>(
if request.issuer.trim().is_empty() || request.client_id.trim().is_empty() {
bail!("ema_auth requires a non-empty enterprise IdP issuer and client ID");
}
validate_ema_oauth_endpoint(request.issuer, "enterprise IdP issuer")?;
let credentials = request.credentials.credentials();
if credentials.url != request.issuer
|| credentials.bound_issuer() != Some(request.issuer)
@@ -85,22 +117,8 @@ async fn resolve_ema_idp_identity_in<K: KeyringStore + Clone + 'static>(
/*has_configured_headers*/ false,
request.redirect_mode,
)?;
let mut manager = AuthorizationManager::new_with_oauth_http_client(
request.issuer.to_string(),
Arc::new(client),
)
.await
.context("failed to create enterprise IdP metadata discovery client")?;
manager.set_allow_missing_issuer(false);
let metadata = manager
.resolve_metadata()
.await
.context("failed to discover enterprise IdP authorization metadata")?
.metadata;
if metadata.issuer.as_deref() != Some(request.issuer) {
bail!("enterprise IdP authorization metadata issuer does not match configuration");
}
validate_ema_oauth_endpoint(&metadata.token_endpoint, "enterprise IdP token endpoint")?;
let (_, metadata) =
resolve_ema_idp_authorization_manager(request.issuer, Arc::new(client)).await?;
for (name, expected) in [
(
"identity_chaining_requested_token_types_supported",
@@ -116,12 +134,6 @@ async fn resolve_ema_idp_identity_in<K: KeyringStore + Clone + 'static>(
);
}
}
validate_ema_public_client_auth(
metadata
.additional_fields
.get("token_endpoint_auth_methods_supported"),
"enterprise IdP",
)?;
let credential_lock =
RefreshCredentialLock::acquire_for_server(&credentials.server_name, &credentials.url)
.await?;

84
codex-rs/rmcp-client/src/ema_identity_tests.rs
View File

@@ -91,6 +91,83 @@ async fn discovery() -> (MockServer, String) {
(server, issuer)
}
#[tokio::test]
async fn enterprise_discovery_shares_idp_checks_but_keeps_login_endpoint_policy() -> Result<()> {
for case in [
"valid",
"missing-discovery",
"issuer",
"missing-issuer",
"token-endpoint",
"public-client",
"authorization-endpoint",
] {
let server = MockServer::start().await;
let issuer = format!("{}/idp", server.uri());
// Login need not advertise ID-JAG exchange capabilities.
let mut metadata = json!({
"issuer": issuer, "authorization_endpoint": format!("{issuer}/authorize"),
"token_endpoint": format!("{issuer}/token"),
"token_endpoint_auth_methods_supported": ["none"],
"grant_types_supported": ["authorization_code"],
});
match case {
"issuer" => metadata["issuer"] = json!("https://other.example"),
"missing-issuer" => {
metadata.as_object_mut().expect("metadata").remove("issuer");
}
"token-endpoint" => metadata["token_endpoint"] = json!("http://unsafe.example/token"),
"public-client" => {
metadata["token_endpoint_auth_methods_supported"] = json!(["client_secret_basic"])
}
"authorization-endpoint" => {
metadata["authorization_endpoint"] = json!("http://unsafe.example/authorize")
}
"valid" | "missing-discovery" => {}
_ => unreachable!(),
}
if case != "missing-discovery" {
Mock::given(method("GET"))
.and(path("/.well-known/oauth-authorization-server/idp"))
.respond_with(ResponseTemplate::new(200).set_body_json(metadata))
.mount(&server)
.await;
}
let client = Arc::new(OAuthHttpClientAdapter::new_with_redirect_mode(
Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new(
OutboundProxyPolicy::ReqwestDefault,
))),
build_default_headers(/*http_headers*/ None, /*env_http_headers*/ None)?,
&issuer,
/*has_configured_headers*/ false,
StreamableHttpRedirectMode::Legacy,
)?);
let shared = resolve_ema_idp_authorization_manager(&issuer, client.clone()).await;
let login = crate::enterprise_oauth_login::resolve_enterprise_authorization_manager(
&issuer, client,
)
.await;
assert_eq!(
(shared.is_ok(), login.is_ok()),
(
matches!(case, "valid" | "authorization-endpoint"),
case == "valid"
),
"{case}",
);
if case == "missing-discovery" {
assert_eq!(
shared
.err()
.expect("reject synthesized metadata")
.to_string(),
"enterprise IdP must publish authorization metadata",
);
}
}
Ok(())
}
const REREAD_TEST_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5);
const REREAD_PANIC_SENTINEL: &str = "credential-panic-payload-sentinel";
@@ -175,10 +252,9 @@ async fn refresh_subject_reread_is_cancellable_and_releases_guard_on_failure() -
entered: Arc::new(Mutex::new(Some(entered_tx))),
release: Arc::new(Mutex::new(release_rx)),
};
let mut identity = Box::pin(resolve_ema_idp_identity_in(
request(&issuer, &snapshot),
&keyring,
));
let request = request(&issuer, &snapshot);
crate::oauth::test_support::warm_http_client(request.http_client.as_ref()).await?;
let mut identity = Box::pin(resolve_ema_idp_identity_in(request, &keyring));
tokio::select! {
result = &mut identity => {
result?;

384
codex-rs/rmcp-client/src/enterprise_oauth_login.rs Normal file
View File

@@ -0,0 +1,384 @@
//! Enterprise OIDC policy and staged, keyring-only credential commits.
//! Browser completion never persists a grant; its owner revalidates authority under the commit lock.
use std::future::Future;
use std::net::IpAddr;
use std::net::Ipv4Addr;
use std::sync::Arc;
use anyhow::Result;
use anyhow::anyhow;
use anyhow::bail;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::HttpClient;
use http::Method;
use http::header::CONTENT_LENGTH;
use http::header::CONTENT_TYPE;
use oauth2::TokenResponse;
use rmcp::transport::AuthorizationManager;
use rmcp::transport::auth::AuthorizationMetadata;
use rmcp::transport::auth::OAuthHttpClient;
use rmcp::transport::auth::OAuthHttpClientFuture;
use rmcp::transport::auth::OAuthHttpRequest;
use tracing::instrument::WithSubscriber;
use url::Host;
use url::Url;
use crate::StoredOAuthTokens;
use crate::ema_auth_policy::validate_ema_oauth_endpoint;
use crate::ema_claims::validate_oidc_identity_assertion;
use crate::ema_identity::resolve_ema_idp_authorization_manager;
use crate::http_client_adapter::StreamableHttpRedirectMode;
use crate::oauth::EnterpriseOAuthGeneration;
use crate::oauth::EnterpriseOAuthGenerationFile;
use crate::oauth::RefreshCredentialLock;
use crate::oauth::delete_oauth_tokens_with_lock_held;
use crate::oauth::save_oauth_tokens_with_lock_held;
use crate::oauth::validate_authorization_server_endpoints;
use crate::oauth_client_registration::McpOAuthClientRegistration;
use crate::perform_oauth_login::OAuthHttpContext;
use crate::perform_oauth_login::OAuthLoginPurpose;
use crate::perform_oauth_login::OauthLoginFlow;
/// An exclusive credential mutation guard. Hold it through primary account logout
/// so a competing process cannot commit between deleting the grant and signing out.
/// Coordination, like ordinary OAuth refresh, is scoped to the same CODEX_HOME.
pub struct EnterpriseOAuthCredentialGuard {
credential_name: String,
issuer: String,
keyring_backend: AuthKeyringBackendKind,
generation_file: EnterpriseOAuthGenerationFile,
_lock: RefreshCredentialLock,
}
impl EnterpriseOAuthCredentialGuard {
pub async fn acquire(
credential_name: &str,
issuer: &str,
keyring_backend: AuthKeyringBackendKind,
) -> Result<Self> {
let lock = RefreshCredentialLock::acquire_for_server(credential_name, issuer)
.with_subscriber(tracing::subscriber::NoSubscriber::default())
.await
.map_err(|_| anyhow!("failed to lock enterprise credentials"))?;
let generation_file = EnterpriseOAuthGenerationFile::open(credential_name, issuer, &lock)
.map_err(|_| anyhow!("failed to open enterprise login generation"))?;
Ok(Self {
credential_name: credential_name.to_owned(),
issuer: issuer.to_owned(),
keyring_backend,
generation_file,
_lock: lock,
})
}
/// Invalidate pending logins and delete the grant, if present. A false return
/// means no grant was stored; earlier login attempts are still invalidated.
pub fn delete_tokens(&self) -> Result<bool> {
// Invalidate attempts that have not stored a grant yet, including other processes.
// Persist before deletion so no successful logout can admit an earlier login.
self.generation_file
.replace()
.map_err(|_| anyhow!("failed to invalidate pending enterprise sign-ins"))?;
// Underlying keyring diagnostics include the account-scoped key. Suppress
// them and discard their error chain, not merely the outer error message.
tracing::subscriber::with_default(tracing::subscriber::NoSubscriber::default(), || {
delete_oauth_tokens_with_lock_held(
&self._lock,
&self.credential_name,
&self.issuer,
OAuthCredentialsStoreMode::Keyring,
self.keyring_backend,
)
})
.map_err(|_| anyhow!("failed to delete enterprise credentials"))
}
}
/// Invalidate earlier enterprise logins across processes, then delete any stored grant.
pub async fn delete_enterprise_oauth_tokens(
credential_name: &str,
issuer: &str,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<bool> {
EnterpriseOAuthCredentialGuard::acquire(credential_name, issuer, keyring_backend_kind)
.await?
.delete_tokens()
}
/// A browser login with no detached persistence worker. Dropping this handle or
/// its wait future unblocks the callback listener and cannot write credentials.
pub struct EnterpriseOAuthLoginHandle {
flow: OauthLoginFlow,
keyring_backend: AuthKeyringBackendKind,
generation: EnterpriseOAuthGeneration,
}
impl EnterpriseOAuthLoginHandle {
pub fn authorization_url(&self) -> String {
self.flow.authorization_url()
}
pub async fn wait(self) -> Result<EnterpriseOAuthCredentials> {
let stored = self
.flow
.complete(/*emit_browser_url*/ false)
.with_subscriber(tracing::subscriber::NoSubscriber::default())
.await
.map_err(|_| anyhow!("enterprise IdP authorization failed"))?;
validate_enterprise_credentials(&stored)?;
Ok(EnterpriseOAuthCredentials {
stored,
keyring_backend: self.keyring_backend,
generation: self.generation,
})
}
}
/// A validated grant that is not yet stored. It intentionally does not expose tokens.
pub struct EnterpriseOAuthCredentials {
stored: StoredOAuthTokens,
keyring_backend: AuthKeyringBackendKind,
generation: EnterpriseOAuthGeneration,
}
impl EnterpriseOAuthCredentials {
/// Revalidate the current attempt, account and configuration under the same
/// exclusive lock used by logout, then persist without an intervening await.
/// Return the authority proof so the caller can retire its attempt while still
/// holding the commit gate, before notifying clients or refreshing runtimes.
pub async fn commit_if<F, Fut, T>(self, is_current: F) -> Result<T>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Option<T>>,
{
let guard = EnterpriseOAuthCredentialGuard::acquire(
&self.stored.server_name,
&self.stored.url,
self.keyring_backend,
)
.await?;
let generation = guard
.generation_file
.current()
.map_err(|_| anyhow!("failed to read enterprise login generation"))?;
if generation.as_ref() != Some(&self.generation) {
bail!("enterprise sign-in was invalidated by logout");
}
let Some(authority) = is_current().await else {
bail!("enterprise sign-in no longer matches the active account or configuration");
};
tracing::subscriber::with_default(tracing::subscriber::NoSubscriber::default(), || {
save_oauth_tokens_with_lock_held(
&guard._lock,
&self.stored.server_name,
&self.stored,
OAuthCredentialsStoreMode::Keyring,
self.keyring_backend,
)
})
.map_err(|_| anyhow!("failed to store enterprise credentials"))?;
Ok(authority)
}
}
/// Start a host-registered enterprise login without launching the browser or
/// persisting a credential. The caller owns both the attempt and its later commit.
pub async fn perform_enterprise_oauth_login_return_url(
request: EnterpriseOAuthLoginRequest<'_>,
) -> Result<EnterpriseOAuthLoginHandle> {
// Capture before discovery or browser setup, then release the lock while the user signs in.
let generation = {
let guard = EnterpriseOAuthCredentialGuard::acquire(
request.credential_name,
request.issuer,
request.keyring_backend_kind,
)
.await?;
match guard.generation_file.current() {
Ok(Some(generation)) => generation,
Ok(None) => guard
.generation_file
.replace()
.map_err(|_| anyhow!("failed to initialize enterprise login generation"))?,
Err(_) => bail!("failed to read enterprise login generation"),
}
};
let flow = OauthLoginFlow::new(
request.credential_name,
request.issuer,
OAuthCredentialsStoreMode::Keyring,
request.keyring_backend_kind,
OAuthHttpContext {
http_headers: None,
env_http_headers: None,
http_client: request.http_client,
redirect_mode: request.redirect_mode,
},
&["openid".to_string(), "offline_access".to_string()],
Some(request.client_id),
OAuthLoginPurpose::EnterpriseIdp,
McpOAuthClientRegistration::Auto,
/*oauth_resource*/ None,
/*launch_browser*/ false,
request.callback_port,
request.callback_url,
/*global_callback_url*/ None,
request.timeout_secs,
)
.with_subscriber(tracing::subscriber::NoSubscriber::default())
.await
.map_err(|_| anyhow!("failed to start enterprise IdP authorization"))?;
Ok(EnterpriseOAuthLoginHandle {
flow,
keyring_backend: request.keyring_backend_kind,
generation,
})
}
pub(crate) fn enterprise_callback_settings(
issuer: &str,
client_id: Option<&str>,
callback_url: Option<&str>,
callback_port: Option<u16>,
) -> Result<(IpAddr, Option<u16>)> {
validate_ema_oauth_endpoint(issuer, "enterprise IdP issuer")?;
if client_id.is_none_or(|client_id| client_id.trim().is_empty()) {
bail!("enterprise IdP login requires its registered client ID");
}
let ip = enterprise_callback_bind_ip(callback_url)?;
let registered_port = callback_url
.map(Url::parse)
.transpose()?
.and_then(|url| url.port());
if callback_port
.zip(registered_port)
.is_some_and(|(configured, registered)| configured != registered)
{
bail!("enterprise IdP callback URL and listener specify different ports");
}
Ok((ip, callback_port.or(registered_port)))
}
fn enterprise_callback_bind_ip(callback_url: Option<&str>) -> Result<IpAddr> {
let Some(callback_url) = callback_url else {
return Ok(Ipv4Addr::LOCALHOST.into());
};
validate_ema_oauth_endpoint(callback_url, "enterprise IdP callback URL")?;
let callback = Url::parse(callback_url)?;
if callback.scheme() == "http" {
match callback.host() {
Some(Host::Domain("localhost")) => return Ok(Ipv4Addr::LOCALHOST.into()),
Some(Host::Ipv4(ip)) if ip.is_loopback() => return Ok(ip.into()),
Some(Host::Ipv6(ip)) if ip.is_loopback() => return Ok(ip.into()),
_ => {}
}
}
bail!("enterprise IdP callback URL must use an HTTP loopback address")
}
fn validate_enterprise_credentials(stored: &StoredOAuthTokens) -> Result<()> {
let credentials = &stored.token_response.0;
if credentials
.refresh_token()
.is_none_or(|refresh_token| refresh_token.secret().trim().is_empty())
{
bail!("enterprise IdP login did not return a refresh token");
}
let assertion = credentials
.extra_fields()
.0
.get("id_token")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| anyhow!("enterprise IdP login did not return an OIDC identity assertion"))?;
validate_oidc_identity_assertion(
assertion,
stored.issuer.as_deref().unwrap_or(&stored.url),
&stored.client_id,
)
.map_err(|_| anyhow!("enterprise IdP returned an invalid OIDC identity assertion"))
}
pub(crate) async fn resolve_enterprise_authorization_manager(
issuer: &str,
http_client: Arc<dyn OAuthHttpClient>,
) -> Result<(AuthorizationManager, AuthorizationMetadata)> {
let (manager, metadata) = resolve_ema_idp_authorization_manager(
issuer,
Arc::new(EnterpriseOAuthHttpClient(http_client)),
)
.await?;
validate_authorization_server_endpoints(&metadata)?;
validate_ema_oauth_endpoint(
&metadata.authorization_endpoint,
"enterprise IdP authorization endpoint",
)?;
Ok((manager, metadata))
}
pub(crate) fn enterprise_authorization_url(auth_url: &str) -> Result<String> {
let mut url = Url::parse(auth_url)?;
let query = url::form_urlencoded::Serializer::new(String::new())
.extend_pairs(
url.query_pairs()
.filter(|(key, _)| key != "resource" && key != "prompt"),
)
.append_pair("prompt", "consent")
.finish();
url.set_query(Some(&query));
Ok(url.to_string())
}
pub(crate) fn without_oauth_resource(encoded: &[u8]) -> String {
url::form_urlencoded::Serializer::new(String::new())
.extend_pairs(url::form_urlencoded::parse(encoded).filter(|(key, _)| key != "resource"))
.finish()
}
/// rmcp supplies a resource indicator for MCP OAuth, but the independent OIDC
/// login must not request the IdP issuer as a protected-resource audience.
pub(crate) struct EnterpriseOAuthHttpClient(pub(crate) Arc<dyn OAuthHttpClient>);
impl OAuthHttpClient for EnterpriseOAuthHttpClient {
fn execute(&self, mut request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> {
if request.request.method() == Method::POST
&& request
.request
.headers()
.get(CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| {
value.split(';').next().is_some_and(|mime| {
mime.trim()
.eq_ignore_ascii_case("application/x-www-form-urlencoded")
})
})
&& url::form_urlencoded::parse(request.request.body())
.any(|(key, value)| key == "grant_type" && value == "authorization_code")
{
let body = without_oauth_resource(request.request.body()).into_bytes();
*request.request.body_mut() = body;
request.request.headers_mut().remove(CONTENT_LENGTH);
}
self.0.execute(request)
}
}
/// Enterprise registration and host-owned login settings. The flow always uses
/// OpenID Connect, strict issuer validation, and keyring-only credential storage.
pub struct EnterpriseOAuthLoginRequest<'a> {
pub credential_name: &'a str,
pub issuer: &'a str,
pub client_id: &'a str,
pub keyring_backend_kind: AuthKeyringBackendKind,
pub callback_port: Option<u16>,
pub callback_url: Option<&'a str>,
pub timeout_secs: Option<i64>,
pub http_client: Arc<dyn HttpClient>,
pub redirect_mode: StreamableHttpRedirectMode,
}
#[cfg(test)]
#[path = "enterprise_oauth_login_tests.rs"]
mod tests;

466
codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs Normal file
View File

@@ -0,0 +1,466 @@
//! Public enterprise login exercises the real callback and keyring adapter in an isolated process.
use std::any::Any;
use std::collections::HashMap;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use codex_exec_server::RouteAwareHttpClient;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use keyring::credential::Credential;
use keyring::credential::CredentialApi;
use keyring::credential::CredentialBuilderApi;
use keyring::credential::CredentialPersistence;
use keyring::mock::MockCredential;
use pretty_assertions::assert_eq;
use serde_json::json;
use sha2::Digest;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tracing_test::traced_test;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
use super::*;
const SECRET: &str = "enterprise-secret-sentinel";
const CREDENTIAL_NAME: &str = "ema-idp:enterprise-secret-sentinel";
#[path = "enterprise_oauth_logout_tests.rs"]
mod logout;
async fn isolated_process(test_name: &str) -> Result<bool> {
const CHILD: &str = "CODEX_ENTERPRISE_LOGIN_TEST_CHILD";
if std::env::var_os(CHILD).is_some() {
return Ok(false);
}
let home = tempfile::tempdir()?;
let output = tokio::process::Command::new(std::env::current_exe()?)
.args(["--exact", test_name, "--nocapture"])
.env(CHILD, "1")
.env("CODEX_HOME", home.path())
.current_dir(home.path())
.output()
.await?;
assert!(
output.status.success(),
"{}\n{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
Ok(true)
}
async fn login(issuer: &str, callback_url: Option<&str>) -> Result<EnterpriseOAuthLoginHandle> {
perform_enterprise_oauth_login_return_url(EnterpriseOAuthLoginRequest {
credential_name: CREDENTIAL_NAME,
issuer,
client_id: "enterprise-client",
keyring_backend_kind: AuthKeyringBackendKind::Direct,
callback_port: None,
callback_url,
timeout_secs: Some(5),
http_client: Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new(
OutboundProxyPolicy::ReqwestDefault,
))),
redirect_mode: StreamableHttpRedirectMode::Legacy,
})
.await
}
async fn metadata(server: &MockServer, issuer: &str) {
Mock::given(method("GET")).and(path("/.well-known/oauth-authorization-server/idp"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"issuer": issuer, "authorization_endpoint": format!("{issuer}/authorize?prompt=none"),
"token_endpoint": format!("{issuer}/token"), "token_endpoint_auth_methods_supported": ["none"],
}))).mount(server).await;
}
async fn callback(authorization_url: &str, issuer: &str, provider_error: bool) -> Result<()> {
let query = Url::parse(authorization_url)?
.query_pairs()
.into_owned()
.collect::<HashMap<_, _>>();
assert_eq!(query.get("prompt").map(String::as_str), Some("consent"));
assert!(!query.contains_key("resource"));
let mut callback = Url::parse(&query["redirect_uri"])?;
let mut pairs = callback.query_pairs_mut();
if provider_error {
pairs
.append_pair("error", SECRET)
.append_pair("error_description", SECRET);
} else {
pairs
.append_pair("code", SECRET)
.append_pair("state", &query["state"])
.append_pair("iss", issuer);
}
drop(pairs);
let port = callback.port().expect("actual listener port");
let mut stream = tokio::net::TcpStream::connect(("127.0.0.1", port)).await?;
stream
.write_all(
format!(
"GET {}?{} HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\nConnection: close\r\n\r\n",
callback.path(),
callback.query().expect("query")
)
.as_bytes(),
)
.await?;
stream.read_to_end(&mut Vec::new()).await?;
Ok(())
}
#[tokio::test]
#[traced_test]
async fn enterprise_callback_errors_and_sdk_logs_exclude_credentials() -> Result<()> {
if isolated_process(
"enterprise_oauth_login::tests::enterprise_callback_errors_and_sdk_logs_exclude_credentials",
)
.await?
{
return Ok(());
}
for (callback_error, token_success) in [(true, false), (false, false), (false, true)] {
let server = MockServer::start().await;
let issuer = format!("{}/idp", server.uri());
metadata(&server, &issuer).await;
let response = if token_success {
ResponseTemplate::new(200).set_body_json(json!({
"access_token": SECRET, "refresh_token": SECRET, "id_token": SECRET, "token_type": "Bearer",
}))
} else {
ResponseTemplate::new(400)
.set_body_json(json!({"error": "invalid_grant", "error_description": SECRET}))
};
Mock::given(method("POST"))
.and(path("/idp/token"))
.respond_with(response)
.expect(u64::from(!callback_error))
.mount(&server)
.await;
let login = login(&issuer, /*callback_url*/ None).await?;
callback(&login.authorization_url(), &issuer, callback_error).await?;
tracing::trace!(target: "rmcp::transport::auth", "SDK trace capture enabled");
let error = login.wait().await.err().expect("reject provider response");
assert!(!format!("{error} {error:?} {error:#}").contains(SECRET));
assert!(logs_contain("SDK trace capture enabled"));
assert!(!logs_contain(SECRET));
}
Ok(())
}
#[test]
fn enterprise_callback_requires_loopback() -> Result<()> {
for (callback, expected) in [
(None, "127.0.0.1"),
(Some("http://localhost/callback"), "127.0.0.1"),
(Some("http://127.0.0.2/callback"), "127.0.0.2"),
(Some("http://[::1]/callback"), "::1"),
] {
assert_eq!(
enterprise_callback_bind_ip(callback)?,
expected.parse::<IpAddr>()?
);
}
for callback in [
"http://0.0.0.0/callback",
"http://[::]/callback",
"https://127.0.0.1/callback",
"http://remote.example/callback",
] {
assert!(enterprise_callback_bind_ip(Some(callback)).is_err());
}
Ok(())
}
#[tokio::test]
#[traced_test]
async fn enterprise_public_api_storage_and_privacy() -> Result<()> {
if isolated_process("enterprise_oauth_login::tests::enterprise_public_api_storage_and_privacy")
.await?
{
return Ok(());
}
let keyring = TestKeyring::default();
keyring::set_default_credential_builder(Box::new(keyring.clone()));
let server = MockServer::start().await;
let issuer = format!("{}/idp", server.uri());
metadata(&server, &issuer).await;
assert!(
login(&issuer, Some(&format!("{}/callback", server.uri())))
.await
.is_err(),
"an occupied registered callback port must not silently move"
);
let assertion = format!(
"{}.{}.signature",
URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#),
URL_SAFE_NO_PAD.encode(
json!({"iss":issuer,"sub":"user","aud":"enterprise-client","exp":4102444800_u64})
.to_string()
)
);
Mock::given(method("POST"))
.and(path("/idp/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"access_token":SECRET,"refresh_token":SECRET,"id_token":assertion,"token_type":"Bearer",
})))
.mount(&server)
.await;
// Callback, exchange and DefaultKeyringStore all run through public entrypoints.
for (login_index, callback_url) in ["http://localhost/callback", "http://127.0.0.1/callback"]
.into_iter()
.enumerate()
{
let login = login(&issuer, Some(callback_url)).await?;
let authorization_url = login.authorization_url();
let authorization_query = Url::parse(&authorization_url)?
.query_pairs()
.into_owned()
.collect::<HashMap<_, _>>();
callback(&authorization_url, &issuer, /*provider_error*/ false).await?;
let credentials = login.wait().await?;
let token_requests = server
.received_requests()
.await
.expect("recorded OAuth requests")
.into_iter()
.filter(|request| {
request.method == http::Method::POST && request.url.path() == "/idp/token"
})
.collect::<Vec<_>>();
assert_eq!(token_requests.len(), login_index + 1);
let token_form = url::form_urlencoded::parse(&token_requests[login_index].body)
.into_owned()
.collect::<HashMap<_, _>>();
assert_eq!(
token_form.get("grant_type").map(String::as_str),
Some("authorization_code")
);
assert!(!token_form.contains_key("resource"));
assert_eq!(
token_form.get("redirect_uri"),
Some(
authorization_query
.get("redirect_uri")
.expect("authorization redirect URI")
)
);
assert_eq!(
authorization_query
.get("code_challenge_method")
.map(String::as_str),
Some("S256")
);
let verifier = token_form.get("code_verifier").expect("PKCE verifier");
let challenge = URL_SAFE_NO_PAD.encode(sha2::Sha256::digest(verifier.as_bytes()));
assert_eq!(authorization_query.get("code_challenge"), Some(&challenge));
assert!(
keyring
.values
.lock()
.unwrap()
.values()
.all(|value| value.get_secret().is_err())
);
let attempt = tokio::sync::Mutex::new(Some("active"));
{
let mut authority = credentials
.commit_if(|| async { Some(attempt.lock().await) })
.await?;
assert!(
attempt.try_lock().is_err(),
"commit retains the attempt gate until the caller retires it"
);
*authority = None;
}
assert_eq!(*attempt.lock().await, None);
let stored = tracing::subscriber::with_default(
tracing::subscriber::NoSubscriber::default(),
|| {
crate::stored_oauth_credentials(
CREDENTIAL_NAME,
&issuer,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)
},
)?
.expect("stored grant");
assert_eq!(
stored
.token_response
.0
.refresh_token()
.map(|token| token.secret().as_str()),
Some(SECRET)
);
assert!(
delete_enterprise_oauth_tokens(
CREDENTIAL_NAME,
&issuer,
AuthKeyringBackendKind::Direct
)
.await?
);
}
// Cancellation while blocked on the actual credential lock cannot leave a
// detached persistence worker that writes after the other process releases it.
let canceled = login(&issuer, /*callback_url*/ None).await?;
callback(
&canceled.authorization_url(),
&issuer,
/*provider_error*/ false,
)
.await?;
let canceled = canceled.wait().await?;
let guard = EnterpriseOAuthCredentialGuard::acquire(
CREDENTIAL_NAME,
&issuer,
AuthKeyringBackendKind::Direct,
)
.await?;
let mut commit = Box::pin(canceled.commit_if(|| async { Some(()) }));
assert!(futures::poll!(&mut commit).is_pending());
drop(commit);
drop(guard);
assert!(
keyring
.values
.lock()
.unwrap()
.values()
.all(|value| value.get_secret().is_err())
);
// Rejected old attempts neither write nor delete a newer grant.
let old = login(&issuer, /*callback_url*/ None).await?;
callback(
&old.authorization_url(),
&issuer,
/*provider_error*/ false,
)
.await?;
let old = old.wait().await?;
let winner = login(&issuer, /*callback_url*/ None).await?;
callback(
&winner.authorization_url(),
&issuer,
/*provider_error*/ false,
)
.await?;
winner
.wait()
.await?
.commit_if(|| async { Some(()) })
.await?;
assert!(old.commit_if(|| async { None::<()> }).await.is_err());
assert!(
keyring
.values
.lock()
.unwrap()
.values()
.any(|value| value.get_secret().is_ok())
);
// Inject raw account identifiers into the actual keyring adapter's error chain.
keyring.fail.store(true, Ordering::SeqCst);
let failed = login(&issuer, /*callback_url*/ None).await?;
callback(
&failed.authorization_url(),
&issuer,
/*provider_error*/ false,
)
.await?;
let save_error = failed
.wait()
.await?
.commit_if(|| async { Some(()) })
.await
.unwrap_err();
let delete_error =
delete_enterprise_oauth_tokens(CREDENTIAL_NAME, &issuer, AuthKeyringBackendKind::Direct)
.await
.unwrap_err();
for error in [save_error, delete_error] {
assert!(!format!("{error} {error:?} {error:#}").contains(SECRET));
}
keyring.fail.store(false, Ordering::SeqCst);
let home = std::path::PathBuf::from(std::env::var("CODEX_HOME")?);
assert!(
!home.join(".credentials.json").exists(),
"enterprise storage never falls back to plaintext"
);
let locks = home.join("mcp-oauth-locks");
std::fs::rename(&locks, home.join("held-locks"))?;
std::fs::write(&locks, b"not a directory")?;
let lock_error =
delete_enterprise_oauth_tokens(CREDENTIAL_NAME, &issuer, AuthKeyringBackendKind::Direct)
.await
.unwrap_err();
assert!(!format!("{lock_error} {lock_error:?} {lock_error:#}").contains(SECRET));
assert!(!logs_contain(SECRET));
Ok(())
}
#[derive(Clone, Default)]
struct TestKeyring {
values: Arc<Mutex<HashMap<String, Arc<MockCredential>>>>,
fail: Arc<AtomicBool>,
}
struct TestCredential(Arc<MockCredential>);
impl CredentialApi for TestCredential {
fn set_secret(&self, secret: &[u8]) -> keyring::Result<()> {
self.0.set_secret(secret)
}
fn get_secret(&self) -> keyring::Result<Vec<u8>> {
self.0.get_secret()
}
fn delete_credential(&self) -> keyring::Result<()> {
self.0.delete_credential()
}
fn as_any(&self) -> &dyn Any {
self
}
}
impl CredentialBuilderApi for TestKeyring {
fn build(
&self,
_target: Option<&str>,
_service: &str,
user: &str,
) -> keyring::Result<Box<Credential>> {
if self.fail.load(Ordering::SeqCst) {
return Err(keyring::Error::Invalid("account".into(), user.into()));
}
Ok(Box::new(TestCredential(
self.values
.lock()
.unwrap()
.entry(user.into())
.or_default()
.clone(),
)))
}
fn as_any(&self) -> &dyn Any {
self
}
fn persistence(&self) -> CredentialPersistence {
CredentialPersistence::ProcessOnly
}
}

195
codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs Normal file
View File

@@ -0,0 +1,195 @@
//! Cross-process logout exercises staged login and the production keyring adapter.
use std::fs;
use std::path::PathBuf;
use pretty_assertions::assert_eq;
use sha2::Digest;
use sha2::Sha256;
use super::*;
const TEST: &str =
"enterprise_oauth_login::tests::logout::logout_invalidates_pending_login_across_processes";
const LOGOUT_ISSUER: &str = "CODEX_ENTERPRISE_LOGOUT_TEST_ISSUER";
#[tokio::test]
async fn logout_invalidates_pending_login_across_processes() -> Result<()> {
if isolated_process(TEST).await? {
return Ok(());
}
let home = PathBuf::from(std::env::var("CODEX_HOME")?);
let keyring = FileKeyring(home.join("test-keyring"));
fs::create_dir_all(&keyring.0)?;
keyring::set_default_credential_builder(Box::new(keyring));
if let Ok(issuer) = std::env::var(LOGOUT_ISSUER) {
delete_enterprise_oauth_tokens(CREDENTIAL_NAME, &issuer, AuthKeyringBackendKind::Direct)
.await?;
return Ok(());
}
let server = MockServer::start().await;
let issuer = format!("{}/idp", server.uri());
metadata(&server, &issuer).await;
let assertion = format!(
"{}.{}.signature",
URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#),
URL_SAFE_NO_PAD.encode(
json!({"iss":issuer,"sub":"user","aud":"enterprise-client","exp":4102444800_u64})
.to_string()
)
);
Mock::given(method("POST"))
.and(path("/idp/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"access_token":SECRET,"refresh_token":SECRET,"id_token":assertion,"token_type":"Bearer",
})))
.mount(&server)
.await;
// The first logout has no grant to delete. The second deletes the fresh grant
// from the previous iteration. Both must invalidate pending callbacks and staged grants.
for _ in 0..2 {
let pending = login(&issuer, /*callback_url*/ None).await?;
let staged = complete_login(&issuer).await?;
let output = tokio::process::Command::new(std::env::current_exe()?)
.args(["--exact", TEST, "--nocapture"])
.env(LOGOUT_ISSUER, &issuer)
.output()
.await?;
assert!(
output.status.success(),
"{}\n{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
assert!(stored(&issuer)?.is_none());
// The logout process has exited; its invalidation must survive that exit.
assert!(staged.commit_if(|| async { Some(()) }).await.is_err());
assert!(stored(&issuer)?.is_none());
// A fresh login after logout remains usable, including for the same account.
complete_login(&issuer)
.await?
.commit_if(|| async { Some(()) })
.await?;
let fresh = stored(&issuer)?.expect("fresh grant");
callback(
&pending.authorization_url(),
&issuer,
/*provider_error*/ false,
)
.await?;
assert!(
pending
.wait()
.await?
.commit_if(|| async { Some(()) })
.await
.is_err()
);
assert_eq!(stored(&issuer)?, Some(fresh));
}
let stale = complete_login(&issuer).await?;
let generation_path = fs::read_dir(home.join("mcp-oauth-locks"))?
.filter_map(std::result::Result::ok)
.map(|entry| entry.path())
.find(|path| {
path.extension()
.is_some_and(|ext| ext == "enterprise-generation")
})
.expect("persistent generation");
fs::write(&generation_path, b"incomplete write")?;
assert!(stale.commit_if(|| async { Some(()) }).await.is_err());
assert!(login(&issuer, /*callback_url*/ None).await.is_err());
// A lost marker never revives a pre-logout attempt or a pre-reset generation.
fs::remove_file(&generation_path)?;
let stale = complete_login(&issuer).await?;
fs::remove_file(&generation_path)?;
complete_login(&issuer)
.await?
.commit_if(|| async { Some(()) })
.await?;
assert!(stale.commit_if(|| async { Some(()) }).await.is_err());
// Metadata failures are logout errors and must leave the stored grant intact.
let before = stored(&issuer)?;
fs::remove_file(&generation_path)?;
fs::create_dir(&generation_path)?;
assert!(
delete_enterprise_oauth_tokens(CREDENTIAL_NAME, &issuer, AuthKeyringBackendKind::Direct)
.await
.is_err()
);
assert_eq!(stored(&issuer)?, before);
assert!(!home.join(".credentials.json").exists());
Ok(())
}
async fn complete_login(issuer: &str) -> Result<EnterpriseOAuthCredentials> {
let handle = login(issuer, /*callback_url*/ None).await?;
callback(
&handle.authorization_url(),
issuer,
/*provider_error*/ false,
)
.await?;
handle.wait().await
}
fn stored(issuer: &str) -> Result<Option<StoredOAuthTokens>> {
crate::stored_oauth_credentials(
CREDENTIAL_NAME,
issuer,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)
}
// Only synthetic fixture credentials are persisted here. Every process still uses
// DefaultKeyringStore and the production credential lock, serializer and logout API.
struct FileKeyring(PathBuf);
struct FileCredential(PathBuf);
impl CredentialBuilderApi for FileKeyring {
fn build(
&self,
_target: Option<&str>,
_service: &str,
user: &str,
) -> keyring::Result<Box<Credential>> {
let name = format!("{:x}", Sha256::digest(user.as_bytes()));
Ok(Box::new(FileCredential(self.0.join(name))))
}
fn as_any(&self) -> &dyn Any {
self
}
}
impl CredentialApi for FileCredential {
fn set_secret(&self, secret: &[u8]) -> keyring::Result<()> {
fs::write(&self.0, secret).map_err(keyring_error)
}
fn get_secret(&self) -> keyring::Result<Vec<u8>> {
fs::read(&self.0).map_err(keyring_error)
}
fn delete_credential(&self) -> keyring::Result<()> {
fs::remove_file(&self.0).map_err(keyring_error)
}
fn as_any(&self) -> &dyn Any {
self
}
}
fn keyring_error(error: std::io::Error) -> keyring::Error {
if error.kind() == std::io::ErrorKind::NotFound {
keyring::Error::NoEntry
} else {
keyring::Error::PlatformFailure(Box::new(error))
}
}

7
codex-rs/rmcp-client/src/lib.rs
View File

@@ -5,6 +5,7 @@ mod ema_auth_policy;
mod ema_claims;
mod ema_exchange;
mod ema_identity;
mod enterprise_oauth_login;
mod event_notification_transport;
mod executor_process_transport;
mod http_client_adapter;
@@ -49,6 +50,12 @@ pub use ema_identity::EmaIdpIdentity;
pub use ema_identity::EmaIdpIdentityRequest;
pub use ema_identity::resolve_ema_idp_identity;
pub use ema_identity::stored_ema_identity_is_usable;
pub use enterprise_oauth_login::EnterpriseOAuthCredentialGuard;
pub use enterprise_oauth_login::EnterpriseOAuthCredentials;
pub use enterprise_oauth_login::EnterpriseOAuthLoginHandle;
pub use enterprise_oauth_login::EnterpriseOAuthLoginRequest;
pub use enterprise_oauth_login::delete_enterprise_oauth_tokens;
pub use enterprise_oauth_login::perform_enterprise_oauth_login_return_url;
pub use event_notification_transport::EventNotificationReceiver;
pub use http_client_adapter::StreamableHttpRedirectMode;
pub use http_headers::with_http_headers_helper;

29
codex-rs/rmcp-client/src/oauth.rs
View File

@@ -18,6 +18,7 @@
mod credential_store;
mod ema_identity;
mod enterprise_generation;
mod issuer_binding;
mod refresh_lock;
mod refresh_transaction;
@@ -76,6 +77,8 @@ use codex_utils_home_dir::find_codex_home;
pub(crate) use self::credential_store::OAuthCredentialStore;
pub(crate) use self::ema_identity::stored_oidc_identity;
pub(crate) use self::enterprise_generation::EnterpriseOAuthGeneration;
pub(crate) use self::enterprise_generation::EnterpriseOAuthGenerationFile;
pub(crate) use self::issuer_binding::validate_authorization_server_endpoints;
pub(crate) use self::issuer_binding::validate_refresh_token_issuer;
pub(crate) use self::refresh_lock::RefreshCredentialLock;
@@ -446,7 +449,18 @@ pub async fn save_oauth_tokens(
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<()> {
let _lock = RefreshCredentialLock::acquire_for_server(server_name, &tokens.url).await?;
let lock = RefreshCredentialLock::acquire_for_server(server_name, &tokens.url).await?;
save_oauth_tokens_with_lock_held(&lock, server_name, tokens, store_mode, keyring_backend_kind)
}
/// Save while retaining the matching credential lock acquired by the caller.
pub(crate) fn save_oauth_tokens_with_lock_held(
_lock: &RefreshCredentialLock,
server_name: &str,
tokens: &StoredOAuthTokens,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<()> {
let keyring_store = DefaultKeyringStore;
match store_mode {
OAuthCredentialsStoreMode::Auto => save_oauth_tokens_with_keyring_with_fallback_to_file(
@@ -592,7 +606,18 @@ pub async fn delete_oauth_tokens(
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<bool> {
let _lock = RefreshCredentialLock::acquire_for_server(server_name, url).await?;
let lock = RefreshCredentialLock::acquire_for_server(server_name, url).await?;
delete_oauth_tokens_with_lock_held(&lock, server_name, url, store_mode, keyring_backend_kind)
}
/// Delete while retaining the matching credential lock acquired by the caller.
pub(crate) fn delete_oauth_tokens_with_lock_held(
_lock: &RefreshCredentialLock,
server_name: &str,
url: &str,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<bool> {
let keyring_store = DefaultKeyringStore;
delete_oauth_tokens_from_keyring_and_file(
&keyring_store,

85
codex-rs/rmcp-client/src/oauth/enterprise_generation.rs Normal file
View File

@@ -0,0 +1,85 @@
//! Persistent, non-secret logout generations for staged enterprise logins.
//! All reads and writes require the credential lock; missing state never admits an old attempt.
use std::fs::File;
use std::fs::OpenOptions;
use std::io::Read;
use std::io::Seek;
use std::io::Write;
use anyhow::Result;
use anyhow::ensure;
use codex_utils_home_dir::find_codex_home;
use oauth2::CsrfToken;
use sha2::Digest;
use sha2::Sha256;
use super::RefreshCredentialLock;
#[derive(PartialEq, Eq)]
pub(crate) struct EnterpriseOAuthGeneration([u8; 32]);
pub(crate) struct EnterpriseOAuthGenerationFile {
file: File,
}
impl EnterpriseOAuthGenerationFile {
pub(crate) fn open(
credential_name: &str,
issuer: &str,
_lock: &RefreshCredentialLock,
) -> Result<Self> {
let key = super::compute_store_key(credential_name, issuer)?;
let name = format!("{:x}.enterprise-generation", Sha256::digest(key.as_bytes()));
let path = find_codex_home()?.join("mcp-oauth-locks").join(name);
let mut options = OpenOptions::new();
options.read(true).write(true).create(true).truncate(false);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
options
.mode(0o600)
.custom_flags(libc::O_NOFOLLOW | libc::O_NONBLOCK);
}
#[cfg(windows)]
{
use std::os::windows::fs::OpenOptionsExt;
use windows_sys::Win32::Storage::FileSystem::FILE_FLAG_OPEN_REPARSE_POINT;
options.custom_flags(FILE_FLAG_OPEN_REPARSE_POINT);
}
let file = options.open(path)?;
ensure!(
file.metadata()?.is_file(),
"invalid enterprise generation file"
);
Ok(Self { file })
}
pub(crate) fn current(&self) -> Result<Option<EnterpriseOAuthGeneration>> {
let mut file = &self.file;
file.rewind()?;
match file.metadata()?.len() {
0 => Ok(None),
32 => {
let mut generation = [0; 32];
file.read_exact(&mut generation)?;
Ok(Some(EnterpriseOAuthGeneration(generation)))
}
_ => anyhow::bail!("invalid enterprise generation file"),
}
}
pub(crate) fn replace(&self) -> Result<EnterpriseOAuthGeneration> {
// Reuse OAuth's randomness without storing any credential. Random initialization
// also prevents an old attempt becoming valid if metadata is removed or truncated.
let random = CsrfToken::new_random_len(/*num_bytes*/ 32);
let generation =
EnterpriseOAuthGeneration(Sha256::digest(random.secret().as_bytes()).into());
let mut file = &self.file;
file.set_len(/*size*/ 0)?;
file.rewind()?;
file.write_all(&generation.0)?;
file.sync_all()?;
Ok(generation)
}
}

23
codex-rs/rmcp-client/src/oauth/test_support.rs
View File

@@ -6,6 +6,29 @@ use std::sync::PoisonError;
use tempfile::tempdir;
/// Build the actual client's Stop-policy route pool before measuring protocol deadlines.
/// Client construction loads platform/custom CA state and can be slow on developer hosts.
/// Use a separate local endpoint so warmup cannot affect the test server's request assertions;
/// the same production HTTP capability still applies all proxy and certificate policy.
pub(crate) async fn warm_http_client(
client: &dyn codex_exec_server::HttpClient,
) -> anyhow::Result<()> {
let server = wiremock::MockServer::start().await;
client
.http_request(codex_exec_server::HttpRequestParams {
method: "GET".to_string(),
url: server.uri(),
headers: Vec::new(),
body: None,
timeout_ms: None,
redirect_policy: codex_exec_server::HttpRedirectPolicy::Stop,
request_id: "test-client-warmup".to_string(),
stream_response: false,
})
.await?;
Ok(())
}
/// Serializes tests that mutate process-wide CODEX_HOME.
///
/// Keep OAuth tests on this one guard instead of defining per-module helpers; otherwise

178
codex-rs/rmcp-client/src/perform_oauth_login.rs
View File

@@ -1,4 +1,5 @@
use std::collections::HashMap;
use std::net::SocketAddr;
use std::string::String;
use std::sync::Arc;
use std::time::Duration;
@@ -23,6 +24,9 @@ use urlencoding::decode;
use crate::StoredOAuthTokens;
use crate::WrappedOAuthTokenResponse;
use crate::enterprise_oauth_login::enterprise_authorization_url;
use crate::enterprise_oauth_login::enterprise_callback_settings;
use crate::enterprise_oauth_login::resolve_enterprise_authorization_manager;
use crate::http_client_adapter::StreamableHttpRedirectMode;
use crate::oauth::compute_expires_at_millis;
use crate::oauth::validate_authorization_server_endpoints;
@@ -30,6 +34,7 @@ use crate::oauth_callback::McpOAuthCallbackMode;
use crate::oauth_callback::append_callback_id_to_redirect_uri;
use crate::oauth_callback::callback_id_from_server_url;
use crate::oauth_callback::callback_mode;
use crate::oauth_callback::resolve_mcp_oauth_callback_url;
use crate::oauth_callback::validate_callback_redirect;
use crate::oauth_client_registration::McpOAuthClientRegistration;
use crate::oauth_client_registration::PreparedOAuthLogin;
@@ -40,11 +45,17 @@ use crate::utils::build_default_headers;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
struct OAuthHttpContext {
http_headers: Option<HashMap<String, String>>,
env_http_headers: Option<HashMap<String, String>>,
http_client: Arc<dyn HttpClient>,
redirect_mode: StreamableHttpRedirectMode,
#[derive(Clone, Copy)]
pub(crate) enum OAuthLoginPurpose {
Mcp,
EnterpriseIdp,
}
pub(crate) struct OAuthHttpContext {
pub(crate) http_headers: Option<HashMap<String, String>>,
pub(crate) env_http_headers: Option<HashMap<String, String>>,
pub(crate) http_client: Arc<dyn HttpClient>,
pub(crate) redirect_mode: StreamableHttpRedirectMode,
}
struct CallbackServerGuard {
@@ -197,6 +208,7 @@ async fn perform_oauth_login_with_browser_output(
http_context,
scopes,
oauth_client_id,
OAuthLoginPurpose::Mcp,
client_registration,
oauth_resource,
/*launch_browser*/ true,
@@ -243,6 +255,7 @@ pub async fn perform_oauth_login_return_url(
http_context,
scopes,
oauth_client_id,
OAuthLoginPurpose::Mcp,
client_registration,
oauth_resource,
/*launch_browser*/ false,
@@ -279,12 +292,15 @@ fn spawn_callback_server(
if let Err(err) = request.respond(response) {
eprintln!("Failed to respond to OAuth callback: {err}");
}
if let Err(err) = tx.send(CallbackResult::Success(OauthCallbackResult {
code,
state,
issuer,
})) {
eprintln!("Failed to send OAuth callback: {err:?}");
if let Err(message) = send_oauth_callback(
tx,
CallbackResult::Success(OauthCallbackResult {
code,
state,
issuer,
}),
) {
eprintln!("{message}");
}
break;
}
@@ -293,8 +309,8 @@ fn spawn_callback_server(
if let Err(err) = request.respond(response) {
eprintln!("Failed to respond to OAuth callback: {err}");
}
if let Err(err) = tx.send(CallbackResult::Error(error)) {
eprintln!("Failed to send OAuth callback error: {err:?}");
if let Err(message) = send_oauth_callback(tx, CallbackResult::Error(error)) {
eprintln!("{message}");
}
break;
}
@@ -323,6 +339,14 @@ enum CallbackResult {
Error(OAuthProviderError),
}
fn send_oauth_callback(
tx: oneshot::Sender<CallbackResult>,
result: CallbackResult,
) -> std::result::Result<(), &'static str> {
tx.send(result)
.map_err(|_| "OAuth callback receiver closed")
}
#[derive(Debug, PartialEq, Eq)]
enum CallbackOutcome {
Success(OauthCallbackResult),
@@ -405,7 +429,7 @@ impl OauthLoginHandle {
}
}
struct OauthLoginFlow {
pub(crate) struct OauthLoginFlow {
auth_url: String,
oauth_state: OAuthState,
authorization_server_issuer: Option<String>,
@@ -502,7 +526,7 @@ fn callback_bind_host(callback_url: Option<&str>) -> &'static str {
impl OauthLoginFlow {
#[allow(clippy::too_many_arguments)]
async fn new(
pub(crate) async fn new(
server_name: &str,
server_url: &str,
store_mode: OAuthCredentialsStoreMode,
@@ -510,6 +534,7 @@ impl OauthLoginFlow {
http_context: OAuthHttpContext,
scopes: &[String],
oauth_client_id: Option<&str>,
purpose: OAuthLoginPurpose,
client_registration: McpOAuthClientRegistration,
oauth_resource: Option<&str>,
launch_browser: bool,
@@ -521,6 +546,18 @@ impl OauthLoginFlow {
const DEFAULT_OAUTH_TIMEOUT_SECS: i64 = 300;
let callback_port = resolve_callback_port(callback_port)?;
let is_enterprise_idp = matches!(purpose, OAuthLoginPurpose::EnterpriseIdp);
let (enterprise_bind_ip, callback_port) = if is_enterprise_idp {
let (ip, port) = enterprise_callback_settings(
server_url,
oauth_client_id,
callback_url,
callback_port,
)?;
(Some(ip), port)
} else {
(None, callback_port)
};
let callback_id = callback_id_from_server_url(server_url)?;
let oauth_client_id = oauth_client_id.filter(|client_id| !client_id.trim().is_empty());
let configured_callback = if oauth_client_id.is_some() {
@@ -556,7 +593,10 @@ impl OauthLoginFlow {
redirect_mode,
)?);
let registered_authorization = if oauth_client_id.is_some() {
Some(resolve_authorization_manager(server_url, Arc::clone(&oauth_http_client)).await?)
Some(
resolve_authorization_manager(server_url, Arc::clone(&oauth_http_client), purpose)
.await?,
)
} else {
None
};
@@ -564,8 +604,8 @@ impl OauthLoginFlow {
.as_ref()
.map(|(_, metadata)| callback_mode(metadata))
.transpose()?;
let use_legacy_fallback = registered_callback_mode
== Some(McpOAuthCallbackMode::CallbackSpecific)
let use_legacy_fallback = !is_enterprise_idp
&& registered_callback_mode == Some(McpOAuthCallbackMode::CallbackSpecific)
&& configured_callback.as_ref().is_some_and(|callback_url| {
callback_url
.path_segments()
@@ -584,18 +624,32 @@ impl OauthLoginFlow {
callback_url
};
let bind_host = callback_bind_host(callback_url);
let bind_ip = match enterprise_bind_ip {
Some(ip) => ip,
None => callback_bind_host(callback_url).parse()?,
};
// Port zero asks the OS for a free ephemeral port; the resolved
// redirect receives that port after the listener has been bound.
let bind_addr = match callback_port {
Some(port) => format!("{bind_host}:{port}"),
None => format!("{bind_host}:0"),
};
let server = Arc::new(Server::http(&bind_addr).map_err(|err| anyhow!(err))?);
let bind_addr = SocketAddr::new(bind_ip, callback_port.unwrap_or(0));
let server = Arc::new(Server::http(bind_addr).map_err(|err| anyhow!(err))?);
let guard = CallbackServerGuard {
server: Arc::clone(&server),
};
let redirect_uri = resolve_redirect_uri(&server, callback_url)?;
let redirect_uri = if is_enterprise_idp {
let listener_port = server
.server_addr()
.to_ip()
.ok_or_else(|| anyhow!("unable to determine enterprise callback listener port"))?
.port();
let mut redirect = Url::parse(&redirect_uri)?;
redirect
.set_port(Some(listener_port))
.map_err(|()| anyhow!("invalid enterprise callback port"))?;
redirect.to_string()
} else {
redirect_uri
};
let scope_refs: Vec<&str> = scopes.iter().map(String::as_str).collect();
let PreparedOAuthLogin {
@@ -605,7 +659,13 @@ impl OauthLoginFlow {
} = if let Some((oauth_client_id, (auth_manager, metadata))) =
oauth_client_id.zip(registered_authorization)
{
let redirect_uri = if callback_url.is_some() && !use_legacy_fallback {
let redirect_uri = if is_enterprise_idp {
resolve_mcp_oauth_callback_url(
server_url,
Some(&redirect_uri),
callback_mode(&metadata)?,
)?
} else if callback_url.is_some() && !use_legacy_fallback {
redirect_uri
} else {
append_callback_id_to_redirect_uri(&redirect_uri, &callback_id)?
@@ -617,6 +677,7 @@ impl OauthLoginFlow {
&redirect_uri,
&callback_id,
oauth_client_id,
purpose,
)
.await?
} else {
@@ -656,11 +717,24 @@ impl OauthLoginFlow {
})
}
fn authorization_url(&self) -> String {
pub(crate) fn authorization_url(&self) -> String {
self.auth_url.clone()
}
async fn finish(mut self, emit_browser_url: bool) -> Result<()> {
async fn finish(self, emit_browser_url: bool) -> Result<()> {
let store_mode = self.store_mode;
let keyring_backend_kind = self.keyring_backend_kind;
let stored = self.complete(emit_browser_url).await?;
save_oauth_tokens(
&stored.server_name,
&stored,
store_mode,
keyring_backend_kind,
)
.await
}
pub(crate) async fn complete(mut self, emit_browser_url: bool) -> Result<StoredOAuthTokens> {
if self.launch_browser {
let server_name = &self.server_name;
let auth_url = &self.auth_url;
@@ -706,25 +780,15 @@ impl OauthLoginFlow {
.context("failed to retrieve OAuth credentials")?;
let credentials = credentials_opt
.ok_or_else(|| anyhow!("OAuth provider did not return credentials"))?;
let expires_at = compute_expires_at_millis(&credentials);
let stored = StoredOAuthTokens {
Ok(StoredOAuthTokens {
server_name: self.server_name.clone(),
url: self.server_url.clone(),
issuer: self.authorization_server_issuer.clone(),
client_id,
token_response: WrappedOAuthTokenResponse(credentials),
expires_at,
};
save_oauth_tokens(
&self.server_name,
&stored,
self.store_mode,
self.keyring_backend_kind,
)
.await?;
Ok(())
})
}
.await;
@@ -733,16 +797,13 @@ impl OauthLoginFlow {
}
fn spawn(self) -> oneshot::Receiver<Result<()>> {
let server_name_for_logging = self.server_name.clone();
let server_name = self.server_name.clone();
let (tx, rx) = oneshot::channel();
tokio::spawn(async move {
let result = self.finish(/*emit_browser_url*/ false).await;
if let Err(err) = &result {
eprintln!(
"Failed to complete OAuth login for '{server_name_for_logging}': {err:#}"
);
eprintln!("Failed to complete OAuth login for '{server_name}': {err:#}");
}
let _ = tx.send(result);
@@ -755,7 +816,11 @@ impl OauthLoginFlow {
async fn resolve_authorization_manager(
server_url: &str,
http_client: Arc<dyn OAuthHttpClient>,
purpose: OAuthLoginPurpose,
) -> Result<(AuthorizationManager, AuthorizationMetadata)> {
if matches!(purpose, OAuthLoginPurpose::EnterpriseIdp) {
return resolve_enterprise_authorization_manager(server_url, http_client).await;
}
let mut auth_manager =
AuthorizationManager::new_with_oauth_http_client(server_url, http_client).await?;
auth_manager.set_allow_missing_issuer(true);
@@ -771,15 +836,21 @@ async fn start_authorization(
redirect_uri: &str,
callback_id: &str,
oauth_client_id: &str,
purpose: OAuthLoginPurpose,
) -> Result<PreparedOAuthLogin> {
let strict_enterprise_idp = matches!(purpose, OAuthLoginPurpose::EnterpriseIdp);
let authorization_server_issuer = metadata.issuer.clone();
validate_callback_redirect(redirect_uri, callback_id, callback_mode(&metadata)?)?;
auth_manager.set_metadata(metadata);
auth_manager.configure_client(
OAuthClientConfig::new(oauth_client_id, redirect_uri)
.with_scopes(scopes.iter().map(|scope| (*scope).to_string()).collect()),
)?;
let client_config = OAuthClientConfig::new(oauth_client_id, redirect_uri)
.with_scopes(scopes.iter().map(|scope| (*scope).to_string()).collect());
auth_manager.configure_client(client_config)?;
let auth_url = auth_manager.get_authorization_url(scopes).await?;
let auth_url = if strict_enterprise_idp {
enterprise_authorization_url(&auth_url)?
} else {
auth_url
};
Ok(PreparedOAuthLogin {
oauth_state: OAuthState::Session(AuthorizationSession::for_scope_upgrade(
@@ -835,6 +906,7 @@ mod tests {
use codex_http_client::OutboundProxyPolicy;
use futures::future::BoxFuture;
use http::HeaderMap;
use oauth2::TokenResponse;
use pretty_assertions::assert_eq;
use serde_json::json;
use tokio::net::TcpListener;
@@ -844,6 +916,7 @@ mod tests {
use super::McpOAuthClientRegistration;
use super::OAuthHttpClientAdapter;
use super::OAuthHttpContext;
use super::OAuthLoginPurpose;
use super::OAuthProviderError;
use super::OauthLoginFlow;
use super::StreamableHttpRedirectMode;
@@ -937,7 +1010,6 @@ mod tests {
Json(json!({
"access_token": "test-access-token",
"token_type": "Bearer",
"refresh_token": "test-refresh-token",
}))
}),
);
@@ -981,7 +1053,7 @@ mod tests {
}
#[tokio::test]
async fn oauth_login_persists_discovered_issuer() -> anyhow::Result<()> {
async fn ordinary_oauth_login_persists_issuer_without_a_refresh_token() -> anyhow::Result<()> {
let _env = TempCodexHome::new();
let (base_url, _registration_requests) = spawn_oauth_metadata_server().await;
let server_url = format!("{base_url}/mcp");
@@ -1000,6 +1072,7 @@ mod tests {
},
&[],
Some("test-client"),
OAuthLoginPurpose::Mcp,
McpOAuthClientRegistration::Auto,
/*oauth_resource*/ None,
/*launch_browser*/ false,
@@ -1033,6 +1106,7 @@ mod tests {
)?
.expect("OAuth login should persist credentials");
assert_eq!(stored.issuer.as_deref(), Some(server_url.as_str()));
assert!(stored.token_response.0.refresh_token().is_none());
Ok(())
}
@@ -1050,6 +1124,7 @@ mod tests {
HeaderMap::new(),
&format!("{base_url}/mcp"),
)),
OAuthLoginPurpose::Mcp,
)
.await
.expect("resolve pre-registered OAuth metadata");
@@ -1060,6 +1135,7 @@ mod tests {
redirect_uri,
"configured-client",
"eci-prd-pub-codex-123",
OAuthLoginPurpose::Mcp,
)
.await
.expect("start pre-registered OAuth authorization");
@@ -1157,6 +1233,7 @@ mod tests {
HeaderMap::new(),
&format!("{issuer}/mcp"),
)),
OAuthLoginPurpose::Mcp,
)
.await
.expect("resolve issuer-aware authorization metadata");
@@ -1167,6 +1244,7 @@ mod tests {
redirect_uri,
"test-callback",
"test-client",
OAuthLoginPurpose::Mcp,
)
.await
.expect("start issuer-aware authorization");