bitrouter-providers 0.27.2

//! Routing shim installer.
//!
//! Generates a small wrapper script at `<shim_dir>/<agent>` (Unix) or
//! `<shim_dir>/<agent>.cmd` (Windows). The shim probes BitRouter health,
//! sets the agent-specific `*_BASE_URL` if reachable, then `exec`s the
//! real binary. When BitRouter is down, the shim falls through to the
//! real binary unchanged — no broken state.
//!
//! Shims are marked with [`SHIM_MARKER`]; install/uninstall refuse to
//! touch a file at the target path that lacks the marker.

use std::net::SocketAddr;
use std::path::{Path, PathBuf};

/// Marker line embedded in every shim. Used to recognise our own files
/// during refresh and uninstall.
pub const SHIM_MARKER: &str = "# bitrouter-shim:v1";

/// Marker for the Windows `.cmd` flavour. Cmd uses `::` for comments,
/// not `#`, so the marker line is different in form but serves the
/// same role.
pub const SHIM_MARKER_WINDOWS: &str = ":: bitrouter-shim:v1";

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Platform {
    Unix,
    Windows,
}

impl Platform {
    pub fn current() -> Self {
        if cfg!(windows) {
            Self::Windows
        } else {
            Self::Unix
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShimAction {
    Created,
    Updated,
    SkippedConflict,
}

#[derive(Debug)]
pub struct ShimEnv {
    pub var: String,
    pub value: String,
}

/// Map an agent id to the `*_BASE_URL` it consumes. Returns `None` for
/// agents we don't have a mapping for — caller should skip rather than
/// guess, since installing a shim that doesn't set anything useful is
/// only blast radius without benefit.
pub fn shim_env_for(agent_id: &str, listen: SocketAddr) -> Option<ShimEnv> {
    let base = format!("http://{listen}");
    let (var, value) = match agent_id {
        "claude" | "claude-acp" | "claude-code-acp" => ("ANTHROPIC_BASE_URL", format!("{base}/v1")),
        "codex" | "codex-acp" => ("OPENAI_BASE_URL", format!("{base}/v1")),
        "gemini" | "gemini-acp" => ("GOOGLE_AI_BASE_URL", format!("{base}/v1beta")),
        _ => return None,
    };
    Some(ShimEnv {
        var: var.to_owned(),
        value,
    })
}

/// Where the shim file lives for a given agent name on the given platform.
pub fn shim_path_for(platform: Platform, shim_dir: &Path, name: &str) -> PathBuf {
    match platform {
        Platform::Unix => shim_dir.join(name),
        Platform::Windows => shim_dir.join(format!("{name}.cmd")),
    }
}

/// Render the shim script body. Pure string; testable on any host.
pub fn render_shim(
    platform: Platform,
    real_binary: &Path,
    listen: SocketAddr,
    env: &ShimEnv,
) -> String {
    match platform {
        Platform::Unix => render_unix_shim(real_binary, listen, env),
        Platform::Windows => render_windows_shim(real_binary, listen, env),
    }
}

fn render_unix_shim(real_binary: &Path, listen: SocketAddr, env: &ShimEnv) -> String {
    let real = shell_single_quote(&real_binary.display().to_string());
    let value = shell_single_quote(&env.value);
    format!(
        "#!/usr/bin/env bash\n\
         {SHIM_MARKER}\n\
         # Auto-generated by `bitrouter init`. Probes BitRouter health and\n\
         # routes through it when reachable; otherwise falls through to the\n\
         # real binary so the command keeps working.\n\
         REAL={real}\n\
         LISTEN={listen}\n\
         if curl -fsS --max-time 1 \"http://${{LISTEN}}/health\" >/dev/null 2>&1; then\n\
         \x20   export {var}={value}\n\
         fi\n\
         exec \"$REAL\" \"$@\"\n",
        var = env.var,
    )
}

fn render_windows_shim(real_binary: &Path, listen: SocketAddr, env: &ShimEnv) -> String {
    format!(
        "@echo off\r\n\
         {SHIM_MARKER_WINDOWS}\r\n\
         :: Auto-generated by `bitrouter init`. Probes BitRouter health and\r\n\
         :: routes through it when reachable; otherwise falls through to the\r\n\
         :: real binary so the command keeps working.\r\n\
         setlocal\r\n\
         set \"REAL={real}\"\r\n\
         set \"LISTEN={listen}\"\r\n\
         curl.exe -fsS --max-time 1 \"http://%LISTEN%/health\" >nul 2>&1\r\n\
         if %ERRORLEVEL% EQU 0 set \"{var}={value}\"\r\n\
         \"%REAL%\" %*\r\n\
         exit /b %ERRORLEVEL%\r\n",
        real = real_binary.display(),
        var = env.var,
        value = env.value,
    )
}

/// Single-quote a string for safe interpolation into a bash script.
fn shell_single_quote(s: &str) -> String {
    let escaped = s.replace('\'', r"'\''");
    format!("'{escaped}'")
}

/// Install (or refresh) a shim. Returns [`ShimAction::SkippedConflict`]
/// if a non-bitrouter file already exists at `shim_path` — the caller
/// can surface this to the user rather than blowing the file away.
pub fn install_shim(
    platform: Platform,
    shim_path: &Path,
    real_binary: &Path,
    listen: SocketAddr,
    env: &ShimEnv,
) -> Result<ShimAction, String> {
    if let Some(parent) = shim_path.parent() {
        std::fs::create_dir_all(parent).map_err(|e| format!("create {}: {e}", parent.display()))?;
    }

    let action = match std::fs::read_to_string(shim_path) {
        Ok(existing) if is_bitrouter_shim(&existing) => ShimAction::Updated,
        Ok(_) => return Ok(ShimAction::SkippedConflict),
        Err(_) => ShimAction::Created,
    };

    let body = render_shim(platform, real_binary, listen, env);
    std::fs::write(shim_path, body).map_err(|e| format!("write {}: {e}", shim_path.display()))?;

    if matches!(platform, Platform::Unix) {
        chmod_executable(shim_path)?;
    }

    Ok(action)
}

#[cfg(unix)]
fn chmod_executable(path: &Path) -> Result<(), String> {
    use std::os::unix::fs::PermissionsExt;
    let mut perms = std::fs::metadata(path)
        .map_err(|e| format!("stat shim: {e}"))?
        .permissions();
    perms.set_mode(0o755);
    std::fs::set_permissions(path, perms).map_err(|e| format!("chmod shim: {e}"))?;
    Ok(())
}

#[cfg(not(unix))]
fn chmod_executable(_path: &Path) -> Result<(), String> {
    // .cmd files are executable by extension via PATHEXT — no mode bit
    // to set on Windows.
    Ok(())
}

/// Remove a shim, but only if it carries our marker. Returns true if a
/// bitrouter shim was removed; false if absent or hand-written by the
/// user.
pub fn uninstall_shim(shim_path: &Path) -> Result<bool, String> {
    let Ok(existing) = std::fs::read_to_string(shim_path) else {
        return Ok(false);
    };
    if !is_bitrouter_shim(&existing) {
        return Ok(false);
    }
    std::fs::remove_file(shim_path).map_err(|e| format!("remove shim: {e}"))?;
    Ok(true)
}

/// True if the file at `shim_path` exists and is a bitrouter shim.
pub fn is_installed(shim_path: &Path) -> bool {
    std::fs::read_to_string(shim_path)
        .map(|s| is_bitrouter_shim(&s))
        .unwrap_or(false)
}

fn is_bitrouter_shim(body: &str) -> bool {
    body.contains(SHIM_MARKER) || body.contains(SHIM_MARKER_WINDOWS)
}

/// Resolve the real binary for an agent name on PATH, **excluding**
/// `exclude_dir` (the directory the shim itself lives in). On Windows,
/// also tries each `PATHEXT` suffix so the locator finds `claude.exe`,
/// `claude.cmd`, etc.
pub fn locate_real_binary(name: &str, exclude_dir: &Path) -> Option<PathBuf> {
    let path_var = std::env::var_os("PATH")?;
    let extensions = path_extensions();
    for dir in std::env::split_paths(&path_var) {
        if dir == exclude_dir {
            continue;
        }
        let bare = dir.join(name);
        if bare.is_file() {
            return Some(bare);
        }
        for ext in &extensions {
            let candidate = dir.join(format!("{name}{ext}"));
            if candidate.is_file() {
                return Some(candidate);
            }
        }
    }
    None
}

#[cfg(windows)]
fn path_extensions() -> Vec<String> {
    std::env::var("PATHEXT")
        .unwrap_or_else(|_| ".COM;.EXE;.BAT;.CMD".to_string())
        .split(';')
        .filter(|s| !s.is_empty())
        .map(|s| s.to_lowercase())
        .collect()
}

#[cfg(not(windows))]
fn path_extensions() -> Vec<String> {
    Vec::new()
}

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

    fn listen() -> SocketAddr {
        "127.0.0.1:8787".parse().expect("static addr parses")
    }

    fn anth_env() -> ShimEnv {
        ShimEnv {
            var: "ANTHROPIC_BASE_URL".to_owned(),
            value: "http://127.0.0.1:8787/v1".to_owned(),
        }
    }

    #[test]
    fn env_mapping_per_agent() {
        let l = listen();
        assert_eq!(
            shim_env_for("claude", l).map(|e| e.var),
            Some("ANTHROPIC_BASE_URL".to_owned())
        );
        assert_eq!(
            shim_env_for("codex", l).map(|e| e.var),
            Some("OPENAI_BASE_URL".to_owned())
        );
        assert_eq!(
            shim_env_for("gemini-acp", l).map(|e| e.var),
            Some("GOOGLE_AI_BASE_URL".to_owned())
        );
        assert!(shim_env_for("unknown-agent", l).is_none());
    }

    #[test]
    fn unix_render_embeds_absolute_real_path_and_marker() {
        let body = render_shim(
            Platform::Unix,
            Path::new("/usr/local/bin/claude"),
            listen(),
            &anth_env(),
        );
        assert!(body.contains(SHIM_MARKER));
        assert!(body.contains("'/usr/local/bin/claude'"));
        assert!(body.contains("export ANTHROPIC_BASE_URL='http://127.0.0.1:8787/v1'"));
        assert!(body.contains("curl -fsS --max-time 1"));
        assert!(body.contains("exec \"$REAL\" \"$@\""));
    }

    #[test]
    fn unix_render_shell_quotes_special_chars() {
        let body = render_shim(
            Platform::Unix,
            Path::new("/path with space/it's/claude"),
            listen(),
            &anth_env(),
        );
        assert!(body.contains(r"'/path with space/it'\''s/claude'"));
    }

    #[test]
    fn windows_render_uses_cmd_idioms() {
        let body = render_shim(
            Platform::Windows,
            Path::new(r"C:\Program Files\Claude\claude.exe"),
            listen(),
            &anth_env(),
        );
        assert!(body.contains("@echo off"));
        assert!(body.contains(SHIM_MARKER_WINDOWS));
        assert!(body.contains("setlocal"));
        assert!(body.contains("curl.exe -fsS --max-time 1"));
        assert!(body.contains(r"C:\Program Files\Claude\claude.exe"));
        assert!(body.contains("ANTHROPIC_BASE_URL=http://127.0.0.1:8787/v1"));
        assert!(body.contains("\r\n"));
        assert!(body.contains("exit /b %ERRORLEVEL%"));
    }

    #[test]
    fn shim_path_extension_matches_platform() {
        let dir = Path::new("/tmp/x");
        assert_eq!(
            shim_path_for(Platform::Unix, dir, "claude"),
            dir.join("claude")
        );
        assert_eq!(
            shim_path_for(Platform::Windows, dir, "claude"),
            dir.join("claude.cmd")
        );
    }

    #[test]
    fn install_creates_executable_shim() -> Result<(), String> {
        let dir = TempDir::new().map_err(|e| e.to_string())?;
        let shim = dir.path().join("claude");
        let action = install_shim(
            Platform::Unix,
            &shim,
            Path::new("/bin/echo"),
            listen(),
            &anth_env(),
        )?;
        assert_eq!(action, ShimAction::Created);
        assert!(shim.exists());
        assert!(is_installed(&shim));

        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let mode = std::fs::metadata(&shim)
                .map_err(|e| e.to_string())?
                .permissions()
                .mode()
                & 0o777;
            assert_eq!(mode, 0o755);
        }
        Ok(())
    }

    #[test]
    fn install_refreshes_existing_bitrouter_shim() -> Result<(), String> {
        let dir = TempDir::new().map_err(|e| e.to_string())?;
        let shim = dir.path().join("claude");
        install_shim(
            Platform::Unix,
            &shim,
            Path::new("/bin/echo"),
            listen(),
            &anth_env(),
        )?;
        let updated_env = ShimEnv {
            var: "ANTHROPIC_BASE_URL".to_owned(),
            value: "http://127.0.0.1:9999/v1".to_owned(),
        };
        let action = install_shim(
            Platform::Unix,
            &shim,
            Path::new("/bin/echo"),
            listen(),
            &updated_env,
        )?;
        assert_eq!(action, ShimAction::Updated);
        let body = std::fs::read_to_string(&shim).map_err(|e| e.to_string())?;
        assert!(body.contains("http://127.0.0.1:9999/v1"));
        Ok(())
    }

    #[test]
    fn install_refuses_to_clobber_foreign_file() -> Result<(), String> {
        let dir = TempDir::new().map_err(|e| e.to_string())?;
        let shim = dir.path().join("claude");
        std::fs::write(&shim, "#!/bin/sh\necho hand-written\n").map_err(|e| e.to_string())?;
        let action = install_shim(
            Platform::Unix,
            &shim,
            Path::new("/bin/echo"),
            listen(),
            &anth_env(),
        )?;
        assert_eq!(action, ShimAction::SkippedConflict);
        let body = std::fs::read_to_string(&shim).map_err(|e| e.to_string())?;
        assert!(body.contains("hand-written"));
        Ok(())
    }

    #[test]
    fn uninstall_only_removes_marked_shim() -> Result<(), String> {
        let dir = TempDir::new().map_err(|e| e.to_string())?;
        let shim = dir.path().join("claude");

        std::fs::write(&shim, "#!/bin/sh\n").map_err(|e| e.to_string())?;
        assert!(!uninstall_shim(&shim)?);
        assert!(shim.exists());

        // Remove the foreign file so install can write a real shim.
        std::fs::remove_file(&shim).map_err(|e| e.to_string())?;
        install_shim(
            Platform::Unix,
            &shim,
            Path::new("/bin/echo"),
            listen(),
            &anth_env(),
        )?;
        assert!(uninstall_shim(&shim)?);
        assert!(!shim.exists());
        Ok(())
    }

    #[cfg(unix)]
    #[test]
    fn shim_falls_through_when_bitrouter_unreachable() -> Result<(), String> {
        use std::os::unix::fs::PermissionsExt;
        use std::process::Command;

        let dir = TempDir::new().map_err(|e| e.to_string())?;
        let shim = dir.path().join("claude");

        let real = dir.path().join("real-claude.sh");
        std::fs::write(
            &real,
            "#!/usr/bin/env bash\n\
             if [ -n \"$ANTHROPIC_BASE_URL\" ]; then echo routed; \
             else echo direct; fi\n",
        )
        .map_err(|e| e.to_string())?;
        std::fs::set_permissions(&real, std::fs::Permissions::from_mode(0o755))
            .map_err(|e| e.to_string())?;

        // Port 1 — nothing listens here. The shim's curl probe must fail
        // and fall through to the real binary unchanged.
        let dead: SocketAddr = "127.0.0.1:1".parse().expect("static addr parses");
        install_shim(Platform::Unix, &shim, &real, dead, &anth_env())?;

        let out = Command::new(&shim)
            .output()
            .map_err(|e| format!("spawn shim: {e}"))?;
        let stdout = String::from_utf8_lossy(&out.stdout);
        assert!(
            stdout.contains("direct"),
            "expected fallback, got: {stdout}"
        );
        Ok(())
    }
}