diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 388a9f452..64413617e 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -207,6 +207,7 @@ pub async fn run_login_with_access_token( &config.codex_home, &access_token, config.cli_auth_credentials_store_mode, + config.forced_chatgpt_workspace_id.as_deref(), Some(&config.chatgpt_base_url), ) .await diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index ff12aaff6..4831d0171 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -108,6 +108,7 @@ async fn login_with_access_token_writes_only_token() { dir.path(), &agent_identity, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, Some(&chatgpt_base_url), ) .await @@ -144,10 +145,12 @@ async fn login_with_access_token_writes_only_personal_access_token() { .mount(&server) .await; let _authapi_guard = EnvVarGuard::set("CODEX_AUTHAPI_BASE_URL", &server.uri()); + let allowed_workspaces = [WORKSPACE_ID_ALLOWED.to_string()]; super::login_with_access_token( dir.path(), "at-login-test", AuthCredentialsStoreMode::File, + Some(&allowed_workspaces), /*chatgpt_base_url*/ None, ) .await @@ -176,6 +179,42 @@ async fn login_with_access_token_writes_only_personal_access_token() { server.verify().await; } +#[tokio::test] +#[serial(codex_auth_env)] +async fn login_with_access_token_rejects_personal_access_token_workspace_mismatch() { + let dir = tempdir().unwrap(); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/user-auth-credential/whoami")) + .and(header("authorization", "Bearer at-workspace-mismatch")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(personal_access_token_whoami(WORKSPACE_ID_DISALLOWED)), + ) + .expect(1) + .mount(&server) + .await; + let _authapi_guard = EnvVarGuard::set("CODEX_AUTHAPI_BASE_URL", &server.uri()); + let allowed_workspaces = [WORKSPACE_ID_ALLOWED.to_string()]; + + let err = super::login_with_access_token( + dir.path(), + "at-workspace-mismatch", + AuthCredentialsStoreMode::File, + Some(&allowed_workspaces), + /*chatgpt_base_url*/ None, + ) + .await + .expect_err("personal access token workspace mismatch should fail"); + + assert_eq!(err.kind(), std::io::ErrorKind::PermissionDenied); + assert!( + !get_auth_file(dir.path()).exists(), + "workspace mismatch should not write auth.json" + ); + server.verify().await; +} + #[tokio::test] #[serial(codex_auth_env)] async fn login_with_access_token_rejects_invalid_personal_access_token() { @@ -193,6 +232,7 @@ async fn login_with_access_token_rejects_invalid_personal_access_token() { dir.path(), "at-invalid-login", AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -214,6 +254,7 @@ async fn login_with_access_token_rejects_invalid_jwt() { dir.path(), "not-a-jwt", AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -244,6 +285,7 @@ async fn login_with_access_token_rejects_unsigned_jwt() { dir.path(), &agent_identity, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, Some(&chatgpt_base_url), ) .await @@ -290,6 +332,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { codex_home.path(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -348,6 +391,7 @@ async fn loads_api_key_from_auth_json() { dir.path(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -427,6 +471,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { codex_home.path(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -828,6 +873,7 @@ async fn load_auth_reads_access_token_from_env() { codex_home.path(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, Some(&chatgpt_base_url), ) .await @@ -872,6 +918,7 @@ async fn load_auth_reads_personal_access_token_from_env() { codex_home.path(), /*enable_codex_api_key_env*/ false, auth_credentials_store_mode, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -900,6 +947,89 @@ async fn load_auth_reads_personal_access_token_from_env() { server.verify().await; } +#[tokio::test] +#[serial(codex_auth_env)] +async fn auth_manager_rejects_env_personal_access_token_workspace_mismatch() { + let codex_home = tempdir().unwrap(); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/user-auth-credential/whoami")) + .and(header("authorization", "Bearer at-env-workspace-mismatch")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(personal_access_token_whoami(WORKSPACE_ID_DISALLOWED)), + ) + .expect(1) + .mount(&server) + .await; + let _authapi_guard = EnvVarGuard::set("CODEX_AUTHAPI_BASE_URL", &server.uri()); + let _access_token_guard = + EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, "at-env-workspace-mismatch"); + + let manager = AuthManager::new_with_workspace_restriction( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ + Some(vec![WORKSPACE_ID_ALLOWED.to_string()]), + /*chatgpt_base_url*/ None, + ) + .await; + + assert_eq!(manager.auth().await, None); + server.verify().await; +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn auth_manager_rejects_stored_personal_access_token_workspace_mismatch() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/user-auth-credential/whoami")) + .and(header( + "authorization", + "Bearer at-stored-workspace-mismatch", + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(personal_access_token_whoami(WORKSPACE_ID_DISALLOWED)), + ) + .expect(4) + .mount(&server) + .await; + let _authapi_guard = EnvVarGuard::set("CODEX_AUTHAPI_BASE_URL", &server.uri()); + let _access_token_guard = remove_access_token_env_var(); + + for auth_credentials_store_mode in [ + AuthCredentialsStoreMode::File, + AuthCredentialsStoreMode::Ephemeral, + ] { + let codex_home = tempdir().unwrap(); + super::login_with_access_token( + codex_home.path(), + "at-stored-workspace-mismatch", + auth_credentials_store_mode, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + ) + .await + .expect("personal access token login should succeed"); + + let manager = AuthManager::new_with_workspace_restriction( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + auth_credentials_store_mode, + /*forced_chatgpt_workspace_id*/ + Some(vec![WORKSPACE_ID_ALLOWED.to_string()]), + /*chatgpt_base_url*/ None, + ) + .await; + + assert_eq!(manager.auth().await, None); + } + server.verify().await; +} + #[tokio::test] #[serial(codex_auth_env)] async fn personal_access_token_does_not_offer_unauthorized_recovery() { @@ -951,6 +1081,7 @@ async fn load_auth_keeps_codex_api_key_env_precedence() { codex_home.path(), /*enable_codex_api_key_env*/ true, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -1020,6 +1151,53 @@ async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() { ); } +#[tokio::test] +#[serial(codex_auth_env)] +async fn enforce_login_restrictions_logs_out_for_personal_access_token_workspace_mismatch() { + let codex_home = tempdir().unwrap(); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/user-auth-credential/whoami")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(personal_access_token_whoami(WORKSPACE_ID_DISALLOWED)), + ) + .expect(2) + .mount(&server) + .await; + let _access_token_guard = remove_access_token_env_var(); + let _authapi_guard = EnvVarGuard::set("CODEX_AUTHAPI_BASE_URL", &server.uri()); + super::login_with_access_token( + codex_home.path(), + "at-workspace-mismatch", + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + ) + .await + .expect("personal access token login should succeed"); + + let config = AuthConfig { + codex_home: codex_home.path().to_path_buf(), + auth_credentials_store_mode: AuthCredentialsStoreMode::File, + forced_login_method: None, + forced_chatgpt_workspace_id: Some(vec![WORKSPACE_ID_ALLOWED.to_string()]), + chatgpt_base_url: None, + }; + + let err = super::enforce_login_restrictions(&config) + .await + .expect_err("expected workspace mismatch to error"); + assert!(err.to_string().contains(&format!( + "current credentials belong to {WORKSPACE_ID_DISALLOWED}" + ))); + assert!( + !codex_home.path().join("auth.json").exists(), + "auth.json should be removed on mismatch" + ); + server.verify().await; +} + #[tokio::test] #[serial(codex_auth_env)] async fn enforce_login_restrictions_allows_matching_workspace() { @@ -1371,6 +1549,7 @@ async fn plan_type_maps_known_plan() { codex_home.path(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -1399,6 +1578,7 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() { codex_home.path(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -1430,6 +1610,7 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() { codex_home.path(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -1461,6 +1642,7 @@ async fn plan_type_maps_unknown_to_unknown() { codex_home.path(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await @@ -1489,6 +1671,7 @@ async fn missing_plan_type_maps_to_unknown() { codex_home.path(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, ) .await diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 383011f24..ececbe271 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -279,6 +279,7 @@ impl CodexAuth { codex_home, /*enable_codex_api_key_env*/ false, auth_credentials_store_mode, + /*forced_chatgpt_workspace_id*/ None, chatgpt_base_url, ) .await @@ -621,11 +622,13 @@ pub async fn login_with_access_token( codex_home: &Path, access_token: &str, auth_credentials_store_mode: AuthCredentialsStoreMode, + forced_chatgpt_workspace_id: Option<&[String]>, chatgpt_base_url: Option<&str>, ) -> std::io::Result<()> { let auth_dot_json = match classify_codex_access_token(access_token) { CodexAccessToken::PersonalAccessToken(access_token) => { - PersonalAccessTokenAuth::load(access_token).await?; + let auth = PersonalAccessTokenAuth::load(access_token).await?; + ensure_personal_access_token_workspace_allowed(forced_chatgpt_workspace_id, &auth)?; AuthDotJson { // Infer PAT auth from the credential field so older Codex builds can still // deserialize auth.json after a rollback. @@ -658,6 +661,14 @@ pub async fn login_with_access_token( save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) } +fn ensure_personal_access_token_workspace_allowed( + expected_workspace_ids: Option<&[String]>, + auth: &PersonalAccessTokenAuth, +) -> std::io::Result<()> { + crate::server::ensure_workspace_account_allowed(expected_workspace_ids, auth.account_id()) + .map_err(|message| std::io::Error::new(std::io::ErrorKind::PermissionDenied, message)) +} + /// Writes an in-memory auth payload for externally managed ChatGPT tokens. pub fn login_with_chatgpt_auth_tokens( codex_home: &Path, @@ -714,6 +725,7 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result< &config.codex_home, /*enable_codex_api_key_env*/ true, config.auth_credentials_store_mode, + /*forced_chatgpt_workspace_id*/ None, config.chatgpt_base_url.as_deref(), ) .await? @@ -754,10 +766,10 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result< if let Some(expected_account_ids) = config.forced_chatgpt_workspace_id.as_deref() { let chatgpt_account_id = match &auth { - CodexAuth::ApiKey(_) - | CodexAuth::PersonalAccessToken(_) - | CodexAuth::BedrockApiKey(_) => return Ok(()), - CodexAuth::AgentIdentity(_) => auth.get_account_id(), + CodexAuth::ApiKey(_) | CodexAuth::BedrockApiKey(_) => return Ok(()), + CodexAuth::AgentIdentity(_) | CodexAuth::PersonalAccessToken(_) => { + auth.get_account_id() + } CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { let token_data = match auth.get_token_data() { Ok(data) => data, @@ -833,6 +845,7 @@ async fn load_auth( codex_home: &Path, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, + forced_chatgpt_workspace_id: Option<&[String]>, chatgpt_base_url: Option<&str>, ) -> std::io::Result> { // API key via env var takes precedence over any other auth method. @@ -854,15 +867,18 @@ async fn load_auth( chatgpt_base_url, ) .await?; + if let CodexAuth::PersonalAccessToken(auth) = &auth { + ensure_personal_access_token_workspace_allowed(forced_chatgpt_workspace_id, auth)?; + } return Ok(Some(auth)); } if let Some(access_token) = read_codex_access_token_from_env() { return match classify_codex_access_token(&access_token) { CodexAccessToken::PersonalAccessToken(access_token) => { - CodexAuth::from_personal_access_token(access_token) - .await - .map(Some) + let auth = PersonalAccessTokenAuth::load(access_token).await?; + ensure_personal_access_token_workspace_allowed(forced_chatgpt_workspace_id, &auth)?; + Ok(Some(CodexAuth::PersonalAccessToken(auth))) } CodexAccessToken::AgentIdentityJwt(jwt) => { CodexAuth::from_agent_identity_jwt(jwt, chatgpt_base_url) @@ -891,6 +907,9 @@ async fn load_auth( chatgpt_base_url, ) .await?; + if let CodexAuth::PersonalAccessToken(auth) = &auth { + ensure_personal_access_token_workspace_allowed(forced_chatgpt_workspace_id, auth)?; + } Ok(Some(auth)) } @@ -1438,11 +1457,29 @@ impl AuthManager { enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, chatgpt_base_url: Option, + ) -> Self { + Self::new_with_workspace_restriction( + codex_home, + enable_codex_api_key_env, + auth_credentials_store_mode, + /*forced_chatgpt_workspace_id*/ None, + chatgpt_base_url, + ) + .await + } + + async fn new_with_workspace_restriction( + codex_home: PathBuf, + enable_codex_api_key_env: bool, + auth_credentials_store_mode: AuthCredentialsStoreMode, + forced_chatgpt_workspace_id: Option>, + chatgpt_base_url: Option, ) -> Self { let managed_auth = load_auth( &codex_home, enable_codex_api_key_env, auth_credentials_store_mode, + forced_chatgpt_workspace_id.as_deref(), chatgpt_base_url.as_deref(), ) .await @@ -1458,7 +1495,7 @@ impl AuthManager { auth_change_tx, enable_codex_api_key_env, auth_credentials_store_mode, - forced_chatgpt_workspace_id: RwLock::new(None), + forced_chatgpt_workspace_id: RwLock::new(forced_chatgpt_workspace_id), chatgpt_base_url, refresh_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), @@ -1658,10 +1695,12 @@ impl AuthManager { } async fn load_auth_from_storage(&self) -> Option { + let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); load_auth( &self.codex_home, self.enable_codex_api_key_env, self.auth_credentials_store_mode, + forced_chatgpt_workspace_id.as_deref(), self.chatgpt_base_url.as_deref(), ) .await @@ -1753,15 +1792,16 @@ impl AuthManager { config: &impl AuthManagerConfig, enable_codex_api_key_env: bool, ) -> Arc { - let auth_manager = Self::shared( - config.codex_home(), - enable_codex_api_key_env, - config.cli_auth_credentials_store_mode(), - Some(config.chatgpt_base_url()), + Arc::new( + Self::new_with_workspace_restriction( + config.codex_home(), + enable_codex_api_key_env, + config.cli_auth_credentials_store_mode(), + config.forced_chatgpt_workspace_id(), + Some(config.chatgpt_base_url()), + ) + .await, ) - .await; - auth_manager.set_forced_chatgpt_workspace_id(config.forced_chatgpt_workspace_id()); - auth_manager } pub fn unauthorized_recovery(self: &Arc) -> UnauthorizedRecovery {