kz_proxy/

lib.rs

//! kz-proxy: run a subprocess with masked secrets and an HTTP proxy that rewrites tokens.
//!
//! The main type is [`Sandbox`]: build it from a [`SandboxConfig`] (with optional [`SecretMapping`]s,
//! [`StringMapping`]s, and connection allow/deny rules), then call [`Sandbox::run`] to execute a shell
//! command with masked env vars and proxied HTTP that rewrites tokens to real secrets.
//!
//! The proxy is implemented with the [hyper](https://github.com/hyperium/hyper) stack so that
//! HTTP parsing, Content-Length, chunked encoding, and CONNECT tunneling follow RFC 7230/9110.

pub mod backend;
mod enforce;
mod runner;
pub use enforce::Enforce;
#[cfg(target_os = "macos")]
mod enforce_docker;
#[cfg(target_os = "linux")]
mod enforce_linux;
mod mitm;

use std::collections::HashMap;
use std::sync::Arc;

use bytes::Bytes;
use http_body_util::{combinators::BoxBody, BodyExt, Full};
use hyper::service::service_fn;
use hyper::{Method, Request, Response};
use hyper_util::rt::TokioIo;
use regex::Regex;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;

/// Mapping from environment variable name to the real secret value.
/// The sandbox will inject a masked token into the subprocess env and
/// the proxy will replace that token with this value in outgoing HTTP requests.
/// Real values must not contain CR, LF, or NUL (validated at run time).
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct SecretMapping {
    pub var: String,
    pub value: String,
}

impl SecretMapping {
    /// Create a secret mapping from env var name to real value.
    pub fn new(var: impl Into<String>, value: impl Into<String>) -> Self {
        Self {
            var: var.into(),
            value: value.into(),
        }
    }
}

/// Mapping from a unique string identifier (token) to the actual value.
/// The proxy will replace occurrences of the token with the value in URIs, headers, and body.
/// Used when the process already uses placeholders; no env injection. Values must not contain CR, LF, or NUL.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct StringMapping {
    pub token: String,
    pub value: String,
}

impl StringMapping {
    /// Create a string mapping from token to value.
    pub fn new(token: impl Into<String>, value: impl Into<String>) -> Self {
        Self {
            token: token.into(),
            value: value.into(),
        }
    }
}

/// Pattern for matching a host in connection allow/deny rules.
#[derive(Clone, Debug)]
pub enum HostPattern {
    /// Exact host string match (e.g. `example.com`).
    Exact(String),
    /// Regex match over the host (e.g. `^api\.example\.com$`). Pattern string is kept for serialization.
    Regex { pattern: String, re: Regex },
}

impl HostPattern {
    /// Create an exact host pattern.
    pub fn exact(host: impl Into<String>) -> Self {
        Self::Exact(host.into())
    }

    /// Create a regex host pattern. Returns an error if the pattern is invalid.
    pub fn regex(pattern: &str) -> Result<Self, regex::Error> {
        let re = Regex::new(pattern)?;
        Ok(Self::Regex {
            pattern: pattern.to_string(),
            re,
        })
    }

    /// Returns true if the given host matches this pattern.
    pub fn matches(&self, host: &str) -> bool {
        match self {
            HostPattern::Exact(s) => s == host,
            HostPattern::Regex { re, .. } => re.is_match(host),
        }
    }
}

/// Allow or deny outbound connections to a host (or hosts matching a regex).
/// Policies are evaluated in order; the first matching policy wins. If no policy matches, the connection is allowed.
#[derive(Clone, Debug)]
pub struct ConnectionPolicy {
    pub pattern: HostPattern,
    /// If true, allow the connection; if false, deny (proxy returns an error to the client).
    pub allow: bool,
}

impl ConnectionPolicy {
    /// Create an allow rule for the given host pattern.
    pub fn allow(pattern: HostPattern) -> Self {
        Self {
            pattern,
            allow: true,
        }
    }

