zed/crates/oauth_callback_server/src/oauth_callback_server.rs

//! Loopback OAuth 2.0 callback server and shared HTML response page.
//!
//! Used by Zed's OAuth-based sign-in flows (e.g. MCP servers, ChatGPT
//! Subscription) to receive the authorization code redirect from the user's
//! browser. The HTML response page rendered to the browser is kept alongside
//! the server so all OAuth callback presentation lives in one place.

/// Generate a styled HTML page for OAuth callback responses.
///
/// Returns a complete HTML document (no HTTP headers) with a centered card
/// layout styled to match Zed's dark theme. The `title` is rendered as a
/// heading and `message` as body text below it.
///
/// When `is_error` is true, a red X icon is shown instead of the green
/// checkmark.
pub fn oauth_callback_page(title: &str, message: &str, is_error: bool) -> String {
    let title = html_escape(title);
    let message = html_escape(message);
    let (icon_bg, icon_svg) = if is_error {
        (
            "#f38ba8",
            r#"<svg viewBox="0 0 24 24"><line x1="18" y1="6" x2="6" y2="18"/><line x1="6" y1="6" x2="18" y2="18"/></svg>"#,
        )
    } else {
        (
            "#a6e3a1",
            r#"<svg viewBox="0 0 24 24"><polyline points="20 6 9 17 4 12"/></svg>"#,
        )
    };
    format!(
        r#"<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>{title} — Zed</title>
<style>
  * {{ margin: 0; padding: 0; box-sizing: border-box; }}
  body {{
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
    background: #1e1e2e;
    color: #cdd6f4;
    display: flex;
    align-items: center;
    justify-content: center;
    min-height: 100vh;
    padding: 1rem;
  }}
  .card {{
    background: #313244;
    border-radius: 12px;
    padding: 2.5rem;
    max-width: 420px;
    width: 100%;
    text-align: center;
    box-shadow: 0 4px 24px rgba(0, 0, 0, 0.3);
  }}
  .icon {{
    width: 48px;
    height: 48px;
    margin: 0 auto 1.5rem;
    background: {icon_bg};
    border-radius: 50%;
    display: flex;
    align-items: center;
    justify-content: center;
  }}
  .icon svg {{
    width: 24px;
    height: 24px;
    stroke: #1e1e2e;
    stroke-width: 3;
    fill: none;
  }}
  h1 {{
    font-size: 1.25rem;
    font-weight: 600;
    margin-bottom: 0.75rem;
    color: #cdd6f4;
  }}
  p {{
    font-size: 0.925rem;
    line-height: 1.5;
    color: #a6adc8;
  }}
  .brand {{
    margin-top: 1.5rem;
    font-size: 0.8rem;
    color: #585b70;
    letter-spacing: 0.05em;
  }}
</style>
</head>
<body>
<div class="card">
  <div class="icon">
    {icon_svg}
  </div>
  <h1>{title}</h1>
  <p>{message}</p>
  <div class="brand">Zed</div>
</div>
</body>
</html>"#,
        title = title,
        message = message,
        icon_bg = icon_bg,
        icon_svg = icon_svg,
    )
}

fn html_escape(input: &str) -> String {
    let mut output = String::with_capacity(input.len());
    for ch in input.chars() {
        match ch {
            '&' => output.push_str("&amp;"),
            '<' => output.push_str("&lt;"),
            '>' => output.push_str("&gt;"),
            '"' => output.push_str("&quot;"),
            '\'' => output.push_str("&#x27;"),
            _ => output.push(ch),
        }
    }
    output
}

#[cfg(not(target_family = "wasm"))]
mod server {
    use super::oauth_callback_page;
    use anyhow::{Context as _, Result, anyhow};
    use std::str::FromStr;
    use std::time::Duration;
    use url::Url;

    /// Parsed OAuth callback parameters from the authorization server redirect.
    pub struct OAuthCallbackParams {
        pub code: String,
        pub state: String,
    }

