lesavka/client/src/relay_transport.rs

use anyhow::{Context, Result, bail};
use std::path::PathBuf;
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity};

const DEFAULT_CLIENT_PKI_DIR: &str = ".config/lesavka/pki";

pub fn endpoint(server_addr: &str) -> Result<Endpoint> {
    enforce_transport_policy(server_addr)?;
    let mut endpoint =
        Channel::from_shared(server_addr.to_string()).context("invalid relay server address")?;
    if is_https(server_addr) {
        endpoint = endpoint
            .tls_config(client_tls_config(server_addr)?)
            .context("configuring relay TLS")?;
    }
    Ok(endpoint)
}

pub async fn connect(server_addr: &str) -> Result<Channel> {
    endpoint(server_addr)?
        .tcp_nodelay(true)
        .connect()
        .await
        .with_context(|| format!("connecting to relay at {server_addr}"))
}

pub fn enforce_transport_policy(server_addr: &str) -> Result<()> {
    if is_https(server_addr) || is_local_http(server_addr) || allow_insecure_transport() {
        return Ok(());
    }

    bail!(
        "refusing insecure relay transport for {server_addr}; use https:// or set LESAVKA_ALLOW_INSECURE=1 for a deliberate lab override"
    )
}

fn client_tls_config(server_addr: &str) -> Result<ClientTlsConfig> {
    let mut tls = ClientTlsConfig::new()
        .domain_name(tls_domain(server_addr))
        .with_enabled_roots();

    let ca_path = env_path("LESAVKA_TLS_CA").or_else(|| default_pki_path("ca.crt"));
    let cert_path = env_path("LESAVKA_TLS_CLIENT_CERT").or_else(|| default_pki_path("client.crt"));
    let key_path = env_path("LESAVKA_TLS_CLIENT_KEY").or_else(|| default_pki_path("client.key"));
    require_tls_enrollment(&ca_path, &cert_path, &key_path)?;

    let ca_path = ca_path.expect("enrollment check guarantees CA path");
    let cert_path = cert_path.expect("enrollment check guarantees cert path");
    let key_path = key_path.expect("enrollment check guarantees key path");

    let ca = std::fs::read(&ca_path)
        .with_context(|| format!("reading TLS CA certificate {}", ca_path.display()))?;
    tls = tls.ca_certificate(Certificate::from_pem(ca));

    let cert = std::fs::read(&cert_path)
        .with_context(|| format!("reading TLS client certificate {}", cert_path.display()))?;
    let key = std::fs::read(&key_path)
        .with_context(|| format!("reading TLS client key {}", key_path.display()))?;
    tls = tls.identity(Identity::from_pem(cert, key));

    Ok(tls)
}

fn require_tls_enrollment(
    ca_path: &Option<PathBuf>,
    cert_path: &Option<PathBuf>,
    key_path: &Option<PathBuf>,
) -> Result<()> {
    if path_exists(ca_path) && path_exists(cert_path) && path_exists(key_path) {
        return Ok(());
    }

    let pki_hint = default_pki_path("")
        .map(|path| path.display().to_string())
        .unwrap_or_else(|| "~/.config/lesavka/pki".to_string());
    bail!(
        "TLS enrollment is missing for the Lesavka relay. Run the server installer first, then run the client installer with the server-generated bundle, or let the client installer auto-fetch it from the configured SSH host. Expected ca.crt, client.crt, and client.key under {pki_hint}."
    )
}

fn path_exists(path: &Option<PathBuf>) -> bool {
    path.as_ref().is_some_and(|path| path.exists())
}

fn tls_domain(server_addr: &str) -> String {
    std::env::var("LESAVKA_TLS_DOMAIN")
        .ok()
        .filter(|value| !value.trim().is_empty())
        .unwrap_or_else(|| {
            host_from_uri(server_addr)
                .unwrap_or("lesavka-server")
                .to_string()
        })
}

fn env_path(name: &str) -> Option<PathBuf> {
    std::env::var_os(name)
        .filter(|value| !value.is_empty())
        .map(PathBuf::from)
}

fn default_pki_path(file_name: &str) -> Option<PathBuf> {
    let home = std::env::var_os("HOME")?;
    Some(
        PathBuf::from(home)
            .join(DEFAULT_CLIENT_PKI_DIR)
            .join(file_name),
    )
}

fn allow_insecure_transport() -> bool {
    std::env::var("LESAVKA_ALLOW_INSECURE")
        .map(|value| {
            let value = value.trim().to_ascii_lowercase();
            !value.is_empty() && value != "0" && value != "false" && value != "no"
        })
        .unwrap_or(false)
}

fn is_https(server_addr: &str) -> bool {
    server_addr
        .get(..8)
        .is_some_and(|scheme| scheme.eq_ignore_ascii_case("https://"))
}

fn is_local_http(server_addr: &str) -> bool {
    let Some(rest) = server_addr.strip_prefix("http://") else {
        return false;
    };
    let host = rest
        .split(['/', '?', '#'])
        .next()
        .unwrap_or_default()
        .rsplit_once('@')
        .map(|(_, host)| host)
        .unwrap_or(rest);
    let host = if host.starts_with('[') {
        host.split(']')
            .next()
            .unwrap_or(host)
            .trim_start_matches('[')
    } else {
        host.split(':').next().unwrap_or(host)
    };
    matches!(host, "localhost" | "::1" | "0.0.0.0") || host.starts_with("127.")
}