    /// Create a deny rule for the given host pattern.
    pub fn deny(pattern: HostPattern) -> Self {
        Self {
            pattern,
            allow: false,
        }
    }
}

/// Configuration for the sandbox: secrets, string mappings, connection allow/deny, and proxy options.
#[derive(Clone, Debug)]
pub struct SandboxConfig {
    /// Env-based secret mappings (env var name → value). Default empty.
    pub secrets: Vec<SecretMapping>,
    /// String token → value mappings. Default empty.
    pub strings: Vec<StringMapping>,
    /// Allow/deny rules for outbound connections (host or host regex). Evaluated in order; first match wins; no match = allow. Default empty.
    pub connections: Vec<ConnectionPolicy>,
    /// If true, CONNECT to private/local addresses (e.g. 127.0.0.1) is allowed. For testing only; default false.
    pub allow_private_connect: bool,
    /// Optional path to PEM file with extra CA cert(s) to trust for upstream (e.g. self-signed server).
    pub upstream_ca: Option<std::path::PathBuf>,
    /// If true, run the child in a sandbox backend (Firecracker on Linux, Docker on macOS) so all traffic is forced through the proxy. Default true.
    pub force_traffic_through_proxy: bool,
    /// Sandbox backend when force_traffic_through_proxy is true. Default: Firecracker on Linux, Docker on macOS.
    pub sandbox_backend: Option<crate::backend::SandboxBackend>,
}

impl Default for SandboxConfig {
    fn default() -> Self {
        Self {
            secrets: Vec::new(),
            strings: Vec::new(),
            connections: Vec::new(),
            allow_private_connect: false,
            upstream_ca: None,
            force_traffic_through_proxy: true,
            sandbox_backend: None,
        }
    }
}

/// Sandbox for running a subprocess with masked secrets and an HTTP proxy that rewrites tokens.
///
/// HTTP requests are forwarded with token replacement. HTTPS (CONNECT) is always handled by MITM:
/// the proxy decrypts, rewrites tokens, and re-encrypts to upstream. The subprocess must trust our CA
/// (we set `SSL_CERT_FILE`). For self-signed or custom upstream servers, set `upstream_ca` on [`SandboxConfig`].
#[derive(Clone, Debug)]
pub struct Sandbox {
    config: SandboxConfig,
}

impl Sandbox {
    /// Create a sandbox from the given config.
    pub fn new(config: SandboxConfig) -> Self {
        Self { config }
    }

    /// Run a shell command in the sandbox: start an HTTP proxy that rewrites masked tokens
    /// and string tokens to real values, set subprocess env with masked tokens (if any) and
    /// HTTP_PROXY/HTTPS_PROXY, then wait for the process to exit.
    pub async fn run(
        &self,
        cmd: &str,
    ) -> Result<std::process::ExitStatus, Box<dyn std::error::Error + Send + Sync>> {
        run_impl(
            cmd,
            self.config.secrets.clone(),
            self.config.strings.clone(),
            self.config.allow_private_connect,
            self.config.upstream_ca.clone(),
            self.config.connections.clone(),
            self.config.force_traffic_through_proxy,
            self.config.sandbox_backend,
        )
        .await
    }
}

/// Returns true if an outbound connection to `host` is allowed by the given policies.
/// If policies is None or empty, returns true (allow all). Otherwise first matching policy wins; no match = allow.
pub(crate) fn connection_allowed(host: &str, policies: Option<&[ConnectionPolicy]>) -> bool {
    let Some(policies) = policies else {
        return true;
    };
    if policies.is_empty() {
        return true;
    }
    for p in policies {
        if p.pattern.matches(host) {
            return p.allow;
        }
    }
    true
}

/// Characters that must not appear in real (or masked) secrets to avoid header/body injection.
const FORBIDDEN_IN_SECRET: &[u8] = b"\r\n\0";

