codex/codex-rs/workload-identity/src/workload_identity_tests.rs

use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;

use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use pretty_assertions::assert_eq;
use tempfile::TempDir;
use url::Url;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;

use super::ACCESS_TOKEN_TYPE;
use super::JWT_BEARER_GRANT_TYPE;
use super::WorkloadIdentityExchange;
use super::WorkloadIdentityToken;
use crate::ApiTokenMetadata;
use crate::ChatgptTokenMetadata;
use crate::WorkloadIdentityConfig;
use crate::WorkloadIdentityError;
use crate::WorkloadIdentityMetadata;

fn assertion_file(assertion: &str) -> (TempDir, PathBuf) {
    let temp_dir = TempDir::new().expect("tempdir");
    let path = temp_dir.path().join("identity-token");
    std::fs::write(&path, assertion).expect("write assertion");
    (temp_dir, path)
}

fn make_exchange(path: PathBuf, server: &MockServer) -> WorkloadIdentityExchange {
    WorkloadIdentityExchange::new(
        WorkloadIdentityConfig::new(
            "idpm_rule_one".to_string(),
            path,
            /*workload_identity_context*/ None,
        )
        .expect("valid config"),
        Url::parse(&format!("{}/oauth/token", server.uri())).expect("valid token URL"),
        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
    )
    .expect("valid exchange")
}

fn success_body(access_token: &str, expires_in: u64) -> serde_json::Value {
    serde_json::json!({
        "access_token": access_token,
        "issued_token_type": ACCESS_TOKEN_TYPE,
        "token_type": "Bearer",
        "expires_in": expires_in,
        "scope": "openid profile email chatgpt.workspace.feature.allow-codex-local-access.access",
        "chatgpt_account_id": "workspace-one",
        "chatgpt_account_user_id": "membership-one",
        "user_id": "user-one",
        "chatgpt_plan_type": "enterprise"
    })
}

fn success(access_token: &str, expires_in: u64) -> ResponseTemplate {
    ResponseTemplate::new(200).set_body_json(success_body(access_token, expires_in))
}

fn api_success_body(access_token: &str) -> serde_json::Value {
    serde_json::json!({
        "access_token": access_token,
        "issued_token_type": ACCESS_TOKEN_TYPE,
        "token_type": "Bearer",
        "expires_in": 600,
        "wif_product": "api",
        "api_organization_id": "org-one",
        "api_project_id": "proj-one",
        "user_id": "service-account-one"
    })
}

#[tokio::test]
async fn exchange_sends_same_contract_for_both_products_and_caches_valid_response() {
    for product in [None, Some("chatgpt"), Some("api")] {
        let (mut body, metadata) = if product == Some("api") {
            (
                api_success_body("sensitive-access-token"),
                WorkloadIdentityMetadata::Api(ApiTokenMetadata {
                    api_organization_id: "org-one".to_string(),
                    api_project_id: "proj-one".to_string(),
                    user_id: "service-account-one".to_string(),
                }),
            )
        } else {
            (
                success_body("sensitive-access-token", /*expires_in*/ 600),
                WorkloadIdentityMetadata::Chatgpt(ChatgptTokenMetadata {
                    chatgpt_account_id: "workspace-one".to_string(),
                    chatgpt_account_user_id: "membership-one".to_string(),
                    chatgpt_plan_type: Some("enterprise".to_string()),
                    scope: "openid profile email chatgpt.workspace.feature.allow-codex-local-access.access"
                        .to_string(),
                    user_id: "user-one".to_string(),
                }),
            )
        };
        if let Some(product) = product {
            body["wif_product"] = product.into();
        }
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/oauth/token"))
            .and(header("content-type", "application/x-www-form-urlencoded"))
            .respond_with(ResponseTemplate::new(200).set_body_json(body))
            .mount(&server)
            .await;
        let (_temp_dir, assertion_path) = assertion_file("assertion-one\n");
        let exchange = make_exchange(assertion_path, &server);

        let expected = WorkloadIdentityToken {
            access_token: "sensitive-access-token".to_string(),
            expires_in: 600,
            metadata,
            version: 1,
        };
        assert_eq!(exchange.resolve().await.expect("exchange"), expected);
        assert_eq!(exchange.resolve().await.expect("cached token"), expected);
        assert!(!format!("{expected:?}").contains("sensitive-access-token"));

        let requests = server.received_requests().await.expect("requests");
        assert_eq!(requests.len(), 1);
        assert_eq!(
            url::form_urlencoded::parse(&requests[0].body)
                .into_owned()
                .collect::<Vec<_>>(),
            vec![
                ("grant_type".to_string(), JWT_BEARER_GRANT_TYPE.to_string()),
                ("assertion".to_string(), "assertion-one".to_string()),
                (
                    "federation_rule_id".to_string(),
                    "idpm_rule_one".to_string()
                ),
            ]
        );
    }
}

