zerogravity/src/standalone.rs

//! Standalone Language Server — spawn and lifecycle management.
//!
//! Launches an isolated LS instance as a child process that the proxy fully owns.
//! The standalone LS shares auth via the main extension server but has its own
//! HTTPS port, data directory, and cascade space. This means the real LS (the
//! one powering the user's coding session) is never touched.

use crate::constants;
use crate::proto;
use std::io::Write;
use std::net::TcpListener;
use std::process::{Child, Command, Stdio};
use tokio::time::{sleep, Duration};
use tracing::{debug, info};

/// Default path to the LS binary.
const LS_BINARY_PATH: &str =
    "/usr/share/antigravity/resources/app/extensions/antigravity/bin/language_server_linux_x64";

/// App root for ANTIGRAVITY_EDITOR_APP_ROOT env var.
const APP_ROOT: &str = "/usr/share/antigravity/resources/app";

/// Data directory for the standalone LS.
const DATA_DIR: &str = "/tmp/antigravity-standalone";

/// System user for UID-scoped iptables isolation.
const LS_USER: &str = "antigravity-ls";

/// A running standalone LS process.
pub struct StandaloneLS {
    child: Child,
    pub port: u16,
    pub csrf: String,
}

/// Config needed from the real (main) LS to bootstrap the standalone one.
pub struct MainLSConfig {
    pub extension_server_port: String,
    pub csrf: String,
}

/// Optional MITM proxy config for the standalone LS.
pub struct StandaloneMitmConfig {
    pub proxy_addr: String,   // e.g. "http://127.0.0.1:8742"
    pub ca_cert_path: String, // path to MITM CA .pem
}

impl StandaloneLS {
    /// Spawn a standalone LS process.
    ///
    /// Discovers the main LS's extension server port and CSRF token,
    /// picks a free port, builds init metadata, and launches the binary.
    ///
    /// If `mitm_config` is provided, sets HTTPS_PROXY and SSL_CERT_FILE
    /// so the LS routes LLM API calls through the MITM proxy.
    pub fn spawn(
        main_config: &MainLSConfig,
        mitm_config: Option<&StandaloneMitmConfig>,
    ) -> Result<Self, String> {
        let port = find_free_port()?;
        let ts = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        // Build init metadata protobuf
        let api_key = format!("standalone-api-key-{ts}");
        let session_id = format!("standalone-session-{ts}");
        let metadata = proto::build_init_metadata(
            &api_key,
            constants::antigravity_version(),
            constants::client_version(),
            &session_id,
            1, // DETECT_AND_USE_PROXY_ENABLED
        );

        // Setup data dir (mode 1777 so both current user and antigravity-ls can write)
        let gemini_dir = format!("{DATA_DIR}/.gemini");
        std::fs::create_dir_all(&gemini_dir)
            .map_err(|e| format!("Failed to create standalone data dir: {e}"))?;
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let _ = std::fs::set_permissions(DATA_DIR, std::fs::Permissions::from_mode(0o1777));
            let _ = std::fs::set_permissions(&gemini_dir, std::fs::Permissions::from_mode(0o1777));
        }

        // LS args — mirrors standalone-ls.sh but with correct params
        let args = vec![
            "-enable_lsp".to_string(),
            "-extension_server_port".to_string(),
            main_config.extension_server_port.clone(),
            "-csrf_token".to_string(),
            main_config.csrf.clone(),
            "-server_port".to_string(),
            port.to_string(),
            "-workspace_id".to_string(),
            format!("standalone_{ts}"),
            "-cloud_code_endpoint".to_string(),
            "https://daily-cloudcode-pa.googleapis.com".to_string(),
            "-app_data_dir".to_string(),
            "antigravity-standalone".to_string(),
            "-gemini_dir".to_string(),
            gemini_dir,
        ];

        info!(port, "Spawning standalone LS");
        debug!(?args, "LS args");

        // Build env vars for the LS process
        let mut env_vars: Vec<(String, String)> = vec![
            ("ANTIGRAVITY_EDITOR_APP_ROOT".into(), APP_ROOT.into()),
        ];