fn validate_secret(value: &str) -> Result<(), String> {
    if value.bytes().any(|b| FORBIDDEN_IN_SECRET.contains(&b)) {
        return Err("secret must not contain CR, LF, or NUL".to_string());
    }
    Ok(())
}

/// Run as the runner (proxy + child) inside the sandbox backend (Docker/Firecracker). Reads RunnerConfig JSON from stdin.
/// On Linux, brings up loopback first. Used when the process is invoked as `blinders --runner` inside the container/VM.
#[cfg(target_os = "linux")]
pub fn run_as_linux_runner(
    config_json: String,
) -> Result<std::process::ExitStatus, Box<dyn std::error::Error + Send + Sync>> {
    let config: runner::RunnerConfig =
        serde_json::from_str(&config_json).map_err(|e| format!("runner config JSON: {}", e))?;
    // In Docker, loopback is already up; only bring it up when needed (e.g. Firecracker VM).
    if !std::path::Path::new("/.dockerenv").exists() {
        enforce_linux::bring_up_loopback()?;
    }
    let connection_policies =
        runner::runner_rules_to_connection_policies(&config.connection_policies)?;
    let upstream_ca = config.upstream_ca.map(std::path::PathBuf::from);
    let rt = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {}", e))?;
    rt.block_on(run_impl_inner(
        &config.cmd,
        config.secret_mappings,
        config.string_mappings,
        config.allow_private_connect,
        upstream_ca,
        connection_policies,
        false, // already in isolate, no need to force again
        &enforce::ENV_ONLY_ENFORCER,
    ))
}

#[allow(clippy::too_many_arguments)]
async fn run_impl(
    cmd: &str,
    secret_mappings: Vec<SecretMapping>,
    string_mappings: Vec<StringMapping>,
    allow_private_connect: bool,
    upstream_ca: Option<std::path::PathBuf>,
    connection_policies: Vec<ConnectionPolicy>,
    force_traffic_through_proxy: bool,
    sandbox_backend: Option<backend::SandboxBackend>,
) -> Result<std::process::ExitStatus, Box<dyn std::error::Error + Send + Sync>> {
    let enforcer = enforce::enforcer_for(sandbox_backend, force_traffic_through_proxy)?;
    if force_traffic_through_proxy {
        if let Some(status) = enforcer.maybe_spawn_runner(
            cmd,
            &secret_mappings,
            &string_mappings,
            allow_private_connect,
            &upstream_ca,
            &connection_policies,
        )? {
            return Ok(status);
        }
    }

    run_impl_inner(
        cmd,
        secret_mappings,
        string_mappings,
        allow_private_connect,
        upstream_ca,
        connection_policies,
        force_traffic_through_proxy,
        enforcer,
    )
    .await
}

