mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-14 04:59:28 +08:00
* feat(kubernetes): support HA gateway rebalancing Signed-off-by: Drew Newberry <anewberry@nvidia.com> * perf(server): cache peer connections, tokens, and owner lookups Every forwarded relay rebuilt its setup from scratch: an owner lookup, a blocking read of the peer token, a TLS connect to the owning replica, and a TokenReview plus Pod GET on the receiving side. Sandbox service routing does this per HTTP request, so the apiserver calls scaled with traffic. Cache all of it on ServerState: - peer channels pooled per endpoint, so relays multiplex over one connection instead of redialing - peer tokens keyed by SHA-256, expiring at min(ttl, token exp) so a hit cannot accept an expired token - owner records for 3s against a 45s ownership TTL, still freshness checked before use Entries are evicted when a relay fails. Also raise HTTP/2 max_concurrent_streams to 1024, since pooling funnels every relay between two replicas onto one connection and hyper's default of 200 sits below the 256 pending-relay budget. Signed-off-by: divesh <dgude@nvidia.com> * perf(server): pool upstream connections for sandbox services Each HTTP request to a sandbox service opened its own supervisor relay, paying a new TCP connection and HTTP/1 handshake every time. Worse, it counted against the 32 in-flight relay cap, so a service handling more than 32 concurrent requests failed outright. Pool idle upstreams per endpoint and port, up to 8 each for 15s. Reuse is safe because the pool only returns a connection hyper reports as ready, and HTTP/1 cannot start a request until the previous body has drained. Upgrades are never pooled since they take the connection over, and a failed send evicts that endpoint. Pruning is bounded per key, with the full sweep limited to once per 30s. Signed-off-by: divesh <dgude@nvidia.com> * fix(server): address HA gateway review findings (#3449) - Let a gateway own supervisor sessions without a peer endpoint. Requiring one whenever the store is PostgreSQL broke every single-instance PostgreSQL deployment, because no sandbox supervisor could connect. A cross-replica request to an owner that advertises no endpoint now fails immediately naming the cause, instead of retrying until the wait timeout. - Close a supervisor session on heartbeat only when another replica owns it, or after renewals fail for the ownership TTL. A database error no longer drops every session heartbeating during an outage. - Clamp owner record ages at zero so a skewed or corrupt stored timestamp cannot produce a negative age. - Bound the cross-object advisory lock with a lock timeout, so a stuck holder fails instead of blocking every mutation in the fleet. - Refuse to start when a peer endpoint is configured on a multi-replica backend but peer authentication is unavailable, and warn when a multi-replica backend has no peer endpoint at all. - Reject a plaintext peer endpoint when the gateway serves TLS. - Skip the sandbox watch poller on single-replica backends, where the local update bus already sees every write. - Rate-limit the peer owner cache sweep so an insert no longer scans the whole map under the lock. - Retry GET and HEAD on a pooled upstream the sandbox closed, instead of returning 502, and drop an emptied endpoint from the pool right away. - Document the gateway peer environment variables and the post-rollout ownership skew operators should expect. Signed-off-by: divesh <dgude@nvidia.com> * fix(server): harden HA supervisor ownership Signed-off-by: Drew Newberry <anewberry@nvidia.com> --------- Signed-off-by: Drew Newberry <anewberry@nvidia.com> Signed-off-by: divesh <dgude@nvidia.com> Co-authored-by: Drew Newberry <anewberry@nvidia.com> Co-authored-by: divesh <dgude@nvidia.com> Co-authored-by: Divesh Chowdary <47188680+FrostGod@users.noreply.github.com>
863 lines
33 KiB
Rust
863 lines
33 KiB
Rust
// 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,
|
|
ExchangeProviderSubjectTokenRequest, ExchangeProviderSubjectTokenResponse, ExecSandboxEvent,
|
|
ExecSandboxInput, ExecSandboxRequest, GatewayMessage, GetGatewayConfigRequest,
|
|
GetGatewayConfigResponse, GetProviderRequest, GetSandboxConfigRequest,
|
|
GetSandboxConfigResponse, GetSandboxProviderEnvironmentRequest,
|
|
GetSandboxProviderEnvironmentResponse, GetSandboxRequest, HealthRequest, HealthResponse,
|
|
IssueSandboxTokenRequest, IssueSandboxTokenResponse, ListProvidersRequest,
|
|
ListProvidersResponse, ListSandboxesRequest, ListSandboxesResponse, PeerRelayFrame,
|
|
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 peer_report_provider_readiness(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ReportProviderReadinessRequest>,
|
|
) -> Result<Response<openshell_core::proto::ReportProviderReadinessResponse>, Status> {
|
|
Err(Status::unimplemented("not used by this test server"))
|
|
}
|
|
|
|
async fn peer_report_endpoint_status(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ReportEndpointStatusRequest>,
|
|
) -> Result<Response<openshell_core::proto::ReportEndpointStatusResponse>, Status> {
|
|
Err(Status::unimplemented("not used by this test server"))
|
|
}
|
|
|
|
async fn peer_get_sandbox_provider_status(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetSandboxProviderStatusRequest>,
|
|
) -> Result<Response<openshell_core::proto::GetSandboxProviderStatusResponse>, Status> {
|
|
Err(Status::unimplemented("not used by this test server"))
|
|
}
|
|
|
|
async fn report_endpoint_status(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ReportEndpointStatusRequest>,
|
|
) -> Result<Response<openshell_core::proto::ReportEndpointStatusResponse>, Status> {
|
|
Ok(Response::new(
|
|
openshell_core::proto::ReportEndpointStatusResponse {},
|
|
))
|
|
}
|
|
|
|
async fn begin_rootfs_tar_staging(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::BeginRootfsTarStagingRequest>,
|
|
) -> Result<Response<openshell_core::proto::BeginRootfsTarStagingResponse>, Status> {
|
|
Err(Status::unimplemented("not used by this test server"))
|
|
}
|
|
|
|
async fn report_main_process_exit(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ReportMainProcessExitRequest>,
|
|
) -> Result<Response<openshell_core::proto::ReportMainProcessExitResponse>, Status> {
|
|
Err(Status::unimplemented("not used by this test server"))
|
|
}
|
|
|
|
async fn finalize_main_process_exit(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::FinalizeMainProcessExitRequest>,
|
|
) -> Result<Response<openshell_core::proto::FinalizeMainProcessExitResponse>, Status> {
|
|
Err(Status::unimplemented("not used by this test server"))
|
|
}
|
|
|
|
async fn get_current_user(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetCurrentUserRequest>,
|
|
) -> Result<Response<openshell_core::proto::GetCurrentUserResponse>, Status> {
|
|
Err(Status::unimplemented("not used by this test server"))
|
|
}
|
|
|
|
async fn health(
|
|
&self,
|
|
_request: tonic::Request<HealthRequest>,
|
|
) -> Result<Response<HealthResponse>, Status> {
|
|
Ok(Response::new(HealthResponse {
|
|
status: ServiceStatus::Healthy.into(),
|
|
version: "test".to_string(),
|
|
}))
|
|
}
|
|
|
|
async fn get_gateway_info(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetGatewayInfoRequest>,
|
|
) -> Result<Response<openshell_core::proto::GetGatewayInfoResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn create_sandbox(
|
|
&self,
|
|
_request: tonic::Request<CreateSandboxRequest>,
|
|
) -> Result<Response<SandboxResponse>, Status> {
|
|
Ok(Response::new(SandboxResponse::default()))
|
|
}
|
|
|
|
async fn create_sandbox_template(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::CreateSandboxTemplateRequest>,
|
|
) -> Result<Response<openshell_core::proto::SandboxTemplateResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn get_sandbox_template(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetSandboxTemplateRequest>,
|
|
) -> Result<Response<openshell_core::proto::SandboxTemplateResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn list_sandbox_templates(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ListSandboxTemplatesRequest>,
|
|
) -> Result<Response<openshell_core::proto::ListSandboxTemplatesResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn delete_sandbox_template(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::DeleteSandboxTemplateRequest>,
|
|
) -> Result<Response<openshell_core::proto::DeleteSandboxTemplateResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn stop_sandbox(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::StopSandboxRequest>,
|
|
) -> Result<Response<SandboxResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn start_sandbox(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::StartSandboxRequest>,
|
|
) -> Result<Response<SandboxResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn get_sandbox(
|
|
&self,
|
|
_request: tonic::Request<GetSandboxRequest>,
|
|
) -> Result<Response<SandboxResponse>, Status> {
|
|
Ok(Response::new(SandboxResponse::default()))
|
|
}
|
|
|
|
async fn list_sandboxes(
|
|
&self,
|
|
_request: tonic::Request<ListSandboxesRequest>,
|
|
) -> Result<Response<ListSandboxesResponse>, Status> {
|
|
Ok(Response::new(ListSandboxesResponse::default()))
|
|
}
|
|
|
|
async fn list_sandbox_providers(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ListSandboxProvidersRequest>,
|
|
) -> Result<Response<openshell_core::proto::ListSandboxProvidersResponse>, Status> {
|
|
Ok(Response::new(
|
|
openshell_core::proto::ListSandboxProvidersResponse::default(),
|
|
))
|
|
}
|
|
|
|
async fn attach_sandbox_provider(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::AttachSandboxProviderRequest>,
|
|
) -> Result<Response<openshell_core::proto::AttachSandboxProviderResponse>, Status> {
|
|
Ok(Response::new(
|
|
openshell_core::proto::AttachSandboxProviderResponse::default(),
|
|
))
|
|
}
|
|
|
|
async fn detach_sandbox_provider(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::DetachSandboxProviderRequest>,
|
|
) -> Result<Response<openshell_core::proto::DetachSandboxProviderResponse>, Status> {
|
|
Ok(Response::new(
|
|
openshell_core::proto::DetachSandboxProviderResponse::default(),
|
|
))
|
|
}
|
|
|
|
async fn delete_sandbox(
|
|
&self,
|
|
_request: tonic::Request<DeleteSandboxRequest>,
|
|
) -> Result<Response<DeleteSandboxResponse>, Status> {
|
|
Ok(Response::new(DeleteSandboxResponse {
|
|
sandbox_id: String::new(),
|
|
outcome: openshell_core::proto::DeletionOutcome::Completed.into(),
|
|
}))
|
|
}
|
|
|
|
async fn get_sandbox_config(
|
|
&self,
|
|
_request: tonic::Request<GetSandboxConfigRequest>,
|
|
) -> Result<Response<GetSandboxConfigResponse>, Status> {
|
|
Ok(Response::new(GetSandboxConfigResponse::default()))
|
|
}
|
|
|
|
async fn get_gateway_config(
|
|
&self,
|
|
_request: tonic::Request<GetGatewayConfigRequest>,
|
|
) -> Result<Response<GetGatewayConfigResponse>, Status> {
|
|
Ok(Response::new(GetGatewayConfigResponse::default()))
|
|
}
|
|
|
|
async fn get_sandbox_provider_status(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetSandboxProviderStatusRequest>,
|
|
) -> Result<Response<openshell_core::proto::GetSandboxProviderStatusResponse>, Status> {
|
|
Err(Status::unimplemented(
|
|
"provider readiness is not exercised by this mock",
|
|
))
|
|
}
|
|
|
|
async fn report_provider_readiness(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ReportProviderReadinessRequest>,
|
|
) -> Result<Response<openshell_core::proto::ReportProviderReadinessResponse>, Status> {
|
|
Err(Status::unimplemented("provider readiness"))
|
|
}
|
|
|
|
async fn get_sandbox_provider_environment(
|
|
&self,
|
|
_request: tonic::Request<GetSandboxProviderEnvironmentRequest>,
|
|
) -> Result<Response<GetSandboxProviderEnvironmentResponse>, Status> {
|
|
Ok(Response::new(
|
|
GetSandboxProviderEnvironmentResponse::default(),
|
|
))
|
|
}
|
|
|
|
async fn create_ssh_session(
|
|
&self,
|
|
_request: tonic::Request<CreateSshSessionRequest>,
|
|
) -> Result<Response<CreateSshSessionResponse>, Status> {
|
|
Ok(Response::new(CreateSshSessionResponse::default()))
|
|
}
|
|
|
|
async fn expose_service(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ExposeServiceRequest>,
|
|
) -> Result<Response<openshell_core::proto::ServiceEndpointResponse>, Status> {
|
|
Ok(Response::new(
|
|
openshell_core::proto::ServiceEndpointResponse::default(),
|
|
))
|
|
}
|
|
|
|
async fn get_service(
|
|
&self,
|
|
_: tonic::Request<openshell_core::proto::GetServiceRequest>,
|
|
) -> Result<Response<openshell_core::proto::ServiceEndpointResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn list_services(
|
|
&self,
|
|
_: tonic::Request<openshell_core::proto::ListServicesRequest>,
|
|
) -> Result<Response<openshell_core::proto::ListServicesResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn delete_service(
|
|
&self,
|
|
_: tonic::Request<openshell_core::proto::DeleteServiceRequest>,
|
|
) -> Result<Response<openshell_core::proto::DeleteServiceResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn revoke_ssh_session(
|
|
&self,
|
|
_request: tonic::Request<RevokeSshSessionRequest>,
|
|
) -> Result<Response<RevokeSshSessionResponse>, Status> {
|
|
Ok(Response::new(RevokeSshSessionResponse::default()))
|
|
}
|
|
|
|
async fn exchange_provider_subject_token(
|
|
&self,
|
|
_request: tonic::Request<ExchangeProviderSubjectTokenRequest>,
|
|
) -> Result<Response<ExchangeProviderSubjectTokenResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn create_provider(
|
|
&self,
|
|
_request: tonic::Request<CreateProviderRequest>,
|
|
) -> Result<Response<ProviderResponse>, Status> {
|
|
Err(Status::unimplemented(
|
|
"create_provider not implemented in test",
|
|
))
|
|
}
|
|
|
|
async fn get_provider(
|
|
&self,
|
|
_request: tonic::Request<GetProviderRequest>,
|
|
) -> Result<Response<ProviderResponse>, Status> {
|
|
Err(Status::unimplemented(
|
|
"get_provider not implemented in test",
|
|
))
|
|
}
|
|
|
|
async fn list_providers(
|
|
&self,
|
|
_request: tonic::Request<ListProvidersRequest>,
|
|
) -> Result<Response<ListProvidersResponse>, Status> {
|
|
Err(Status::unimplemented(
|
|
"list_providers not implemented in test",
|
|
))
|
|
}
|
|
|
|
async fn list_provider_profiles(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ListProviderProfilesRequest>,
|
|
) -> Result<Response<openshell_core::proto::ListProviderProfilesResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn get_provider_profile(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetProviderProfileRequest>,
|
|
) -> Result<Response<openshell_core::proto::ProviderProfileResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn import_provider_profiles(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ImportProviderProfilesRequest>,
|
|
) -> Result<Response<openshell_core::proto::ImportProviderProfilesResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn update_provider_profiles(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::UpdateProviderProfilesRequest>,
|
|
) -> Result<Response<openshell_core::proto::UpdateProviderProfilesResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn lint_provider_profiles(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::LintProviderProfilesRequest>,
|
|
) -> Result<Response<openshell_core::proto::LintProviderProfilesResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn delete_provider_profile(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::DeleteProviderProfileRequest>,
|
|
) -> Result<Response<openshell_core::proto::DeleteProviderProfileResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn update_provider(
|
|
&self,
|
|
_request: tonic::Request<UpdateProviderRequest>,
|
|
) -> Result<Response<ProviderResponse>, Status> {
|
|
Err(Status::unimplemented(
|
|
"update_provider not implemented in test",
|
|
))
|
|
}
|
|
|
|
async fn get_provider_refresh_status(
|
|
&self,
|
|
_: tonic::Request<openshell_core::proto::GetProviderRefreshStatusRequest>,
|
|
) -> Result<Response<openshell_core::proto::GetProviderRefreshStatusResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn configure_provider_refresh(
|
|
&self,
|
|
_: tonic::Request<openshell_core::proto::ConfigureProviderRefreshRequest>,
|
|
) -> Result<Response<openshell_core::proto::ConfigureProviderRefreshResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn rotate_provider_credential(
|
|
&self,
|
|
_: tonic::Request<openshell_core::proto::RotateProviderCredentialRequest>,
|
|
) -> Result<Response<openshell_core::proto::RotateProviderCredentialResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn delete_provider_refresh(
|
|
&self,
|
|
_: tonic::Request<openshell_core::proto::DeleteProviderRefreshRequest>,
|
|
) -> Result<Response<openshell_core::proto::DeleteProviderRefreshResponse>, Status> {
|
|
Err(Status::unimplemented("unused"))
|
|
}
|
|
|
|
async fn delete_provider(
|
|
&self,
|
|
_request: tonic::Request<DeleteProviderRequest>,
|
|
) -> Result<Response<DeleteProviderResponse>, Status> {
|
|
Err(Status::unimplemented(
|
|
"delete_provider not implemented in test",
|
|
))
|
|
}
|
|
|
|
type WatchSandboxStream = ReceiverStream<Result<SandboxStreamEvent, Status>>;
|
|
type ExecSandboxStream = ReceiverStream<Result<ExecSandboxEvent, Status>>;
|
|
type ConnectSupervisorStream = ReceiverStream<Result<GatewayMessage, Status>>;
|
|
|
|
async fn watch_sandbox(
|
|
&self,
|
|
_request: tonic::Request<WatchSandboxRequest>,
|
|
) -> Result<Response<Self::WatchSandboxStream>, Status> {
|
|
let (_tx, rx) = mpsc::channel(1);
|
|
Ok(Response::new(ReceiverStream::new(rx)))
|
|
}
|
|
|
|
async fn exec_sandbox(
|
|
&self,
|
|
_request: tonic::Request<ExecSandboxRequest>,
|
|
) -> Result<Response<Self::ExecSandboxStream>, Status> {
|
|
let (_tx, rx) = mpsc::channel(1);
|
|
Ok(Response::new(ReceiverStream::new(rx)))
|
|
}
|
|
|
|
type ExecSandboxInteractiveStream = ReceiverStream<Result<ExecSandboxEvent, Status>>;
|
|
|
|
async fn exec_sandbox_interactive(
|
|
&self,
|
|
_request: tonic::Request<tonic::Streaming<ExecSandboxInput>>,
|
|
) -> Result<Response<Self::ExecSandboxInteractiveStream>, Status> {
|
|
let (_tx, rx) = mpsc::channel(1);
|
|
Ok(Response::new(ReceiverStream::new(rx)))
|
|
}
|
|
|
|
async fn update_config(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::UpdateConfigRequest>,
|
|
) -> Result<Response<openshell_core::proto::UpdateConfigResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn get_sandbox_policy_status(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetSandboxPolicyStatusRequest>,
|
|
) -> Result<Response<openshell_core::proto::GetSandboxPolicyStatusResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn list_sandbox_policies(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ListSandboxPoliciesRequest>,
|
|
) -> Result<Response<openshell_core::proto::ListSandboxPoliciesResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn report_sandbox_configuration(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ReportSandboxConfigurationRequest>,
|
|
) -> Result<Response<openshell_core::proto::ReportSandboxConfigurationResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn report_policy_status(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ReportPolicyStatusRequest>,
|
|
) -> Result<Response<openshell_core::proto::ReportPolicyStatusResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn get_sandbox_logs(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetSandboxLogsRequest>,
|
|
) -> Result<Response<openshell_core::proto::GetSandboxLogsResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn push_sandbox_logs(
|
|
&self,
|
|
_request: tonic::Request<tonic::Streaming<openshell_core::proto::PushSandboxLogsRequest>>,
|
|
) -> Result<Response<openshell_core::proto::PushSandboxLogsResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn submit_policy_analysis(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::SubmitPolicyAnalysisRequest>,
|
|
) -> Result<Response<openshell_core::proto::SubmitPolicyAnalysisResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn get_draft_policy(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetDraftPolicyRequest>,
|
|
) -> Result<Response<openshell_core::proto::GetDraftPolicyResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn approve_draft_chunk(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ApproveDraftChunkRequest>,
|
|
) -> Result<Response<openshell_core::proto::ApproveDraftChunkResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn reject_draft_chunk(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::RejectDraftChunkRequest>,
|
|
) -> Result<Response<openshell_core::proto::RejectDraftChunkResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn approve_all_draft_chunks(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ApproveAllDraftChunksRequest>,
|
|
) -> Result<Response<openshell_core::proto::ApproveAllDraftChunksResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn edit_draft_chunk(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::EditDraftChunkRequest>,
|
|
) -> Result<Response<openshell_core::proto::EditDraftChunkResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn undo_draft_chunk(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::UndoDraftChunkRequest>,
|
|
) -> Result<Response<openshell_core::proto::UndoDraftChunkResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn clear_draft_chunks(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ClearDraftChunksRequest>,
|
|
) -> Result<Response<openshell_core::proto::ClearDraftChunksResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn get_draft_history(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetDraftHistoryRequest>,
|
|
) -> Result<Response<openshell_core::proto::GetDraftHistoryResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn issue_sandbox_token(
|
|
&self,
|
|
_request: tonic::Request<IssueSandboxTokenRequest>,
|
|
) -> Result<Response<IssueSandboxTokenResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn refresh_sandbox_token(
|
|
&self,
|
|
_request: tonic::Request<RefreshSandboxTokenRequest>,
|
|
) -> Result<Response<RefreshSandboxTokenResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn connect_supervisor(
|
|
&self,
|
|
_request: tonic::Request<tonic::Streaming<SupervisorMessage>>,
|
|
) -> Result<Response<Self::ConnectSupervisorStream>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
type RelayStreamStream = ReceiverStream<Result<RelayFrame, Status>>;
|
|
|
|
async fn relay_stream(
|
|
&self,
|
|
_request: tonic::Request<tonic::Streaming<RelayFrame>>,
|
|
) -> Result<Response<Self::RelayStreamStream>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
type PeerRelayStream = ReceiverStream<Result<PeerRelayFrame, Status>>;
|
|
|
|
async fn peer_relay(
|
|
&self,
|
|
_request: tonic::Request<tonic::Streaming<PeerRelayFrame>>,
|
|
) -> Result<Response<Self::PeerRelayStream>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
type ForwardTcpStream =
|
|
std::pin::Pin<Box<dyn tokio_stream::Stream<Item = Result<TcpForwardFrame, Status>> + Send>>;
|
|
|
|
async fn forward_tcp(
|
|
&self,
|
|
_request: tonic::Request<tonic::Streaming<TcpForwardFrame>>,
|
|
) -> Result<Response<Self::ForwardTcpStream>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn create_workspace(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::CreateWorkspaceRequest>,
|
|
) -> Result<Response<openshell_core::proto::CreateWorkspaceResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn get_workspace(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::GetWorkspaceRequest>,
|
|
) -> Result<Response<openshell_core::proto::GetWorkspaceResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn list_workspaces(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ListWorkspacesRequest>,
|
|
) -> Result<Response<openshell_core::proto::ListWorkspacesResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn delete_workspace(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::DeleteWorkspaceRequest>,
|
|
) -> Result<Response<openshell_core::proto::DeleteWorkspaceResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn add_workspace_member(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::AddWorkspaceMemberRequest>,
|
|
) -> Result<Response<openshell_core::proto::AddWorkspaceMemberResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn remove_workspace_member(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::RemoveWorkspaceMemberRequest>,
|
|
) -> Result<Response<openshell_core::proto::RemoveWorkspaceMemberResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
|
|
async fn list_workspace_members(
|
|
&self,
|
|
_request: tonic::Request<openshell_core::proto::ListWorkspaceMembersRequest>,
|
|
) -> Result<Response<openshell_core::proto::ListWorkspaceMembersResponse>, Status> {
|
|
Err(Status::unimplemented("not implemented in test"))
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// TLS / PKI helpers (used by TLS integration tests)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// 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<u8>,
|
|
pub server_cert_pem: Vec<u8>,
|
|
pub server_key_pem: Vec<u8>,
|
|
pub client_cert_pem: Vec<u8>,
|
|
pub client_key_pem: Vec<u8>,
|
|
}
|
|
|
|
/// 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::<String>::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::<String>::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::<String>::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::<String>::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<Store> {
|
|
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::<Result<Vec<CertificateDer<'static>>, _>>()
|
|
.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<u8>,
|
|
client_cert_pem: Vec<u8>,
|
|
client_key_pem: Vec<u8>,
|
|
) -> OpenShellClient<Channel> {
|
|
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)
|
|
}
|