OpenShell/crates/openshell-cli/src/tls.rs

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use miette::{IntoDiagnostic, Result, WrapErr};
use openshell_core::proto::inference_client::InferenceClient;
use openshell_core::proto::open_shell_client::OpenShellClient;
use rustls::{
    RootCertStore,
    client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
    pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime},
};
use std::collections::HashMap;
use std::future::Future;
use std::io::Cursor;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::OnceLock;
use std::time::Duration;
use tokio::sync::Mutex;
use tonic::service::interceptor::InterceptedService;
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity};
use tracing::debug;

/// Concrete gRPC client type used by all commands.
pub type GrpcClient = OpenShellClient<InterceptedService<Channel, EdgeAuthInterceptor>>;
/// Concrete inference client type.
pub type GrpcInferenceClient = InferenceClient<InterceptedService<Channel, EdgeAuthInterceptor>>;

#[derive(Clone, Debug, Default)]
pub struct TlsOptions {
    ca: Option<PathBuf>,
    cert: Option<PathBuf>,
    key: Option<PathBuf>,
    /// Gateway name for resolving default cert directory.
    gateway_name: Option<String>,
    /// Edge auth bearer token — when set, disables mTLS client certs and
    /// injects authentication headers on every gRPC request instead.
    pub edge_token: Option<String>,
    /// OIDC bearer token — when set, injects `authorization: Bearer <token>`
    /// on every gRPC request. Takes precedence over `edge_token`.
    pub oidc_token: Option<String>,
    /// Skip TLS certificate verification for gateway connections.
    pub gateway_insecure: bool,
}

impl TlsOptions {
    pub fn new(ca: Option<PathBuf>, cert: Option<PathBuf>, key: Option<PathBuf>) -> Self {
        Self {
            ca,
            cert,
            key,
            gateway_name: None,
            edge_token: None,
            oidc_token: None,
            gateway_insecure: false,
        }
    }

    pub fn has_any(&self) -> bool {
        self.ca.is_some() || self.cert.is_some() || self.key.is_some()
    }

    /// Return the gateway name, if set.
    pub fn gateway_name(&self) -> Option<&str> {
        self.gateway_name.as_deref()
    }

    /// Set the gateway name for cert directory resolution.
    #[must_use]
    pub fn with_gateway_name(&self, name: &str) -> Self {
        Self {
            gateway_name: Some(name.to_string()),
            ..self.clone()
        }
    }

    #[must_use]
    pub fn with_default_paths(&self, server: &str) -> Self {
        let base = self
            .gateway_name
            .as_deref()
            .and_then(tls_dir_for_gateway)
            .or_else(|| default_tls_dir(server));
        Self {
            ca: self
                .ca
                .clone()
                .or_else(|| base.as_ref().map(|dir| dir.join("ca.crt"))),
            cert: self
                .cert
                .clone()
                .or_else(|| base.as_ref().map(|dir| dir.join("tls.crt"))),
            key: self
                .key
                .clone()
                .or_else(|| base.as_ref().map(|dir| dir.join("tls.key"))),
            gateway_name: self.gateway_name.clone(),
            ..self.clone()
        }
    }

    /// Returns `true` when using bearer token auth (edge or OIDC).
    pub fn is_bearer_auth(&self) -> bool {
        self.edge_token.is_some() || self.oidc_token.is_some()
    }
}

pub struct TlsMaterials {
    ca: Vec<u8>,
    cert: Vec<u8>,
    key: Vec<u8>,
}

/// Resolve the TLS cert directory for a known gateway name.
fn tls_dir_for_gateway(name: &str) -> Option<PathBuf> {
    let safe_name = sanitize_name(name);
    let base = xdg_config_dir().ok()?.join("openshell").join("gateways");
    Some(base.join(safe_name).join("mtls"))
}

/// Fallback TLS directory resolution from a server URL.
///
/// Used when no gateway name is set (e.g., `SshProxy` which receives a raw URL).
fn default_tls_dir(server: &str) -> Option<PathBuf> {
    let mut name = std::env::var("OPENSHELL_GATEWAY")
        .ok()
        .filter(|value| !value.trim().is_empty());

    if name.is_none()
        && let Ok(uri) = server.parse::<hyper::Uri>()
        && let Some(host) = uri.host()
    {
        name = Some(
            if host == "127.0.0.1" || host.eq_ignore_ascii_case("localhost") {
                "openshell".to_string()
            } else {
                host.to_string()
            },
        );
    }

    let name = name.unwrap_or_else(|| "openshell".to_string());
    let safe_name = sanitize_name(&name);
    let base = xdg_config_dir().ok()?.join("openshell").join("gateways");
    Some(base.join(safe_name).join("mtls"))
}