#[allow(clippy::too_many_arguments)]
async fn run_impl_inner(
    cmd: &str,
    secret_mappings: Vec<SecretMapping>,
    string_mappings: Vec<StringMapping>,
    allow_private_connect: bool,
    upstream_ca: Option<std::path::PathBuf>,
    connection_policies: Vec<ConnectionPolicy>,
    force_traffic_through_proxy: bool,
    enforcer: &'static dyn enforce::Enforce,
) -> Result<std::process::ExitStatus, Box<dyn std::error::Error + Send + Sync>> {
    if secret_mappings.is_empty() && string_mappings.is_empty() {
        return Err("sandbox requires at least one secret mapping or string mapping".into());
    }

    for m in &secret_mappings {
        validate_secret(&m.value).map_err(|e| format!("{}: {}", m.var, e))?;
    }
    for m in &string_mappings {
        validate_secret(&m.value).map_err(|e| format!("{}: {}", m.token, e))?;
    }

    let mut env_vars_with_masked: Vec<(String, String)> = Vec::with_capacity(secret_mappings.len());
    let mut proxy_map: HashMap<String, String> =
        HashMap::with_capacity(secret_mappings.len() + string_mappings.len());

    for m in &secret_mappings {
        let masked = format!(
            "{}-{}",
            m.var.to_lowercase().replace('_', "-"),
            uuid::Uuid::new_v4()
        );
        env_vars_with_masked.push((m.var.clone(), masked.clone()));
        proxy_map.insert(masked, m.value.clone());
    }
    for m in &string_mappings {
        proxy_map.insert(m.token.clone(), m.value.clone());
    }

    // Deterministic replacement order: longest token first so substrings are not partially replaced.
    let replacement_order: Vec<(String, String)> = {
        let mut v: Vec<_> = proxy_map.into_iter().collect();
        v.sort_by(|a, b| b.0.len().cmp(&a.0.len()));
        v
    };
    let token_map = Arc::new(replacement_order);

    let (mitm_config, ssl_cert_file) = {
        let (config, ca_pem) =
            mitm::MitmConfig::new(upstream_ca).map_err(|e| format!("MITM config: {}", e))?;
        let temp = std::env::temp_dir().join(format!("blinders-ca-{}.pem", uuid::Uuid::new_v4()));
        std::fs::write(&temp, &ca_pem).map_err(|e| format!("write CA cert: {}", e))?;
        (Arc::new(config), temp)
    };

    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let port = listener.local_addr()?.port();
    let proxy_url = format!("http://127.0.0.1:{}", port);

    let (shutdown_tx, mut shutdown_rx) = mpsc::channel::<()>(1);
    let server_handle = tokio::spawn(async move {
        loop {
            tokio::select! {
                _ = shutdown_rx.recv() => break,
                accept_result = listener.accept() => {
                    let (stream, _) = match accept_result {
                        Ok(x) => x,
                        Err(_) => continue,
                    };
                    let token_map = Arc::clone(&token_map);
                    let allow_private = allow_private_connect;
                    let mitm_config = Arc::clone(&mitm_config);
                    let connection_policies = connection_policies.clone();
                    tokio::spawn(async move {
                        let io = TokioIo::new(stream);
                        let service = service_fn(move |req| {
                            let token_map = Arc::clone(&token_map);
                            let mitm_config = Arc::clone(&mitm_config);
                            let connection_policies = connection_policies.clone();
                            async move { proxy_handler(req, token_map, allow_private, mitm_config, connection_policies).await }
                        });
                        let conn = hyper::server::conn::http1::Builder::new()
                            .serve_connection(io, service)
                            .with_upgrades();
                        if let Err(e) = conn.await {
                            eprintln!("proxy connection error: {}", e);
                        }
                    });
                }
            }
        }
    });

    let cmd = cmd.to_string();
    let env_vars_with_masked = env_vars_with_masked;
    let exit_status = tokio::task::spawn_blocking(move || {
        enforcer.run_child(
            &cmd,
            &proxy_url,
            &env_vars_with_masked,
            &ssl_cert_file,
            force_traffic_through_proxy,
        )
    })
    .await
    .map_err(|e| format!("subprocess join: {}", e))??;

    let _ = shutdown_tx.send(()).await;
    let _ = server_handle.await;

    Ok(exit_status)
}

/// Re-exported for the MITM module when feature "mitm" is enabled.
pub(crate) type BoxBodyType = BoxBody<Bytes, hyper::Error>;

fn full_body(chunk: Bytes) -> BoxBodyType {
    Full::new(chunk).map_err(|never| match never {}).boxed()
}

pub(crate) fn bad_request(msg: &str) -> Response<BoxBodyType> {
    Response::builder()
        .status(http::StatusCode::BAD_REQUEST)
        .body(full_body(Bytes::from(msg.to_string())))
        .unwrap()
}

pub(crate) fn bad_gateway(msg: &str) -> Response<BoxBodyType> {
    Response::builder()
        .status(http::StatusCode::BAD_GATEWAY)
        .body(full_body(Bytes::from(msg.to_string())))
        .unwrap()
}