    /// Configuration for the loopback OAuth callback server.
    ///
    /// OAuth servers compare `redirect_uri` against a per-client allow-list using
    /// exact string matching (RFC 6749 §3.1.2), so the `host`, `preferred_port`,
    /// and `path` here must match what's registered for the OAuth client_id.
    #[derive(Clone, Copy)]
    pub struct OAuthCallbackServerConfig {
        /// Host portion of the redirect URI (typically `127.0.0.1` or `localhost`).
        pub host: &'static str,
        /// Preferred port. Use `0` for an OS-assigned ephemeral port.
        pub preferred_port: u16,
        /// Optional fallback port if `preferred_port` is unavailable. Only used
        /// when `preferred_port` is non-zero.
        pub fallback_port: Option<u16>,
        /// Callback path on the redirect URI (e.g. `/callback`, `/auth/callback`).
        pub path: &'static str,
    }

    impl Default for OAuthCallbackServerConfig {
        fn default() -> Self {
            Self {
                host: "127.0.0.1",
                preferred_port: 0,
                fallback_port: None,
                path: "/callback",
            }
        }
    }

    impl OAuthCallbackParams {
        /// Parse the query string from a callback URL like
        /// `http://127.0.0.1:<port>/callback?code=...&state=...`.
        pub fn parse_query(query: &str) -> Result<Self> {
            let mut code: Option<String> = None;
            let mut state: Option<String> = None;
            let mut error: Option<String> = None;
            let mut error_description: Option<String> = None;

            for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
                match key.as_ref() {
                    "code" => {
                        if !value.is_empty() {
                            code = Some(value.into_owned());
                        }
                    }
                    "state" => {
                        if !value.is_empty() {
                            state = Some(value.into_owned());
                        }
                    }
                    "error" => {
                        if !value.is_empty() {
                            error = Some(value.into_owned());
                        }
                    }
                    "error_description" => {
                        if !value.is_empty() {
                            error_description = Some(value.into_owned());
                        }
                    }
                    _ => {}
                }
            }

            if let Some(error_code) = error {
                anyhow::bail!(
                    "OAuth authorization failed: {} ({})",
                    error_code,
                    error_description.as_deref().unwrap_or("no description")
                );
            }

            let code = code.ok_or_else(|| anyhow!("missing 'code' parameter in OAuth callback"))?;
            let state =
                state.ok_or_else(|| anyhow!("missing 'state' parameter in OAuth callback"))?;