fn sanitize_name(value: &str) -> String {
    value
        .chars()
        .map(|ch| {
            if ch.is_ascii_alphanumeric() || matches!(ch, '.' | '-' | '_') {
                ch
            } else {
                '_'
            }
        })
        .collect()
}

fn xdg_config_dir() -> Result<PathBuf> {
    openshell_core::paths::xdg_config_dir()
}

pub fn require_tls_materials(server: &str, tls: &TlsOptions) -> Result<TlsMaterials> {
    let resolved = tls.with_default_paths(server);
    let default_hint = default_tls_dir(server).map_or_else(String::new, |dir| {
        format!(" or place certs in {}", dir.display())
    });
    let ca_path = resolved
        .ca
        .as_ref()
        .ok_or_else(|| miette::miette!("TLS CA is required for https endpoints{default_hint}"))?;
    let cert_path = resolved.cert.as_ref().ok_or_else(|| {
        miette::miette!("TLS client cert is required for https endpoints{default_hint}")
    })?;
    let key_path = resolved.key.as_ref().ok_or_else(|| {
        miette::miette!("TLS client key is required for https endpoints{default_hint}")
    })?;

    let ca = std::fs::read(ca_path)
        .into_diagnostic()
        .wrap_err_with(|| format!("failed to read TLS CA from {}", ca_path.display()))?;
    let cert = std::fs::read(cert_path)
        .into_diagnostic()
        .wrap_err_with(|| format!("failed to read TLS cert from {}", cert_path.display()))?;
    let key = std::fs::read(key_path)
        .into_diagnostic()
        .wrap_err_with(|| format!("failed to read TLS key from {}", key_path.display()))?;

    Ok(TlsMaterials { ca, cert, key })
}

fn load_private_key(pem: &[u8]) -> Result<PrivateKeyDer<'static>> {
    let mut cursor = Cursor::new(pem);
    let key = rustls_pemfile::private_key(&mut cursor)
        .into_diagnostic()?
        .ok_or_else(|| miette::miette!("no private key found in TLS key PEM"))?;
    Ok(key)
}

pub fn build_rustls_config(materials: &TlsMaterials) -> Result<rustls::ClientConfig> {
    let mut roots = RootCertStore::empty();
    let mut ca_cursor = Cursor::new(&materials.ca);
    let ca_certs = rustls_pemfile::certs(&mut ca_cursor)
        .collect::<Result<Vec<CertificateDer<'static>>, _>>()
        .into_diagnostic()?;
    for cert in ca_certs {
        roots.add(cert).into_diagnostic()?;
    }

    let mut cert_cursor = Cursor::new(&materials.cert);
    let cert_chain = rustls_pemfile::certs(&mut cert_cursor)
        .collect::<Result<Vec<CertificateDer<'static>>, _>>()
        .into_diagnostic()?;
    let key = load_private_key(&materials.key)?;

    rustls::ClientConfig::builder()
        .with_root_certificates(roots)
        .with_client_auth_cert(cert_chain, key)
        .into_diagnostic()
}

pub fn build_tonic_tls_config(materials: &TlsMaterials) -> ClientTlsConfig {
    let ca_cert = Certificate::from_pem(materials.ca.clone());
    let identity = Identity::from_pem(materials.cert.clone(), materials.key.clone());
    ClientTlsConfig::new()
        .ca_certificate(ca_cert)
        .identity(identity)
}

#[derive(Debug)]
struct InsecureServerCertVerifier;

impl ServerCertVerifier for InsecureServerCertVerifier {
    fn verify_server_cert(
        &self,
        _end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &ServerName<'_>,
        _ocsp_response: &[u8],
        _now: UnixTime,
    ) -> Result<ServerCertVerified, rustls::Error> {
        Ok(ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &rustls::DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &rustls::DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        rustls::crypto::ring::default_provider()
            .signature_verification_algorithms
            .supported_schemes()
    }
}

#[derive(Clone)]
struct InsecureTlsConnector {
    tls_connector: tokio_rustls::TlsConnector,
}

impl tower::Service<hyper::Uri> for InsecureTlsConnector {
    type Response = hyper_util::rt::TokioIo<tokio_rustls::client::TlsStream<tokio::net::TcpStream>>;
    type Error = Box<dyn std::error::Error + Send + Sync>;
    type Future =
        std::pin::Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(
        &mut self,
        _cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        std::task::Poll::Ready(Ok(()))
    }

