magi-code 0.77.1

Repository-aware CLI coding agent for terminal work
Documentation
use super::*;

#[derive(Clone)]
pub(crate) struct McpOAuthLoginOutcome {
    pub(crate) authorization_url: String,
    pub(crate) redirect_uri: String,
    pub(crate) state: String,
    pub(crate) verifier: String,
    pub(crate) client_id: String,
    pub(crate) client_secret: Option<String>,
    pub(crate) metadata: AuthorizationServerMetadata,
    pub(crate) resource: String,
    pub(crate) authorization_server: String,
}

impl fmt::Debug for McpOAuthLoginOutcome {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("McpOAuthLoginOutcome")
            .field("authorization_url", &sanitize_url(&self.authorization_url))
            .field("redirect_uri", &self.redirect_uri)
            .field("state", &"[REDACTED]")
            .field("verifier", &"[REDACTED]")
            .field("client_id", &self.client_id)
            .field(
                "client_secret",
                &self.client_secret.as_ref().map(|_| "[REDACTED]"),
            )
            .field("metadata", &self.metadata)
            .field("resource", &self.resource)
            .field("authorization_server", &self.authorization_server)
            .finish()
    }
}
pub(super) fn generate_code_verifier() -> String {
    random_urlsafe(64)
}

pub(super) fn code_challenge(verifier: &str) -> String {
    let digest = Sha256::digest(verifier.as_bytes());
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest)
}

pub(super) fn generate_state() -> String {
    random_urlsafe(32)
}

pub(super) fn authorization_url(
    authorization_endpoint: &str,
    client_id: &str,
    redirect_uri: &str,
    challenge: &str,
    state: &str,
    resource: &str,
    scopes: &[String],
) -> String {
    let mut pairs = vec![
        ("response_type", "code"),
        ("client_id", client_id),
        ("redirect_uri", redirect_uri),
        ("code_challenge", challenge),
        ("code_challenge_method", "S256"),
        ("state", state),
        ("resource", resource),
    ];
    let scope = scopes.join(" ");
    if !scope.is_empty() {
        pairs.push(("scope", scope.as_str()));
    }
    let query = pairs
        .into_iter()
        .map(|(key, value)| format!("{}={}", pct(key), pct(value)))
        .collect::<Vec<_>>()
        .join("&");
    format!("{authorization_endpoint}?{query}")
}
pub(crate) fn prepare_login(
    server_url: &str,
    oauth: &McpOAuthConfig,
    redirect_uri: &str,
    client: &reqwest::blocking::Client,
) -> McpResult<McpOAuthLoginOutcome> {
    let protected = discover_protected_resource(server_url, client)?;
    let auth_server =
        select_authorization_server(oauth.authorization_server.as_deref(), protected.as_ref())?;
    let metadata = discover_authorization_server(&auth_server, client)?;
    let registration = if oauth.client_id.is_none() {
        match metadata.registration_endpoint.as_deref() {
            Some(endpoint) => Some(register_client(
                endpoint,
                vec![redirect_uri.to_string()],
                client,
            )?),
            None => {
                return Err(McpError::Config(
                    "server does not support dynamic client registration; configure client_id"
                        .to_string(),
                ));
            }
        }
    } else {
        None
    };
    let client_id = oauth
        .client_id
        .clone()
        .or_else(|| registration.as_ref().map(|r| r.client_id.clone()))
        .ok_or_else(|| McpError::Config("MCP OAuth client_id missing".to_string()))?;
    let client_secret = registration.and_then(|r| r.client_secret);
    let verifier = generate_code_verifier();
    let challenge = code_challenge(&verifier);
    let state = generate_state();
    let resource = protected
        .as_ref()
        .map(|p| p.resource.clone())
        .unwrap_or_else(|| server_url.to_string());
    let authorization_url = authorization_url(
        &metadata.authorization_endpoint,
        &client_id,
        redirect_uri,
        &challenge,
        &state,
        &resource,
        &oauth.scopes,
    );
    Ok(McpOAuthLoginOutcome {
        authorization_url,
        redirect_uri: redirect_uri.to_string(),
        state,
        verifier,
        client_id,
        client_secret,
        metadata,
        resource,
        authorization_server: auth_server,
    })
}
pub(crate) fn bind_callback_listener() -> McpResult<(TcpListener, String)> {
    let listener = TcpListener::bind("127.0.0.1:0").map_err(|_| {
        McpError::Transport("could not bind MCP OAuth callback on 127.0.0.1".to_string())
    })?;
    listener
        .set_nonblocking(true)
        .map_err(McpError::transport)?;
    let addr = listener.local_addr().map_err(McpError::transport)?;
    Ok((listener, format!("http://{addr}{CALLBACK_PATH}")))
}

