codex/codex-rs/cli/tests/mcp_login.rs

//! Exercises terminal OAuth pasteback and its shared callback validation through the CLI.

use std::collections::BTreeMap;
use std::process::Output;
use std::process::Stdio;
use std::time::Duration;

use anyhow::Context;
use anyhow::Result;
use anyhow::bail;
use codex_core::config::load_global_mcp_servers;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use tempfile::TempDir;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::net::TcpStream;
use tokio::process::Child;
use tokio::process::ChildStdout;
use tokio::process::Command;
use tokio::time::timeout;
use url::Url;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;

const TEST_TIMEOUT: Duration = Duration::from_secs(30);

struct Fixture {
    server: MockServer,
    home: TempDir,
    issuer: String,
}

impl Fixture {
    async fn new(client_secret: Option<&'static str>) -> Result<Self> {
        Self::with_scopes(client_secret, &["mcp.read"]).await
    }

    async fn with_scopes(client_secret: Option<&'static str>, scopes: &[&str]) -> Result<Self> {
        let server = MockServer::start().await;
        let home = TempDir::new()?;
        let issuer = format!("{}/mcp", server.uri());
        let saved_client_secret = client_secret
            .map(|secret| format!("client_secret = \"{secret}\"\n"))
            .unwrap_or_default();
        let saved_scopes = if scopes.is_empty() {
            String::new()
        } else {
            format!("scopes = {scopes:?}\n")
        };
        std::fs::write(
            home.path().join("config.toml"),
            format!(
                "mcp_oauth_credentials_store = \"file\"\n\
             [mcp_servers.manual]\nurl = \"{issuer}\"\n{saved_scopes}\
             [mcp_servers.manual.oauth]\nclient_id = \"registered-client\"\n\
             {saved_client_secret}"
            ),
        )?;
        let fixture = Self {
            server,
            home,
            issuer,
        };
        fixture.mount_oauth(client_secret).await;
        Ok(fixture)
    }

    async fn mount_oauth(&self, client_secret: Option<&'static str>) {
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-authorization-server/mcp"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "issuer": self.issuer,
                "authorization_endpoint": format!("{}/authorize", self.server.uri()),
                "token_endpoint": format!("{}/token", self.server.uri()),
                "response_types_supported": ["code"],
                "code_challenge_methods_supported": ["S256"],
                "scopes_supported": ["mcp.read"],
                "authorization_response_iss_parameter_supported": true,
                "token_endpoint_auth_methods_supported": [if client_secret.is_some() {
                    "client_secret_post"
                } else {
                    "none"
                }],
            })))
            .mount(&self.server)
            .await;
        Mock::given(method("POST"))
            .and(path("/token"))
            .and(move |request: &wiremock::Request| {
                let body: BTreeMap<_, _> = url::form_urlencoded::parse(&request.body)
                    .into_owned()
                    .collect();
                body.get("client_secret").map(String::as_str) == client_secret
            })
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "access_token": "mock-access-token",
                "token_type": "Bearer",
                "refresh_token": "mock-refresh-token",
                "scope": "mcp.read",
            })))
            .mount(&self.server)
            .await;
    }

    async fn mount_resource_metadata(&self, scopes: &[&str], challenge_scope: Option<&str>) {
        let mut challenge = format!(
            "Bearer resource_metadata=\"{}/resource-metadata\"",
            self.server.uri()
        );
        if let Some(scope) = challenge_scope {
            challenge.push_str(&format!(", scope=\"{scope}\""));
        }
        Mock::given(method("GET"))
            .and(path("/mcp"))
            .respond_with(ResponseTemplate::new(401).insert_header("www-authenticate", challenge))
            .mount(&self.server)
            .await;
        Mock::given(method("GET"))
            .and(path("/resource-metadata"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "resource": self.issuer,
                "authorization_servers": [self.issuer],
                "scopes_supported": scopes,
            })))
            .mount(&self.server)
            .await;
    }

    async fn start(&self) -> Result<Login> {
        let mut child = Command::new(codex_utils_cargo_bin::cargo_bin("codex")?)
            .kill_on_drop(true)
            .current_dir(self.home.path())
            .env("CODEX_HOME", self.home.path())
            .env("NO_PROXY", "127.0.0.1,localhost")
            .env("no_proxy", "127.0.0.1,localhost")
            .args(["mcp", "login", "manual", "--no-browser"])
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()?;
        let mut stdout = BufReader::new(child.stdout.take().context("missing stdout")?);
        let mut transcript = String::new();
        let authorization_url = match read_authorization_url(&mut stdout, &mut transcript).await {
            Ok(authorization_url) => authorization_url,
            Err(error) => {
                let mut stderr = String::new();
                if let Some(stderr_pipe) = child.stderr.take() {
                    let _ = timeout(
                        Duration::from_secs(1),
                        stderr_pipe.take(/*limit*/ 4096).read_to_string(&mut stderr),
                    )
                    .await;
                }
                return Err(error.context(format!("MCP login stderr: {stderr}")));
            }
        };
        Ok(Login {
            child,
            stdout,
            transcript,
            authorization_url,
        })
    }
}