/// Replace all occurrences of `from` with `to` in `buf` (byte-wise).
fn replace_bytes(buf: &[u8], from: &[u8], to: &[u8]) -> Vec<u8> {
    if from.is_empty() || buf.is_empty() {
        return buf.to_vec();
    }
    let mut out = Vec::with_capacity(buf.len());
    let mut i = 0;
    while i <= buf.len().saturating_sub(from.len()) {
        if buf[i..].starts_with(from) {
            out.extend_from_slice(to);
            i += from.len();
        } else {
            out.push(buf[i]);
            i += 1;
        }
    }
    out.extend_from_slice(&buf[i..]);
    out
}

/// Apply token replacements (longest first) to a byte buffer.
pub(crate) fn replace_tokens_in_bytes(
    buf: &[u8],
    replacement_order: &[(String, String)],
) -> Vec<u8> {
    let mut current = buf.to_vec();
    for (masked, real) in replacement_order {
        current = replace_bytes(&current, masked.as_bytes(), real.as_bytes());
    }
    current
}

/// Apply token replacements to a header value string; returns None if result would be invalid.
pub(crate) fn replace_tokens_in_header_value(
    value: &str,
    replacement_order: &[(String, String)],
) -> Option<String> {
    let mut s = value.to_string();
    for (masked, real) in replacement_order {
        s = s.replace(masked, real);
    }
    if s.bytes().any(|b| FORBIDDEN_IN_SECRET.contains(&b)) {
        return None;
    }
    Some(s)
}

/// Check if the authority host is a private/local address (SSRF mitigation).
fn is_private_authority(authority: &str) -> bool {
    let (host, _port) = authority.split_once(':').unwrap_or((authority, "80"));
    let host = host.trim_start_matches('[').trim_end_matches(']');
    if host == "localhost" || host.is_empty() {
        return true;
    }
    if let Ok(ip) = host.parse::<std::net::IpAddr>() {
        return is_private_ip(ip);
    }
    false
}

fn is_private_ip(ip: std::net::IpAddr) -> bool {
    match ip {
        std::net::IpAddr::V4(a) => {
            a.is_loopback()
                || a.is_private()
                || a.is_link_local()
                || a.is_broadcast()
                || a.is_documentation()
        }
        std::net::IpAddr::V6(a) => a.is_loopback() || a.is_unspecified(),
    }
}

async fn proxy_handler(
    req: Request<hyper::body::Incoming>,
    token_map: Arc<Vec<(String, String)>>,
    allow_private_connect: bool,
    mitm_config: Arc<mitm::MitmConfig>,
    connection_policies: Vec<ConnectionPolicy>,
) -> Result<Response<BoxBodyType>, hyper::Error> {
    if req.method() == Method::CONNECT {
        let authority = match req.uri().authority() {
            Some(a) => a.to_string(),
            None => return Ok(bad_request("CONNECT must include authority (host:port)")),
        };
        let host = req
            .uri()
            .authority()
            .map(|a| a.host().to_string())
            .unwrap_or_default();
        if !connection_allowed(host.as_str(), Some(&connection_policies)) {
            return Ok(bad_request("CONNECT to this host is not allowed by policy"));
        }
        if !allow_private_connect && is_private_authority(&authority) {
            return Ok(bad_request("CONNECT to private/local address not allowed"));
        }
        return mitm::handle_connect_mitm(req, token_map, mitm_config, connection_policies).await;
    }
    handle_forward(req, token_map, connection_policies).await
}

