// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 //! Shared helpers for openshell-server integration tests. //! //! Include with `mod common;` at the top of each integration test file. //! Items may not be used by every test file; the blanket `#[allow]` prevents //! spurious dead-code warnings. #![allow(dead_code)] use hyper_util::{ rt::{TokioExecutor, TokioIo}, server::conn::auto::Builder, }; use openshell_core::proto::{ CreateProviderRequest, CreateSandboxRequest, CreateSshSessionRequest, CreateSshSessionResponse, DeleteProviderRequest, DeleteProviderResponse, DeleteSandboxRequest, DeleteSandboxResponse, ExecSandboxEvent, ExecSandboxInput, ExecSandboxRequest, GatewayMessage, GetGatewayConfigRequest, GetGatewayConfigResponse, GetProviderRequest, GetSandboxConfigRequest, GetSandboxConfigResponse, GetSandboxProviderEnvironmentRequest, GetSandboxProviderEnvironmentResponse, GetSandboxRequest, HealthRequest, HealthResponse, IssueSandboxTokenRequest, IssueSandboxTokenResponse, ListProvidersRequest, ListProvidersResponse, ListSandboxesRequest, ListSandboxesResponse, ProviderResponse, RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RelayFrame, RevokeSshSessionRequest, RevokeSshSessionResponse, SandboxResponse, SandboxStreamEvent, ServiceStatus, SupervisorMessage, TcpForwardFrame, UpdateProviderRequest, WatchSandboxRequest, open_shell_client::OpenShellClient, open_shell_server::{OpenShell, OpenShellServer}, }; use openshell_server::{MultiplexedService, Store, TlsAcceptor, health_router}; use rcgen::{CertificateParams, IsCa, KeyPair}; use rustls::RootCertStore; use rustls::pki_types::CertificateDer; use rustls_pemfile::certs; use std::io::Write; use std::net::SocketAddr; use std::sync::Arc; use tempfile::tempdir; use tokio::net::TcpListener; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use tonic::transport::{Channel, ClientTlsConfig, Endpoint}; use tonic::{Response, Status}; // --------------------------------------------------------------------------- // Minimal OpenShell stub: all methods return defaults or Unimplemented. // --------------------------------------------------------------------------- #[derive(Clone, Default)] pub struct TestOpenShell; #[tonic::async_trait] impl OpenShell for TestOpenShell { async fn health( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(HealthResponse { status: ServiceStatus::Healthy.into(), version: "test".to_string(), })) } async fn get_gateway_info( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn create_sandbox( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(SandboxResponse::default())) } async fn get_sandbox( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(SandboxResponse::default())) } async fn list_sandboxes( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(ListSandboxesResponse::default())) } async fn list_sandbox_providers( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new( openshell_core::proto::ListSandboxProvidersResponse::default(), )) } async fn attach_sandbox_provider( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new( openshell_core::proto::AttachSandboxProviderResponse::default(), )) } async fn detach_sandbox_provider( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new( openshell_core::proto::DetachSandboxProviderResponse::default(), )) } async fn delete_sandbox( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(DeleteSandboxResponse { deleted: true })) } async fn get_sandbox_config( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(GetSandboxConfigResponse::default())) } async fn get_gateway_config( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(GetGatewayConfigResponse::default())) } async fn get_sandbox_provider_environment( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new( GetSandboxProviderEnvironmentResponse::default(), )) } async fn create_ssh_session( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(CreateSshSessionResponse::default())) } async fn expose_service( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new( openshell_core::proto::ServiceEndpointResponse::default(), )) } async fn get_service( &self, _: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn list_services( &self, _: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn delete_service( &self, _: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn revoke_ssh_session( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(RevokeSshSessionResponse::default())) } async fn create_provider( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented( "create_provider not implemented in test", )) } async fn get_provider( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented( "get_provider not implemented in test", )) } async fn list_providers( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented( "list_providers not implemented in test", )) } async fn list_provider_profiles( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn get_provider_profile( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn import_provider_profiles( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn update_provider_profiles( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn lint_provider_profiles( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn delete_provider_profile( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn update_provider( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented( "update_provider not implemented in test", )) } async fn get_provider_refresh_status( &self, _: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn configure_provider_refresh( &self, _: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn rotate_provider_credential( &self, _: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn delete_provider_refresh( &self, _: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn delete_provider( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented( "delete_provider not implemented in test", )) } type WatchSandboxStream = ReceiverStream>; type ExecSandboxStream = ReceiverStream>; type ConnectSupervisorStream = ReceiverStream>; async fn watch_sandbox( &self, _request: tonic::Request, ) -> Result, Status> { let (_tx, rx) = mpsc::channel(1); Ok(Response::new(ReceiverStream::new(rx))) } async fn exec_sandbox( &self, _request: tonic::Request, ) -> Result, Status> { let (_tx, rx) = mpsc::channel(1); Ok(Response::new(ReceiverStream::new(rx))) } type ExecSandboxInteractiveStream = ReceiverStream>; async fn exec_sandbox_interactive( &self, _request: tonic::Request>, ) -> Result, Status> { let (_tx, rx) = mpsc::channel(1); Ok(Response::new(ReceiverStream::new(rx))) } async fn update_config( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn get_sandbox_policy_status( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn list_sandbox_policies( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn report_policy_status( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn get_sandbox_logs( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn push_sandbox_logs( &self, _request: tonic::Request>, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn submit_policy_analysis( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn get_draft_policy( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn approve_draft_chunk( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn reject_draft_chunk( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn approve_all_draft_chunks( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn edit_draft_chunk( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn undo_draft_chunk( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn clear_draft_chunks( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn get_draft_history( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn issue_sandbox_token( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn refresh_sandbox_token( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn connect_supervisor( &self, _request: tonic::Request>, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } type RelayStreamStream = ReceiverStream>; async fn relay_stream( &self, _request: tonic::Request>, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } type ForwardTcpStream = std::pin::Pin> + Send>>; async fn forward_tcp( &self, _request: tonic::Request>, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn create_workspace( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn get_workspace( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn list_workspaces( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn delete_workspace( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn add_workspace_member( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn remove_workspace_member( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn list_workspace_members( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } } // --------------------------------------------------------------------------- // TLS / PKI helpers (used by TLS integration tests) // --------------------------------------------------------------------------- /// Initialise the rustls crypto provider (idempotent). pub fn install_rustls_provider() { let _ = rustls::crypto::ring::default_provider().install_default(); } /// PKI bundle: CA cert, server cert+key, client cert+key (all PEM). #[allow(clippy::struct_field_names)] pub struct PkiBundle { pub ca_cert_pem: Vec, pub server_cert_pem: Vec, pub server_key_pem: Vec, pub client_cert_pem: Vec, pub client_key_pem: Vec, } /// Generate a full PKI: CA → server cert (for `localhost`) + client cert. /// Returns a `TempDir` that must be kept alive while the paths are in use. pub fn generate_pki() -> (tempfile::TempDir, PkiBundle) { // Generate CA let mut ca_params = CertificateParams::new(Vec::::new()).expect("failed to create CA params"); ca_params.is_ca = IsCa::Ca(rcgen::BasicConstraints::Unconstrained); ca_params .distinguished_name .push(rcgen::DnType::CommonName, "test-ca"); let ca_key = KeyPair::generate().expect("failed to generate CA key"); let ca_cert = ca_params .self_signed(&ca_key) .expect("failed to sign CA cert"); // Generate server cert signed by CA let server_params = CertificateParams::new(vec!["localhost".to_string()]) .expect("failed to create server params"); let server_key = KeyPair::generate().expect("failed to generate server key"); let server_cert = server_params .signed_by(&server_key, &ca_cert, &ca_key) .expect("failed to sign server cert"); // Generate client cert signed by CA let mut client_params = CertificateParams::new(Vec::::new()).expect("failed to create client params"); client_params .distinguished_name .push(rcgen::DnType::CommonName, "test-client"); let client_key = KeyPair::generate().expect("failed to generate client key"); let client_cert = client_params .signed_by(&client_key, &ca_cert, &ca_key) .expect("failed to sign client cert"); let dir = tempdir().expect("failed to create tempdir"); let write_file = |name: &str, data: &[u8]| { let path = dir.path().join(name); std::fs::File::create(&path) .and_then(|mut f| f.write_all(data)) .expect("failed to write file"); }; write_file("ca.pem", ca_cert.pem().as_bytes()); write_file("server-cert.pem", server_cert.pem().as_bytes()); write_file("server-key.pem", server_key.serialize_pem().as_bytes()); write_file("client-cert.pem", client_cert.pem().as_bytes()); write_file("client-key.pem", client_key.serialize_pem().as_bytes()); let bundle = PkiBundle { ca_cert_pem: ca_cert.pem().into_bytes(), server_cert_pem: server_cert.pem().into_bytes(), server_key_pem: server_key.serialize_pem().into_bytes(), client_cert_pem: client_cert.pem().into_bytes(), client_key_pem: client_key.serialize_pem().into_bytes(), }; (dir, bundle) } /// Start a TLS-wrapped test server using the given `TlsAcceptor`. /// Returns the bound address and a task handle (abort to stop). pub async fn start_test_server( tls_acceptor: TlsAcceptor, ) -> (SocketAddr, tokio::task::JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let grpc_service = OpenShellServer::new(TestOpenShell); let http_service = health_router(test_health_store().await); let service = MultiplexedService::new(grpc_service, http_service); let handle = tokio::spawn(async move { loop { let Ok((stream, _)) = listener.accept().await else { continue; }; let svc = service.clone(); let tls = tls_acceptor.clone(); tokio::spawn(async move { let Ok(tls_stream) = tls.acceptor().accept(stream).await else { return; }; let _ = Builder::new(TokioExecutor::new()) .serve_connection(TokioIo::new(tls_stream), svc) .await; }); } }); (addr, handle) } /// Rogue PKI bundle: client cert + key not signed by the server's CA. pub struct RoguePkiBundle { pub client_cert_pem: String, pub client_key_pem: String, } /// Generate a rogue CA and a client certificate signed by that CA. /// /// Used to verify that the server rejects mTLS connections from clients whose /// certificate chain does not trace back to the trusted CA. pub fn generate_rogue_pki() -> RoguePkiBundle { let mut rogue_ca_params = CertificateParams::new(Vec::::new()).expect("failed to create rogue CA params"); rogue_ca_params.is_ca = IsCa::Ca(rcgen::BasicConstraints::Unconstrained); rogue_ca_params .distinguished_name .push(rcgen::DnType::CommonName, "rogue-ca"); let rogue_ca_key = KeyPair::generate().expect("failed to generate rogue CA key"); let rogue_ca_cert = rogue_ca_params .self_signed(&rogue_ca_key) .expect("failed to sign rogue CA cert"); let mut rogue_client_params = CertificateParams::new(Vec::::new()).expect("failed to create rogue client params"); rogue_client_params .distinguished_name .push(rcgen::DnType::CommonName, "rogue-client"); let rogue_client_key = KeyPair::generate().expect("failed to generate rogue client key"); let rogue_client_cert = rogue_client_params .signed_by(&rogue_client_key, &rogue_ca_cert, &rogue_ca_key) .expect("failed to sign rogue client cert"); RoguePkiBundle { client_cert_pem: rogue_client_cert.pem(), client_key_pem: rogue_client_key.serialize_pem(), } } /// Build an in-memory store sufficient for wiring `health_router` in tests /// where the persistence layer itself is not under test. pub async fn test_health_store() -> Arc { Arc::new( Store::connect("sqlite::memory:") .await .expect("connect in-memory sqlite store for tests"), ) } /// Parse PEM cert bytes into a `RootCertStore`. pub fn build_tls_root(cert_pem: &[u8]) -> RootCertStore { let mut roots = RootCertStore::empty(); let mut cursor = std::io::Cursor::new(cert_pem); let parsed = certs(&mut cursor) .collect::>, _>>() .expect("failed to parse cert pem"); for cert in parsed { roots.add(cert).expect("failed to add cert"); } roots } /// Build a gRPC client with mTLS (CA + client cert). pub async fn grpc_client_mtls( addr: SocketAddr, ca_pem: Vec, client_cert_pem: Vec, client_key_pem: Vec, ) -> OpenShellClient { let ca_cert = tonic::transport::Certificate::from_pem(ca_pem); let identity = tonic::transport::Identity::from_pem(client_cert_pem, client_key_pem); let tls = ClientTlsConfig::new() .ca_certificate(ca_cert) .identity(identity) .domain_name("localhost"); let endpoint = Endpoint::from_shared(format!("https://localhost:{}", addr.port())) .expect("invalid endpoint") .tls_config(tls) .expect("failed to set tls"); let channel = endpoint.connect().await.expect("failed to connect"); OpenShellClient::new(channel) }