// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 use openshell_bootstrap::load_last_sandbox; use openshell_cli::run; use openshell_cli::tls::TlsOptions; use openshell_core::proto::open_shell_server::{OpenShell, OpenShellServer}; use openshell_core::proto::{ CreateProviderRequest, CreateSandboxRequest, CreateSshSessionRequest, CreateSshSessionResponse, DeleteProviderRequest, DeleteProviderResponse, DeleteSandboxRequest, DeleteSandboxResponse, ExecSandboxEvent, ExecSandboxRequest, GatewayMessage, GetGatewayConfigRequest, GetGatewayConfigResponse, GetProviderRequest, GetSandboxConfigRequest, GetSandboxConfigResponse, GetSandboxProviderEnvironmentRequest, GetSandboxProviderEnvironmentResponse, GetSandboxRequest, HealthRequest, HealthResponse, ListProvidersRequest, ListProvidersResponse, ListSandboxesRequest, ListSandboxesResponse, PlatformEvent, ProviderResponse, RevokeSshSessionRequest, RevokeSshSessionResponse, Sandbox, SandboxPhase, SandboxResponse, SandboxStreamEvent, ServiceStatus, SupervisorMessage, UpdateProviderRequest, WatchSandboxRequest, sandbox_stream_event, }; use rcgen::{ BasicConstraints, Certificate, CertificateParams, ExtendedKeyUsagePurpose, IsCa, KeyPair, }; use std::collections::HashMap; use std::fs; use std::os::unix::fs::PermissionsExt; use std::sync::Arc; use tempfile::TempDir; use tokio::net::TcpListener; use tokio::sync::{Mutex, mpsc}; use tokio_stream::wrappers::TcpListenerStream; use tonic::transport::{Certificate as TlsCertificate, Identity, Server, ServerTlsConfig}; use tonic::{Response, Status}; static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); struct SavedVar { key: &'static str, original: Option, } struct EnvVarGuard { vars: Vec, _lock: std::sync::MutexGuard<'static, ()>, } #[allow(unsafe_code)] impl EnvVarGuard { fn set(pairs: &[(&'static str, String)]) -> 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); } } } } } #[derive(Clone, Default)] struct SandboxState { deleted_names: Arc>>>, } #[derive(Clone, Default)] struct TestOpenShell { state: SandboxState, } #[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 create_sandbox( &self, request: tonic::Request, ) -> Result, Status> { let name = request.into_inner().name; let sandbox_name = if name.is_empty() { "test-sandbox".to_string() } else { name }; Ok(Response::new(SandboxResponse { sandbox: Some(Sandbox { metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { id: format!("id-{sandbox_name}"), name: sandbox_name, created_at_ms: 0, labels: HashMap::new(), }), phase: SandboxPhase::Provisioning as i32, ..Sandbox::default() }), })) } async fn get_sandbox( &self, request: tonic::Request, ) -> Result, Status> { let name = request.into_inner().name; Ok(Response::new(SandboxResponse { sandbox: Some(Sandbox { metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { id: format!("id-{name}"), name, created_at_ms: 0, labels: HashMap::new(), }), phase: SandboxPhase::Ready as i32, ..Sandbox::default() }), })) } async fn list_sandboxes( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(ListSandboxesResponse::default())) } async fn delete_sandbox( &self, request: tonic::Request, ) -> Result, Status> { let request = request.into_inner(); self.state .deleted_names .lock() .await .push(vec![request.name]); 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> { let sandbox_id = request.into_inner().sandbox_id; Ok(Response::new(CreateSshSessionResponse { sandbox_id, token: "test-token".to_string(), gateway_scheme: "https".to_string(), gateway_host: "localhost".to_string(), gateway_port: 443, connect_path: "/connect/ssh".to_string(), ..CreateSshSessionResponse::default() })) } 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> { Ok(Response::new(ProviderResponse::default())) } async fn get_provider( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::not_found("provider not found")) } async fn list_providers( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(ListProvidersResponse::default())) } async fn update_provider( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(ProviderResponse::default())) } async fn delete_provider( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(DeleteProviderResponse { deleted: true })) } type WatchSandboxStream = tokio_stream::wrappers::ReceiverStream>; type ExecSandboxStream = tokio_stream::wrappers::ReceiverStream>; type ConnectSupervisorStream = tokio_stream::wrappers::ReceiverStream>; async fn watch_sandbox( &self, request: tonic::Request, ) -> Result, Status> { let sandbox_id = request.into_inner().id; let (tx, rx) = mpsc::channel(4); tokio::spawn(async move { let provisioning = Sandbox { metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { id: sandbox_id.clone(), name: sandbox_id.trim_start_matches("id-").to_string(), created_at_ms: 0, labels: HashMap::new(), }), phase: SandboxPhase::Provisioning as i32, ..Sandbox::default() }; let ready = Sandbox { phase: SandboxPhase::Ready as i32, ..provisioning.clone() }; let _ = tx .send(Ok(SandboxStreamEvent { payload: Some(sandbox_stream_event::Payload::Sandbox(provisioning)), })) .await; let _ = tx .send(Ok(SandboxStreamEvent { payload: Some(sandbox_stream_event::Payload::Event(PlatformEvent { reason: "Scheduled".to_string(), message: "Sandbox scheduled".to_string(), ..PlatformEvent::default() })), })) .await; let _ = tx .send(Ok(SandboxStreamEvent { payload: Some(sandbox_stream_event::Payload::Sandbox(ready)), })) .await; }); Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new( rx, ))) } async fn exec_sandbox( &self, _request: tonic::Request, ) -> Result, Status> { let (_tx, rx) = mpsc::channel(1); Ok(Response::new(tokio_stream::wrappers::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 connect_supervisor( &self, _request: tonic::Request>, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } type RelayStreamStream = tokio_stream::wrappers::ReceiverStream< Result, >; async fn relay_stream( &self, _request: tonic::Request>, ) -> Result, tonic::Status> { Err(tonic::Status::unimplemented("not implemented in test")) } } fn install_rustls_provider() { let _ = rustls::crypto::ring::default_provider().install_default(); } fn build_ca() -> (Certificate, KeyPair) { let key_pair = KeyPair::generate().unwrap(); let mut params = CertificateParams::new(Vec::::new()).unwrap(); params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); let cert = params.self_signed(&key_pair).unwrap(); (cert, key_pair) } 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()) } fn build_client_cert(ca: &Certificate, ca_key: &KeyPair) -> (String, String) { let key_pair = KeyPair::generate().unwrap(); let mut params = CertificateParams::new(Vec::::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()) } struct TestServer { endpoint: String, tls: TlsOptions, openshell: TestOpenShell, _dir: TempDir, } async fn run_server() -> TestServer { install_rustls_provider(); let (ca, ca_key) = build_ca(); let (server_cert, server_key) = build_server_cert(&ca, &ca_key); let (client_cert, client_key) = build_client_cert(&ca, &ca_key); let ca_cert = ca.pem(); let identity = Identity::from_pem(server_cert, server_key); let client_ca = TlsCertificate::from_pem(ca_cert.clone()); let tls_config = ServerTlsConfig::new() .identity(identity) .client_ca_root(client_ca); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let incoming = TcpListenerStream::new(listener); let openshell = TestOpenShell::default(); let svc_openshell = openshell.clone(); tokio::spawn(async move { Server::builder() .tls_config(tls_config) .unwrap() .add_service(OpenShellServer::new(svc_openshell)) .serve_with_incoming(incoming) .await .unwrap(); }); let dir = tempfile::tempdir().unwrap(); let ca_path = dir.path().join("ca.crt"); let cert_path = dir.path().join("tls.crt"); let key_path = dir.path().join("tls.key"); fs::write(&ca_path, ca_cert).unwrap(); fs::write(&cert_path, client_cert).unwrap(); fs::write(&key_path, client_key).unwrap(); let tls = TlsOptions::new(Some(ca_path), Some(cert_path), Some(key_path)); let endpoint = format!("https://localhost:{}", addr.port()); TestServer { endpoint, tls, openshell, _dir: dir, } } fn install_fake_ssh(dir: &TempDir) -> std::path::PathBuf { let ssh_path = dir.path().join("ssh"); fs::write(&ssh_path, "#!/bin/sh\nexit 0\n").unwrap(); let mut perms = fs::metadata(&ssh_path).unwrap().permissions(); perms.set_mode(0o755); fs::set_permissions(&ssh_path, perms).unwrap(); ssh_path } fn test_env(fake_ssh_dir: &TempDir, xdg_dir: &TempDir) -> EnvVarGuard { let path = format!( "{}:{}", fake_ssh_dir.path().display(), std::env::var("PATH").unwrap_or_default() ); EnvVarGuard::set(&[ ("PATH", path), ( "XDG_CONFIG_HOME", xdg_dir.path().to_str().unwrap().to_string(), ), ("HOME", xdg_dir.path().to_str().unwrap().to_string()), ]) } async fn deleted_names(server: &TestServer) -> Vec> { server.openshell.state.deleted_names.lock().await.clone() } fn test_tls(server: &TestServer) -> TlsOptions { server.tls.with_gateway_name("openshell") } #[tokio::test] async fn sandbox_create_keeps_command_sessions_by_default() { let server = run_server().await; let fake_ssh_dir = tempfile::tempdir().unwrap(); let xdg_dir = tempfile::tempdir().unwrap(); let _env = test_env(&fake_ssh_dir, &xdg_dir); let tls = test_tls(&server); install_fake_ssh(&fake_ssh_dir); run::sandbox_create( &server.endpoint, Some("default-command"), None, "openshell", None, true, false, None, None, None, &[], None, None, &["echo".to_string(), "OK".to_string()], Some(false), Some(false), Some(false), &HashMap::new(), &tls, ) .await .expect("sandbox create should succeed"); assert!(deleted_names(&server).await.is_empty()); assert_eq!( load_last_sandbox("openshell").as_deref(), Some("default-command"), "default sandboxes should be persisted as last-used" ); } #[tokio::test] async fn sandbox_create_deletes_command_sessions_with_no_keep() { let server = run_server().await; let fake_ssh_dir = tempfile::tempdir().unwrap(); let xdg_dir = tempfile::tempdir().unwrap(); let _env = test_env(&fake_ssh_dir, &xdg_dir); let tls = test_tls(&server); install_fake_ssh(&fake_ssh_dir); run::sandbox_create( &server.endpoint, Some("ephemeral-command"), None, "openshell", None, false, false, None, None, None, &[], None, None, &["echo".to_string(), "OK".to_string()], Some(false), Some(false), Some(false), &HashMap::new(), &tls, ) .await .expect("sandbox create should succeed"); assert_eq!( deleted_names(&server).await, vec![vec!["ephemeral-command".to_string()]] ); assert_eq!( load_last_sandbox("openshell"), None, "no-keep sandboxes should not be persisted as last-used" ); } #[tokio::test] async fn sandbox_create_deletes_shell_sessions_with_no_keep() { let server = run_server().await; let fake_ssh_dir = tempfile::tempdir().unwrap(); let xdg_dir = tempfile::tempdir().unwrap(); let _env = test_env(&fake_ssh_dir, &xdg_dir); let tls = test_tls(&server); install_fake_ssh(&fake_ssh_dir); run::sandbox_create( &server.endpoint, Some("ephemeral-shell"), None, "openshell", None, false, false, None, None, None, &[], None, None, &[], Some(true), Some(false), Some(false), &HashMap::new(), &tls, ) .await .expect("sandbox create shell should succeed"); assert_eq!( deleted_names(&server).await, vec![vec!["ephemeral-shell".to_string()]] ); assert_eq!( load_last_sandbox("openshell"), None, "no-keep shell sessions should not be persisted as last-used" ); } #[tokio::test] async fn sandbox_create_keeps_sandbox_with_hidden_keep_flag() { let server = run_server().await; let fake_ssh_dir = tempfile::tempdir().unwrap(); let xdg_dir = tempfile::tempdir().unwrap(); let _env = test_env(&fake_ssh_dir, &xdg_dir); let tls = test_tls(&server); install_fake_ssh(&fake_ssh_dir); run::sandbox_create( &server.endpoint, Some("persistent-keep"), None, "openshell", None, true, false, None, None, None, &[], None, None, &["echo".to_string(), "OK".to_string()], Some(false), Some(false), Some(false), &HashMap::new(), &tls, ) .await .expect("sandbox create should succeed"); assert!(deleted_names(&server).await.is_empty()); assert_eq!( load_last_sandbox("openshell").as_deref(), Some("persistent-keep"), "persistent sandboxes should remain selectable as last-used" ); } #[tokio::test] async fn sandbox_create_keeps_sandbox_with_forwarding() { let server = run_server().await; let fake_ssh_dir = tempfile::tempdir().unwrap(); let xdg_dir = tempfile::tempdir().unwrap(); let _env = test_env(&fake_ssh_dir, &xdg_dir); let tls = test_tls(&server); install_fake_ssh(&fake_ssh_dir); run::sandbox_create( &server.endpoint, Some("persistent-forward"), None, "openshell", None, false, false, None, None, None, &[], None, Some(openshell_core::forward::ForwardSpec::new(8080)), &["echo".to_string(), "OK".to_string()], Some(false), Some(false), Some(false), &HashMap::new(), &tls, ) .await .expect("sandbox create with forward should succeed"); assert!(deleted_names(&server).await.is_empty()); }