pub(crate) fn capture_loopback_or_manual_code(
    listener: Option<TcpListener>,
    expected_state: &str,
    timeout: Duration,
    cancel: &AtomicBool,
    manual_rx: Option<&Receiver<String>>,
) -> McpResult<String> {
    let deadline = Instant::now() + timeout;
    loop {
        if cancel.load(Ordering::SeqCst) {
            return Err(McpError::Transport(
                "MCP OAuth login cancelled; token unchanged".to_string(),
            ));
        }
        if let Some(rx) = manual_rx {
            match rx.try_recv() {
                Ok(input) => return parse_manual_fallback_input(&input, expected_state),
                Err(std::sync::mpsc::TryRecvError::Empty)
                | Err(std::sync::mpsc::TryRecvError::Disconnected) => {}
            }
        }
        if let Some(listener) = &listener {
            match listener.accept() {
                Ok((mut stream, _)) => return handle_callback_stream(&mut stream, expected_state),
                Err(error) if error.kind() == io::ErrorKind::WouldBlock => {}
                Err(error) => {
                    return Err(McpError::Transport(format!(
                        "MCP OAuth callback failed: {error}"
                    )));
                }
            }
        }
        if Instant::now() >= deadline {
            return Err(McpError::Transport(
                "OAuth login timed out (no callback received within 5 minutes)".to_string(),
            ));
        }
        std::thread::sleep(Duration::from_millis(50));
    }
}

fn handle_callback_stream(stream: &mut TcpStream, expected_state: &str) -> McpResult<String> {
    stream
        .set_read_timeout(Some(CALLBACK_STREAM_TIMEOUT))
        .map_err(McpError::transport)?;
    stream
        .set_write_timeout(Some(CALLBACK_STREAM_TIMEOUT))
        .map_err(McpError::transport)?;
    let mut buf = [0_u8; 4096];
    let n = stream
        .read(&mut buf)
        .map_err(|_| McpError::Transport("MCP OAuth callback read failed".to_string()))?;
    let request = String::from_utf8_lossy(&buf[..n]);
    let first = request.lines().next().unwrap_or_default();
    let result = parse_callback_request_line(first, expected_state);
    let (status, body) = if result.is_ok() {
        (
            "200 OK",
            "MCP OAuth login complete. You can close this tab.",
        )
    } else {
        (
            "400 Bad Request",
            "MCP OAuth login failed. Return to your terminal.",
        )
    };
    let response = format!(
        "HTTP/1.1 {status}\r\ncontent-type: text/html\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
        body.len()
    );
    let _ = stream.write_all(response.as_bytes());
    result
}

pub(super) fn parse_callback_request_line(line: &str, expected_state: &str) -> McpResult<String> {
    let Some(target) = line
        .strip_prefix("GET ")
        .and_then(|rest| rest.split_whitespace().next())
    else {
        return Err(McpError::Transport(
            "MCP OAuth callback was malformed".to_string(),
        ));
    };
    parse_redirect_target(target, expected_state)
}