async fn handle_forward(
    req: Request<hyper::body::Incoming>,
    token_map: Arc<Vec<(String, String)>>,
    connection_policies: Vec<ConnectionPolicy>,
) -> Result<Response<BoxBodyType>, hyper::Error> {
    let (parts, body) = req.into_parts();
    let uri = parts.uri.clone();
    let uri_str = uri.to_string();
    let modified_uri_bytes = replace_tokens_in_bytes(uri_str.as_bytes(), &token_map);
    let modified_uri = match String::from_utf8(modified_uri_bytes) {
        Ok(s) => s.parse().unwrap_or(uri),
        Err(_) => uri.clone(),
    };
    let host = match modified_uri.host() {
        Some(h) => h.to_string(),
        None => return Ok(bad_request("Request URI has no host")),
    };
    if !connection_allowed(host.as_str(), Some(&connection_policies)) {
        return Ok(bad_request(
            "Connection to this host is not allowed by policy",
        ));
    }
    let port = modified_uri.port_u16().unwrap_or(80);

    let body_bytes = body.collect().await?.to_bytes();
    let modified_body = replace_tokens_in_bytes(&body_bytes, &token_map);

    let mut new_headers = http::HeaderMap::new();
    for (name, value) in parts.headers.iter() {
        if name == http::header::CONNECTION
            || name.as_str().eq_ignore_ascii_case("proxy-connection")
            || name == http::header::TRANSFER_ENCODING
        {
            continue;
        }
        let value_str = match value.to_str() {
            Ok(s) => s,
            Err(_) => continue,
        };
        let new_value = replace_tokens_in_header_value(value_str, &token_map);
        if let Some(v) = new_value {
            if let Ok(hv) = v.parse() {
                new_headers.insert(name.clone(), hv);
            }
        }
    }
    new_headers.insert(
        http::header::CONTENT_LENGTH,
        http::HeaderValue::from_str(&modified_body.len().to_string()).unwrap(),
    );

    let body = http_body_util::Full::new(bytes::Bytes::from(modified_body));
    let new_req = match Request::builder()
        .method(parts.method)
        .uri(&modified_uri)
        .body(body)
    {
        Ok(r) => r,
        Err(_) => return Ok(bad_gateway("Invalid request")),
    };
    let (mut new_parts, new_body) = new_req.into_parts();
    new_parts.headers = new_headers;
    let new_req = Request::from_parts(new_parts, new_body);

    let stream = match TcpStream::connect((host.as_str(), port)).await {
        Ok(s) => s,
        Err(e) => {
            eprintln!("proxy connect error: {}", e);
            return Ok(bad_gateway("Upstream connection failed"));
        }
    };
    let io = TokioIo::new(stream);
    let (mut sender, conn) = match hyper::client::conn::http1::Builder::new()
        .handshake(io)
        .await
    {
        Ok(x) => x,
        Err(e) => {
            eprintln!("proxy handshake error: {}", e);
            return Ok(bad_gateway("Upstream handshake failed"));
        }
    };
    tokio::spawn(async move {
        let _ = conn.await;
    });

    let resp = match sender.send_request(new_req).await {
        Ok(r) => r,
        Err(e) => {
            eprintln!("proxy send error: {}", e);
            return Ok(bad_gateway("Upstream request failed"));
        }
    };
    let (resp_parts, resp_body) = resp.into_parts();
    let resp_body = resp_body.boxed();
    let resp = Response::from_parts(resp_parts, resp_body);
    Ok(resp)
}

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

    #[test]
    fn secret_mapping_creation() {
        let m = SecretMapping {
            var: "API_KEY".to_string(),
            value: "secret123".to_string(),
        };
        assert_eq!(m.var, "API_KEY");
        assert_eq!(m.value, "secret123");
    }

    #[test]
    fn sandbox_new_default_config() {
        let _sandbox = Sandbox::new(SandboxConfig::default());
    }

    #[test]
    fn sandbox_new_one_mapping() {
        let m = SecretMapping {
            var: "X".to_string(),
            value: "y".to_string(),
        };
        let config = SandboxConfig {
            secrets: vec![m],
            ..SandboxConfig::default()
        };
        let _sandbox = Sandbox::new(config);
    }

    #[test]
    fn sandbox_new_multiple_mappings() {
        let mappings = vec![
            SecretMapping {
                var: "A".to_string(),
                value: "a".to_string(),
            },
            SecretMapping {
                var: "B".to_string(),
                value: "b".to_string(),
            },
        ];
        let config = SandboxConfig {
            secrets: mappings,
            ..SandboxConfig::default()
        };
        let _sandbox = Sandbox::new(config);
    }

    #[tokio::test]
    async fn sandbox_clone() {
        let m = SecretMapping {
            var: "K".to_string(),
            value: "v".to_string(),
        };
        let config = SandboxConfig {
            secrets: vec![m],
            force_traffic_through_proxy: false,
            ..SandboxConfig::default()
        };
        let sandbox = Sandbox::new(config);
        let cloned = sandbox.clone();
        let status = cloned.run("true").await;
        assert!(status.is_ok());
        assert!(status.unwrap().success());
    }

    #[tokio::test]
    async fn run_requires_at_least_one_mapping() {
        let config = SandboxConfig {
            force_traffic_through_proxy: false,
            ..SandboxConfig::default()
        };
        let sandbox = Sandbox::new(config);
        let result = sandbox.run("true").await;
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(
            err.to_string()
                .contains("at least one secret mapping or string mapping"),
            "expected message about mapping, got: {}",
            err
        );
    }

    #[tokio::test]
    async fn run_with_string_mapping_only() {
        let config = SandboxConfig {
            strings: vec![StringMapping {
                token: "__TOKEN__".to_string(),
                value: "replaced".to_string(),
            }],
            force_traffic_through_proxy: false,
            ..SandboxConfig::default()
        };
        let sandbox = Sandbox::new(config);
        let result = sandbox.run("true").await;
        assert!(result.is_ok());
        assert!(result.unwrap().success());
    }

    #[tokio::test]
    async fn run_returns_exit_status_success() {
        let config = SandboxConfig {
            secrets: vec![SecretMapping {
                var: "X".to_string(),
                value: "x".to_string(),
            }],
            force_traffic_through_proxy: false,
            ..SandboxConfig::default()
        };
        let sandbox = Sandbox::new(config);
        let result = sandbox.run("true").await;
        assert!(result.is_ok(), "{:?}", result.err());
        assert!(result.unwrap().success());
    }

    #[tokio::test]
    async fn run_returns_exit_status_failure() {
        let config = SandboxConfig {
            secrets: vec![SecretMapping {
                var: "X".to_string(),
                value: "x".to_string(),
            }],
            force_traffic_through_proxy: false,
            ..SandboxConfig::default()
        };
        let sandbox = Sandbox::new(config);
        let result = sandbox.run("false").await;
        assert!(result.is_ok());
        assert!(!result.unwrap().success());
    }

    #[tokio::test]
    async fn run_forward_exit_code() {
        let config = SandboxConfig {
            secrets: vec![SecretMapping {
                var: "X".to_string(),
                value: "x".to_string(),
            }],
            force_traffic_through_proxy: false,
            ..SandboxConfig::default()
        };
        let sandbox = Sandbox::new(config);
        let result = sandbox.run("sh -c 'exit 42'").await;
        assert!(result.is_ok());
        assert_eq!(result.unwrap().code(), Some(42));
    }

    #[tokio::test]
    async fn run_command_sees_masked_env() {
        let config = SandboxConfig {
            secrets: vec![SecretMapping {
                var: "API_KEY".to_string(),
                value: "real-secret".to_string(),
            }],
            force_traffic_through_proxy: false,
            ..SandboxConfig::default()
        };
        let sandbox = Sandbox::new(config);
        let result = sandbox.run("sh -c 'v=$API_KEY; if [ \"$v\" = \"real-secret\" ]; then exit 1; fi; case \"$v\" in api-key-*) exit 0;; *) exit 2;; esac'").await;
        assert!(result.is_ok(), "run failed");
        assert_eq!(
            result.unwrap().code(),
            Some(0),
            "expected masked token pattern api-key-*"
        );
    }

    #[tokio::test]
    async fn run_sets_http_proxy_env() {
        let config = SandboxConfig {
            secrets: vec![SecretMapping {
                var: "X".to_string(),
                value: "x".to_string(),
            }],
            force_traffic_through_proxy: false,
            ..SandboxConfig::default()
        };
        let sandbox = Sandbox::new(config);
        let result = sandbox
            .run("sh -c 'case \"$HTTP_PROXY\" in http://127.0.0.1*) exit 0;; *) exit 1;; esac'")
            .await;
        assert!(result.is_ok());
        assert_eq!(result.unwrap().code(), Some(0));
        let result = sandbox
            .run("sh -c 'case \"$HTTPS_PROXY\" in http://127.0.0.1*) exit 0;; *) exit 1;; esac'")
            .await;
        assert!(result.is_ok());
        assert_eq!(result.unwrap().code(), Some(0));
    }

    #[test]
    fn replace_bytes_basic() {
        let buf = b"hello world";
        let out = replace_bytes(buf, b"o", b"X");
        assert_eq!(out.as_slice(), b"hellX wXrld");
    }

    #[test]
    fn replace_tokens_longest_first() {
        let order = vec![
            ("api-key-long".to_string(), "real-long".to_string()),
            ("api-key".to_string(), "real-short".to_string()),
        ];
        let buf = b"prefix api-key-long suffix";
        let out = replace_tokens_in_bytes(buf, &order);
        assert_eq!(out.as_slice(), b"prefix real-long suffix");
    }

    #[test]
    fn is_private_authority_blocks_localhost() {
        assert!(is_private_authority("localhost:443"));
        assert!(is_private_authority("127.0.0.1:8080"));
        assert!(is_private_authority("10.0.0.1:80"));
        assert!(!is_private_authority("example.com:443"));
    }

    #[test]
    fn validate_secret_rejects_crlf() {
        assert!(validate_secret("ok").is_ok());
        assert!(validate_secret("no\rcr").is_err());
        assert!(validate_secret("no\nlf").is_err());
        assert!(validate_secret("no\0nul").is_err());
    }

    #[test]
    fn string_mapping_creation() {
        let m = StringMapping {
            token: "__API_KEY__".to_string(),
            value: "secret123".to_string(),
        };
        assert_eq!(m.token, "__API_KEY__");
        assert_eq!(m.value, "secret123");
    }

    #[test]
    fn replace_tokens_in_bytes_string_mapping() {
        let order = vec![("__TOKEN__".to_string(), "real-value".to_string())];
        let buf = b"Bearer __TOKEN__";
        let out = replace_tokens_in_bytes(buf, &order);
        assert_eq!(out.as_slice(), b"Bearer real-value");
    }

    #[test]
    fn connection_allowed_no_policies() {
        assert!(connection_allowed("example.com", None));
        assert!(connection_allowed("evil.com", Some(&[])));
    }

    #[test]
    fn connection_allowed_first_match_wins() {
        let policies = vec![
            ConnectionPolicy::deny(HostPattern::exact("blocked.com")),
            ConnectionPolicy::allow(HostPattern::exact("blocked.com")),
        ];
        assert!(!connection_allowed("blocked.com", Some(&policies)));
    }

    #[test]
    fn connection_allowed_regex() {
        let policies = vec![
            ConnectionPolicy::deny(HostPattern::regex(r"^internal\.").unwrap()),
            ConnectionPolicy::allow(HostPattern::exact("api.example.com")),
        ];
        assert!(!connection_allowed("internal.service", Some(&policies)));
        assert!(connection_allowed("api.example.com", Some(&policies)));
        assert!(connection_allowed("other.com", Some(&policies))); // no match = allow
    }
}