#[tokio::test]
async fn exchange_rejects_unknown_products_and_incomplete_api_identity() {
    let mut invalid_responses = Vec::new();
    for product in [serde_json::json!("unknown"), serde_json::Value::Null] {
        let mut body = success_body("access-one", /*expires_in*/ 600);
        body["wif_product"] = product;
        invalid_responses.push(body);
    }
    for field in [
        "wif_product",
        "api_organization_id",
        "api_project_id",
        "user_id",
    ] {
        let mut body = api_success_body("access-one");
        body.as_object_mut().expect("response object").remove(field);
        invalid_responses.push(body);
        let mut body = api_success_body("access-one");
        body[field] = " ".into();
        invalid_responses.push(body);
    }
    let (_temp_dir, assertion_path) = assertion_file("assertion-one");
    for body in invalid_responses {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/oauth/token"))
            .respond_with(ResponseTemplate::new(200).set_body_json(body))
            .expect(1)
            .mount(&server)
            .await;
        let exchange = make_exchange(assertion_path.clone(), &server);
        assert!(matches!(
            exchange.resolve().await,
            Err(WorkloadIdentityError::InvalidExchangeResponse)
        ));
    }
}

#[tokio::test]
async fn exchange_forwards_optional_workload_context_without_parsing() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/oauth/token"))
        .and(header("content-type", "application/x-www-form-urlencoded"))
        .respond_with(success("sensitive-access-token", /*expires_in*/ 600))
        .mount(&server)
        .await;
    let (_temp_dir, assertion_path) = assertion_file("assertion-one");
    let context = "server-validates-this-raw-value";
    let config = WorkloadIdentityConfig::new(
        "idpm_rule_one".to_string(),
        assertion_path,
        /*workload_identity_context*/ Some(context.to_string()),
    )
    .expect("valid config");
    let exchange = WorkloadIdentityExchange::new(
        config,
        Url::parse(&format!("{}/oauth/token", server.uri())).expect("valid token URL"),
        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
    )
    .expect("valid exchange");

    exchange.resolve().await.expect("exchange");

    let requests = server.received_requests().await.expect("requests");
    assert_eq!(requests.len(), 1);
    assert_eq!(
        url::form_urlencoded::parse(&requests[0].body)
            .into_owned()
            .collect::<Vec<_>>(),
        vec![
            ("grant_type".to_string(), JWT_BEARER_GRANT_TYPE.to_string()),
            ("assertion".to_string(), "assertion-one".to_string()),
            (
                "federation_rule_id".to_string(),
                "idpm_rule_one".to_string()
            ),
            ("workload_identity_context".to_string(), context.to_string()),
        ]
    );
}

#[tokio::test]
async fn concurrent_resolve_and_rejected_token_refresh_are_single_flight() {
    let server = MockServer::start().await;
    let calls = Arc::new(AtomicUsize::new(0));
    Mock::given(method("POST"))
        .respond_with({
            let calls = Arc::clone(&calls);
            move |_request: &wiremock::Request| {
                let call = calls.fetch_add(1, Ordering::SeqCst) + 1;
                success(&format!("access-{call}"), /*expires_in*/ 600)
                    .set_delay(Duration::from_millis(50))
            }
        })
        .mount(&server)
        .await;
    let (_temp_dir, assertion_path) = assertion_file("assertion-one");
    let exchange = Arc::new(make_exchange(assertion_path.clone(), &server));

    let resolves = (0..8)
        .map(|_| {
            let exchange = Arc::clone(&exchange);
            tokio::spawn(async move { exchange.resolve().await })
        })
        .collect::<Vec<_>>();
    let mut initial = None;
    for resolve in resolves {
        let token = resolve.await.expect("join resolve").expect("resolve");
        assert_eq!(token.access_token, "access-1");
        initial = Some(token);
    }
    tokio::fs::write(&assertion_path, "assertion-two\n")
        .await
        .expect("rotate assertion");
    let version = initial.expect("initial token").version();
    let refreshes = (0..8)
        .map(|_| {
            let exchange = Arc::clone(&exchange);
            tokio::spawn(async move { exchange.refresh(version).await })
        })
        .collect::<Vec<_>>();
    for refresh in refreshes {
        assert_eq!(
            refresh
                .await
                .expect("join refresh")
                .expect("refresh")
                .access_token,
            "access-2"
        );
    }

    assert_eq!(calls.load(Ordering::SeqCst), 2);
    assert_eq!(
        server
            .received_requests()
            .await
            .expect("requests")
            .iter()
            .map(|request| {
                url::form_urlencoded::parse(&request.body)
                    .find(|(name, _)| name == "assertion")
                    .map(|(_, value)| value.into_owned())
                    .expect("assertion field")
            })
            .collect::<Vec<_>>(),
        vec!["assertion-one", "assertion-two"]
    );
}

