OpenShell/crates/openshell-cli/tests/helpers/mod.rs

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

//! Shared helpers for CLI integration tests.
//!
//! Include this module from a test file with:
//! ```ignore
//! mod helpers;
//! ```

#[macro_export]
macro_rules! unimplemented_sandbox_template_rpcs {
    () => {
        fn create_sandbox_template<'life0, 'async_trait>(
            &'life0 self,
            _request: tonic::Request<openshell_core::proto::CreateSandboxTemplateRequest>,
        ) -> std::pin::Pin<
            Box<
                dyn std::future::Future<
                        Output = Result<
                            tonic::Response<openshell_core::proto::SandboxTemplateResponse>,
                            tonic::Status,
                        >,
                    > + Send
                    + 'async_trait,
            >,
        >
        where
            'life0: 'async_trait,
            Self: 'async_trait,
        {
            Box::pin(async { Err(tonic::Status::unimplemented("unused")) })
        }

        fn get_sandbox_template<'life0, 'async_trait>(
            &'life0 self,
            _request: tonic::Request<openshell_core::proto::GetSandboxTemplateRequest>,
        ) -> std::pin::Pin<
            Box<
                dyn std::future::Future<
                        Output = Result<
                            tonic::Response<openshell_core::proto::SandboxTemplateResponse>,
                            tonic::Status,
                        >,
                    > + Send
                    + 'async_trait,
            >,
        >
        where
            'life0: 'async_trait,
            Self: 'async_trait,
        {
            Box::pin(async { Err(tonic::Status::unimplemented("unused")) })
        }

        fn list_sandbox_templates<'life0, 'async_trait>(
            &'life0 self,
            _request: tonic::Request<openshell_core::proto::ListSandboxTemplatesRequest>,
        ) -> std::pin::Pin<
            Box<
                dyn std::future::Future<
                        Output = Result<
                            tonic::Response<openshell_core::proto::ListSandboxTemplatesResponse>,
                            tonic::Status,
                        >,
                    > + Send
                    + 'async_trait,
            >,
        >
        where
            'life0: 'async_trait,
            Self: 'async_trait,
        {
            Box::pin(async { Err(tonic::Status::unimplemented("unused")) })
        }

        fn delete_sandbox_template<'life0, 'async_trait>(
            &'life0 self,
            _request: tonic::Request<openshell_core::proto::DeleteSandboxTemplateRequest>,
        ) -> std::pin::Pin<
            Box<
                dyn std::future::Future<
                        Output = Result<
                            tonic::Response<openshell_core::proto::DeleteSandboxTemplateResponse>,
                            tonic::Status,
                        >,
                    > + Send
                    + 'async_trait,
            >,
        >
        where
            'life0: 'async_trait,
            Self: 'async_trait,
        {
            Box::pin(async { Err(tonic::Status::unimplemented("unused")) })
        }
    };
}

use rcgen::{
    BasicConstraints, Certificate, CertificateParams, ExtendedKeyUsagePurpose, IsCa, KeyPair,
};

// ── EnvVarGuard ──────────────────────────────────────────────────────────────

/// Global mutex that serialises tests which mutate environment variables so
/// concurrent threads don't clobber each other's state.
static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

struct SavedVar {
    key: &'static str,
    original: Option<String>,
}

/// RAII guard that acquires `ENV_LOCK` and restores all modified environment
/// variables on drop.
pub struct EnvVarGuard {
    vars: Vec<SavedVar>,
    _lock: std::sync::MutexGuard<'static, ()>,
}

#[allow(dead_code, unsafe_code)]
impl EnvVarGuard {
    /// Acquire the global env-var lock and atomically set one or more
    /// environment variables.  All variables are restored to their prior
    /// state (or removed) when the guard is dropped.
    pub fn set(pairs: &[(&'static str, &str)]) -> Self {
        let lock = ENV_LOCK
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let mut vars = Vec::with_capacity(pairs.len());
        for &(key, value) in pairs {
            let original = std::env::var(key).ok();
            unsafe {
                std::env::set_var(key, value);
            }
            vars.push(SavedVar { key, original });
        }
        Self { vars, _lock: lock }
    }
}

#[allow(unsafe_code)]
impl Drop for EnvVarGuard {
    fn drop(&mut self) {
        for var in &self.vars {
            if let Some(value) = &var.original {
                unsafe {
                    std::env::set_var(var.key, value);
                }
            } else {
                unsafe {
                    std::env::remove_var(var.key);
                }
            }
        }
        // _lock drops here, releasing the mutex
    }
}

// ── TLS helpers ──────────────────────────────────────────────────────────────

/// Generate a self-signed CA certificate and its key pair.
#[allow(dead_code)]
pub fn build_ca() -> (Certificate, KeyPair) {
    let key_pair = KeyPair::generate().unwrap();
    let mut params = CertificateParams::new(Vec::<String>::new()).unwrap();
    params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
    let cert = params.self_signed(&key_pair).unwrap();
    (cert, key_pair)
}

/// Generate a server certificate signed by `ca`, valid for `localhost`.
///
/// Returns `(cert_pem, key_pem)`.
#[allow(dead_code)]
pub fn build_server_cert(ca: &Certificate, ca_key: &KeyPair) -> (String, String) {
    let key_pair = KeyPair::generate().unwrap();
    let mut params = CertificateParams::new(vec!["localhost".to_string()]).unwrap();
    params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth];
    let cert = params.signed_by(&key_pair, ca, ca_key).unwrap();
    (cert.pem(), key_pair.serialize_pem())
}

/// Generate a client authentication certificate signed by `ca`.
///
/// Returns `(cert_pem, key_pem)`.
#[allow(dead_code)]
pub fn build_client_cert(ca: &Certificate, ca_key: &KeyPair) -> (String, String) {
    let key_pair = KeyPair::generate().unwrap();
    let mut params = CertificateParams::new(Vec::<String>::new()).unwrap();
    params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ClientAuth];
    let cert = params.signed_by(&key_pair, ca, ca_key).unwrap();
    (cert.pem(), key_pair.serialize_pem())
}