asimov-cli 25.4.0

// This is free and unencumbered software released into the public domain.

//! Upstream connection establishment, optionally through a proxy.
//!
//! The proxy is configured through the conventional environment variables,
//! consulted in this order: `https_proxy`, `HTTPS_PROXY`, `all_proxy`,
//! `ALL_PROXY`. (Since the upstream endpoint is always HTTPS, `http_proxy`
//! does not apply.) The `no_proxy`/`NO_PROXY` exclusion list is honored.
//!
//! Supported proxy URL schemes:
//!
//! - `http://` — HTTP proxy (tunneled with `CONNECT`)
//! - `https://` — HTTPS proxy (TLS to the proxy itself, then `CONNECT`)
//! - `socks5://` — SOCKS5 proxy (target DNS resolved locally)
//! - `socks5h://` — SOCKS5 proxy (target DNS resolved by the proxy)

use crate::BoxError;
use base64::Engine as _;

/// How to reach the upstream server.
#[derive(Clone, Debug)]
pub enum ProxyConfig {
    /// Connect directly to the target.
    Direct,

    /// Tunnel through an HTTP(S) proxy using `CONNECT`.
    HttpConnect {
        host: String,
        port: u16,
        /// Whether to speak TLS to the proxy itself (an `https://` proxy).
        tls: bool,
        /// Pre-encoded `Proxy-Authorization: Basic` credentials.
        basic_auth: Option<String>,
    },

    /// Tunnel through a SOCKS5 proxy.
    Socks5 {
        host: String,
        port: u16,
        auth: Option<(String, String)>,
        /// Whether to resolve target DNS through the proxy (`socks5h://`).
        remote_dns: bool,
    },
    //
    // TODO: a `Tor` variant tunneling through the Tor network by means of
    // the `arti-client` crate (its `DataStream` implements `AsyncRead` +
    // `AsyncWrite`, so it slots into `ProxyStream::new()` directly).
}

impl ProxyConfig {
    /// Determines the proxy configuration for `target_host` from the
    /// conventional environment variables.
    pub fn from_env(target_host: &str) -> Result<Self, BoxError> {
        if no_proxy_matches(target_host) {
            return Ok(Self::Direct);
        }
        match env_var(&["https_proxy", "HTTPS_PROXY", "all_proxy", "ALL_PROXY"]) {
            Some(input) => Self::parse(&input),
            None => Ok(Self::Direct),
        }
    }

    /// Parses a proxy URL such as `http://user:pass@host:port` or
    /// `socks5h://host:port`.
    pub fn parse(input: &str) -> Result<Self, BoxError> {
        // A bare `host:port` is conventionally an HTTP proxy:
        let url = if input.contains("://") {
            url::Url::parse(input)
        } else {
            url::Url::parse(&format!("http://{}", input))
        }
        .map_err(|err| format!("invalid proxy URL `{}`: {}", input, err))?;

        let host = url
            .host_str()
            .ok_or_else(|| format!("proxy URL `{}` is missing a host", input))?
            .to_string();

        let username = url.username();
        let password = url.password();

        match url.scheme() {
            "http" | "https" => {
                let tls = url.scheme() == "https";
                let port = url.port().unwrap_or(if tls { 443 } else { 80 });
                let basic_auth = if !username.is_empty() || password.is_some() {
                    let credentials = format!("{}:{}", username, password.unwrap_or_default());
                    Some(base64::engine::general_purpose::STANDARD.encode(credentials))
                } else {
                    None
                };
                Ok(Self::HttpConnect {
                    host,
                    port,
                    tls,
                    basic_auth,
                })
            },

            "socks5" | "socks5h" => {
                let port = url.port().unwrap_or(1080);
                let auth = if !username.is_empty() {
                    Some((
                        username.to_string(),
                        password.unwrap_or_default().to_string(),
                    ))
                } else {
                    None
                };
                Ok(Self::Socks5 {
                    host,
                    port,
                    auth,
                    remote_dns: url.scheme() == "socks5h",
                })
            },

            scheme => Err(format!("unsupported proxy scheme: {}", scheme).into()),
        }
    }
}