pub(super) fn parse_manual_fallback_input(input: &str, expected_state: &str) -> McpResult<String> {
    let trimmed = input.trim();
    if trimmed.is_empty() {
        return Err(McpError::Transport(
            "manual MCP OAuth fallback was empty".to_string(),
        ));
    }
    if trimmed.contains('?') || trimmed.starts_with("http://") || trimmed.starts_with("https://") {
        let target = trimmed
            .split_once("://")
            .and_then(|(_, rest)| rest.find('/').map(|idx| &rest[idx..]))
            .unwrap_or(trimmed);
        return parse_redirect_target(target, expected_state);
    }
    if trimmed.contains(char::is_whitespace) || trimmed.contains('&') || trimmed.contains('=') {
        return Err(McpError::Transport(
            "manual MCP OAuth fallback was malformed".to_string(),
        ));
    }
    Ok(trimmed.to_string())
}

fn parse_redirect_target(target: &str, expected_state: &str) -> McpResult<String> {
    let (path, query) = target.split_once('?').unwrap_or((target, ""));
    if path != CALLBACK_PATH {
        return Err(McpError::Transport(
            "MCP OAuth callback used an unexpected path".to_string(),
        ));
    }
    let params = parse_query(query)?;
    if params.iter().any(|(key, _)| key == "error") {
        return Err(McpError::Transport(
            "MCP OAuth provider rejected login".to_string(),
        ));
    }
    let state = params
        .iter()
        .find(|(key, _)| key == "state")
        .map(|(_, value)| value.as_str())
        .unwrap_or_default();
    if state != expected_state {
        return Err(McpError::Transport(
            "OAuth state mismatch — possible CSRF attack or stale login attempt".to_string(),
        ));
    }
    params
        .into_iter()
        .find(|(key, value)| key == "code" && !value.is_empty())
        .map(|(_, value)| value)
        .ok_or_else(|| {
            McpError::Transport("OAuth callback received without authorization code".to_string())
        })
}

pub(super) fn parse_query(query: &str) -> McpResult<Vec<(String, String)>> {
    query
        .split('&')
        .filter(|part| !part.is_empty())
        .map(|part| {
            let (k, v) = part.split_once('=').unwrap_or((part, ""));
            Ok((decode_pct(k)?, decode_pct(v)?))
        })
        .collect()
}

fn decode_pct(input: &str) -> McpResult<String> {
    let mut out = Vec::new();
    let bytes = input.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i] == b'%' {
            if i + 2 >= bytes.len() {
                return Err(McpError::Transport(
                    "MCP OAuth callback query contains malformed percent escape".to_string(),
                ));
            }
            let high = hex_digit(bytes[i + 1]).ok_or_else(|| {
                McpError::Transport(
                    "MCP OAuth callback query contains malformed percent escape".to_string(),
                )
            })?;
            let low = hex_digit(bytes[i + 2]).ok_or_else(|| {
                McpError::Transport(
                    "MCP OAuth callback query contains malformed percent escape".to_string(),
                )
            })?;
            out.push((high << 4) | low);
            i += 3;
        } else {
            out.push(if bytes[i] == b'+' { b' ' } else { bytes[i] });
            i += 1;
        }
    }
    String::from_utf8(out).map_err(|_| {
        McpError::Transport("MCP OAuth callback query contains invalid UTF-8".to_string())
    })
}

fn hex_digit(byte: u8) -> Option<u8> {
    match byte {
        b'0'..=b'9' => Some(byte - b'0'),
        b'a'..=b'f' => Some(byte - b'a' + 10),
        b'A'..=b'F' => Some(byte - b'A' + 10),
        _ => None,
    }
}

fn random_urlsafe(bytes: usize) -> String {
    let mut out = Vec::new();
    while out.len() < bytes {
        out.extend_from_slice(uuid::Uuid::new_v4().as_bytes());
    }
    out.truncate(bytes);
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(out)
}

fn pct(input: &str) -> String {
    let mut out = String::new();
    for b in input.bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                out.push(char::from(b))
            }
            _ => out.push_str(&format!("%{b:02X}")),
        }
    }
    out
}