        // If MITM is enabled, add SSL + proxy env vars
        if let Some(mitm) = mitm_config {
            // Go's SSL_CERT_FILE replaces the entire system cert pool, so we
            // need a combined bundle: system CAs + our MITM CA
            // Write to /tmp — accessible by antigravity-ls user
            // (user's ~/.config/ is not traversable by other UIDs)
            let combined_ca_path = "/tmp/antigravity-mitm-combined-ca.pem".to_string();
            let system_ca = std::fs::read_to_string("/etc/ssl/certs/ca-certificates.crt")
                .unwrap_or_default();
            let mitm_ca = std::fs::read_to_string(&mitm.ca_cert_path)
                .map_err(|e| format!("Failed to read MITM CA cert: {e}"))?;
            std::fs::write(&combined_ca_path, format!("{system_ca}\n{mitm_ca}"))
                .map_err(|e| format!("Failed to write combined CA bundle: {e}"))?;
            // Make readable by antigravity-ls user
            #[cfg(unix)]
            {
                use std::os::unix::fs::PermissionsExt;
                let _ = std::fs::set_permissions(
                    &combined_ca_path,
                    std::fs::Permissions::from_mode(0o644),
                );
            }

            info!(
                proxy = %mitm.proxy_addr,
                ca = %combined_ca_path,
                "Setting MITM env vars on standalone LS (combined CA bundle)"
            );
            env_vars.push(("SSL_CERT_FILE".into(), combined_ca_path));
            env_vars.push(("SSL_CERT_DIR".into(), "/dev/null".into()));
            env_vars.push(("NODE_EXTRA_CA_CERTS".into(), mitm.ca_cert_path.clone()));
        }

        // Check if 'antigravity-ls' user exists for UID-scoped iptables isolation
        let use_sudo = has_ls_user();

        let mut cmd = if use_sudo {
            info!("Using UID isolation: spawning LS as 'antigravity-ls' user");
            // Build: sudo -n -u antigravity-ls -- /usr/bin/env VAR=val ... LS_BINARY args...
            let mut c = Command::new("sudo");
            c.args(["-n", "-u", LS_USER, "--", "/usr/bin/env"]);
            // Pass env vars as key=value args to /usr/bin/env
            for (k, v) in &env_vars {
                c.arg(format!("{k}={v}"));
            }
            c.arg(LS_BINARY_PATH);
            c.args(&args);
            c
        } else {
            debug!("No 'antigravity-ls' user found, spawning LS as current user");
            let mut c = Command::new(LS_BINARY_PATH);
            c.args(&args);
            for (k, v) in &env_vars {
                c.env(k, v);
            }
            c
        };

        cmd.stdin(Stdio::piped())
            .stdout(Stdio::null())
            .stderr(Stdio::null());

        let mut child = cmd
            .spawn()
            .map_err(|e| format!("Failed to spawn LS binary: {e}"))?;

        // Feed init metadata via stdin, then close it
        if let Some(mut stdin) = child.stdin.take() {
            stdin
                .write_all(&metadata)
                .map_err(|e| format!("Failed to write init metadata to stdin: {e}"))?;
            // stdin drops here → EOF
        }

        info!(pid = child.id(), port, "Standalone LS spawned");

