OpenShell/crates/openshell-driver-podman/src/main.rs

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

use clap::Parser;
use miette::{IntoDiagnostic, Result};
use std::future::Future;
use std::net::SocketAddr;
use std::num::{NonZeroI64, NonZeroU64};
use std::path::PathBuf;
use tracing::info;

use openshell_core::proto::compute::v1::compute_driver_server::ComputeDriverServer;
use openshell_core::{AppArmorProfile, ImagePullPolicy, VERSION};
use openshell_driver_podman::config::{DEFAULT_NETWORK_NAME, DEFAULT_PODMAN_STOP_TIMEOUT_SECS};
use openshell_driver_podman::{ComputeDriverService, PodmanComputeConfig, PodmanComputeDriver};

#[derive(Parser)]
#[command(name = "openshell-driver-podman")]
#[command(version = VERSION)]
struct Args {
    /// Public compute-driver Unix socket used by an external gateway.
    #[arg(long, env = "OPENSHELL_COMPUTE_DRIVER_SOCKET")]
    bind_socket: Option<PathBuf>,

    #[arg(
        long,
        env = "OPENSHELL_COMPUTE_DRIVER_BIND",
        default_value = "127.0.0.1:50061"
    )]
    bind_address: SocketAddr,

    #[arg(long, env = "OPENSHELL_LOG_LEVEL", default_value = "info")]
    log_level: String,

    #[arg(long, env = "OPENSHELL_OTLP_ENDPOINT")]
    otlp_endpoint: Option<String>,

    #[arg(long, env = "OPENSHELL_GATEWAY_NAME")]
    gateway_name: Option<String>,

    /// Path to the Podman API Unix socket.
    #[arg(long, env = "OPENSHELL_PODMAN_SOCKET")]
    podman_socket: Option<PathBuf>,

    #[arg(long, env = "OPENSHELL_SANDBOX_IMAGE")]
    sandbox_image: Option<String>,

    #[arg(
        long,
        env = "OPENSHELL_SANDBOX_IMAGE_PULL_POLICY",
        default_value_t = ImagePullPolicy::IfNotPresent
    )]
    sandbox_image_pull_policy: ImagePullPolicy,

    #[arg(long, env = "OPENSHELL_GRPC_ENDPOINT")]
    grpc_endpoint: Option<String>,

    /// Port the gateway server is listening on.
    ///
    /// Used when `--grpc-endpoint` is not set to auto-detect the endpoint
    /// that sandbox containers dial back to.
    #[arg(
        long,
        env = "OPENSHELL_GATEWAY_PORT",
        default_value_t = openshell_core::config::DEFAULT_SERVER_PORT
    )]
    gateway_port: u16,

    /// Host gateway IP used for sandbox host aliases.
    ///
    /// Empty uses Podman's `host-gateway` resolver.
    #[arg(long, env = "OPENSHELL_PODMAN_HOST_GATEWAY_IP")]
    host_gateway_ip: Option<String>,

    #[arg(
        long,
        env = "OPENSHELL_SANDBOX_SSH_SOCKET_PATH",
        default_value = openshell_core::container_paths::SSH_SOCKET_PATH
    )]
    sandbox_ssh_socket_path: String,

    /// Podman network name retained for driver-managed resources.
    #[arg(long, env = "OPENSHELL_NETWORK_NAME", default_value = DEFAULT_NETWORK_NAME)]
    network_name: String,

    /// Container stop timeout in seconds (SIGTERM → SIGKILL).
    #[arg(long, env = "OPENSHELL_STOP_TIMEOUT", default_value_t = DEFAULT_PODMAN_STOP_TIMEOUT_SECS)]
    stop_timeout: u32,

    /// Container cgroup PID limit for sandbox containers. Omit to use
    /// `OpenShell`'s 2048-process default.
    #[arg(long, env = "OPENSHELL_SANDBOX_PIDS_LIMIT", default_value = "2048")]
    sandbox_pids_limit: Option<NonZeroI64>,

    /// Health check interval in seconds. Omit it in gateway TOML to disable
    /// health checks; the standalone driver keeps its prior 10-second default.
    #[arg(
        long,
        env = "OPENSHELL_HEALTH_CHECK_INTERVAL_SECS",
        default_value = "10"
    )]
    health_check_interval_secs: Option<NonZeroU64>,

    /// OCI image containing the `openshell-sandbox` runtime binary.
    #[arg(long, env = "OPENSHELL_SANDBOX_RUNTIME_IMAGE")]
    sandbox_runtime_image: Option<String>,

    /// OCI image containing the `openshell-supervisor` control binary.
    #[arg(long, env = "OPENSHELL_SUPERVISOR_IMAGE")]
    supervisor_image: Option<String>,

    /// Host path to the CA certificate for sandbox mTLS.
    #[arg(long, env = "OPENSHELL_PODMAN_TLS_CA")]
    podman_tls_ca: Option<PathBuf>,

    /// Host path to the client certificate for sandbox mTLS.
    #[arg(long, env = "OPENSHELL_PODMAN_TLS_CERT")]
    podman_tls_cert: Option<PathBuf>,

    /// Host path to the client private key for sandbox mTLS.
    #[arg(long, env = "OPENSHELL_PODMAN_TLS_KEY")]
    podman_tls_key: Option<PathBuf>,

    /// Host UNIX socket projected into supervisors for provider SPIFFE token exchange.
    #[arg(long, env = "OPENSHELL_PROVIDER_SPIFFE_WORKLOAD_API_SOCKET")]
    provider_spiffe_workload_api_socket: Option<PathBuf>,

    /// `AppArmor` model: `RuntimeDefault`, `Unconfined`, or `Localhost/<profile>`.
    #[arg(long, env = "OPENSHELL_APP_ARMOR_PROFILE")]
    app_armor_profile: Option<AppArmorProfile>,

    /// Corporate forward proxy URL for the supervisor's upstream TLS dials,
    /// in explicit `http://host:port` form (scheme and port required).
    /// Credentials must not be embedded in the URL; use
    /// `--sandbox-proxy-auth-file` instead.
    #[arg(long, env = "OPENSHELL_SANDBOX_HTTPS_PROXY")]
    sandbox_https_proxy: Option<String>,

    /// Comma-separated `NO_PROXY` list injected alongside the proxy URL.
    #[arg(long, env = "OPENSHELL_SANDBOX_NO_PROXY")]
    sandbox_no_proxy: Option<String>,

    /// Path to a file containing the corporate proxy credentials as
    /// `user:pass`. Delivered to the supervisor through a root-only secret
    /// mount so the credentials never appear in config or container metadata.
    #[arg(long, env = "OPENSHELL_SANDBOX_PROXY_AUTH_FILE")]
    sandbox_proxy_auth_file: Option<String>,

    /// Explicit acknowledgement (`true`) that the proxy credential is sent
    /// as cleartext Basic auth over the plain-TCP connection to the http://
    /// proxy. Required when `--sandbox-proxy-auth-file` is set.
    #[arg(long, env = "OPENSHELL_SANDBOX_PROXY_AUTH_ALLOW_INSECURE")]
    sandbox_proxy_auth_allow_insecure: Option<bool>,

    /// Send the destination hostname in CONNECT requests to the corporate
    /// proxy instead of a validated IP. Only for proxies whose ACLs filter
    /// on hostnames: the proxy then resolves the name itself, so sandbox
    /// SSRF/`allowed_ips` validation no longer binds the connection.
    #[arg(long, env = "OPENSHELL_SANDBOX_PROXY_CONNECT_BY_HOSTNAME")]
    sandbox_proxy_connect_by_hostname: Option<bool>,

    /// Path (on the gateway host) to a PEM CA bundle trusted for the corporate
    /// proxy: the TLS handshake with an `https://` proxy and, for
    /// TLS-intercepting proxies, re-signed upstream certificates. Bind-mounted
    /// read-only into the sandbox. Only meaningful with `--sandbox-https-proxy`.
    #[arg(long, env = "OPENSHELL_SANDBOX_PROXY_CA_BUNDLE")]
    sandbox_proxy_ca_bundle: Option<String>,

    /// User namespace mode for sandbox containers (e.g. `auto`).
    /// When unset, containers use the default user namespace.
    #[arg(long, env = "OPENSHELL_PODMAN_USERNS")]
    userns: Option<String>,

    /// Explicit UID mappings for `userns = "private"`.
    /// Each entry is `"container_id:host_id:size"`.
    #[arg(long = "uidmap")]
    uidmap: Vec<String>,

    /// Explicit GID mappings for `userns = "private"`.
    /// Each entry is `"container_id:host_id:size"`.
    #[arg(long = "gidmap")]
    gidmap: Vec<String>,

    /// Allow sandbox requests to attach host bind mounts.
    #[arg(long, env = "OPENSHELL_ENABLE_BIND_MOUNTS", default_value_t = false)]
    enable_bind_mounts: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    let args = Args::parse();
    let _tracing = openshell_otel::install_driver_tracing(
        openshell_driver_podman::otel_tracing::TRACING,
        openshell_otel::DriverTracingConfig {
            endpoint: args.otlp_endpoint.as_deref(),
            gateway_name: args.gateway_name.as_deref(),
            service_version: VERSION,
            log_level: &args.log_level,
        },
    );

    let driver = PodmanComputeDriver::new(PodmanComputeConfig {
        socket_path: args.podman_socket,
        default_image: args.sandbox_image.unwrap_or_default(),
        image_pull_policy: args.sandbox_image_pull_policy,
        grpc_endpoint: args.grpc_endpoint.unwrap_or_default(),
        gateway_port: args.gateway_port,
        host_gateway_ip: args
            .host_gateway_ip
            .unwrap_or_else(PodmanComputeConfig::default_host_gateway_ip),
        ssh_socket_path: args.sandbox_ssh_socket_path,
        network_name: args.network_name,
        stop_timeout_secs: args.stop_timeout,
        sandbox_runtime_image: args
            .sandbox_runtime_image
            .unwrap_or_else(openshell_core::config::default_sandbox_runtime_image),
        supervisor_image: args
            .supervisor_image
            .unwrap_or_else(openshell_core::config::default_supervisor_image),
        guest_tls_ca: args.podman_tls_ca,
        guest_tls_cert: args.podman_tls_cert,
        guest_tls_key: args.podman_tls_key,
        provider_spiffe_workload_api_socket: args.provider_spiffe_workload_api_socket,
        app_armor_profile: args.app_armor_profile,
        sandbox_pids_limit: args.sandbox_pids_limit,
        health_check_interval_secs: args.health_check_interval_secs,
        https_proxy: args.sandbox_https_proxy,
        no_proxy: args.sandbox_no_proxy,
        proxy_auth_file: args.sandbox_proxy_auth_file,
        proxy_auth_allow_insecure: args.sandbox_proxy_auth_allow_insecure,
        proxy_connect_by_hostname: args.sandbox_proxy_connect_by_hostname,
        proxy_ca_bundle: args.sandbox_proxy_ca_bundle,
        userns: args.userns,
        uidmap: args.uidmap,
        gidmap: args.gidmap,
        enable_bind_mounts: args.enable_bind_mounts,
    })
    .await
    .into_diagnostic()?;

    let service = ComputeDriverServer::new(ComputeDriverService::new(driver));
    if let Some(socket_path) = args.bind_socket {
        let listener = openshell_core::external_driver_socket::bind_private(&socket_path)
            .map_err(|err| miette::miette!("{err}"))?;
        let _cleanup =
            openshell_core::external_driver_socket::SocketCleanup::new(socket_path.clone());
        info!(socket = %socket_path.display(), "Starting Podman compute driver");
        tonic::transport::Server::builder()
            .layer(openshell_otel::compute_driver_rpc_layer())
            .add_service(service)
            .serve_with_incoming_shutdown(
                openshell_core::external_driver_socket::SameUidUnixIncoming::new(listener),
                shutdown_signal(),
            )
            .await
            .into_diagnostic()
    } else {
        info!(address = %args.bind_address, "Starting Podman compute driver");
        tonic::transport::Server::builder()
            .layer(openshell_otel::compute_driver_rpc_layer())
            .add_service(service)
            .serve_with_shutdown(args.bind_address, shutdown_signal())
            .await
            .into_diagnostic()
    }
}