async fn read_authorization_url(
    stdout: &mut BufReader<ChildStdout>,
    transcript: &mut String,
) -> Result<Url> {
    timeout(TEST_TIMEOUT, async {
        loop {
            let mut line = String::new();
            if stdout.read_line(&mut line).await? == 0 {
                bail!("login exited before printing its authorization URL: {transcript}");
            }
            transcript.push_str(&line);
            if line.starts_with("http://") {
                return Ok(Url::parse(line.trim())?);
            }
        }
    })
    .await
    .context("authorization URL timeout")?
}

struct Login {
    child: Child,
    stdout: BufReader<ChildStdout>,
    transcript: String,
    authorization_url: Url,
}

impl Login {
    fn callback(&self, issuer: &str) -> Result<Url> {
        let params: BTreeMap<_, _> = self.authorization_url.query_pairs().into_owned().collect();
        let mut callback = Url::parse(&params["redirect_uri"])?;
        callback
            .query_pairs_mut()
            .append_pair("code", "mock-code")
            .append_pair("state", &params["state"])
            .append_pair("iss", issuer);
        Ok(callback)
    }

    async fn finish(mut self) -> Result<Output> {
        let (mut output, _) = timeout(TEST_TIMEOUT, async {
            tokio::try_join!(
                self.child.wait_with_output(),
                self.stdout.read_to_string(&mut self.transcript),
            )
        })
        .await
        .context("login completion timeout")??;
        output.stdout = self.transcript.into_bytes();
        Ok(output)
    }
}

#[tokio::test]
async fn no_browser_exchanges_pasted_callback_and_saves_credentials() -> Result<()> {
    let fixture = Fixture::new(/*client_secret*/ None).await?;
    assert_no_browser_exchanges_pasted_callback_and_saves_credentials(
        fixture, /*client_secret*/ None,
    )
    .await
}

#[tokio::test]
async fn no_browser_prefers_challenge_scopes_for_preregistered_client() -> Result<()> {
    let fixture = Fixture::with_scopes(/*client_secret*/ None, &[]).await?;
    fixture
        .mount_resource_metadata(&["mcp.read", "mcp.write"], Some("mcp.read"))
        .await;
    // The authorization server advertises neither scopes nor dynamic registration.
    Mock::given(method("GET"))
        .and(path("/.well-known/oauth-authorization-server/mcp"))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "issuer": fixture.issuer,
            "authorization_endpoint": format!("{}/authorize", fixture.server.uri()),
            "token_endpoint": format!("{}/token", fixture.server.uri()),
            "response_types_supported": ["code"],
            "code_challenge_methods_supported": ["S256"],
        })))
        .with_priority(/*priority*/ 1)
        .mount(&fixture.server)
        .await;

    assert_no_browser_exchanges_pasted_callback_and_saves_credentials(
        fixture, /*client_secret*/ None,
    )
    .await
}

#[tokio::test]
async fn no_browser_exchanges_pasted_callback_with_client_secret() -> Result<()> {
    let fixture = Fixture::new(Some("headless-client-secret")).await?;
    assert_no_browser_exchanges_pasted_callback_and_saves_credentials(
        fixture,
        Some("headless-client-secret"),
    )
    .await
}