        Ok(StandaloneLS {
            child,
            port,
            csrf: main_config.csrf.clone(),
        })
    }

    /// Wait for the standalone LS to be ready (accepting TCP connections).
    ///
    /// Retries up to `max_attempts` times with a 1-second delay between each.
    pub async fn wait_ready(&self, max_attempts: u32) -> Result<(), String> {
        info!(port = self.port, "Waiting for standalone LS to be ready...");

        for attempt in 1..=max_attempts {
            sleep(Duration::from_secs(1)).await;

            // Simple TCP connect check — if the LS is listening, it's ready
            match tokio::net::TcpStream::connect(format!("127.0.0.1:{}", self.port)).await {
                Ok(_) => {
                    info!(attempt, "Standalone LS is ready (accepting connections)");
                    return Ok(());
                }
                Err(e) => {
                    debug!(attempt, error = %e, "LS not ready yet");
                }
            }
        }

        Err(format!(
            "Standalone LS failed to become ready after {max_attempts} attempts on port {}",
            self.port
        ))
    }

    /// Check if the child process is still running.
    #[allow(dead_code)]
    pub fn is_alive(&mut self) -> bool {
        matches!(self.child.try_wait(), Ok(None))
    }

    /// Kill the standalone LS process.
    pub fn kill(&mut self) {
        info!("Killing standalone LS");
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

impl Drop for StandaloneLS {
    fn drop(&mut self) {
        self.kill();
    }
}

/// Discover only the extension_server_port and csrf_token from the running main LS.
///
/// This does NOT discover the HTTPS port — we don't need to talk to the real LS,
/// only steal its extension server connection info.
pub fn discover_main_ls_config() -> Result<MainLSConfig, String> {
    let pid = find_main_ls_pid()?;

    let cmdline = std::fs::read(format!("/proc/{pid}/cmdline"))
        .map_err(|e| format!("Can't read cmdline for PID {pid}: {e}"))?;
    let args: Vec<&[u8]> = cmdline.split(|&b| b == 0).collect();

    let mut csrf = String::new();
    let mut ext_port = String::new();

    for (i, arg) in args.iter().enumerate() {
        if let Ok(s) = std::str::from_utf8(arg) {
            match s {
                "--csrf_token" | "-csrf_token" => {
                    if let Some(next) = args.get(i + 1) {
                        if let Ok(val) = std::str::from_utf8(next) {
                            csrf = val.to_string();
                        }
                    }
                }
                "--extension_server_port" | "-extension_server_port" => {
                    if let Some(next) = args.get(i + 1) {
                        if let Ok(val) = std::str::from_utf8(next) {
                            ext_port = val.to_string();
                        }
                    }
                }
                _ => {}
            }
        }
    }

    if csrf.is_empty() {
        return Err("Could not find CSRF token from main LS".to_string());
    }
    if ext_port.is_empty() {
        return Err("Could not find extension_server_port from main LS".to_string());
    }

    info!(
        pid,
        ext_port,
        csrf_len = csrf.len(),
        "Discovered main LS config"
    );

    Ok(MainLSConfig {
        extension_server_port: ext_port,
        csrf,
    })
}

/// Find the PID of the main (real) LS process.
///
/// Checks `/proc/<pid>/exe` to ensure we find the actual LS binary,
/// not bash scripts that happen to mention `language_server_linux` in their args.
fn find_main_ls_pid() -> Result<String, String> {
    let proc = std::path::Path::new("/proc");
    if !proc.exists() {
        return Err("No /proc filesystem".to_string());
    }

    let entries = std::fs::read_dir(proc)
        .map_err(|e| format!("Cannot read /proc: {e}"))?;

    for entry in entries.flatten() {
        let name = entry.file_name();
        let name_str = name.to_string_lossy();
        // Only numeric dirs (PIDs)
        if !name_str.chars().all(|c| c.is_ascii_digit()) {
            continue;
        }
        let exe_link = entry.path().join("exe");
        if let Ok(target) = std::fs::read_link(&exe_link) {
            let target_str = target.to_string_lossy().to_string();
            let target_clean = target_str.trim_end_matches(" (deleted)");
            // Must be the actual LS binary, not a bash script
            if target_clean.contains("language_server_linux")
                || target_clean.contains("antigravity-language-server")
            {
                return Ok(name_str.to_string());
            }
        }
    }

    Err("No main LS process found — Antigravity must be running".to_string())
}

/// Find a free TCP port by binding to port 0.
fn find_free_port() -> Result<u16, String> {
    let listener =
        TcpListener::bind("127.0.0.1:0").map_err(|e| format!("Failed to bind for port: {e}"))?;
    listener
        .local_addr()
        .map(|a| a.port())
        .map_err(|e| format!("Failed to get port: {e}"))
}

/// Check if the dedicated LS system user exists.
///
/// When the user exists, the proxy spawns the LS as that UID so iptables
/// can scope the :443 redirect to only the standalone LS process.
fn has_ls_user() -> bool {
    Command::new("id")
        .args(["-u", LS_USER])
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_find_free_port() {
        let port = find_free_port().unwrap();
        assert!(port > 0);
        // Port should be available — try binding to it
        let listener = TcpListener::bind(format!("127.0.0.1:{port}"));
        assert!(listener.is_ok(), "Port {port} should be free");
    }
}