    fn call(&mut self, uri: hyper::Uri) -> Self::Future {
        let tls_connector = self.tls_connector.clone();
        Box::pin(async move {
            let host = uri.host().unwrap_or("localhost").to_string();
            let port = uri.port_u16().unwrap_or(443);
            let addr = format!("{host}:{port}");
            let tcp = tokio::net::TcpStream::connect(addr).await?;
            let server_name = ServerName::try_from(host)?;
            let tls_stream = tls_connector.connect(server_name, tcp).await?;
            Ok(hyper_util::rt::TokioIo::new(tls_stream))
        })
    }
}

pub fn build_insecure_rustls_config() -> Result<rustls::ClientConfig> {
    let config = rustls::ClientConfig::builder()
        .dangerous()
        .with_custom_certificate_verifier(std::sync::Arc::new(InsecureServerCertVerifier))
        .with_no_client_auth();
    Ok(config)
}

/// Tunnel proxy addresses keyed by upstream endpoint + token.
///
/// Each distinct edge-authenticated gateway gets its own local proxy instead of
/// reusing the first gateway touched in the current process.
static EDGE_TUNNEL_ADDRS: OnceLock<Mutex<HashMap<(String, String), SocketAddr>>> = OnceLock::new();

async fn edge_tunnel_addr(server: &str, token: &str) -> Result<SocketAddr> {
    let key = (server.to_string(), token.to_string());
    let registry = EDGE_TUNNEL_ADDRS.get_or_init(|| Mutex::new(HashMap::new()));

    {
        let addrs = registry.lock().await;
        if let Some(addr) = addrs.get(&key).copied() {
            return Ok(addr);
        }
    }

    let proxy = crate::edge_tunnel::start_tunnel_proxy(server, token).await?;
    debug!(
        local_addr = %proxy.local_addr,
        server,
        "edge tunnel proxy started, routing gRPC through local proxy"
    );

    let mut addrs = registry.lock().await;
    Ok(*addrs.entry(key).or_insert(proxy.local_addr))
}

pub async fn build_channel(server: &str, tls: &TlsOptions) -> Result<Channel> {
    if server.starts_with("http://") {
        let endpoint = Endpoint::from_shared(server.to_string())
            .into_diagnostic()?
            .connect_timeout(Duration::from_secs(10))
            .http2_adaptive_window(true)
            .http2_keep_alive_interval(Duration::from_secs(10))
            .keep_alive_while_idle(true);
        return endpoint.connect().await.into_diagnostic();
    }

    // When Cloudflare edge bearer auth is active and the server is HTTPS,
    // route traffic through a local WebSocket tunnel proxy instead.
    // OIDC tokens bypass the tunnel — they connect directly.
    if tls.edge_token.is_some() && server.starts_with("https://") {
        let token = tls
            .edge_token
            .as_deref()
            .ok_or_else(|| miette::miette!("edge token required for tunnel"))?;
        let local_addr = edge_tunnel_addr(server, token).await?;

        // Connect to the local tunnel proxy over plaintext HTTP/2.
        let local_url = format!("http://{local_addr}");
        let endpoint = Endpoint::from_shared(local_url)
            .into_diagnostic()?
            .connect_timeout(Duration::from_secs(10))
            .http2_adaptive_window(true)
            .http2_keep_alive_interval(Duration::from_secs(10))
            .keep_alive_while_idle(true);
        return endpoint.connect().await.into_diagnostic();
    }

    if tls.gateway_insecure && server.starts_with("https://") {
        tracing::warn!("TLS certificate verification is disabled — do not use in production");
        let rustls_config = build_insecure_rustls_config()?;
        let tls_connector = tokio_rustls::TlsConnector::from(std::sync::Arc::new(rustls_config));
        let connector = InsecureTlsConnector { tls_connector };
        // Use http:// so tonic does not layer its own TLS on top — our
        // connector performs TLS with the insecure config.
        let http_uri = server.replacen("https://", "http://", 1);
        let endpoint = Endpoint::from_shared(http_uri)
            .into_diagnostic()?
            .connect_timeout(Duration::from_secs(10))
            .http2_keep_alive_interval(Duration::from_secs(10))
            .keep_alive_while_idle(true);
        return endpoint
            .connect_with_connector(connector)
            .await
            .into_diagnostic();
    }

    let mut endpoint = Endpoint::from_shared(server.to_string())
        .into_diagnostic()?
        .connect_timeout(Duration::from_secs(10))
        .http2_adaptive_window(true)
        .http2_keep_alive_interval(Duration::from_secs(10))
        .keep_alive_while_idle(true);

    let tls_config = if tls.oidc_token.is_some() {
        // OIDC bearer auth over HTTPS: use mTLS certs for the transport layer
        // when available (server may still require client certs), and layer
        // the Bearer token on top via the interceptor.
        require_tls_materials(server, tls).map_or_else(
            |_| {
                let resolved = tls.with_default_paths(server);
                resolved
                    .ca
                    .as_ref()
                    .and_then(|ca_path| std::fs::read(ca_path).ok())
                    .map_or_else(
                        || ClientTlsConfig::new().with_enabled_roots(),
                        |ca_pem| {
                            ClientTlsConfig::new().ca_certificate(Certificate::from_pem(ca_pem))
                        },
                    )
            },
            |materials| build_tonic_tls_config(&materials),
        )
    } else if tls.edge_token.is_some() {
        // Edge bearer mode — routed through tunnel above; if we reach here
        // the server is not HTTPS so connect plaintext.
        return endpoint.connect().await.into_diagnostic();
    } else {
        // Standard mTLS: private CA + client cert.
        let materials = require_tls_materials(server, tls)?;
        build_tonic_tls_config(&materials)
    };
    endpoint = endpoint.tls_config(tls_config).into_diagnostic()?;
    endpoint.connect().await.into_diagnostic()
}

/// Build a gRPC [`OpenShellClient`].
///
/// When `tls.edge_token` is set, the returned client is wrapped with an
/// interceptor that injects authentication headers on every request.
/// Otherwise, standard mTLS is used (interceptor is a no-op).
pub async fn grpc_client(server: &str, tls: &TlsOptions) -> Result<GrpcClient> {
    let channel = build_channel(server, tls).await?;
    let interceptor = EdgeAuthInterceptor::maybe_from(tls)?;
    Ok(OpenShellClient::with_interceptor(channel, interceptor))
}

/// Interceptor that injects authentication headers into every outgoing gRPC request.
///
/// Supports OIDC Bearer tokens (standard `authorization` header) and
/// Cloudflare Access tokens (custom headers). When no token is set, acts
/// as a no-op. OIDC takes precedence over edge tokens.
#[derive(Clone)]
#[allow(clippy::struct_field_names)]
pub struct EdgeAuthInterceptor {
    /// Standard `authorization: Bearer <token>` for OIDC.
    bearer_value: Option<tonic::metadata::MetadataValue<tonic::metadata::Ascii>>,
    /// CF-specific `Cf-Access-Jwt-Assertion` header.
    header_value: Option<tonic::metadata::MetadataValue<tonic::metadata::Ascii>>,
    /// CF-specific `Cookie: CF_Authorization=<token>` header.
    cookie_value: Option<tonic::metadata::MetadataValue<tonic::metadata::Ascii>>,
}

impl EdgeAuthInterceptor {
    /// Create an interceptor from [`TlsOptions`].  Returns a no-op interceptor
    /// when no auth token is configured.
    pub fn maybe_from(tls: &TlsOptions) -> Result<Self> {
        // OIDC bearer token takes precedence.
        if let Some(ref token) = tls.oidc_token {
            let bearer: tonic::metadata::MetadataValue<tonic::metadata::Ascii> =
                format!("Bearer {token}")
                    .parse()
                    .map_err(|_| miette::miette!("invalid OIDC token value"))?;
            return Ok(Self {
                bearer_value: Some(bearer),
                header_value: None,
                cookie_value: None,
            });
        }

        let (header_value, cookie_value) = match tls.edge_token.as_deref() {
            Some(t) => {
                let hv: tonic::metadata::MetadataValue<tonic::metadata::Ascii> = t
                    .parse()
                    .map_err(|_| miette::miette!("invalid edge token value"))?;
                let cv: tonic::metadata::MetadataValue<tonic::metadata::Ascii> =
                    format!("CF_Authorization={t}")
                        .parse()
                        .map_err(|_| miette::miette!("invalid edge token value for cookie"))?;
                (Some(hv), Some(cv))
            }
            None => (None, None),
        };
        Ok(Self {
            bearer_value: None,
            header_value,
            cookie_value,
        })
    }
}

impl tonic::service::Interceptor for EdgeAuthInterceptor {
    fn call(
        &mut self,
        mut req: tonic::Request<()>,
    ) -> std::result::Result<tonic::Request<()>, tonic::Status> {
        if let Some(ref val) = self.bearer_value {
            req.metadata_mut().insert("authorization", val.clone());
        }
        if let Some(ref val) = self.header_value {
            req.metadata_mut()
                .insert("cf-access-jwt-assertion", val.clone());
        }
        if let Some(ref val) = self.cookie_value {
            req.metadata_mut().insert("cookie", val.clone());
        }
        Ok(req)
    }
}

pub async fn grpc_inference_client(server: &str, tls: &TlsOptions) -> Result<GrpcInferenceClient> {
    let channel = build_channel(server, tls).await?;
    let interceptor = EdgeAuthInterceptor::maybe_from(tls)?;
    Ok(InferenceClient::with_interceptor(channel, interceptor))
}