#[tokio::test]
async fn add_persists_client_secret_for_subsequent_login_without_disclosure() -> Result<()> {
    const CLIENT_SECRET: &str = "cli-client-secret";
    let fixture = Fixture::new(Some(CLIENT_SECRET)).await?;
    std::fs::write(
        fixture.home.path().join("config.toml"),
        "mcp_oauth_credentials_store = \"file\"\n",
    )?;
    // Stop automatic login before browser launch, then retry with valid metadata.
    Mock::given(method("GET"))
        .and(path("/.well-known/oauth-authorization-server/mcp"))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "issuer": fixture.issuer,
            "authorization_endpoint": "not-a-valid-authorization-url",
            "token_endpoint": format!("{}/token", fixture.server.uri()),
            "authorization_response_iss_parameter_supported": true,
        })))
        .with_priority(/*priority*/ 1)
        .mount(&fixture.server)
        .await;
    let output = timeout(
        TEST_TIMEOUT,
        Command::new(codex_utils_cargo_bin::cargo_bin("codex")?)
            .kill_on_drop(true)
            .current_dir(fixture.home.path())
            .env("CODEX_HOME", fixture.home.path())
            .env("NO_PROXY", "127.0.0.1,localhost")
            .env("no_proxy", "127.0.0.1,localhost")
            .args([
                "mcp",
                "add",
                "manual",
                "--url",
                &fixture.issuer,
                "--oauth-client-id",
                "registered-client",
                "--oauth-client-secret",
                CLIENT_SECRET,
            ])
            .output(),
    )
    .await??;
    assert!(!output.status.success());
    let stdout = String::from_utf8(output.stdout)?;
    let stderr = String::from_utf8(output.stderr)?;
    assert!(stdout.contains("Detected OAuth support. Starting OAuth flow"));
    assert!(!stdout.contains(CLIENT_SECRET));
    assert!(!stderr.contains(CLIENT_SECRET));
    let servers = load_global_mcp_servers(fixture.home.path()).await?;
    assert_eq!(servers["manual"].oauth_client_secret(), Some(CLIENT_SECRET));
    // Reset removes the priority override without relying on wiremock's scoped indices.
    fixture.server.reset().await;
    fixture.mount_oauth(Some(CLIENT_SECRET)).await;
    fixture
        .mount_resource_metadata(&["mcp.read"], /*challenge_scope*/ None)
        .await;

    assert_no_browser_exchanges_pasted_callback_and_saves_credentials(fixture, Some(CLIENT_SECRET))
        .await
}