#[tokio::test]
async fn rejected_token_waiting_on_proactive_fallback_still_forces_refresh() {
    let server = MockServer::start().await;
    let calls = Arc::new(AtomicUsize::new(0));
    Mock::given(method("POST"))
        .respond_with({
            let calls = Arc::clone(&calls);
            move |_request: &wiremock::Request| match calls.fetch_add(1, Ordering::SeqCst) {
                0 => success("access-one", /*expires_in*/ 600),
                1 => ResponseTemplate::new(503).set_delay(Duration::from_millis(200)),
                2.. => success("access-three", /*expires_in*/ 600),
            }
        })
        .mount(&server)
        .await;
    let (_temp_dir, assertion_path) = assertion_file("assertion-one");
    let exchange = Arc::new(make_exchange(assertion_path, &server));
    let initial = exchange.resolve().await.expect("initial exchange");
    exchange
        .state
        .lock()
        .await
        .cached
        .as_mut()
        .expect("cached token")
        .refresh_at = std::time::Instant::now();

    let proactive = tokio::spawn({
        let exchange = Arc::clone(&exchange);
        async move { exchange.resolve().await }
    });
    tokio::time::timeout(Duration::from_secs(1), async {
        while calls.load(Ordering::SeqCst) < 2 {
            tokio::task::yield_now().await;
        }
    })
    .await
    .expect("proactive exchange started");

    let forced = exchange.refresh(initial.version());
    tokio::pin!(forced);
    assert!(
        tokio::time::timeout(Duration::from_millis(20), forced.as_mut())
            .await
            .is_err(),
        "forced refresh should wait for the proactive exchange"
    );
    let fallback = proactive
        .await
        .expect("join proactive refresh")
        .expect("cached fallback");
    assert_eq!(fallback.access_token, initial.access_token);

    let refreshed = forced.await.expect("forced refresh");
    assert_eq!(refreshed.access_token, "access-three");
    assert_ne!(refreshed.version(), initial.version());
    assert_eq!(calls.load(Ordering::SeqCst), 3);
}

#[tokio::test]
async fn transient_proactive_refresh_failure_uses_still_valid_token() {
    let server = MockServer::start().await;
    let calls = Arc::new(AtomicUsize::new(0));
    Mock::given(method("POST"))
        .respond_with({
            let calls = Arc::clone(&calls);
            move |_request: &wiremock::Request| {
                if calls.fetch_add(1, Ordering::SeqCst) == 0 {
                    success("access-one", /*expires_in*/ 4)
                } else {
                    ResponseTemplate::new(503).set_body_string("sensitive server detail")
                }
            }
        })
        .mount(&server)
        .await;
    let (_temp_dir, assertion_path) = assertion_file("assertion-one");
    let exchange = make_exchange(assertion_path, &server);
    let initial = exchange.resolve().await.expect("initial exchange");
    exchange
        .state
        .lock()
        .await
        .cached
        .as_mut()
        .expect("cached token")
        .refresh_at = std::time::Instant::now();

    let fallback = exchange.resolve().await.expect("cached fallback");
    assert_eq!(fallback.access_token, initial.access_token);
    assert_eq!(fallback.version(), initial.version());
    assert_eq!(exchange.resolve().await.expect("delayed retry"), fallback);
    assert_eq!(calls.load(Ordering::SeqCst), 2);
}