/// Returns the first nonempty environment variable among `names`.
fn env_var(names: &[&str]) -> Option<String> {
    names
        .iter()
        .find_map(|name| std::env::var(name).ok().filter(|value| !value.is_empty()))
}

/// Checks whether `host` is excluded from proxying by `no_proxy`/`NO_PROXY`.
fn no_proxy_matches(host: &str) -> bool {
    let Some(no_proxy) = env_var(&["no_proxy", "NO_PROXY"]) else {
        return false;
    };
    no_proxy_list_matches(&no_proxy, host)
}

fn no_proxy_list_matches(no_proxy: &str, host: &str) -> bool {
    no_proxy.split(',').any(|entry| {
        let entry = entry.trim().trim_start_matches('.');
        !entry.is_empty()
            && (entry == "*"
                || host.eq_ignore_ascii_case(entry)
                || (host.len() > entry.len()
                    && host[..host.len() - entry.len()].ends_with('.')
                    && host[host.len() - entry.len()..].eq_ignore_ascii_case(entry)))
    })
}

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

    #[test]
    fn parse_http_proxy() {
        let config = ProxyConfig::parse("http://proxy.example:3128").unwrap();
        let ProxyConfig::HttpConnect {
            host,
            port,
            tls,
            basic_auth,
        } = config
        else {
            panic!("expected HttpConnect, got: {:?}", config)
        };
        assert_eq!(host, "proxy.example");
        assert_eq!(port, 3128);
        assert!(!tls);
        assert!(basic_auth.is_none());
    }

    #[test]
    fn parse_bare_host_port_as_http_proxy() {
        let config = ProxyConfig::parse("proxy.example:8080").unwrap();
        assert!(matches!(
            config,
            ProxyConfig::HttpConnect {
                tls: false,
                port: 8080,
                ..
            }
        ));
    }

    #[test]
    fn parse_https_proxy_with_auth() {
        let config = ProxyConfig::parse("https://user:pass@proxy.example").unwrap();
        let ProxyConfig::HttpConnect {
            host,
            port,
            tls,
            basic_auth,
        } = config
        else {
            panic!("expected HttpConnect, got: {:?}", config)
        };
        assert_eq!(host, "proxy.example");
        assert_eq!(port, 443);
        assert!(tls);
        assert_eq!(basic_auth.as_deref(), Some("dXNlcjpwYXNz")); // "user:pass"
    }

    #[test]
    fn parse_socks5_proxy() {
        let config = ProxyConfig::parse("socks5://127.0.0.1").unwrap();
        assert!(matches!(
            config,
            ProxyConfig::Socks5 {
                port: 1080,
                remote_dns: false,
                auth: None,
                ..
            }
        ));

        let config = ProxyConfig::parse("socks5h://user:pass@127.0.0.1:9050").unwrap();
        let ProxyConfig::Socks5 {
            port,
            remote_dns,
            auth,
            ..
        } = config
        else {
            panic!("expected Socks5, got: {:?}", config)
        };
        assert_eq!(port, 9050);
        assert!(remote_dns);
        assert_eq!(auth, Some(("user".to_string(), "pass".to_string())));
    }

    #[test]
    fn parse_unsupported_scheme() {
        assert!(ProxyConfig::parse("ftp://proxy.example").is_err());
    }

    #[test]
    fn no_proxy_matching() {
        assert!(no_proxy_list_matches("*", "openrouter.ai"));
        assert!(no_proxy_list_matches("openrouter.ai", "openrouter.ai"));
        assert!(no_proxy_list_matches(".openrouter.ai", "api.openrouter.ai"));
        assert!(no_proxy_list_matches(
            "example.com, openrouter.ai",
            "openrouter.ai"
        ));
        assert!(!no_proxy_list_matches("example.com", "openrouter.ai"));
        assert!(!no_proxy_list_matches("router.ai", "openrouter.ai"));
    }
}