async fn assert_no_browser_exchanges_pasted_callback_and_saves_credentials(
    fixture: Fixture,
    client_secret: Option<&'static str>,
) -> Result<()> {
    let mut login = fixture.start().await?;
    let authorization_url = login.authorization_url.to_string();
    let params: BTreeMap<_, _> = login.authorization_url.query_pairs().into_owned().collect();
    assert_eq!(params.get("scope").map(String::as_str), Some("mcp.read"));
    let callback = login.callback(&fixture.issuer)?;
    login
        .child
        .stdin
        .as_mut()
        .context("missing stdin")?
        .write_all(format!("{callback}\n").as_bytes())
        .await?;
    let output = login.finish().await?;
    assert!(
        output.status.success(),
        "{}",
        String::from_utf8_lossy(&output.stderr)
    );
    let stdout = String::from_utf8(output.stdout)?;
    let stderr = String::from_utf8(output.stderr)?;
    if let Some(client_secret) = client_secret {
        assert!(!stdout.contains(client_secret));
        assert!(!stderr.contains(client_secret));
    }
    insta::allow_duplicates! {
        insta::assert_snapshot!(stdout.replace(&authorization_url, "[AUTHORIZATION_URL]"), @r"
        Authorize the MCP server by opening this URL in your browser:
        [AUTHORIZATION_URL]

        After signing in, copy the full URL from your browser's address bar.
        If the callback page cannot load, paste that URL here anyway.
        Successfully logged in to MCP server 'manual'.
        ");
    }
    assert!(stderr.contains("Callback URL (input hidden): "));
    let requests = fixture
        .server
        .received_requests()
        .await
        .context("missing requests")?;
    let tokens: Vec<_> = requests
        .iter()
        .filter(|r| r.url.path() == "/token")
        .collect();
    assert_eq!(tokens.len(), 1);
    let mut exchange: BTreeMap<_, _> = url::form_urlencoded::parse(&tokens[0].body)
        .into_owned()
        .collect();
    let verifier = exchange
        .remove("code_verifier")
        .context("missing PKCE verifier")?;
    assert!((43..=128).contains(&verifier.len()));
    assert_eq!(params["code_challenge_method"], "S256");
    assert!(!params["code_challenge"].is_empty());
    let mut expected_exchange = BTreeMap::from([
        ("client_id".to_string(), "registered-client".to_string()),
        ("code".to_string(), "mock-code".to_string()),
        ("grant_type".to_string(), "authorization_code".to_string()),
        ("redirect_uri".to_string(), params["redirect_uri"].clone()),
        ("resource".to_string(), fixture.issuer.clone()),
    ]);
    if let Some(client_secret) = client_secret {
        expected_exchange.insert("client_secret".to_string(), client_secret.to_string());
    }
    assert_eq!(exchange, expected_exchange);
    let credentials: BTreeMap<String, Value> = serde_json::from_slice(&std::fs::read(
        fixture.home.path().join(".credentials.json"),
    )?)?;
    assert_eq!(
        credentials.into_values().collect::<Vec<_>>(),
        vec![json!({
            "server_name": "manual", "server_url": fixture.issuer, "issuer": fixture.issuer,
            "client_id": "registered-client", "access_token": "mock-access-token",
            "refresh_token": "mock-refresh-token", "expires_at": null, "scopes": ["mcp.read"],
        })]
    );
    Ok(())
}

#[tokio::test]
async fn no_browser_rejects_invalid_callbacks_without_exchanging_tokens() -> Result<()> {
    for (field, value) in [
        ("state", Some("wrong-state")),
        ("iss", None),
        ("iss", Some("https://secret-issuer-marker.example")),
        ("eof", None),
    ] {
        let fixture = Fixture::new(/*client_secret*/ None).await?;
        let mut login = fixture.start().await?;
        let mut callback = login.callback(&fixture.issuer)?;
        let mut pairs: BTreeMap<_, _> = callback.query_pairs().into_owned().collect();
        pairs.remove(field);
        if let Some(value) = value {
            pairs.insert(field.to_string(), value.to_string());
        }
        callback.set_query(None);
        callback.query_pairs_mut().extend_pairs(pairs);
        if field == "eof" {
            drop(login.child.stdin.take());
        } else {
            login
                .child
                .stdin
                .as_mut()
                .context("missing stdin")?
                .write_all(format!("{callback}\n").as_bytes())
                .await?;
        }
        let output = login.finish().await?;
        assert!(
            !output.status.success(),
            "accepted invalid {field}: {value:?}"
        );
        let stderr = String::from_utf8(output.stderr)?;
        for secret in ["mock-code", "secret-issuer-marker"] {
            assert!(
                !stderr.contains(secret),
                "callback error leaked {secret}: {stderr}"
            );
        }
        assert!(!fixture.home.path().join(".credentials.json").exists());
        assert!(
            fixture
                .server
                .received_requests()
                .await
                .context("missing requests")?
                .iter()
                .all(|request| request.url.path() != "/token")
        );
    }
    Ok(())
}

#[tokio::test]
async fn no_browser_accepts_http_callback_while_terminal_input_remains_open() -> Result<()> {
    let fixture = Fixture::new(/*client_secret*/ None).await?;
    let mut login = fixture.start().await?;
    let callback = login.callback(&fixture.issuer)?;
    let stdin = login.child.stdin.take().context("missing stdin")?;
    let mut stream = TcpStream::connect((
        "127.0.0.1",
        callback.port().context("missing callback port")?,
    ))
    .await?;
    stream
        .write_all(
            format!(
                "GET {}?{} HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n",
                callback.path(),
                callback.query().context("missing callback query")?
            )
            .as_bytes(),
        )
        .await?;
    let output = login.finish().await?;
    drop(stdin);
    assert!(
        output.status.success(),
        "{}",
        String::from_utf8_lossy(&output.stderr)
    );
    assert!(fixture.home.path().join(".credentials.json").exists());
    Ok(())
}

#[tokio::test]
async fn no_browser_retries_rejected_discovered_scopes_with_pasteback() -> Result<()> {
    let fixture = Fixture::with_scopes(/*client_secret*/ None, &[]).await?;
    fixture
        .mount_resource_metadata(&["mcp.read"], /*challenge_scope*/ None)
        .await;
    let mut login = fixture.start().await?;
    assert!(
        login
            .authorization_url
            .query_pairs()
            .any(|(key, value)| key == "scope" && value == "mcp.read")
    );
    let mut callback = login.callback(&fixture.issuer)?;
    let mut params: BTreeMap<_, _> = callback.query_pairs().into_owned().collect();
    params.remove("code");
    params.insert("error".to_string(), "invalid_scope".to_string());
    callback.set_query(None);
    callback.query_pairs_mut().extend_pairs(params);
    login
        .child
        .stdin
        .as_mut()
        .context("missing stdin")?
        .write_all(format!("{callback}\n").as_bytes())
        .await?;
    login.authorization_url =
        read_authorization_url(&mut login.stdout, &mut login.transcript).await?;
    assert!(
        !login
            .authorization_url
            .query_pairs()
            .any(|(key, _)| key == "scope")
    );
    let callback = login.callback(&fixture.issuer)?;
    login
        .child
        .stdin
        .as_mut()
        .context("missing stdin")?
        .write_all(format!("{callback}\n").as_bytes())
        .await?;
    let output = login.finish().await?;
    assert!(
        output.status.success(),
        "{}",
        String::from_utf8_lossy(&output.stderr)
    );
    assert!(String::from_utf8(output.stdout)?.contains("Retrying without scopes"));
    let requests = fixture
        .server
        .received_requests()
        .await
        .context("missing requests")?;
    assert_eq!(
        requests
            .iter()
            .filter(|request| request.url.path() == "/token")
            .count(),
        1
    );
    Ok(())
}

#[cfg(unix)]
#[tokio::test]
async fn no_browser_sigint_cancels_token_exchange_after_pasted_or_http_callback() -> Result<()> {
    for callback_delivery in ["paste", "http"] {
        let fixture = Fixture::new(/*client_secret*/ None).await?;
        let request_started = std::sync::Arc::new(tokio::sync::Notify::new());
        let notify = std::sync::Arc::clone(&request_started);
        Mock::given(method("POST"))
            .and(path("/token"))
            .respond_with(move |_: &wiremock::Request| {
                notify.notify_one();
                ResponseTemplate::new(200)
                    .set_delay(Duration::from_secs(/*secs*/ 30))
                    .set_body_json(
                        json!({"access_token": "mock-access-token", "token_type": "Bearer"}),
                    )
            })
            .with_priority(/*priority*/ 1)
            .expect(1)
            .mount(&fixture.server)
            .await;
        let mut login = fixture.start().await?;
        let callback = login.callback(&fixture.issuer)?;
        // Keep the input pipe open after an HTTP callback, including while waiting for exit.
        let mut stdin = login.child.stdin.take().context("missing stdin")?;
        let mut http_callback = None;
        if callback_delivery == "http" {
            let mut stream = TcpStream::connect((
                "127.0.0.1",
                callback.port().context("missing callback port")?,
            ))
            .await?;
            stream
                .write_all(
                    format!(
                        "GET {}?{} HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n",
                        callback.path(),
                        callback.query().context("missing callback query")?,
                    )
                    .as_bytes(),
                )
                .await?;
            http_callback = Some(stream);
        } else {
            stdin.write_all(format!("{callback}\n").as_bytes()).await?;
        }
        timeout(TEST_TIMEOUT, request_started.notified())
            .await
            .context("token exchange never started")?;
        let child_pid = i32::try_from(login.child.id().context("login exited before SIGINT")?)?;
        // SAFETY: this positive PID belongs to the live child process owned by this test.
        assert_eq!(unsafe { libc::kill(child_pid, libc::SIGINT) }, 0);
        let output = timeout(Duration::from_secs(/*secs*/ 5), login.finish())
            .await
            .with_context(|| {
                format!("SIGINT did not cancel {callback_delivery} callback token exchange")
            })??;
        drop((stdin, http_callback));
        assert!(
            !output.status.success(),
            "SIGINT login unexpectedly succeeded"
        );
        assert!(String::from_utf8(output.stderr)?.contains("OAuth login cancelled"));
        assert!(!fixture.home.path().join(".credentials.json").exists());
    }
    Ok(())
}