#[test]
fn configuration_requires_an_absolute_file_and_secure_token_url() {
    assert!(matches!(
        WorkloadIdentityConfig::new(
            "idpm_rule_one".to_string(),
            PathBuf::from("relative.jwt"),
            /*workload_identity_context*/ None,
        ),
        Err(WorkloadIdentityError::AssertionFileMustBeAbsolute)
    ));

    let (_temp_dir, assertion_path) = assertion_file("assertion-one");
    let config = WorkloadIdentityConfig::new(
        "idpm_rule_one".to_string(),
        assertion_path,
        /*workload_identity_context*/ None,
    )
    .expect("valid config");
    assert!(matches!(
        WorkloadIdentityExchange::new(
            config,
            Url::parse("http://auth.example.com/oauth/token").expect("parse URL"),
            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
        ),
        Err(WorkloadIdentityError::InvalidTokenUrl)
    ));
}

#[tokio::test]
async fn exchange_rejects_oversized_assertions_and_incomplete_responses() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
            "access_token": "access-one",
            "issued_token_type": ACCESS_TOKEN_TYPE,
            "token_type": "Bearer",
            "expires_in": 600,
            "scope": "openid",
            "chatgpt_account_id": "workspace-one",
            "user_id": "user-one"
        })))
        .mount(&server)
        .await;
    let (_temp_dir, assertion_path) = assertion_file(&"x".repeat(16 * 1024 + 1));
    let exchange = make_exchange(assertion_path.clone(), &server);
    assert!(matches!(
        exchange.resolve().await,
        Err(WorkloadIdentityError::AssertionTooLarge)
    ));

    tokio::fs::write(&assertion_path, "valid-assertion")
        .await
        .expect("replace assertion");
    assert!(matches!(
        exchange.resolve().await,
        Err(WorkloadIdentityError::InvalidExchangeResponse)
    ));
}

#[tokio::test]
async fn managed_exchange_denies_before_reading_assertion_or_sending() {
    let server = MockServer::start().await;
    let (_home, assertion_path) = assertion_file("assertion");
    std::fs::remove_file(&assertion_path).unwrap();
    let controller = codex_http_client::NetworkPolicyController::default();
    let exchange = WorkloadIdentityExchange::new(
        WorkloadIdentityConfig::new(
            "rule".into(),
            assertion_path,
            /*workload_identity_context*/ None,
        )
        .unwrap(),
        Url::parse(&format!("{}/oauth/token", server.uri())).unwrap(),
        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
            .with_network_policy(controller.policy()),
    )
    .unwrap();
    assert!(matches!(
        exchange.resolve().await,
        Err(WorkloadIdentityError::Policy(
            codex_http_client::NetworkPolicyDenied::Unavailable
        ))
    ));
    assert!(server.received_requests().await.unwrap().is_empty());
    controller.publish(
        controller.policy().revision(),
        codex_http_client::DestinationPolicy::Restricted {
            allowed_hosts: Default::default(),
        },
    );
    assert!(matches!(
        exchange.resolve().await,
        Err(WorkloadIdentityError::Policy(
            codex_http_client::NetworkPolicyDenied::Destination
        ))
    ));
    assert!(server.received_requests().await.unwrap().is_empty());
}

#[tokio::test]
async fn managed_exchange_revokes_an_active_token_response() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .respond_with(
            success("token", /*expires_in*/ 600).set_delay(Duration::from_secs(/*secs*/ 30)),
        )
        .expect(/*r*/ 1)
        .mount(&server)
        .await;
    let (_home, assertion_path) = assertion_file("assertion");
    let controller = codex_http_client::NetworkPolicyController::default();
    controller.publish(
        controller.policy().revision(),
        codex_http_client::DestinationPolicy::Unrestricted,
    );
    let exchange = WorkloadIdentityExchange::new(
        WorkloadIdentityConfig::new(
            "rule".into(),
            assertion_path,
            /*workload_identity_context*/ None,
        )
        .unwrap(),
        Url::parse(&format!("{}/oauth/token", server.uri())).unwrap(),
        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
            .with_network_policy(controller.policy()),
    )
    .unwrap();
    let revoke = async {
        while server.received_requests().await.unwrap().is_empty() {
            tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await;
        }
        controller.publish(
            controller.policy().revision(),
            codex_http_client::DestinationPolicy::Restricted {
                allowed_hosts: Default::default(),
            },
        );
    };
    let (result, ()) = tokio::time::timeout(Duration::from_secs(/*secs*/ 5), async {
        tokio::join!(exchange.resolve(), revoke)
    })
    .await
    .unwrap();
    assert!(matches!(
        result,
        Err(WorkloadIdentityError::Policy(
            codex_http_client::NetworkPolicyDenied::Revoked
        ))
    ));
}