async fn select_shutdown_signal(
    ctrl_c: impl Future<Output = ()>,
    terminate: impl Future<Output = ()>,
) {
    tokio::select! {
        () = ctrl_c => {}
        () = terminate => {}
    }
}

async fn ctrl_c_signal() {
    if let Err(error) = tokio::signal::ctrl_c().await {
        tracing::warn!(%error, "Failed to install Ctrl-C signal handler");
        std::future::pending::<()>().await;
    }
}

#[cfg(unix)]
async fn terminate_signal() {
    let Ok(mut signal) = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
    else {
        tracing::warn!("Failed to install SIGTERM signal handler");
        std::future::pending::<()>().await;
        return;
    };
    let _ = signal.recv().await;
}

async fn shutdown_signal() {
    #[cfg(unix)]
    select_shutdown_signal(ctrl_c_signal(), terminate_signal()).await;

    #[cfg(not(unix))]
    ctrl_c_signal().await;

    info!("Received shutdown signal, draining in-flight requests");
}

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

    #[tokio::test]
    async fn shutdown_completes_when_termination_signal_arrives() {
        select_shutdown_signal(std::future::pending(), std::future::ready(())).await;
    }

    #[test]
    fn accepts_gateway_otlp_configuration() {
        let args = Args::try_parse_from([
            "openshell-driver-podman",
            "--otlp-endpoint",
            "http://collector.internal:4317",
            "--gateway-name",
            "production-us-west",
        ])
        .expect("OTLP configuration should be accepted");

        assert_eq!(
            args.otlp_endpoint.as_deref(),
            Some("http://collector.internal:4317")
        );
        assert_eq!(args.gateway_name.as_deref(), Some("production-us-west"));
    }

    #[test]
    fn standalone_defaults_preserve_health_checks_and_reject_zero_limits() {
        let defaults = Args::try_parse_from(["openshell-driver-podman"])
            .expect("standalone driver defaults should parse");
        assert_eq!(
            defaults.health_check_interval_secs.map(NonZeroU64::get),
            Some(10)
        );
        assert_eq!(
            defaults.sandbox_pids_limit.map(NonZeroI64::get),
            Some(openshell_core::config::DEFAULT_SANDBOX_PIDS_LIMIT)
        );

        for flag in ["--sandbox-pids-limit", "--health-check-interval-secs"] {
            let result = Args::try_parse_from(["openshell-driver-podman", flag, "0"]);
            assert!(result.is_err(), "zero must be rejected for {flag}");
            let error = result.err().expect("error was asserted above");
            assert!(error.to_string().contains("invalid value"), "flag: {flag}");
        }
    }
}