            Ok(Self { code, state })
        }
    }

    /// How long to wait for the browser to complete the OAuth flow before giving
    /// up and releasing the loopback port.
    const OAUTH_CALLBACK_TIMEOUT: Duration = Duration::from_secs(2 * 60);

    /// Start a loopback HTTP server to receive the OAuth authorization callback.
    ///
    /// Binds to an ephemeral loopback port. Returns `(redirect_uri, callback_future)`.
    /// The caller should use the redirect URI in the authorization request, open
    /// the browser, then await the future to receive the callback.
    pub fn start_oauth_callback_server() -> Result<(
        String,
        futures::channel::oneshot::Receiver<Result<OAuthCallbackParams>>,
    )> {
        start_oauth_callback_server_with_config(OAuthCallbackServerConfig::default())
    }

    /// Start a loopback HTTP server with custom host/port/path.
    ///
    /// Use this when the OAuth client requires a specific redirect URI that the
    /// default ephemeral-port `http://127.0.0.1:<port>/callback` doesn't match.
    pub fn start_oauth_callback_server_with_config(
        config: OAuthCallbackServerConfig,
    ) -> Result<(
        String,
        futures::channel::oneshot::Receiver<Result<OAuthCallbackParams>>,
    )> {
        let server = bind_callback_server(&config)?;
        let port = server
            .server_addr()
            .to_ip()
            .ok_or_else(|| anyhow!("server not bound to a TCP address"))?
            .port();

        let redirect_uri = format!("http://{}:{}{}", config.host, port, config.path);
        let expected_path = config.path;

        let (tx, rx) = futures::channel::oneshot::channel();

        std::thread::spawn(move || {
            let deadline = std::time::Instant::now() + OAUTH_CALLBACK_TIMEOUT;

            loop {
                if tx.is_canceled() {
                    return;
                }
                let remaining = deadline.saturating_duration_since(std::time::Instant::now());
                if remaining.is_zero() {
                    return;
                }

                let timeout = remaining.min(Duration::from_millis(500));
                let Some(request) = (match server.recv_timeout(timeout) {
                    Ok(req) => req,
                    Err(_) => {
                        let _ = tx.send(Err(anyhow!("OAuth callback server I/O error")));
                        return;
                    }
                }) else {
                    continue;
                };

                let raw_url = request.url().to_string();
                let raw_path = raw_url.split('?').next().unwrap_or(&raw_url);
                if raw_path == CANCEL_PATH {
                    let response = tiny_http::Response::from_string("Cancelled")
                        .with_status_code(200)
                        .with_header(
                            tiny_http::Header::from_str("Content-Type: text/plain")
                                .expect("failed to construct response header"),
                        )
                        .with_header(
                            tiny_http::Header::from_str("Connection: close")
                                .expect("failed to construct response header"),
                        );
                    if let Err(err) = request.respond(response) {
                        log::error!("Failed to send OAuth cancel response: {}", err);
                    }
                    let _ = tx.send(Err(anyhow!(
                        "OAuth callback server was cancelled by another sign-in attempt"
                    )));
                    return;
                }

                let result = handle_oauth_callback_request(&request, expected_path);

                let (status_code, body) = match &result {
                    Ok(_) => (
                        200,
                        oauth_callback_page(
                            "Authorization Successful",
                            "You can close this tab and return to Zed.",
                            false,
                        ),
                    ),
                    Err(err) => {
                        log::error!("OAuth callback error: {}", err);
                        (
                            400,
                            oauth_callback_page(
                                "Authorization Failed",
                                "Something went wrong. Please try again from Zed.",
                                true,
                            ),
                        )
                    }
                };

                let response = tiny_http::Response::from_string(body)
                    .with_status_code(status_code)
                    .with_header(
                        tiny_http::Header::from_str("Content-Type: text/html")
                            .expect("failed to construct response header"),
                    )
                    .with_header(
                        tiny_http::Header::from_str("Keep-Alive: timeout=0,max=0")
                            .expect("failed to construct response header"),
                    );
                if let Err(err) = request.respond(response) {
                    log::error!("Failed to send OAuth callback response: {}", err);
                }

                let _ = tx.send(result);
                return;
            }
        });

        Ok((redirect_uri, rx))
    }

    fn handle_oauth_callback_request(
        request: &tiny_http::Request,
        expected_path: &str,
    ) -> Result<OAuthCallbackParams> {
        let url = Url::parse(&format!("http://localhost{}", request.url()))
            .context("malformed callback request URL")?;

        if url.path() != expected_path {
            anyhow::bail!("unexpected path in OAuth callback: {}", url.path());
        }

        let query = url
            .query()
            .ok_or_else(|| anyhow!("OAuth callback has no query string"))?;
        OAuthCallbackParams::parse_query(query)
    }

    /// Callback path reserved for evicting a previously-running OAuth callback
    /// server bound to the same port. Always handled, regardless of `config.path`.
    const CANCEL_PATH: &str = "/cancel";

    const BIND_MAX_ATTEMPTS: u32 = 10;
    const BIND_RETRY_DELAY: Duration = Duration::from_millis(200);
    const CANCEL_REQUEST_TIMEOUT: Duration = Duration::from_secs(2);

    fn bind_callback_server(config: &OAuthCallbackServerConfig) -> Result<tiny_http::Server> {
        // Ephemeral ports always succeed; skip the cancel-retry dance entirely.
        if config.preferred_port == 0 {
            let addr = format!("{}:0", config.host);
            return tiny_http::Server::http(&addr).map_err(|err| {
                anyhow!(err).context(format!(
                    "Failed to bind loopback listener for OAuth callback on {addr}"
                ))
            });
        }

        match try_bind_with_cancel(config.host, config.preferred_port) {
            Ok(server) => Ok(server),
            Err(primary_err) => {
                let Some(fallback_port) = config.fallback_port else {
                    return Err(primary_err.context(format!(
                        "Failed to bind loopback listener for OAuth callback on {}:{}",
                        config.host, config.preferred_port,
                    )));
                };
                log::warn!(
                    "OAuth callback port {}:{} unavailable; falling back to port {}",
                    config.host,
                    config.preferred_port,
                    fallback_port,
                );
                try_bind_with_cancel(config.host, fallback_port).map_err(|fallback_err| {
                    fallback_err.context(format!(
                        "Failed to bind loopback listener for OAuth callback on {}:{} or {}:{}",
                        config.host, config.preferred_port, config.host, fallback_port,
                    ))
                })
            }
        }
    }

    /// Attempts to bind to a fixed `host:port`. On `AddrInUse`, sends a single
    /// `GET /cancel` to the existing listener (to evict a previous OAuth flow
    /// from this or a compatible client) and retries.
    fn try_bind_with_cancel(host: &'static str, port: u16) -> Result<tiny_http::Server> {
        let addr = format!("{host}:{port}");
        let mut cancel_attempted = false;
        let mut last_err: Option<anyhow::Error> = None;

        for _ in 0..BIND_MAX_ATTEMPTS {
            match tiny_http::Server::http(&addr) {
                Ok(server) => return Ok(server),
                Err(err) => {
                    let is_addr_in_use = err
                        .downcast_ref::<std::io::Error>()
                        .map(|io_err| io_err.kind() == std::io::ErrorKind::AddrInUse)
                        .unwrap_or(false);

                    if !is_addr_in_use {
                        return Err(anyhow!(err).context(format!(
                            "Failed to bind loopback listener for OAuth callback on {addr}"
                        )));
                    }

                    if !cancel_attempted {
                        cancel_attempted = true;
                        if let Err(cancel_err) = send_cancel_request(host, port) {
                            log::warn!(
                                "Failed to cancel previous OAuth callback server on {addr}: {cancel_err}"
                            );
                        }
                    }

                    last_err = Some(anyhow!(err));
                    std::thread::sleep(BIND_RETRY_DELAY);
                }
            }
        }

        Err(last_err
            .unwrap_or_else(|| anyhow!("unknown bind error"))
            .context(format!(
                "OAuth callback port {addr} remained in use after {BIND_MAX_ATTEMPTS} attempts"
            )))
    }

    /// Sends `GET /cancel` to a listener on `host:port`, asking it to shut down.
    ///
    /// Best-effort: errors here are surfaced to the caller for logging but do
    /// not block the subsequent rebind attempt.
    fn send_cancel_request(host: &str, port: u16) -> std::io::Result<()> {
        use std::io::{Read as _, Write as _};
        use std::net::{TcpStream, ToSocketAddrs as _};

        let addr = format!("{host}:{port}")
            .to_socket_addrs()?
            .next()
            .ok_or_else(|| {
                std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    format!("could not resolve {host}:{port}"),
                )
            })?;
        let mut stream = TcpStream::connect_timeout(&addr, CANCEL_REQUEST_TIMEOUT)?;
        stream.set_read_timeout(Some(CANCEL_REQUEST_TIMEOUT))?;
        stream.set_write_timeout(Some(CANCEL_REQUEST_TIMEOUT))?;

        stream.write_all(b"GET /cancel HTTP/1.1\r\n")?;
        stream.write_all(format!("Host: {host}:{port}\r\n").as_bytes())?;
        stream.write_all(b"Connection: close\r\n\r\n")?;

        // Drain the response so the server can close cleanly. We don't care
        // about the body; errors here are harmless.
        let mut buf = [0u8; 64];
        let _ = stream.read(&mut buf);
        Ok(())
    }
}

#[cfg(not(target_family = "wasm"))]
pub use server::{
    OAuthCallbackParams, OAuthCallbackServerConfig, start_oauth_callback_server,
    start_oauth_callback_server_with_config,
};