fn host_from_uri(server_addr: &str) -> Option<&str> {
    let rest = server_addr.split_once("://")?.1;
    let authority = rest.split(['/', '?', '#']).next().unwrap_or_default();
    let authority = authority
        .rsplit_once('@')
        .map(|(_, host)| host)
        .unwrap_or(authority);
    if authority.starts_with('[') {
        return authority
            .split(']')
            .next()
            .map(|host| host.trim_start_matches('['))
            .filter(|host| !host.is_empty());
    }
    authority.split(':').next().filter(|host| !host.is_empty())
}

#[cfg(test)]
mod tests {
    use super::{endpoint, enforce_transport_policy, host_from_uri, is_local_http, tls_domain};
    use temp_env::with_vars;
    use tempfile::tempdir;

    #[test]
    fn localhost_http_is_allowed_for_tunnels() {
        with_vars([("LESAVKA_ALLOW_INSECURE", None::<&str>)], || {
            assert!(enforce_transport_policy("http://127.0.0.1:50051").is_ok());
            assert!(enforce_transport_policy("http://localhost:50051").is_ok());
            assert!(enforce_transport_policy("http://[::1]:50051").is_ok());
        });
    }

    #[test]
    fn public_http_requires_deliberate_override() {
        with_vars([("LESAVKA_ALLOW_INSECURE", None::<&str>)], || {
            let err = enforce_transport_policy("http://38.28.125.112:50051")
                .expect_err("public http must be rejected");
            assert!(
                err.to_string()
                    .contains("refusing insecure relay transport")
            );
        });
        with_vars([("LESAVKA_ALLOW_INSECURE", Some("1"))], || {
            assert!(enforce_transport_policy("http://38.28.125.112:50051").is_ok());
        });
    }

    #[test]
    fn tls_domain_can_be_overridden_for_ip_backed_servers() {
        with_vars([("LESAVKA_TLS_DOMAIN", None::<&str>)], || {
            assert_eq!(tls_domain("https://38.28.125.112:50051"), "38.28.125.112");
        });
        with_vars([("LESAVKA_TLS_DOMAIN", Some("lesavka-server"))], || {
            assert_eq!(tls_domain("https://38.28.125.112:50051"), "lesavka-server");
        });
    }

    #[test]
    fn endpoint_reports_invalid_local_uri_after_security_policy_passes() {
        let err = endpoint("http://[::1").expect_err("malformed local URI should fail");
        assert!(err.to_string().contains("invalid relay server address"));
    }

    #[test]
    fn local_http_policy_handles_ipv4_ipv6_and_auth_shapes() {
        assert!(is_local_http("http://user:pass@127.0.0.1:50051/path"));
        assert!(is_local_http("http://[::1]:50051"));
        assert!(is_local_http("http://0.0.0.0:50051"));
        assert!(!is_local_http("https://127.0.0.1:50051"));
        assert!(!is_local_http("http://38.28.125.112:50051"));
    }

    #[test]
    fn host_parser_covers_ipv6_auth_and_missing_scheme() {
        assert_eq!(
            host_from_uri("https://user@lesavka-server:50051"),
            Some("lesavka-server")
        );
        assert_eq!(host_from_uri("https://[::1]:50051"), Some("::1"));
        assert_eq!(host_from_uri("lesavka-server:50051"), None);
    }

    #[test]
    fn https_endpoint_reads_default_and_explicit_pki_paths_when_present() {
        let dir = tempdir().expect("pki dir");
        let pki = dir.path().join(".config/lesavka/pki");
        std::fs::create_dir_all(&pki).expect("create pki");
        std::fs::write(pki.join("ca.crt"), b"not a real ca").expect("ca");
        std::fs::write(pki.join("client.crt"), b"not a real cert").expect("cert");
        std::fs::write(pki.join("client.key"), b"not a real key").expect("key");

        with_vars(
            [
                ("HOME", Some(dir.path().to_string_lossy().as_ref())),
                ("LESAVKA_TLS_CA", None::<&str>),
                ("LESAVKA_TLS_CLIENT_CERT", None::<&str>),
                ("LESAVKA_TLS_CLIENT_KEY", None::<&str>),
            ],
            || {
                let err = endpoint("https://lesavka-server:50051")
                    .expect_err("fake PEM should be rejected after file discovery");
                assert!(err.to_string().contains("configuring relay TLS"));
            },
        );

        with_vars(
            [
                ("HOME", None::<&str>),
                (
                    "LESAVKA_TLS_CA",
                    Some(pki.join("ca.crt").to_string_lossy().as_ref()),
                ),
                (
                    "LESAVKA_TLS_CLIENT_CERT",
                    Some(pki.join("client.crt").to_string_lossy().as_ref()),
                ),
                (
                    "LESAVKA_TLS_CLIENT_KEY",
                    Some(pki.join("client.key").to_string_lossy().as_ref()),
                ),
            ],
            || {
                let err = endpoint("https://lesavka-server:50051")
                    .expect_err("fake explicit PEM should be rejected after file discovery");
                assert!(err.to_string().contains("configuring relay TLS"));
            },
        );
    }

    #[test]
    fn https_endpoint_fails_fast_without_client_enrollment() {
        let dir = tempdir().expect("empty pki dir");
        with_vars(
            [
                ("HOME", Some(dir.path().to_string_lossy().as_ref())),
                ("LESAVKA_TLS_CA", None::<&str>),
                ("LESAVKA_TLS_CLIENT_CERT", None::<&str>),
                ("LESAVKA_TLS_CLIENT_KEY", None::<&str>),
                ("LESAVKA_TLS_DOMAIN", None::<&str>),
            ],
            || {
                let err = endpoint("https://38.28.125.112:50051")
                    .expect_err("missing mTLS enrollment should be explicit");
                assert!(err.to_string().contains("TLS enrollment is missing"));
                assert!(err.to_string().contains("client.key"));
            },
        );
    }
}