From d3cc825fe94185aac5e671f06100f4e1847ca6e2 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Fri, 17 Jul 2026 16:03:06 +0100 Subject: [PATCH 01/10] fix(python): make proto generation install grpcio tools Signed-off-by: Gordon Sim --- tasks/python.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tasks/python.toml b/tasks/python.toml index e04bf27bb4..22e4e1f3c3 100644 --- a/tasks/python.toml +++ b/tasks/python.toml @@ -221,7 +221,7 @@ description = "Generate Python protobuf stubs from .proto files" run = """ #!/usr/bin/env bash set -euo pipefail -uv run --frozen python -m grpc_tools.protoc \ +uv run --frozen --group dev python -m grpc_tools.protoc \ -Iproto \ --python_out=python/openshell/_proto \ --pyi_out=python/openshell/_proto \ @@ -232,7 +232,7 @@ uv run --frozen python -m grpc_tools.protoc \ proto/options.proto \ proto/sandbox.proto # Fix absolute imports in generated stubs to use package-relative imports -uv run --frozen python - <<'PY' +uv run --frozen --group dev python - <<'PY' from pathlib import Path import re From 95f9c70ee9366bc0cdbab574d067f8bce3793e7a Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Fri, 17 Jul 2026 16:04:40 +0100 Subject: [PATCH 02/10] feat(k8s): bootstrap supervisors through pod registration stream Signed-off-by: Gordon Sim --- architecture/compute-runtimes.md | 9 + architecture/gateway.md | 28 +-- .../tests/ensure_providers_integration.rs | 11 ++ .../openshell-cli/tests/mtls_integration.rs | 11 ++ .../tests/provider_commands_integration.rs | 11 ++ .../sandbox_create_lifecycle_integration.rs | 11 ++ .../sandbox_name_fallback_integration.rs | 11 ++ crates/openshell-core/src/grpc_client.rs | 36 ++-- crates/openshell-core/src/sandbox_env.rs | 7 +- crates/openshell-driver-kubernetes/README.md | 5 +- .../openshell-driver-kubernetes/src/config.rs | 6 +- .../openshell-driver-kubernetes/src/driver.rs | 13 +- .../openshell-driver-kubernetes/src/main.rs | 4 +- crates/openshell-sdk/tests/client_mock.rs | 10 + .../src/auth/authenticator.rs | 2 +- crates/openshell-server/src/auth/k8s_sa.rs | 44 +++-- crates/openshell-server/src/auth/principal.rs | 4 +- .../openshell-server/src/auth/sandbox_jwt.rs | 2 +- .../src/auth/sandbox_methods.rs | 3 + crates/openshell-server/src/config_file.rs | 2 +- crates/openshell-server/src/grpc/auth_rpc.rs | 181 +++++++++++++----- crates/openshell-server/src/grpc/mod.rs | 33 +++- crates/openshell-server/src/grpc/sandbox.rs | 2 +- crates/openshell-server/src/lib.rs | 4 +- crates/openshell-server/src/multiplex.rs | 9 +- crates/openshell-server/tests/common/mod.rs | 18 +- .../tests/supervisor_relay_integration.rs | 10 +- proto/openshell.proto | 30 +++ 28 files changed, 388 insertions(+), 129 deletions(-) diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index f211f906d6..7145e90ad3 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -125,6 +125,15 @@ Driver-controlled environment variables must override sandbox image or template values for sandbox ID, sandbox name, gateway endpoint, relay socket path, TLS paths, and command metadata. +Kubernetes supervisors authenticate to the gateway in two stages. They first +call `RegisterSupervisorPod` with the projected ServiceAccount token; the +gateway validates the pod-bound token and live Agent Sandbox owner state, then +activates already-bound cold pods by streaming back a gateway-minted sandbox +JWT. The supervisor installs that JWT in memory and starts the normal +`ConnectSupervisor` session as the activated sandbox. The gateway does not dial +pod IPs or require an inbound activation port; activation is supervisor +initiated over the existing outbound gRPC connection. + Kubernetes can run the supervisor in the default combined topology or in a sidecar topology. Combined mode keeps network and process supervision in the agent container. Sidecar mode runs network enforcement, the proxy, and gateway diff --git a/architecture/gateway.md b/architecture/gateway.md index c8b323ea13..defb0019c2 100644 --- a/architecture/gateway.md +++ b/architecture/gateway.md @@ -155,24 +155,28 @@ token is reported as connected but unauthenticated. Sandbox supervisor RPCs authenticate with explicit sandbox credentials; mTLS does not grant sandbox identity. Kubernetes deployments use the gateway-minted JWT bootstrap path: the supervisor starts with a projected -ServiceAccount token, exchanges it for a gateway-minted sandbox JWT, and uses -that JWT on subsequent gateway RPCs. +ServiceAccount token, registers the pod with the gateway, receives a +gateway-minted sandbox JWT after activation, and uses that JWT on subsequent +gateway RPCs. User-facing mutations are authorized by role policy when OIDC or edge identity is enabled. Sandbox secrets are gateway-signed JWTs bound to a single sandbox ID. Docker, Podman, and VM drivers deliver the initial token through supervisor-only runtime material; Kubernetes supervisors exchange a projected ServiceAccount -token through `IssueSandboxToken`. The gateway validates that projected token -with Kubernetes `TokenReview`, requires the configured sandbox service account, -checks the returned pod binding against the live pod UID, and verifies the pod's -controlling `Sandbox` ownerReference against the live Sandbox CR UID and -sandbox-id label before minting the gateway JWT. The bootstrap path accepts -both `agents.x-k8s.io/v1beta1` ownerReferences from newer Agent Sandbox -controllers and `agents.x-k8s.io/v1alpha1` ownerReferences from existing -deployments. Supervisors renew gateway JWTs in memory before expiry only while -the sandbox record still exists. Older tokens are not server-revoked; shared -deployments bound replay exposure with short `gateway_jwt.ttl_secs` lifetimes. +token through `RegisterSupervisorPod`. The gateway validates that projected +token with Kubernetes `TokenReview`, requires the configured sandbox service +account, checks the returned pod binding against the live pod UID, and verifies +the pod's controlling `Sandbox` ownerReference against the live Sandbox CR UID +and sandbox-id label before sending an activation message with the gateway JWT. +The bootstrap path accepts both `agents.x-k8s.io/v1beta1` ownerReferences from +newer Agent Sandbox controllers and `agents.x-k8s.io/v1alpha1` ownerReferences +from existing deployments. `IssueSandboxToken` remains as a compatibility shim +for older supervisor images and mints through the same already-bound pod +activation logic. Supervisors renew gateway JWTs in memory before expiry only +while the sandbox record still exists. Older tokens are not server-revoked; +shared deployments bound replay exposure with short `gateway_jwt.ttl_secs` +lifetimes. The config default is `gateway_jwt.ttl_secs = 0` for local single-player Docker, Podman, and VM gateways; those tokens carry `exp = 0` and do not expire. Kubernetes and other diff --git a/crates/openshell-cli/tests/ensure_providers_integration.rs b/crates/openshell-cli/tests/ensure_providers_integration.rs index 883c8c4446..3e56df00c2 100644 --- a/crates/openshell-cli/tests/ensure_providers_integration.rs +++ b/crates/openshell-cli/tests/ensure_providers_integration.rs @@ -564,6 +564,17 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = tokio_stream::wrappers::ReceiverStream< + Result, + >; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-cli/tests/mtls_integration.rs b/crates/openshell-cli/tests/mtls_integration.rs index 622e3c1170..b5da73ecbb 100644 --- a/crates/openshell-cli/tests/mtls_integration.rs +++ b/crates/openshell-cli/tests/mtls_integration.rs @@ -449,6 +449,17 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = tokio_stream::wrappers::ReceiverStream< + Result, + >; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-cli/tests/provider_commands_integration.rs b/crates/openshell-cli/tests/provider_commands_integration.rs index 53304b57c5..94e102c261 100644 --- a/crates/openshell-cli/tests/provider_commands_integration.rs +++ b/crates/openshell-cli/tests/provider_commands_integration.rs @@ -967,6 +967,17 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = tokio_stream::wrappers::ReceiverStream< + Result, + >; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs b/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs index fda3d529f3..046f54c3d4 100644 --- a/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs +++ b/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs @@ -637,6 +637,17 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = tokio_stream::wrappers::ReceiverStream< + Result, + >; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-cli/tests/sandbox_name_fallback_integration.rs b/crates/openshell-cli/tests/sandbox_name_fallback_integration.rs index 699749f0d1..8e8e65dbed 100644 --- a/crates/openshell-cli/tests/sandbox_name_fallback_integration.rs +++ b/crates/openshell-cli/tests/sandbox_name_fallback_integration.rs @@ -537,6 +537,17 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = tokio_stream::wrappers::ReceiverStream< + Result, + >; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-core/src/grpc_client.rs b/crates/openshell-core/src/grpc_client.rs index 070704fb0a..c27ff4dbb7 100644 --- a/crates/openshell-core/src/grpc_client.rs +++ b/crates/openshell-core/src/grpc_client.rs @@ -12,8 +12,8 @@ //! Podman / VM drivers write this to a bundle file at sandbox-create //! time). //! 3. `OPENSHELL_K8S_SA_TOKEN_FILE` — projected `ServiceAccount` JWT; the -//! supervisor exchanges it for a gateway JWT via `IssueSandboxToken` -//! once at startup. +//! supervisor registers its pod and receives a gateway JWT via +//! `RegisterSupervisorPod` once at startup. //! //! The resolved bearer credential is held in process memory thereafter and //! injected on every outbound call by [`AuthInterceptor`]. @@ -24,11 +24,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use crate::proto::{ DenialSummary, GetDraftPolicyRequest, GetInferenceBundleRequest, GetInferenceBundleResponse, - GetSandboxConfigRequest, GetSandboxProviderEnvironmentRequest, IssueSandboxTokenRequest, - NetworkActivitySummary, PolicyChunk, PolicySource, PolicyStatus, RefreshSandboxTokenRequest, - ReportPolicyStatusRequest, SandboxPolicy as ProtoSandboxPolicy, SubmitPolicyAnalysisRequest, - SubmitPolicyAnalysisResponse, UpdateConfigRequest, inference_client::InferenceClient, - open_shell_client::OpenShellClient, + GetSandboxConfigRequest, GetSandboxProviderEnvironmentRequest, NetworkActivitySummary, + PolicyChunk, PolicySource, PolicyStatus, RefreshSandboxTokenRequest, + RegisterSupervisorPodRequest, ReportPolicyStatusRequest, SandboxPolicy as ProtoSandboxPolicy, + SubmitPolicyAnalysisRequest, SubmitPolicyAnalysisResponse, UpdateConfigRequest, + inference_client::InferenceClient, open_shell_client::OpenShellClient, }; use crate::sandbox_env; use miette::{IntoDiagnostic, Result, WrapErr}; @@ -242,7 +242,7 @@ async fn token_slot(endpoint: &str, plain_channel: &Channel) -> Result<(TokenSlo /// /// `endpoint` is logged on errors but never used for transport here; the /// actual network call lives inside this function only on the K8s -/// bootstrap path, which uses `plain_channel` to call `IssueSandboxToken` +/// bootstrap path, which uses `plain_channel` to call `RegisterSupervisorPod` /// once before the steady-state Bearer-authenticated channel is built. async fn acquire_sandbox_token(endpoint: &str, plain_channel: &Channel) -> Result { if let Ok(t) = std::env::var(sandbox_env::SANDBOX_TOKEN) @@ -295,7 +295,7 @@ async fn acquire_k8s_sandbox_token( .wrap_err_with(|| format!("failed to read K8s SA token from {sa_path}"))? .trim() .to_string(); - info!(endpoint = %endpoint, "exchanging K8s ServiceAccount token for sandbox JWT"); + info!(endpoint = %endpoint, "registering K8s supervisor pod for sandbox activation"); // The bootstrap exchange uses a one-off interceptor pinned to the // SA token; the resulting gateway JWT becomes the value in the // shared `TOKEN_SLOT` once `connect_channel` returns. @@ -308,11 +308,23 @@ async fn acquire_k8s_sandbox_token( let bootstrap = InterceptedService::new(plain_channel.clone(), interceptor); let mut client = OpenShellClient::new(bootstrap); let resp = client - .issue_sandbox_token(IssueSandboxTokenRequest {}) + .register_supervisor_pod(RegisterSupervisorPodRequest {}) .await .into_diagnostic() - .wrap_err("IssueSandboxToken bootstrap exchange failed")?; - Ok(resp.into_inner().token) + .wrap_err("RegisterSupervisorPod bootstrap stream failed")?; + let mut stream = resp.into_inner(); + let activation = stream + .message() + .await + .into_diagnostic() + .wrap_err("RegisterSupervisorPod activation stream failed")? + .ok_or_else(|| miette::miette!("RegisterSupervisorPod stream closed before activation"))?; + info!( + sandbox_id = %activation.sandbox_id, + sandbox_name = %activation.sandbox_name, + "received supervisor pod activation" + ); + Ok(activation.token) } /// Build an authenticated channel for direct external use (e.g. the diff --git a/crates/openshell-core/src/sandbox_env.rs b/crates/openshell-core/src/sandbox_env.rs index f066580636..08da8b946c 100644 --- a/crates/openshell-core/src/sandbox_env.rs +++ b/crates/openshell-core/src/sandbox_env.rs @@ -87,9 +87,10 @@ pub const USER_ENVIRONMENT: &str = "OPENSHELL_USER_ENVIRONMENT"; /// Path to the projected `ServiceAccount` JWT (Kubernetes driver). /// -/// Used to bootstrap a gateway-minted JWT via `IssueSandboxToken`. Kubelet -/// writes and rotates this file; the supervisor exchanges its contents -/// for a gateway JWT at startup and on refresh. +/// Used to register the supervisor pod and receive a gateway-minted JWT via +/// `RegisterSupervisorPod`. Kubelet writes and rotates this file; the +/// supervisor presents its contents at startup and when rebootstrap is needed +/// after refresh authentication failure. pub const K8S_SA_TOKEN_FILE: &str = "OPENSHELL_K8S_SA_TOKEN_FILE"; /// Filesystem path to the SPIFFE Workload API UNIX socket used for provider diff --git a/crates/openshell-driver-kubernetes/README.md b/crates/openshell-driver-kubernetes/README.md index 96e54ad448..ce8a88d55e 100644 --- a/crates/openshell-driver-kubernetes/README.md +++ b/crates/openshell-driver-kubernetes/README.md @@ -45,8 +45,9 @@ values must override image-provided environment variables. Sandbox pods run as `service_account_name` and keep `automountServiceAccountToken: false`. The only Kubernetes token exposed to the supervisor is an explicit, audience-bound projected token mounted at -`/var/run/secrets/openshell/token` for the one-shot `IssueSandboxToken` -bootstrap exchange. +`/var/run/secrets/openshell/token` for the `RegisterSupervisorPod` bootstrap +stream. The gateway validates the pod-bound token and activates already-bound +cold pods by returning a gateway-minted sandbox JWT on that stream. The gateway uses the supervisor relay for connect, exec, and file sync. Sandbox pods do not need direct external ingress for SSH. diff --git a/crates/openshell-driver-kubernetes/src/config.rs b/crates/openshell-driver-kubernetes/src/config.rs index 1eeaac8396..b9b43de6fe 100644 --- a/crates/openshell-driver-kubernetes/src/config.rs +++ b/crates/openshell-driver-kubernetes/src/config.rs @@ -272,9 +272,9 @@ pub struct KubernetesComputeConfig { /// Empty string (default) = omit the field, using the cluster default. pub default_runtime_class_name: String, /// Lifetime (seconds) of the projected `ServiceAccount` token kubelet - /// writes into each sandbox pod. Used only for the one-shot - /// `IssueSandboxToken` bootstrap exchange — the gateway-minted JWT - /// that follows has its own TTL set via `gateway_jwt.ttl_secs`. + /// writes into each sandbox pod. Used only for the + /// `RegisterSupervisorPod` bootstrap stream — the gateway-minted JWT that + /// follows has its own TTL set via `gateway_jwt.ttl_secs`. /// /// Kubelet enforces a minimum of 600 seconds; the supervisor uses /// this token within a few seconds of pod start, so any value at diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index c784f10db9..3c69b80a83 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -2199,7 +2199,7 @@ struct SandboxPodParams<'a> { workspace_default_storage_size: &'a str, default_runtime_class_name: &'a str, /// Lifetime (seconds) of the projected `ServiceAccount` token used - /// for the bootstrap `IssueSandboxToken` exchange. + /// for the bootstrap `RegisterSupervisorPod` stream. sa_token_ttl_secs: i64, provider_spiffe_enabled: bool, provider_spiffe_workload_api_socket_path: &'a str, @@ -2425,7 +2425,7 @@ fn sandbox_template_to_k8s_with_validated_config( } // Carry the sandbox UUID as a pod annotation so the gateway can resolve // a projected SA token claim (pod name + uid) back to a sandbox identity - // when the supervisor calls `IssueSandboxToken` at startup. The gateway + // when the supervisor calls `RegisterSupervisorPod` at startup. The gateway // also verifies the pod's controlling Sandbox ownerReference against the // live CR before accepting this annotation. Its K8s Role does NOT grant // `patch pods`, so this annotation is effectively immutable post-create. @@ -2633,9 +2633,10 @@ fn sandbox_template_to_k8s_with_validated_config( } // Projected ServiceAccountToken volume — kubelet writes a short-lived // audience-bound JWT into /var/run/secrets/openshell/token and rotates - // it automatically. The supervisor exchanges this for a gateway-minted - // JWT via `IssueSandboxToken` once at startup. In sidecar topology both - // supervisor containers run with the sandbox GID and need group-read access. + // it automatically. The supervisor presents this to `RegisterSupervisorPod` + // and receives a gateway-minted JWT on the activation stream. In sidecar + // topology both supervisor containers run with the sandbox GID and need + // group-read access. let sa_token_default_mode = match params.topology { SupervisorTopology::Combined => 0o400, SupervisorTopology::Sidecar => 0o440, @@ -2975,7 +2976,7 @@ fn apply_required_env( } // Projected ServiceAccount token written by kubelet (see the volume // definition in `sandbox_template_to_k8s`). The supervisor reads this - // and exchanges it for a gateway-minted JWT via `IssueSandboxToken`. + // and receives a gateway-minted JWT via `RegisterSupervisorPod`. upsert_env( env, openshell_core::sandbox_env::K8S_SA_TOKEN_FILE, diff --git a/crates/openshell-driver-kubernetes/src/main.rs b/crates/openshell-driver-kubernetes/src/main.rs index c733b8a45b..e6e9aac33e 100644 --- a/crates/openshell-driver-kubernetes/src/main.rs +++ b/crates/openshell-driver-kubernetes/src/main.rs @@ -107,8 +107,8 @@ struct Args { app_armor_profile: Option, /// Lifetime (seconds) of the projected `ServiceAccount` token - /// kubelet writes into each sandbox pod for the `IssueSandboxToken` - /// bootstrap exchange. Kubelet enforces a minimum of 600s; the + /// kubelet writes into each sandbox pod for the `RegisterSupervisorPod` + /// bootstrap stream. Kubelet enforces a minimum of 600s; the /// gateway clamps values outside `[600, 86400]`. Default 3600. #[arg(long, env = "OPENSHELL_K8S_SA_TOKEN_TTL_SECS", default_value_t = 3600)] sa_token_ttl_secs: i64, diff --git a/crates/openshell-sdk/tests/client_mock.rs b/crates/openshell-sdk/tests/client_mock.rs index d0512d26e1..5b24032b17 100644 --- a/crates/openshell-sdk/tests/client_mock.rs +++ b/crates/openshell-sdk/tests/client_mock.rs @@ -243,6 +243,16 @@ impl OpenShell for TestOpenShell { Ok(Response::new(proto::CreateSshSessionResponse::default())) } + type RegisterSupervisorPodStream = + tokio_stream::wrappers::ReceiverStream>; + + async fn register_supervisor_pod( + &self, + _: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("unused")) + } + async fn expose_service( &self, _: tonic::Request, diff --git a/crates/openshell-server/src/auth/authenticator.rs b/crates/openshell-server/src/auth/authenticator.rs index f5d5c7b2af..bc52729224 100644 --- a/crates/openshell-server/src/auth/authenticator.rs +++ b/crates/openshell-server/src/auth/authenticator.rs @@ -15,7 +15,7 @@ //! Live authenticators slotting into the chain: //! - [`super::sandbox_jwt::SandboxJwtAuthenticator`] — gateway-minted JWTs //! - [`super::k8s_sa::K8sServiceAccountAuthenticator`] — K8s projected SA -//! tokens (path-scoped to `IssueSandboxToken`) +//! tokens (path-scoped to Kubernetes bootstrap RPCs) //! - [`super::oidc::OidcAuthenticator`] — user OIDC Bearer tokens use super::principal::Principal; use async_trait::async_trait; diff --git a/crates/openshell-server/src/auth/k8s_sa.rs b/crates/openshell-server/src/auth/k8s_sa.rs index eed0e5f083..c1b578f4b3 100644 --- a/crates/openshell-server/src/auth/k8s_sa.rs +++ b/crates/openshell-server/src/auth/k8s_sa.rs @@ -3,13 +3,13 @@ //! Kubernetes `ServiceAccount` bootstrap authenticator. //! -//! Path-scoped to `IssueSandboxToken`. Validates a projected SA token +//! Path-scoped to the Kubernetes supervisor bootstrap RPCs. Validates a projected SA token //! presented by a sandbox pod, reads the pod's `openshell.io/sandbox-id` //! annotation, verifies the pod is controlled by the corresponding Sandbox CR, //! and returns a [`Principal::Sandbox`] with -//! [`SandboxIdentitySource::K8sServiceAccount`]. The `IssueSandboxToken` handler -//! then mints a gateway-signed JWT for that sandbox id; subsequent gRPC calls -//! from the supervisor use the gateway-minted JWT validated by +//! [`SandboxIdentitySource::K8sServiceAccount`]. The bootstrap handlers then +//! mint a gateway-signed JWT for that sandbox id; subsequent gRPC calls from +//! the supervisor use the gateway-minted JWT validated by //! [`super::sandbox_jwt::SandboxJwtAuthenticator`]. //! //! This is the only authenticator that talks to the K8s apiserver. It is @@ -30,9 +30,13 @@ use std::sync::Arc; use tonic::Status; use tracing::{debug, info, warn}; -/// gRPC method path that this authenticator accepts. All other paths fall -/// through (return `Ok(None)`) so a gateway-minted JWT is required there. +/// Legacy gRPC method path accepted by this authenticator. All other +/// non-bootstrap paths fall through (return `Ok(None)`) so a gateway-minted JWT +/// is required there. pub const ISSUE_SANDBOX_TOKEN_PATH: &str = "/openshell.v1.OpenShell/IssueSandboxToken"; +/// Supervisor pod registration path accepted by this authenticator for the +/// phase-1 cold Kubernetes activation flow. +pub const REGISTER_SUPERVISOR_POD_PATH: &str = "/openshell.v1.OpenShell/RegisterSupervisorPod"; /// Pod annotation that binds a sandbox pod to its UUID. Set by the /// Kubernetes compute driver at pod-create time. The gateway accepts this @@ -95,9 +99,9 @@ impl Authenticator for K8sServiceAccountAuthenticator { headers: &http::HeaderMap, path: &str, ) -> Result, Status> { - // Scope: only the bootstrap RPC. Other paths fall through so the - // SandboxJwtAuthenticator (or OIDC) handles them. - if path != ISSUE_SANDBOX_TOKEN_PATH { + // Scope: only Kubernetes bootstrap RPCs. Other paths fall through so + // the SandboxJwtAuthenticator (or OIDC) handles them. + if path != ISSUE_SANDBOX_TOKEN_PATH && path != REGISTER_SUPERVISOR_POD_PATH { return Ok(None); } @@ -881,7 +885,7 @@ mod tests { } #[tokio::test] - async fn authenticates_on_issue_path_only() { + async fn authenticates_on_bootstrap_paths_only() { let resolved = ResolvedK8sIdentity { sandbox_id: "sandbox-a".to_string(), pod_name: "openshell-sandbox-a".to_string(), @@ -906,6 +910,22 @@ mod tests { _ => panic!("expected sandbox principal"), } + let on_register = auth + .authenticate(&bearer_headers("sa-jwt"), REGISTER_SUPERVISOR_POD_PATH) + .await + .unwrap() + .expect("expected principal"); + match on_register { + Principal::Sandbox(p) => { + assert_eq!(p.sandbox_id, "sandbox-a"); + assert!(matches!( + p.source, + SandboxIdentitySource::K8sServiceAccount { .. } + )); + } + _ => panic!("expected sandbox principal"), + } + let off_issue = auth .authenticate( &bearer_headers("sa-jwt"), @@ -915,11 +935,11 @@ mod tests { .unwrap(); assert!( off_issue.is_none(), - "K8s SA authenticator must be scoped to IssueSandboxToken" + "K8s SA authenticator must be scoped to bootstrap paths" ); assert_eq!( fake.seen_tokens.lock().unwrap().len(), - 1, + 2, "off-path call must not consult the apiserver" ); } diff --git a/crates/openshell-server/src/auth/principal.rs b/crates/openshell-server/src/auth/principal.rs index 1d4cb7276c..e32f01c2e0 100644 --- a/crates/openshell-server/src/auth/principal.rs +++ b/crates/openshell-server/src/auth/principal.rs @@ -70,7 +70,7 @@ pub enum SandboxIdentitySource { /// Per-sandbox client certificate. Reserved for channel-bound sandbox /// identity. BootstrapCert { fingerprint: String }, - /// K8s `ServiceAccount` token used to bootstrap a gateway-minted JWT - /// via `IssueSandboxToken`. Populated only on that one RPC path. + /// K8s `ServiceAccount` token used to bootstrap a gateway-minted JWT. + /// Populated only on Kubernetes bootstrap RPC paths. K8sServiceAccount { pod_name: String, pod_uid: String }, } diff --git a/crates/openshell-server/src/auth/sandbox_jwt.rs b/crates/openshell-server/src/auth/sandbox_jwt.rs index 39f5982ca0..46dab8e112 100644 --- a/crates/openshell-server/src/auth/sandbox_jwt.rs +++ b/crates/openshell-server/src/auth/sandbox_jwt.rs @@ -8,7 +8,7 @@ //! supervisor-to-gateway gRPC calls. This module implements both sides of the //! gateway-controlled token: //! - [`SandboxJwtIssuer`] mints fresh tokens (called from -//! `handle_create_sandbox` and the `IssueSandboxToken` RPC). +//! `handle_create_sandbox` and the Kubernetes bootstrap RPCs). //! - [`SandboxJwtAuthenticator`] validates tokens on inbound requests and //! produces a [`Principal::Sandbox`] with [`SandboxIdentitySource::BootstrapJwt`]. //! diff --git a/crates/openshell-server/src/auth/sandbox_methods.rs b/crates/openshell-server/src/auth/sandbox_methods.rs index b90841d85a..a0d07ce1df 100644 --- a/crates/openshell-server/src/auth/sandbox_methods.rs +++ b/crates/openshell-server/src/auth/sandbox_methods.rs @@ -26,6 +26,9 @@ mod tests { "/openshell.v1.OpenShell/ConnectSupervisor" )); assert!(is_sandbox_callable("/openshell.v1.OpenShell/RelayStream")); + assert!(is_sandbox_callable( + "/openshell.v1.OpenShell/RegisterSupervisorPod" + )); assert!(is_sandbox_callable( "/openshell.v1.OpenShell/GetSandboxConfig" )); diff --git a/crates/openshell-server/src/config_file.rs b/crates/openshell-server/src/config_file.rs index c4e0cbc959..3271c86881 100644 --- a/crates/openshell-server/src/config_file.rs +++ b/crates/openshell-server/src/config_file.rs @@ -129,7 +129,7 @@ pub struct GatewayFileSection { #[serde(default)] pub enable_user_namespaces: Option, /// Lifetime (seconds) of the projected `ServiceAccount` token kubelet - /// writes for the `IssueSandboxToken` bootstrap exchange. Driver + /// writes for the `RegisterSupervisorPod` bootstrap stream. Driver /// clamps to `[600, 86400]`. #[serde(default)] pub sa_token_ttl_secs: Option, diff --git a/crates/openshell-server/src/grpc/auth_rpc.rs b/crates/openshell-server/src/grpc/auth_rpc.rs index f4cb6a872a..b23f8ab132 100644 --- a/crates/openshell-server/src/grpc/auth_rpc.rs +++ b/crates/openshell-server/src/grpc/auth_rpc.rs @@ -3,8 +3,9 @@ //! Authentication-related RPC handlers. //! -//! Hosts the two sandbox-identity RPCs: -//! - `IssueSandboxToken` — bootstrap exchange (K8s SA token → gateway JWT) +//! Hosts the sandbox-identity RPCs: +//! - `RegisterSupervisorPod` — Kubernetes pod registration and activation +//! - `IssueSandboxToken` — legacy bootstrap compatibility shim //! - `RefreshSandboxToken` — renew a still-valid gateway JWT //! //! Both end in a fresh gateway-signed JWT minted by @@ -12,20 +13,55 @@ //! until their own `exp` and are bounded by the configured short TTL. use crate::ServerState; -use crate::auth::principal::{Principal, SandboxIdentitySource}; +use crate::auth::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; use openshell_core::proto::{ - IssueSandboxTokenRequest, IssueSandboxTokenResponse, RefreshSandboxTokenRequest, - RefreshSandboxTokenResponse, Sandbox, + IssueSandboxTokenRequest, IssueSandboxTokenResponse, PodActivationMessage, + RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RegisterSupervisorPodRequest, Sandbox, }; use std::sync::Arc; +use std::{pin::Pin, result::Result as StdResult}; +use tokio_stream::Stream; use tonic::{Request, Response, Status}; use tracing::{debug, info, warn}; +pub type RegisterSupervisorPodStream = + Pin> + Send + 'static>>; + #[allow(clippy::result_large_err, clippy::unused_async)] pub async fn handle_issue_sandbox_token( state: &Arc, request: Request, ) -> Result, Status> { + // Compatibility shim for older supervisor images. New Kubernetes + // supervisors should use RegisterSupervisorPod so warm-pool activation can + // later remain pending on the same bootstrap stream. + let sandbox = require_k8s_bootstrap_sandbox(request.extensions(), "IssueSandboxToken")?; + let activation = mint_cold_pod_activation(state, &sandbox, "IssueSandboxToken").await?; + Ok(Response::new(IssueSandboxTokenResponse { + token: activation.token, + expires_at_ms: activation.token_expires_at_ms, + })) +} + +#[allow(clippy::result_large_err, clippy::unused_async)] +pub async fn handle_register_supervisor_pod( + state: &Arc, + request: Request, +) -> Result, Status> { + let sandbox = require_k8s_bootstrap_sandbox(request.extensions(), "RegisterSupervisorPod")?; + let activation = mint_cold_pod_activation(state, &sandbox, "RegisterSupervisorPod").await?; + info!( + sandbox_id = %activation.sandbox_id, + "activated registered supervisor pod" + ); + Ok(Response::new(Box::pin(tokio_stream::once(Ok(activation))))) +} + +#[allow(clippy::result_large_err, clippy::unused_async)] +pub async fn handle_refresh_sandbox_token( + state: &Arc, + request: Request, +) -> Result, Status> { let principal = request .extensions() .get::() @@ -34,99 +70,112 @@ pub async fn handle_issue_sandbox_token( let Principal::Sandbox(sandbox) = principal else { return Err(Status::permission_denied( - "IssueSandboxToken requires a sandbox principal", + "RefreshSandboxToken requires a sandbox principal", )); }; - // Only the bootstrap K8s ServiceAccount path can mint a fresh gateway JWT - // via this RPC. Sandboxes already holding a gateway JWT use - // `RefreshSandboxToken` instead. - if !matches!( - sandbox.source, - SandboxIdentitySource::K8sServiceAccount { .. } - ) { + // Only callers already holding a gateway-minted JWT may refresh; the K8s + // bootstrap path must use RegisterSupervisorPod. + let SandboxIdentitySource::BootstrapJwt { .. } = &sandbox.source else { debug!( sandbox_id = %sandbox.sandbox_id, - "IssueSandboxToken rejected: non-bootstrap principal source" + "RefreshSandboxToken rejected: non-gateway-JWT principal source" ); return Err(Status::permission_denied( - "this principal cannot mint a sandbox token; use RefreshSandboxToken", + "this principal cannot refresh; use RegisterSupervisorPod for bootstrap", )); - } + }; let issuer = state.sandbox_jwt_issuer.as_ref().ok_or_else(|| { warn!( sandbox_id = %sandbox.sandbox_id, - "IssueSandboxToken called but sandbox JWT issuer is not configured" + "RefreshSandboxToken called but sandbox JWT issuer is not configured" ); Status::unavailable("sandbox JWT minting is not configured on this gateway") })?; - ensure_sandbox_exists(state, &sandbox.sandbox_id).await?; + load_sandbox(state, &sandbox.sandbox_id).await?; let minted = issuer.mint(&sandbox.sandbox_id)?; info!( sandbox_id = %sandbox.sandbox_id, - "issued gateway sandbox JWT" + "renewed gateway sandbox JWT" ); - Ok(Response::new(IssueSandboxTokenResponse { + + Ok(Response::new(RefreshSandboxTokenResponse { token: minted.token, expires_at_ms: minted.expires_at_ms, })) } -#[allow(clippy::result_large_err, clippy::unused_async)] -pub async fn handle_refresh_sandbox_token( - state: &Arc, - request: Request, -) -> Result, Status> { - let principal = request - .extensions() +fn require_k8s_bootstrap_sandbox( + extensions: &tonic::Extensions, + rpc_name: &'static str, +) -> Result { + let principal = extensions .get::() .cloned() .ok_or_else(|| Status::unauthenticated("missing principal"))?; let Principal::Sandbox(sandbox) = principal else { - return Err(Status::permission_denied( - "RefreshSandboxToken requires a sandbox principal", - )); + return Err(Status::permission_denied(format!( + "{rpc_name} requires a sandbox principal" + ))); }; - // Only callers already holding a gateway-minted JWT may refresh; the - // K8s bootstrap path must use `IssueSandboxToken`. - let SandboxIdentitySource::BootstrapJwt { .. } = &sandbox.source else { + if !matches!( + sandbox.source, + SandboxIdentitySource::K8sServiceAccount { .. } + ) { debug!( sandbox_id = %sandbox.sandbox_id, - "RefreshSandboxToken rejected: non-gateway-JWT principal source" + rpc = rpc_name, + "bootstrap RPC rejected: non-K8s ServiceAccount principal source" ); return Err(Status::permission_denied( - "this principal cannot refresh; use IssueSandboxToken for bootstrap", + "this principal cannot mint a sandbox token; use RefreshSandboxToken", )); - }; + } + Ok(sandbox) +} + +async fn mint_cold_pod_activation( + state: &Arc, + sandbox: &SandboxPrincipal, + rpc_name: &'static str, +) -> Result { let issuer = state.sandbox_jwt_issuer.as_ref().ok_or_else(|| { warn!( sandbox_id = %sandbox.sandbox_id, - "RefreshSandboxToken called but sandbox JWT issuer is not configured" + rpc = rpc_name, + "bootstrap RPC called but sandbox JWT issuer is not configured" ); Status::unavailable("sandbox JWT minting is not configured on this gateway") })?; - ensure_sandbox_exists(state, &sandbox.sandbox_id).await?; - + let record = load_sandbox(state, &sandbox.sandbox_id).await?; let minted = issuer.mint(&sandbox.sandbox_id)?; + let sandbox_name = record + .metadata + .as_ref() + .map_or_else(String::new, |m| m.name.clone()); info!( sandbox_id = %sandbox.sandbox_id, - "renewed gateway sandbox JWT" + rpc = rpc_name, + "issued gateway sandbox JWT for cold pod activation" ); - Ok(Response::new(RefreshSandboxTokenResponse { + Ok(PodActivationMessage { + sandbox_id: sandbox.sandbox_id.clone(), + sandbox_name, token: minted.token, - expires_at_ms: minted.expires_at_ms, - })) + token_expires_at_ms: minted.expires_at_ms, + startup_metadata: std::collections::HashMap::default(), + }) } -async fn ensure_sandbox_exists(state: &Arc, sandbox_id: &str) -> Result<(), Status> { +async fn load_sandbox(state: &Arc, sandbox_id: &str) -> Result { if sandbox_id.is_empty() { return Err(Status::invalid_argument("sandbox_id is required")); } @@ -136,9 +185,7 @@ async fn ensure_sandbox_exists(state: &Arc, sandbox_id: &str) -> Re .get_message::(sandbox_id) .await .map_err(|e| Status::internal(format!("fetch sandbox failed: {e}")))? - .ok_or_else(|| Status::not_found("sandbox not found"))?; - - Ok(()) + .ok_or_else(|| Status::not_found("sandbox not found")) } #[cfg(test)] @@ -159,6 +206,7 @@ mod tests { use openshell_core::proto::{Sandbox, SandboxPhase, SandboxSpec}; use std::collections::HashMap; use std::time::Duration; + use tokio_stream::StreamExt; async fn state_with_issuer() -> Arc { let mat = generate_jwt_key().expect("jwt key"); @@ -273,6 +321,39 @@ mod tests { assert!(resp.expires_at_ms > 0); } + #[tokio::test] + async fn register_supervisor_pod_returns_immediate_activation_for_existing_sandbox() { + use crate::auth::principal::SandboxIdentitySource; + + let state = state_with_issuer().await; + let mut req = Request::new(RegisterSupervisorPodRequest {}); + req.extensions_mut() + .insert(Principal::Sandbox(SandboxPrincipal { + sandbox_id: "sandbox-a".to_string(), + source: SandboxIdentitySource::K8sServiceAccount { + pod_name: "pod-a".to_string(), + pod_uid: "uid-a".to_string(), + }, + trust_domain: Some("openshell".to_string()), + })); + + let mut stream = handle_register_supervisor_pod(&state, req) + .await + .expect("register OK") + .into_inner(); + let activation = stream + .next() + .await + .expect("activation message") + .expect("activation OK"); + assert_eq!(activation.sandbox_id, "sandbox-a"); + assert_eq!(activation.sandbox_name, "sandbox-a"); + assert!(!activation.token.is_empty()); + assert!(activation.token_expires_at_ms > 0); + assert!(activation.startup_metadata.is_empty()); + assert!(stream.next().await.is_none()); + } + #[tokio::test] async fn issue_rejects_missing_sandbox() { use crate::auth::principal::SandboxIdentitySource; @@ -316,9 +397,9 @@ mod tests { #[tokio::test] async fn refresh_rejects_k8s_sa_principal() { - // K8s SA-bootstrap principals must use IssueSandboxToken, not - // RefreshSandboxToken — the refresh path assumes a still-valid - // gateway-minted JWT exists. + // K8s SA-bootstrap principals must use RegisterSupervisorPod, not + // RefreshSandboxToken. The refresh path assumes a still-valid + // gateway-minted JWT already exists. use crate::auth::principal::SandboxIdentitySource; let state = state_with_issuer().await; let mut req = Request::new(RefreshSandboxTokenRequest {}); diff --git a/crates/openshell-server/src/grpc/mod.rs b/crates/openshell-server/src/grpc/mod.rs index ecd96ea90d..ac0f3bd7f2 100644 --- a/crates/openshell-server/src/grpc/mod.rs +++ b/crates/openshell-server/src/grpc/mod.rs @@ -39,16 +39,17 @@ use openshell_core::proto::{ ListSandboxProvidersRequest, ListSandboxProvidersResponse, ListSandboxesRequest, ListSandboxesResponse, ListServicesRequest, ListServicesResponse, ListWorkspaceMembersRequest, ListWorkspaceMembersResponse, ListWorkspacesRequest, ListWorkspacesResponse, - ProviderProfileResponse, ProviderResponse, PushSandboxLogsRequest, PushSandboxLogsResponse, - RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RejectDraftChunkRequest, - RejectDraftChunkResponse, RelayFrame, RemoveWorkspaceMemberRequest, - RemoveWorkspaceMemberResponse, ReportPolicyStatusRequest, ReportPolicyStatusResponse, - RevokeSshSessionRequest, RevokeSshSessionResponse, RotateProviderCredentialRequest, - RotateProviderCredentialResponse, SandboxResponse, SandboxStreamEvent, ServiceEndpointResponse, - ServiceStatus, SubmitPolicyAnalysisRequest, SubmitPolicyAnalysisResponse, SupervisorMessage, - TcpForwardFrame, UndoDraftChunkRequest, UndoDraftChunkResponse, UpdateConfigRequest, - UpdateConfigResponse, UpdateProviderProfilesRequest, UpdateProviderProfilesResponse, - UpdateProviderRequest, WatchSandboxRequest, open_shell_server::OpenShell, + PodActivationMessage, ProviderProfileResponse, ProviderResponse, PushSandboxLogsRequest, + PushSandboxLogsResponse, RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, + RegisterSupervisorPodRequest, RejectDraftChunkRequest, RejectDraftChunkResponse, RelayFrame, + RemoveWorkspaceMemberRequest, RemoveWorkspaceMemberResponse, ReportPolicyStatusRequest, + ReportPolicyStatusResponse, RevokeSshSessionRequest, RevokeSshSessionResponse, + RotateProviderCredentialRequest, RotateProviderCredentialResponse, SandboxResponse, + SandboxStreamEvent, ServiceEndpointResponse, ServiceStatus, SubmitPolicyAnalysisRequest, + SubmitPolicyAnalysisResponse, SupervisorMessage, TcpForwardFrame, UndoDraftChunkRequest, + UndoDraftChunkResponse, UpdateConfigRequest, UpdateConfigResponse, + UpdateProviderProfilesRequest, UpdateProviderProfilesResponse, UpdateProviderRequest, + WatchSandboxRequest, open_shell_server::OpenShell, }; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; @@ -687,6 +688,18 @@ impl OpenShell for OpenShellService { auth_rpc::handle_issue_sandbox_token(&self.state, request).await } + type RegisterSupervisorPodStream = Pin< + Box> + Send + 'static>, + >; + + #[rpc_auth(auth = "sandbox")] + async fn register_supervisor_pod( + &self, + request: Request, + ) -> Result, Status> { + auth_rpc::handle_register_supervisor_pod(&self.state, request).await + } + #[rpc_auth(auth = "sandbox")] async fn refresh_sandbox_token( &self, diff --git a/crates/openshell-server/src/grpc/sandbox.rs b/crates/openshell-server/src/grpc/sandbox.rs index 57bf421b88..dfab3ad9f8 100644 --- a/crates/openshell-server/src/grpc/sandbox.rs +++ b/crates/openshell-server/src/grpc/sandbox.rs @@ -219,7 +219,7 @@ async fn handle_create_sandbox_inner( })?; // Mint the gateway JWT for singleplayer drivers. K8s sandboxes skip - // this mint and bootstrap via `IssueSandboxToken` at supervisor + // this mint and bootstrap via `RegisterSupervisorPod` at supervisor // startup; identifying "is this K8s?" lives in the compute layer, so // we mint unconditionally here when the issuer is configured and let // the K8s driver simply ignore the field. diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index a75534d397..ff371125b8 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -136,7 +136,7 @@ pub struct ServerState { pub oidc_cache: Option>, /// Gateway-minted sandbox JWT issuer. `None` when `config.gateway_jwt` - /// is not configured; in that mode `IssueSandboxToken` returns + /// is not configured; in that mode Kubernetes bootstrap RPCs return /// `Status::unavailable`. Populated at startup from the on-disk key /// material that `certgen` writes. pub sandbox_jwt_issuer: Option>, @@ -147,7 +147,7 @@ pub struct ServerState { pub sandbox_jwt_authenticator: Option>, /// Optional K8s `ServiceAccount` authenticator that backs the - /// `IssueSandboxToken` bootstrap path. Only present when the gateway + /// Kubernetes supervisor bootstrap paths. Only present when the gateway /// runs in-cluster. pub k8s_sa_authenticator: Option>, diff --git a/crates/openshell-server/src/multiplex.rs b/crates/openshell-server/src/multiplex.rs index 83e27d27cb..3c4c105dfc 100644 --- a/crates/openshell-server/src/multiplex.rs +++ b/crates/openshell-server/src/multiplex.rs @@ -694,10 +694,10 @@ where /// Assemble the authenticator chain for the gateway. /// /// Chain order (first-match-wins): -/// 1. `K8sServiceAccountAuthenticator` (path-scoped to `IssueSandboxToken`) -/// — exchanges a projected SA token for a `Principal::Sandbox` so the -/// `IssueSandboxToken` handler can mint a gateway JWT. No-op on every -/// other path; only present when the gateway runs in-cluster. +/// 1. `K8sServiceAccountAuthenticator` (path-scoped to Kubernetes bootstrap +/// RPCs) — resolves a projected SA token to a `Principal::Sandbox` so the +/// bootstrap handler can mint a gateway JWT. No-op on every other path; +/// only present when the gateway runs in-cluster. /// 2. `SandboxJwtAuthenticator` — validates gateway-minted JWTs. Recognized /// via a distinctive `kid` so non-matching Bearer tokens fall through. /// 3. `OidcAuthenticator` — validates user Bearer tokens against the @@ -2126,6 +2126,7 @@ mod tests { "/openshell.v1.OpenShell/ConnectSupervisor", "/openshell.v1.OpenShell/RelayStream", "/openshell.v1.OpenShell/IssueSandboxToken", + "/openshell.v1.OpenShell/RegisterSupervisorPod", "/openshell.v1.OpenShell/RefreshSandboxToken", "/openshell.inference.v1.Inference/GetInferenceBundle", ] { diff --git a/crates/openshell-server/tests/common/mod.rs b/crates/openshell-server/tests/common/mod.rs index 93beeacf15..3352d8d991 100644 --- a/crates/openshell-server/tests/common/mod.rs +++ b/crates/openshell-server/tests/common/mod.rs @@ -21,10 +21,11 @@ use openshell_core::proto::{ 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, + ListProvidersResponse, ListSandboxesRequest, ListSandboxesResponse, PodActivationMessage, + ProviderResponse, RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, + RegisterSupervisorPodRequest, RelayFrame, RevokeSshSessionRequest, RevokeSshSessionResponse, + SandboxResponse, SandboxStreamEvent, ServiceStatus, SupervisorMessage, TcpForwardFrame, + UpdateProviderRequest, WatchSandboxRequest, open_shell_client::OpenShellClient, open_shell_server::{OpenShell, OpenShellServer}, }; @@ -448,6 +449,15 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = ReceiverStream>; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-server/tests/supervisor_relay_integration.rs b/crates/openshell-server/tests/supervisor_relay_integration.rs index 721baacd88..134e58b840 100644 --- a/crates/openshell-server/tests/supervisor_relay_integration.rs +++ b/crates/openshell-server/tests/supervisor_relay_integration.rs @@ -23,7 +23,8 @@ use hyper_util::{ server::conn::auto::Builder, }; use openshell_core::proto::{ - GatewayMessage, RelayFrame, RelayInit, SupervisorMessage, TcpForwardFrame, + GatewayMessage, PodActivationMessage, RelayFrame, RelayInit, SupervisorMessage, + TcpForwardFrame, open_shell_client::OpenShellClient, open_shell_server::{OpenShell, OpenShellServer}, }; @@ -413,6 +414,13 @@ impl OpenShell for RelayGateway { ) -> Result, Status> { Err(Status::unimplemented("unused")) } + type RegisterSupervisorPodStream = ReceiverStream>; + async fn register_supervisor_pod( + &self, + _: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("unused")) + } async fn refresh_sandbox_token( &self, _: tonic::Request, diff --git a/proto/openshell.proto b/proto/openshell.proto index 1b447a17e3..030fa8be48 100644 --- a/proto/openshell.proto +++ b/proto/openshell.proto @@ -240,6 +240,14 @@ service OpenShell { // and never call this RPC. rpc IssueSandboxToken(IssueSandboxTokenRequest) returns (IssueSandboxTokenResponse); + // Register a Kubernetes supervisor pod and wait for gateway activation. + // + // The initial trial activates already-bound cold pods immediately. Later + // warm-pool stages keep this stream pending until a SandboxClaim adopts the + // registered pod. + rpc RegisterSupervisorPod(RegisterSupervisorPodRequest) + returns (stream PodActivationMessage); + // Renew the calling sandbox's gateway JWT. Older tokens remain valid // until their own expiry; deployments should keep token TTLs short to // bound replay exposure. The supervisor calls this from a background @@ -291,6 +299,28 @@ message IssueSandboxTokenResponse { int64 expires_at_ms = 2; } +// RegisterSupervisorPod request. Empty body; identity is established by the +// authentication credentials carried in the request headers (a projected +// Kubernetes ServiceAccount JWT in the K8s driver path). +message RegisterSupervisorPodRequest {} + +// Activation sent by the gateway once a registered supervisor pod is bound to +// an OpenShell sandbox identity. +message PodActivationMessage { + // OpenShell sandbox UUID the supervisor should use for ConnectSupervisor. + string sandbox_id = 1; + // Human-readable sandbox name, if known. + string sandbox_name = 2; + // Gateway-minted JWT bound to the sandbox UUID. + string token = 3; + // Absolute expiry of the issued token, milliseconds since the epoch. 0 means + // the token is non-expiring. + int64 token_expires_at_ms = 4; + // Reserved for future activation-time metadata. The phase-1 cold path leaves + // this empty. + map startup_metadata = 5; +} + // RefreshSandboxToken request. Empty body; the calling principal must // already be a sandbox principal (i.e. the request carries a still-valid // gateway-minted JWT in its Authorization header). From d3c13522ce03d5337ac92f4f1236ee4885eacae7 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Mon, 20 Jul 2026 19:15:56 +0100 Subject: [PATCH 03/10] feat(server): add warm supervisor pod registration groundwork Signed-off-by: Gordon Sim --- crates/openshell-server-macros/src/lib.rs | 11 +- crates/openshell-server/src/auth/guard.rs | 7 +- crates/openshell-server/src/auth/k8s_sa.rs | 175 ++++++----- .../openshell-server/src/auth/method_authz.rs | 15 +- crates/openshell-server/src/auth/principal.rs | 27 ++ .../src/auth/sandbox_methods.rs | 6 +- crates/openshell-server/src/grpc/auth_rpc.rs | 103 +++++-- crates/openshell-server/src/grpc/mod.rs | 2 +- crates/openshell-server/src/grpc/policy.rs | 3 + crates/openshell-server/src/inference.rs | 5 +- crates/openshell-server/src/lib.rs | 9 + crates/openshell-server/src/multiplex.rs | 82 +++++- .../src/supervisor_pod_registration.rs | 276 ++++++++++++++++++ 13 files changed, 608 insertions(+), 113 deletions(-) create mode 100644 crates/openshell-server/src/supervisor_pod_registration.rs diff --git a/crates/openshell-server-macros/src/lib.rs b/crates/openshell-server-macros/src/lib.rs index a698ae6623..d085080f48 100644 --- a/crates/openshell-server-macros/src/lib.rs +++ b/crates/openshell-server-macros/src/lib.rs @@ -54,6 +54,7 @@ struct RpcAuth { enum AuthMode { Unauthenticated, Sandbox, + PodRegistration, Bearer, Dual, } @@ -116,11 +117,11 @@ impl RpcAuth { }; match mode { - AuthMode::Unauthenticated | AuthMode::Sandbox => { + AuthMode::Unauthenticated | AuthMode::Sandbox | AuthMode::PodRegistration => { if let Some(ref s) = scope { return Err(Error::new( s.span(), - "`scope` is only valid for `auth = \"bearer\"` or `auth = \"dual\"` (sandbox principals don't carry scopes)", + "`scope` is only valid for `auth = \"bearer\"` or `auth = \"dual\"` (sandbox and pod-registration principals don't carry scopes)", )); } if role.is_some() { @@ -154,12 +155,13 @@ fn parse_auth_mode(value: &LitStr) -> Result { match value.value().as_str() { "unauthenticated" => Ok(AuthMode::Unauthenticated), "sandbox" => Ok(AuthMode::Sandbox), + "pod_registration" => Ok(AuthMode::PodRegistration), "bearer" => Ok(AuthMode::Bearer), "dual" => Ok(AuthMode::Dual), other => Err(Error::new( value.span(), format!( - "invalid auth mode `{other}`; expected one of `unauthenticated`, `sandbox`, `bearer`, `dual`" + "invalid auth mode `{other}`; expected one of `unauthenticated`, `sandbox`, `pod_registration`, `bearer`, `dual`" ), )), } @@ -288,6 +290,9 @@ fn expand(args: &AuthzArgs, item: &mut ItemImpl) -> Result { quote! { crate::auth::method_authz::AuthMode::Sandbox } } + AuthMode::PodRegistration => { + quote! { crate::auth::method_authz::AuthMode::PodRegistration } + } AuthMode::Bearer => quote! { crate::auth::method_authz::AuthMode::Bearer }, AuthMode::Dual => quote! { crate::auth::method_authz::AuthMode::Dual }, }; diff --git a/crates/openshell-server/src/auth/guard.rs b/crates/openshell-server/src/auth/guard.rs index edcd6bc013..481cfed2ff 100644 --- a/crates/openshell-server/src/auth/guard.rs +++ b/crates/openshell-server/src/auth/guard.rs @@ -20,6 +20,8 @@ use tracing::info; /// scope at the router level). /// - [`Principal::Sandbox`] must reference the same canonical UUID it /// was authenticated with. +/// - [`Principal::K8sPod`] is rejected — pod registration identity is not a +/// sandbox-scoped credential. /// - [`Principal::Anonymous`] is rejected — sandbox-class methods are /// never anonymously callable. /// @@ -44,6 +46,9 @@ pub fn ensure_sandbox_scope(principal: &Principal, claimed_sandbox_id: &str) -> )) } } + Principal::K8sPod(_) => Err(Status::permission_denied( + "sandbox-scoped methods require a sandbox principal", + )), Principal::Anonymous => Err(Status::unauthenticated( "sandbox-scoped methods require an authenticated caller", )), @@ -84,7 +89,7 @@ pub fn ensure_sandbox_principal_scope( ensure_sandbox_scope(principal, claimed_sandbox_id)?; Ok(p.clone()) } - Principal::User(_) => Err(Status::permission_denied( + Principal::User(_) | Principal::K8sPod(_) => Err(Status::permission_denied( "supervisor RPCs require a sandbox principal", )), Principal::Anonymous => Err(Status::unauthenticated( diff --git a/crates/openshell-server/src/auth/k8s_sa.rs b/crates/openshell-server/src/auth/k8s_sa.rs index c1b578f4b3..d288ebed91 100644 --- a/crates/openshell-server/src/auth/k8s_sa.rs +++ b/crates/openshell-server/src/auth/k8s_sa.rs @@ -6,17 +6,16 @@ //! Path-scoped to the Kubernetes supervisor bootstrap RPCs. Validates a projected SA token //! presented by a sandbox pod, reads the pod's `openshell.io/sandbox-id` //! annotation, verifies the pod is controlled by the corresponding Sandbox CR, -//! and returns a [`Principal::Sandbox`] with -//! [`SandboxIdentitySource::K8sServiceAccount`]. The bootstrap handlers then -//! mint a gateway-signed JWT for that sandbox id; subsequent gRPC calls from -//! the supervisor use the gateway-minted JWT validated by -//! [`super::sandbox_jwt::SandboxJwtAuthenticator`]. +//! and returns either a sandbox principal for the legacy token RPC or a +//! registration-scoped pod principal for `RegisterSupervisorPod`. Warm pods may +//! register before they are bound to a sandbox, so registration identity must +//! not grant sandbox-scoped RPC access. //! //! This is the only authenticator that talks to the K8s apiserver. It is //! optional — the gateway boots without it in singleplayer deployments. use super::authenticator::Authenticator; -use super::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; +use super::principal::{Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal}; use async_trait::async_trait; use k8s_openapi::api::{ authentication::v1::{TokenReview, TokenReviewSpec, TokenReviewStatus, UserInfo}, @@ -54,24 +53,16 @@ const SANDBOX_ID_LABEL: &str = "openshell.ai/sandbox-id"; const POD_NAME_EXTRA: &str = "authentication.kubernetes.io/pod-name"; const POD_UID_EXTRA: &str = "authentication.kubernetes.io/pod-uid"; -/// Resolved identity extracted from a validated SA token + pod lookup. -#[derive(Debug, Clone)] -pub struct ResolvedK8sIdentity { - pub sandbox_id: String, - pub pod_name: String, - pub pod_uid: String, -} - /// Apiserver-facing operations the authenticator depends on. Split out so /// tests can fake the apiserver without standing up a kube cluster. #[async_trait] pub trait K8sIdentityResolver: Send + Sync + 'static { /// Validate `token` via `TokenReview` (`aud == openshell-gateway`), - /// extract the pod name/uid, then `GET` the pod and read - /// `openshell.io/sandbox-id`. Returns `Ok(None)` when the token is - /// well-formed but does not authenticate (e.g. wrong audience); returns - /// `Err` for transport/server errors. - async fn resolve(&self, token: &str) -> Result, Status>; + /// extract the pod name/uid, then `GET` the pod and owning Sandbox CR. + /// Returns `Ok(None)` when the token is well-formed but does not + /// authenticate (e.g. wrong audience); returns `Err` for + /// transport/server errors. + async fn resolve(&self, token: &str) -> Result, Status>; } /// Authenticator wrapper around a [`K8sIdentityResolver`]. @@ -118,18 +109,20 @@ impl Authenticator for K8sServiceAccountAuthenticator { return Ok(None); }; - if resolved.sandbox_id.is_empty() { + if path == REGISTER_SUPERVISOR_POD_PATH { + return Ok(Some(Principal::K8sPod(resolved))); + } + + let sandbox_id = resolved.sandbox_id.clone().ok_or_else(|| { warn!( pod = %resolved.pod_name, - "pod missing openshell.io/sandbox-id annotation; rejecting" + "pod missing openshell.io/sandbox-id annotation; rejecting legacy token issue" ); - return Err(Status::permission_denied( - "pod is not bound to a sandbox identity", - )); - } + Status::permission_denied("pod is not bound to a sandbox identity") + })?; Ok(Some(Principal::Sandbox(SandboxPrincipal { - sandbox_id: resolved.sandbox_id, + sandbox_id, source: SandboxIdentitySource::K8sServiceAccount { pod_name: resolved.pod_name, pod_uid: resolved.pod_uid, @@ -222,7 +215,7 @@ impl LiveK8sResolver { #[async_trait] impl K8sIdentityResolver for LiveK8sResolver { - async fn resolve(&self, token: &str) -> Result, Status> { + async fn resolve(&self, token: &str) -> Result, Status> { let review = TokenReview { metadata: ObjectMeta::default(), spec: TokenReviewSpec { @@ -295,7 +288,7 @@ impl K8sIdentityResolver for LiveK8sResolver { return Err(Status::permission_denied("SA token pod UID mismatch")); } - let sandbox_id = pod_sandbox_id(&pod)?; + let sandbox_id = pod_sandbox_id(&pod); let owner = sandbox_owner_reference(&pod)?; let sandbox_cr = self.get_sandbox_cr_for_owner(&owner).await.map_err(|e| { @@ -317,12 +310,17 @@ impl K8sIdentityResolver for LiveK8sResolver { ); return Err(Status::permission_denied("sandbox owner not found")); }; - validate_sandbox_owner_reference(&owner, &sandbox_id, &sandbox_cr)?; + validate_sandbox_owner_reference(&owner, &sandbox_cr)?; + if let Some(ref sandbox_id) = sandbox_id { + validate_sandbox_owner_binding(&owner, sandbox_id, &sandbox_cr)?; + } - Ok(Some(ResolvedK8sIdentity { + Ok(Some(RegisteredPodIdentity { sandbox_id, pod_name: identity.pod_name, pod_uid: identity.pod_uid, + sandbox_owner_name: owner.name, + sandbox_owner_uid: owner.uid, })) } } @@ -393,20 +391,13 @@ fn user_extra_one(user: &UserInfo, key: &str) -> Result { } #[allow(clippy::result_large_err)] -fn pod_sandbox_id(pod: &Pod) -> Result { - let sandbox_id = pod - .metadata +fn pod_sandbox_id(pod: &Pod) -> Option { + pod.metadata .annotations .as_ref() .and_then(|a| a.get(SANDBOX_ID_ANNOTATION)) .cloned() - .unwrap_or_default(); - if sandbox_id.is_empty() { - return Err(Status::permission_denied( - "pod is not bound to a sandbox identity", - )); - } - Ok(sandbox_id) + .filter(|sandbox_id| !sandbox_id.is_empty()) } #[allow(clippy::result_large_err)] @@ -478,7 +469,6 @@ fn should_try_next_sandbox_api_version(err: &KubeError) -> bool { #[allow(clippy::result_large_err)] fn validate_sandbox_owner_reference( owner: &SandboxOwnerReference, - sandbox_id: &str, sandbox_cr: &DynamicObject, ) -> Result<(), Status> { let actual_uid = sandbox_cr.metadata.uid.as_deref().unwrap_or_default(); @@ -492,6 +482,15 @@ fn validate_sandbox_owner_reference( return Err(Status::permission_denied("sandbox owner UID mismatch")); } + Ok(()) +} + +#[allow(clippy::result_large_err)] +fn validate_sandbox_owner_binding( + owner: &SandboxOwnerReference, + sandbox_id: &str, + sandbox_cr: &DynamicObject, +) -> Result<(), Status> { let actual_sandbox_id = sandbox_cr .metadata .labels @@ -521,12 +520,12 @@ pub mod test_support { /// Fake resolver for unit tests. Returns the configured outcome on /// every call and records the tokens it observed. pub struct FakeResolver { - pub outcome: Result, Status>, + pub outcome: Result, Status>, pub seen_tokens: Mutex>, } impl FakeResolver { - pub fn returning(outcome: Result, Status>) -> Self { + pub fn returning(outcome: Result, Status>) -> Self { Self { outcome, seen_tokens: Mutex::new(Vec::new()), @@ -536,7 +535,7 @@ pub mod test_support { #[async_trait] impl K8sIdentityResolver for FakeResolver { - async fn resolve(&self, token: &str) -> Result, Status> { + async fn resolve(&self, token: &str) -> Result, Status> { self.seen_tokens.lock().unwrap().push(token.to_string()); match &self.outcome { Ok(opt) => Ok(opt.clone()), @@ -655,6 +654,16 @@ mod tests { } } + fn registered_pod(sandbox_id: Option<&str>) -> RegisteredPodIdentity { + RegisteredPodIdentity { + sandbox_id: sandbox_id.map(str::to_string), + pod_name: "openshell-sandbox-a".to_string(), + pod_uid: "uid-a".to_string(), + sandbox_owner_name: "sandbox-owner-a".to_string(), + sandbox_owner_uid: "cr-uid-a".to_string(), + } + } + fn sandbox_cr(name: &str, uid: &str, sandbox_id: &str) -> DynamicObject { let sandbox_gvk = GroupVersionKind::gvk(SANDBOX_API_GROUP, SANDBOX_API_VERSION_V1BETA1, SANDBOX_KIND); @@ -769,15 +778,12 @@ mod tests { } #[test] - fn pod_sandbox_id_requires_annotation() { + fn pod_sandbox_id_is_optional_for_warm_registration() { assert_eq!( - pod_sandbox_id(&pod_with_sandbox_id(Some("sandbox-id-a"))).unwrap(), - "sandbox-id-a" + pod_sandbox_id(&pod_with_sandbox_id(Some("sandbox-id-a"))).as_deref(), + Some("sandbox-id-a") ); - - let err = pod_sandbox_id(&pod_with_sandbox_id(None)) - .expect_err("missing sandbox-id annotation must fail"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); + assert!(pod_sandbox_id(&pod_with_sandbox_id(None)).is_none()); } #[test] @@ -863,34 +869,40 @@ mod tests { } #[test] - fn validate_sandbox_owner_reference_requires_matching_cr_uid_and_label() { + fn validate_sandbox_owner_reference_requires_matching_cr_uid() { let owner = SandboxOwnerReference { api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), name: "sandbox-a".to_string(), uid: "cr-uid-a".to_string(), }; let cr = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-a"); - validate_sandbox_owner_reference(&owner, "sandbox-id-a", &cr) - .expect("matching CR should be accepted"); + validate_sandbox_owner_reference(&owner, &cr).expect("matching CR should be accepted"); let wrong_uid = sandbox_cr("sandbox-a", "cr-uid-b", "sandbox-id-a"); - let err = validate_sandbox_owner_reference(&owner, "sandbox-id-a", &wrong_uid) + let err = validate_sandbox_owner_reference(&owner, &wrong_uid) .expect_err("wrong CR UID must fail"); assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + #[test] + fn validate_sandbox_owner_binding_requires_matching_label() { + let owner = SandboxOwnerReference { + api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), + name: "sandbox-a".to_string(), + uid: "cr-uid-a".to_string(), + }; + let cr = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-a"); + validate_sandbox_owner_binding(&owner, "sandbox-id-a", &cr) + .expect("matching label should be accepted"); let wrong_label = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-b"); - let err = validate_sandbox_owner_reference(&owner, "sandbox-id-a", &wrong_label) + let err = validate_sandbox_owner_binding(&owner, "sandbox-id-a", &wrong_label) .expect_err("wrong sandbox-id label must fail"); assert_eq!(err.code(), tonic::Code::PermissionDenied); } #[tokio::test] async fn authenticates_on_bootstrap_paths_only() { - let resolved = ResolvedK8sIdentity { - sandbox_id: "sandbox-a".to_string(), - pod_name: "openshell-sandbox-a".to_string(), - pod_uid: "uid-a".to_string(), - }; + let resolved = registered_pod(Some("sandbox-a")); let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); let auth = K8sServiceAccountAuthenticator::new(fake.clone()); @@ -916,14 +928,12 @@ mod tests { .unwrap() .expect("expected principal"); match on_register { - Principal::Sandbox(p) => { - assert_eq!(p.sandbox_id, "sandbox-a"); - assert!(matches!( - p.source, - SandboxIdentitySource::K8sServiceAccount { .. } - )); + Principal::K8sPod(p) => { + assert_eq!(p.sandbox_id.as_deref(), Some("sandbox-a")); + assert_eq!(p.pod_name, "openshell-sandbox-a"); + assert_eq!(p.pod_uid, "uid-a"); } - _ => panic!("expected sandbox principal"), + _ => panic!("expected pod registration principal"), } let off_issue = auth @@ -970,12 +980,8 @@ mod tests { } #[tokio::test] - async fn pod_without_annotation_is_rejected() { - let resolved = ResolvedK8sIdentity { - sandbox_id: String::new(), - pod_name: "stray-pod".to_string(), - pod_uid: "uid".to_string(), - }; + async fn issue_token_rejects_unbound_registered_pod() { + let resolved = registered_pod(None); let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); let auth = K8sServiceAccountAuthenticator::new(fake); let err = auth @@ -985,6 +991,27 @@ mod tests { assert_eq!(err.code(), tonic::Code::PermissionDenied); } + #[tokio::test] + async fn register_supervisor_pod_accepts_unbound_registered_pod() { + let resolved = registered_pod(None); + let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); + let auth = K8sServiceAccountAuthenticator::new(fake); + let principal = auth + .authenticate(&bearer_headers("sa-jwt"), REGISTER_SUPERVISOR_POD_PATH) + .await + .expect("unbound warm pod registration should authenticate") + .expect("expected principal"); + + match principal { + Principal::K8sPod(p) => { + assert!(p.sandbox_id.is_none()); + assert_eq!(p.pod_name, "openshell-sandbox-a"); + assert_eq!(p.sandbox_owner_uid, "cr-uid-a"); + } + _ => panic!("expected pod registration principal"), + } + } + #[tokio::test] async fn resolver_error_propagates() { let fake = Arc::new(FakeResolver::returning(Err(Status::unavailable( diff --git a/crates/openshell-server/src/auth/method_authz.rs b/crates/openshell-server/src/auth/method_authz.rs index ec8dc5bca3..c53ca8024d 100644 --- a/crates/openshell-server/src/auth/method_authz.rs +++ b/crates/openshell-server/src/auth/method_authz.rs @@ -34,6 +34,9 @@ pub enum AuthMode { /// Only callable by a `Principal::Sandbox` (gateway-minted sandbox JWT). /// See `auth/sandbox_jwt.rs`. Sandbox, + /// Only callable by a Kubernetes pod principal during supervisor pod + /// registration. This does not grant sandbox-scoped RPC access. + PodRegistration, /// Bearer (OIDC) authentication required. Bearer, /// Either sandbox principal or Bearer; scope and role apply on @@ -109,6 +112,16 @@ pub fn is_sandbox_callable(method: &str) -> bool { ) } +/// `true` if the method is callable by a Kubernetes pod registration +/// principal. +#[must_use] +pub fn is_pod_registration_callable(method: &str) -> bool { + matches!( + lookup(method).map(|m| m.mode), + Some(AuthMode::PodRegistration) + ) +} + /// `true` if the method is callable by a `Principal::User` (`bearer` or /// `dual` auth mode). /// @@ -119,7 +132,7 @@ pub fn is_sandbox_callable(method: &str) -> bool { #[must_use] pub fn is_user_callable(method: &str) -> bool { match lookup(method).map(|m| m.mode) { - Some(AuthMode::Sandbox | AuthMode::Unauthenticated) => false, + Some(AuthMode::Sandbox | AuthMode::PodRegistration | AuthMode::Unauthenticated) => false, Some(AuthMode::Bearer | AuthMode::Dual) | None => true, } } diff --git a/crates/openshell-server/src/auth/principal.rs b/crates/openshell-server/src/auth/principal.rs index e32f01c2e0..d395805250 100644 --- a/crates/openshell-server/src/auth/principal.rs +++ b/crates/openshell-server/src/auth/principal.rs @@ -28,6 +28,12 @@ pub enum Principal { /// sandbox UUID. The wrapped `sandbox_id` MUST match any sandbox referenced /// in the request body for sandbox-class methods. Sandbox(#[allow(dead_code)] SandboxPrincipal), + /// Kubernetes pod authenticated for supervisor pod registration only. + /// + /// This is intentionally not a sandbox principal: warm pods can be valid + /// Kubernetes pods before they are claimed by a sandbox. The router only + /// allows this principal to call `RegisterSupervisorPod`. + K8sPod(RegisteredPodIdentity), /// Truly unauthenticated caller (health probes, reflection). Sandbox-class /// and user-class methods reject this variant. #[allow(dead_code)] @@ -41,6 +47,27 @@ pub struct UserPrincipal { pub identity: Identity, } +/// Kubernetes pod identity validated by `TokenReview` and live pod lookup. +/// +/// `sandbox_id` is present only when the pod is already bound to a concrete +/// `OpenShell` sandbox. Unbound warm pods carry the owning Sandbox CR metadata +/// but must not receive a gateway sandbox JWT until a later claim binds the +/// exact pod UID to a sandbox. +#[derive(Debug, Clone)] +#[allow(dead_code)] +pub struct RegisteredPodIdentity { + /// Live pod name from the `TokenReview` binding. + pub pod_name: String, + /// Live pod UID from the `TokenReview` binding and pod lookup. + pub pod_uid: String, + /// `OpenShell` sandbox UUID from the pod annotation, when already bound. + pub sandbox_id: Option, + /// Owning Sandbox CR name from the pod ownerReference. + pub sandbox_owner_name: String, + /// Owning Sandbox CR UID from the pod ownerReference and live CR lookup. + pub sandbox_owner_uid: String, +} + /// Sandbox caller — bound to one specific sandbox UUID. /// /// `sandbox_id` and `source` are consumed by the router and handler guards. diff --git a/crates/openshell-server/src/auth/sandbox_methods.rs b/crates/openshell-server/src/auth/sandbox_methods.rs index a0d07ce1df..2566c29ab5 100644 --- a/crates/openshell-server/src/auth/sandbox_methods.rs +++ b/crates/openshell-server/src/auth/sandbox_methods.rs @@ -26,9 +26,6 @@ mod tests { "/openshell.v1.OpenShell/ConnectSupervisor" )); assert!(is_sandbox_callable("/openshell.v1.OpenShell/RelayStream")); - assert!(is_sandbox_callable( - "/openshell.v1.OpenShell/RegisterSupervisorPod" - )); assert!(is_sandbox_callable( "/openshell.v1.OpenShell/GetSandboxConfig" )); @@ -54,6 +51,9 @@ mod tests { assert!(!is_sandbox_callable( "/openshell.inference.v1.Inference/GetInferenceRoute" )); + assert!(!is_sandbox_callable( + "/openshell.v1.OpenShell/RegisterSupervisorPod" + )); assert!(!is_sandbox_callable( "/openshell.inference.v1.Inference/SetInferenceRoute" )); diff --git a/crates/openshell-server/src/grpc/auth_rpc.rs b/crates/openshell-server/src/grpc/auth_rpc.rs index b23f8ab132..52d26a0eba 100644 --- a/crates/openshell-server/src/grpc/auth_rpc.rs +++ b/crates/openshell-server/src/grpc/auth_rpc.rs @@ -13,7 +13,9 @@ //! until their own `exp` and are bounded by the configured short TTL. use crate::ServerState; -use crate::auth::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; +use crate::auth::principal::{ + Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal, +}; use openshell_core::proto::{ IssueSandboxTokenRequest, IssueSandboxTokenResponse, PodActivationMessage, RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RegisterSupervisorPodRequest, Sandbox, @@ -36,7 +38,8 @@ pub async fn handle_issue_sandbox_token( // supervisors should use RegisterSupervisorPod so warm-pool activation can // later remain pending on the same bootstrap stream. let sandbox = require_k8s_bootstrap_sandbox(request.extensions(), "IssueSandboxToken")?; - let activation = mint_cold_pod_activation(state, &sandbox, "IssueSandboxToken").await?; + let activation = + mint_cold_pod_activation(state, &sandbox.sandbox_id, "IssueSandboxToken").await?; Ok(Response::new(IssueSandboxTokenResponse { token: activation.token, expires_at_ms: activation.token_expires_at_ms, @@ -48,13 +51,21 @@ pub async fn handle_register_supervisor_pod( state: &Arc, request: Request, ) -> Result, Status> { - let sandbox = require_k8s_bootstrap_sandbox(request.extensions(), "RegisterSupervisorPod")?; - let activation = mint_cold_pod_activation(state, &sandbox, "RegisterSupervisorPod").await?; - info!( - sandbox_id = %activation.sandbox_id, - "activated registered supervisor pod" - ); - Ok(Response::new(Box::pin(tokio_stream::once(Ok(activation))))) + let pod = require_registered_pod(request.extensions())?; + if let Some(sandbox_id) = pod.sandbox_id.as_deref() { + let activation = + mint_cold_pod_activation(state, sandbox_id, "RegisterSupervisorPod").await?; + info!( + sandbox_id = %activation.sandbox_id, + pod = %pod.pod_name, + pod_uid = %pod.pod_uid, + "activated bound supervisor pod" + ); + return Ok(Response::new(Box::pin(tokio_stream::once(Ok(activation))))); + } + + let stream = state.supervisor_pod_registrations.register_pending(pod)?; + Ok(Response::new(Box::pin(stream))) } #[allow(clippy::result_large_err, clippy::unused_async)] @@ -140,34 +151,53 @@ fn require_k8s_bootstrap_sandbox( Ok(sandbox) } +fn require_registered_pod(extensions: &tonic::Extensions) -> Result { + if let Some(pod) = extensions.get::().cloned() { + return Ok(pod); + } + + let principal = extensions + .get::() + .cloned() + .ok_or_else(|| Status::unauthenticated("missing principal"))?; + + let Principal::K8sPod(pod) = principal else { + return Err(Status::permission_denied( + "RegisterSupervisorPod requires a Kubernetes pod registration principal", + )); + }; + + Ok(pod) +} + async fn mint_cold_pod_activation( state: &Arc, - sandbox: &SandboxPrincipal, + sandbox_id: &str, rpc_name: &'static str, ) -> Result { let issuer = state.sandbox_jwt_issuer.as_ref().ok_or_else(|| { warn!( - sandbox_id = %sandbox.sandbox_id, + sandbox_id = %sandbox_id, rpc = rpc_name, "bootstrap RPC called but sandbox JWT issuer is not configured" ); Status::unavailable("sandbox JWT minting is not configured on this gateway") })?; - let record = load_sandbox(state, &sandbox.sandbox_id).await?; - let minted = issuer.mint(&sandbox.sandbox_id)?; + let record = load_sandbox(state, sandbox_id).await?; + let minted = issuer.mint(sandbox_id)?; let sandbox_name = record .metadata .as_ref() .map_or_else(String::new, |m| m.name.clone()); info!( - sandbox_id = %sandbox.sandbox_id, + sandbox_id = %sandbox_id, rpc = rpc_name, "issued gateway sandbox JWT for cold pod activation" ); Ok(PodActivationMessage { - sandbox_id: sandbox.sandbox_id.clone(), + sandbox_id: sandbox_id.to_string(), sandbox_name, token: minted.token, token_expires_at_ms: minted.expires_at_ms, @@ -273,6 +303,16 @@ mod tests { }) } + fn registered_pod(sandbox_id: Option<&str>) -> RegisteredPodIdentity { + RegisteredPodIdentity { + pod_name: "pod-a".to_string(), + pod_uid: "uid-a".to_string(), + sandbox_id: sandbox_id.map(str::to_string), + sandbox_owner_name: "sandbox-owner-a".to_string(), + sandbox_owner_uid: "owner-uid-a".to_string(), + } + } + #[tokio::test] async fn refresh_returns_new_token() { let state = state_with_issuer().await; @@ -323,19 +363,10 @@ mod tests { #[tokio::test] async fn register_supervisor_pod_returns_immediate_activation_for_existing_sandbox() { - use crate::auth::principal::SandboxIdentitySource; - let state = state_with_issuer().await; let mut req = Request::new(RegisterSupervisorPodRequest {}); req.extensions_mut() - .insert(Principal::Sandbox(SandboxPrincipal { - sandbox_id: "sandbox-a".to_string(), - source: SandboxIdentitySource::K8sServiceAccount { - pod_name: "pod-a".to_string(), - pod_uid: "uid-a".to_string(), - }, - trust_domain: Some("openshell".to_string()), - })); + .insert(Principal::K8sPod(registered_pod(Some("sandbox-a")))); let mut stream = handle_register_supervisor_pod(&state, req) .await @@ -354,6 +385,28 @@ mod tests { assert!(stream.next().await.is_none()); } + #[tokio::test] + async fn register_supervisor_pod_keeps_unbound_warm_pod_pending() { + let state = state_with_issuer().await; + let mut req = Request::new(RegisterSupervisorPodRequest {}); + req.extensions_mut() + .insert(Principal::K8sPod(registered_pod(None))); + + let mut stream = handle_register_supervisor_pod(&state, req) + .await + .expect("register OK") + .into_inner(); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 1); + assert!( + tokio::time::timeout(Duration::from_millis(20), stream.next()) + .await + .is_err(), + "unbound warm pod must wait for later activation" + ); + drop(stream); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); + } + #[tokio::test] async fn issue_rejects_missing_sandbox() { use crate::auth::principal::SandboxIdentitySource; diff --git a/crates/openshell-server/src/grpc/mod.rs b/crates/openshell-server/src/grpc/mod.rs index ac0f3bd7f2..23331fbcf0 100644 --- a/crates/openshell-server/src/grpc/mod.rs +++ b/crates/openshell-server/src/grpc/mod.rs @@ -692,7 +692,7 @@ impl OpenShell for OpenShellService { Box> + Send + 'static>, >; - #[rpc_auth(auth = "sandbox")] + #[rpc_auth(auth = "pod_registration")] async fn register_supervisor_pod( &self, request: Request, diff --git a/crates/openshell-server/src/grpc/policy.rs b/crates/openshell-server/src/grpc/policy.rs index 0703771846..66261e3d5e 100644 --- a/crates/openshell-server/src/grpc/policy.rs +++ b/crates/openshell-server/src/grpc/policy.rs @@ -1252,6 +1252,9 @@ async fn resolve_sandbox_by_name_for_principal( Ok(sandbox) } Principal::User(_) => sandbox.ok_or_else(|| Status::not_found("sandbox not found")), + Principal::K8sPod(_) => Err(Status::permission_denied( + "sandbox-scoped methods require a sandbox principal", + )), Principal::Anonymous => Err(Status::unauthenticated( "sandbox-scoped methods require an authenticated caller", )), diff --git a/crates/openshell-server/src/inference.rs b/crates/openshell-server/src/inference.rs index 97217c2565..7fe80b36fa 100644 --- a/crates/openshell-server/src/inference.rs +++ b/crates/openshell-server/src/inference.rs @@ -930,7 +930,10 @@ fn authorize_inference_bundle( ) -> Result { match principal { Some(crate::auth::principal::Principal::Sandbox(s)) => Ok(s.sandbox_id.clone()), - Some(crate::auth::principal::Principal::User(_)) => Err(Status::permission_denied( + Some( + crate::auth::principal::Principal::User(_) + | crate::auth::principal::Principal::K8sPod(_), + ) => Err(Status::permission_denied( "GetInferenceBundle requires a sandbox principal", )), Some(crate::auth::principal::Principal::Anonymous) | None => Err(Status::unauthenticated( diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index ff371125b8..55c13654fb 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -43,6 +43,7 @@ mod sandbox_index; mod sandbox_watch; mod service_routing; mod ssh_sessions; +mod supervisor_pod_registration; pub mod supervisor_session; mod telemetry; #[cfg(any(test, feature = "test-support"))] @@ -132,6 +133,11 @@ pub struct ServerState { /// Validated built-in and operator-registered supervisor middleware. pub middleware_registry: Arc, + /// Pending Kubernetes supervisor pod registrations awaiting warm-pool + /// activation. + pub(crate) supervisor_pod_registrations: + Arc, + /// OIDC JWKS cache for JWT validation. `None` when OIDC is not configured. pub oidc_cache: Option>, @@ -207,6 +213,9 @@ impl ServerState { settings_mutex: tokio::sync::Mutex::new(()), supervisor_sessions, middleware_registry: Arc::new(MiddlewareRegistry::default()), + supervisor_pod_registrations: Arc::new( + supervisor_pod_registration::SupervisorPodRegistrationRegistry::new(), + ), oidc_cache, sandbox_jwt_issuer: None, sandbox_jwt_authenticator: None, diff --git a/crates/openshell-server/src/multiplex.rs b/crates/openshell-server/src/multiplex.rs index 3c4c105dfc..e40a3f831a 100644 --- a/crates/openshell-server/src/multiplex.rs +++ b/crates/openshell-server/src/multiplex.rs @@ -471,6 +471,22 @@ fn gateway_principal_fields(principal: &Principal) -> BTreeMap { fields.insert("trust_domain".to_string(), trust_domain.clone()); } } + Principal::K8sPod(pod) => { + fields.insert("kind".to_string(), "k8s_pod".to_string()); + fields.insert("pod_name".to_string(), pod.pod_name.clone()); + fields.insert("pod_uid".to_string(), pod.pod_uid.clone()); + fields.insert( + "sandbox_owner_name".to_string(), + pod.sandbox_owner_name.clone(), + ); + fields.insert( + "sandbox_owner_uid".to_string(), + pod.sandbox_owner_uid.clone(), + ); + if let Some(sandbox_id) = &pod.sandbox_id { + fields.insert("sandbox_id".to_string(), sandbox_id.clone()); + } + } Principal::Anonymous => { fields.insert("kind".to_string(), "anonymous".to_string()); } @@ -695,8 +711,9 @@ where /// /// Chain order (first-match-wins): /// 1. `K8sServiceAccountAuthenticator` (path-scoped to Kubernetes bootstrap -/// RPCs) — resolves a projected SA token to a `Principal::Sandbox` so the -/// bootstrap handler can mint a gateway JWT. No-op on every other path; +/// RPCs) — resolves a projected SA token to either a legacy +/// `Principal::Sandbox` for `IssueSandboxToken` or a registration-scoped +/// pod principal for `RegisterSupervisorPod`. No-op on every other path; /// only present when the gateway runs in-cluster. /// 2. `SandboxJwtAuthenticator` — validates gateway-minted JWTs. Recognized /// via a distinctive `kid` so non-matching Bearer tokens fall through. @@ -743,6 +760,7 @@ fn build_authenticator_chain(state: &ServerState) -> Option /// `Principal::User` is gated by the RBAC `AuthzPolicy`. /// `Principal::Sandbox` is gated by a supervisor-method allowlist, then /// handlers enforce same-sandbox scope on request bodies. +/// `Principal::K8sPod` is gated by the pod-registration method only. #[derive(Clone)] pub struct AuthGrpcRouter { inner: S, @@ -881,6 +899,14 @@ where ))); } } + Principal::K8sPod(ref pod) => { + if !crate::auth::method_authz::is_pod_registration_callable(&path) { + return Ok(status_response(tonic::Status::permission_denied( + "Kubernetes pod principals may only register supervisor pods", + ))); + } + req.extensions_mut().insert(pod.clone()); + } Principal::Anonymous => { return Ok(status_response(tonic::Status::unauthenticated( "anonymous callers may not call authenticated methods", @@ -1816,7 +1842,8 @@ mod tests { use crate::auth::authenticator::test_support::MockAuthenticator; use crate::auth::identity::{Identity, IdentityProvider}; use crate::auth::principal::{ - Principal, SandboxIdentitySource, SandboxPrincipal, UserPrincipal, + Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal, + UserPrincipal, }; use http_body_util::Full; use std::sync::Arc; @@ -1906,6 +1933,16 @@ mod tests { }) } + fn pod_registration_principal() -> Principal { + Principal::K8sPod(RegisteredPodIdentity { + pod_name: "pod-a".to_string(), + pod_uid: "pod-uid-a".to_string(), + sandbox_id: None, + sandbox_owner_name: "sandbox-owner-a".to_string(), + sandbox_owner_uid: "owner-uid-a".to_string(), + }) + } + #[tokio::test] async fn mtls_peer_identity_fills_missing_principal_when_enabled() { let mock = Arc::new(MockAuthenticator::returning(Ok(None))); @@ -2094,6 +2131,43 @@ mod tests { )); } + #[tokio::test] + async fn pod_registration_principal_can_only_register_supervisor_pod() { + let mock = Arc::new(MockAuthenticator::returning(Ok(Some( + pod_registration_principal(), + )))); + let chain = AuthenticatorChain::new(vec![mock]); + let (recorder, seen) = PrincipalRecorder::new(); + let mut router = AuthGrpcRouter::new(recorder, Some(chain), None); + + let res = router + .call(empty_request( + "/openshell.v1.OpenShell/RegisterSupervisorPod", + )) + .await + .unwrap(); + + assert_eq!(res.status(), 200); + assert!(matches!( + seen.lock().unwrap().as_ref(), + Some(Principal::K8sPod(_)) + )); + + let mock = Arc::new(MockAuthenticator::returning(Ok(Some( + pod_registration_principal(), + )))); + let chain = AuthenticatorChain::new(vec![mock]); + let (recorder, seen) = PrincipalRecorder::new(); + let mut router = AuthGrpcRouter::new(recorder, Some(chain), None); + let res = router + .call(empty_request("/openshell.v1.OpenShell/GetSandboxConfig")) + .await + .unwrap(); + + assert!(seen.lock().unwrap().is_none()); + assert_eq!(grpc_status(&res).as_deref(), Some("7")); + } + /// A user principal — even one carrying `openshell:all` and the /// admin role — must not reach a `sandbox`-annotated method. The /// router enforces this from the per-handler auth-mode declarations @@ -2126,7 +2200,6 @@ mod tests { "/openshell.v1.OpenShell/ConnectSupervisor", "/openshell.v1.OpenShell/RelayStream", "/openshell.v1.OpenShell/IssueSandboxToken", - "/openshell.v1.OpenShell/RegisterSupervisorPod", "/openshell.v1.OpenShell/RefreshSandboxToken", "/openshell.inference.v1.Inference/GetInferenceBundle", ] { @@ -2151,6 +2224,7 @@ mod tests { "/openshell.v1.OpenShell/CreateProvider", "/openshell.v1.OpenShell/ApproveDraftChunk", "/openshell.inference.v1.Inference/GetInferenceRoute", + "/openshell.v1.OpenShell/RegisterSupervisorPod", "/openshell.inference.v1.Inference/SetInferenceRoute", ] { let mock = Arc::new(MockAuthenticator::returning(Ok(Some(sandbox_principal())))); diff --git a/crates/openshell-server/src/supervisor_pod_registration.rs b/crates/openshell-server/src/supervisor_pod_registration.rs new file mode 100644 index 0000000000..1c847c43d4 --- /dev/null +++ b/crates/openshell-server/src/supervisor_pod_registration.rs @@ -0,0 +1,276 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Pending supervisor pod registrations for warm-pool activation. +//! +//! Cold pods already bound to a sandbox are activated immediately by the +//! `RegisterSupervisorPod` handler. Warm pods can register before claim +//! assignment; the future claim controller will activate the stored stream once +//! it binds the exact pod UID to a sandbox. + +use crate::auth::principal::RegisteredPodIdentity; +use openshell_core::proto::PodActivationMessage; +use std::collections::HashMap; +use std::pin::Pin; +use std::result::Result as StdResult; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; +use std::time::Instant; +use tokio::sync::mpsc; +use tokio_stream::Stream; +use tokio_stream::wrappers::ReceiverStream; +use tonic::Status; +use tracing::{debug, info, warn}; + +#[derive(Debug, Default)] +pub struct SupervisorPodRegistrationRegistry { + inner: Mutex, + next_session_id: AtomicU64, +} + +#[derive(Debug, Default)] +struct Inner { + pending_by_pod_uid: HashMap, +} + +#[derive(Debug)] +struct PendingRegistration { + identity: RegisteredPodIdentity, + sender: mpsc::Sender>, + session_id: u64, + registered_at: Instant, +} + +impl SupervisorPodRegistrationRegistry { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + #[allow(clippy::result_large_err)] + pub fn register_pending( + self: &Arc, + identity: RegisteredPodIdentity, + ) -> Result { + if identity.pod_uid.is_empty() { + return Err(Status::permission_denied("registered pod UID is required")); + } + + let (sender, receiver) = mpsc::channel(1); + let session_id = self.next_session_id.fetch_add(1, Ordering::Relaxed); + let pod_uid = identity.pod_uid.clone(); + let pod_name = identity.pod_name.clone(); + let sandbox_owner = identity.sandbox_owner_name.clone(); + + let replaced = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + inner + .pending_by_pod_uid + .insert( + pod_uid.clone(), + PendingRegistration { + identity, + sender, + session_id, + registered_at: Instant::now(), + }, + ) + .is_some() + }; + + if replaced { + info!( + pod = %pod_name, + pod_uid = %pod_uid, + sandbox_owner = %sandbox_owner, + "replaced duplicate pending supervisor pod registration" + ); + } else { + info!( + pod = %pod_name, + pod_uid = %pod_uid, + sandbox_owner = %sandbox_owner, + "registered warm supervisor pod pending activation" + ); + } + + Ok(PendingRegistrationStream { + registry: self.clone(), + pod_uid, + session_id, + inner: ReceiverStream::new(receiver), + }) + } + + #[allow(clippy::result_large_err)] + #[allow(dead_code)] + pub fn activate( + &self, + pod_uid: &str, + activation: PodActivationMessage, + ) -> Result { + let pending = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + inner.pending_by_pod_uid.remove(pod_uid) + }; + + let Some(pending) = pending else { + debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to activate"); + return Ok(false); + }; + + let pod_name = pending.identity.pod_name.clone(); + pending.sender.try_send(Ok(activation)).map_err(|_| { + warn!( + pod = %pod_name, + pod_uid = %pod_uid, + "pending supervisor pod registration stream closed before activation" + ); + Status::unavailable("registered pod stream closed before activation") + })?; + info!( + pod = %pod_name, + pod_uid = %pod_uid, + pending_ms = pending.registered_at.elapsed().as_millis(), + "activated pending supervisor pod registration" + ); + Ok(true) + } + + #[must_use] + #[allow(dead_code)] + pub fn pending_count(&self) -> usize { + self.inner + .lock() + .expect("pending pod registry poisoned") + .pending_by_pod_uid + .len() + } + + fn remove_if_session(&self, pod_uid: &str, session_id: u64) { + let removed = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner + .pending_by_pod_uid + .get(pod_uid) + .is_some_and(|pending| pending.session_id == session_id) + { + inner.pending_by_pod_uid.remove(pod_uid) + } else { + None + } + }; + + if let Some(pending) = removed { + debug!( + pod = %pending.identity.pod_name, + pod_uid = %pod_uid, + pending_ms = pending.registered_at.elapsed().as_millis(), + "removed pending supervisor pod registration" + ); + } + } +} + +pub struct PendingRegistrationStream { + registry: Arc, + pod_uid: String, + session_id: u64, + inner: ReceiverStream>, +} + +impl Stream for PendingRegistrationStream { + type Item = StdResult; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.inner).poll_next(cx) + } +} + +impl Drop for PendingRegistrationStream { + fn drop(&mut self) { + self.registry + .remove_if_session(&self.pod_uid, self.session_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + use tokio_stream::StreamExt; + + fn pod_identity(pod_uid: &str) -> RegisteredPodIdentity { + RegisteredPodIdentity { + pod_name: "warm-pod-a".to_string(), + pod_uid: pod_uid.to_string(), + sandbox_id: None, + sandbox_owner_name: "sandbox-owner-a".to_string(), + sandbox_owner_uid: "owner-uid-a".to_string(), + } + } + + #[test] + fn dropping_stream_removes_pending_registration() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let stream = registry + .register_pending(pod_identity("pod-uid-a")) + .expect("register pending"); + assert_eq!(registry.pending_count(), 1); + drop(stream); + assert_eq!(registry.pending_count(), 0); + } + + #[test] + fn duplicate_registration_replaces_prior_stream() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let old = registry + .register_pending(pod_identity("pod-uid-a")) + .expect("register old"); + let new = registry + .register_pending(pod_identity("pod-uid-a")) + .expect("register new"); + + assert_eq!(registry.pending_count(), 1); + drop(old); + assert_eq!( + registry.pending_count(), + 1, + "old stream drop must not remove replacement" + ); + drop(new); + assert_eq!(registry.pending_count(), 0); + } + + #[tokio::test] + async fn activation_sends_message_and_removes_pending_registration() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let mut stream = registry + .register_pending(pod_identity("pod-uid-a")) + .expect("register pending"); + + let activation = PodActivationMessage { + sandbox_id: "sandbox-a".to_string(), + sandbox_name: "sandbox-a".to_string(), + token: "token-a".to_string(), + token_expires_at_ms: 123, + startup_metadata: HashMap::default(), + }; + assert!( + registry + .activate("pod-uid-a", activation) + .expect("activate") + ); + assert_eq!(registry.pending_count(), 0); + + let received = stream + .next() + .await + .expect("activation message") + .expect("activation OK"); + assert_eq!(received.sandbox_id, "sandbox-a"); + assert!(stream.next().await.is_none()); + } +} From 07a8eeba2926e4555bae07a6d3dac5b090843d81 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Mon, 20 Jul 2026 19:58:56 +0100 Subject: [PATCH 04/10] feat(server): scaffold claim-driven warm pod activation Signed-off-by: Gordon Sim --- crates/openshell-server/src/grpc/auth_rpc.rs | 57 +-- crates/openshell-server/src/lib.rs | 1 + .../src/supervisor_pod_registration.rs | 180 ++++++++- .../src/warm_pod_activation.rs | 381 ++++++++++++++++++ 4 files changed, 551 insertions(+), 68 deletions(-) create mode 100644 crates/openshell-server/src/warm_pod_activation.rs diff --git a/crates/openshell-server/src/grpc/auth_rpc.rs b/crates/openshell-server/src/grpc/auth_rpc.rs index 52d26a0eba..7bfdcc820d 100644 --- a/crates/openshell-server/src/grpc/auth_rpc.rs +++ b/crates/openshell-server/src/grpc/auth_rpc.rs @@ -16,9 +16,10 @@ use crate::ServerState; use crate::auth::principal::{ Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal, }; +use crate::warm_pod_activation::{load_sandbox, mint_pod_activation}; use openshell_core::proto::{ IssueSandboxTokenRequest, IssueSandboxTokenResponse, PodActivationMessage, - RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RegisterSupervisorPodRequest, Sandbox, + RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RegisterSupervisorPodRequest, }; use std::sync::Arc; use std::{pin::Pin, result::Result as StdResult}; @@ -38,8 +39,7 @@ pub async fn handle_issue_sandbox_token( // supervisors should use RegisterSupervisorPod so warm-pool activation can // later remain pending on the same bootstrap stream. let sandbox = require_k8s_bootstrap_sandbox(request.extensions(), "IssueSandboxToken")?; - let activation = - mint_cold_pod_activation(state, &sandbox.sandbox_id, "IssueSandboxToken").await?; + let activation = mint_pod_activation(state, &sandbox.sandbox_id, "IssueSandboxToken").await?; Ok(Response::new(IssueSandboxTokenResponse { token: activation.token, expires_at_ms: activation.token_expires_at_ms, @@ -53,8 +53,7 @@ pub async fn handle_register_supervisor_pod( ) -> Result, Status> { let pod = require_registered_pod(request.extensions())?; if let Some(sandbox_id) = pod.sandbox_id.as_deref() { - let activation = - mint_cold_pod_activation(state, sandbox_id, "RegisterSupervisorPod").await?; + let activation = mint_pod_activation(state, sandbox_id, "RegisterSupervisorPod").await?; info!( sandbox_id = %activation.sandbox_id, pod = %pod.pod_name, @@ -170,54 +169,6 @@ fn require_registered_pod(extensions: &tonic::Extensions) -> Result, - sandbox_id: &str, - rpc_name: &'static str, -) -> Result { - let issuer = state.sandbox_jwt_issuer.as_ref().ok_or_else(|| { - warn!( - sandbox_id = %sandbox_id, - rpc = rpc_name, - "bootstrap RPC called but sandbox JWT issuer is not configured" - ); - Status::unavailable("sandbox JWT minting is not configured on this gateway") - })?; - - let record = load_sandbox(state, sandbox_id).await?; - let minted = issuer.mint(sandbox_id)?; - let sandbox_name = record - .metadata - .as_ref() - .map_or_else(String::new, |m| m.name.clone()); - info!( - sandbox_id = %sandbox_id, - rpc = rpc_name, - "issued gateway sandbox JWT for cold pod activation" - ); - - Ok(PodActivationMessage { - sandbox_id: sandbox_id.to_string(), - sandbox_name, - token: minted.token, - token_expires_at_ms: minted.expires_at_ms, - startup_metadata: std::collections::HashMap::default(), - }) -} - -async fn load_sandbox(state: &Arc, sandbox_id: &str) -> Result { - if sandbox_id.is_empty() { - return Err(Status::invalid_argument("sandbox_id is required")); - } - - state - .store - .get_message::(sandbox_id) - .await - .map_err(|e| Status::internal(format!("fetch sandbox failed: {e}")))? - .ok_or_else(|| Status::not_found("sandbox not found")) -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index 55c13654fb..caecdb80d6 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -52,6 +52,7 @@ mod tls; #[cfg(test)] pub(crate) mod tls_test_utils; pub mod tracing_bus; +pub(crate) mod warm_pod_activation; mod ws_tunnel; use metrics_exporter_prometheus::PrometheusBuilder; diff --git a/crates/openshell-server/src/supervisor_pod_registration.rs b/crates/openshell-server/src/supervisor_pod_registration.rs index 1c847c43d4..a8eb9eafa0 100644 --- a/crates/openshell-server/src/supervisor_pod_registration.rs +++ b/crates/openshell-server/src/supervisor_pod_registration.rs @@ -32,6 +32,7 @@ pub struct SupervisorPodRegistrationRegistry { #[derive(Debug, Default)] struct Inner { pending_by_pod_uid: HashMap, + activated_pod_uid: HashMap, } #[derive(Debug)] @@ -65,6 +66,11 @@ impl SupervisorPodRegistrationRegistry { let replaced = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_pod_uid.contains_key(&pod_uid) { + return Err(Status::already_exists( + "supervisor pod has already been activated", + )); + } inner .pending_by_pod_uid .insert( @@ -104,38 +110,103 @@ impl SupervisorPodRegistrationRegistry { } #[allow(clippy::result_large_err)] - #[allow(dead_code)] - pub fn activate( + pub(crate) fn pending_identity(&self, pod_uid: &str) -> Result { + let inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_pod_uid.contains_key(pod_uid) { + return Err(Status::already_exists( + "supervisor pod has already been activated", + )); + } + inner + .pending_by_pod_uid + .get(pod_uid) + .map(|pending| pending.identity.clone()) + .ok_or_else(|| Status::not_found("pending supervisor pod registration not found")) + } + + #[allow(clippy::result_large_err)] + pub(crate) fn activate( &self, pod_uid: &str, activation: PodActivationMessage, - ) -> Result { + ) -> Result<(), Status> { + let pending = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_pod_uid.contains_key(pod_uid) { + return Err(Status::already_exists( + "supervisor pod has already been activated", + )); + } + let Some(pending) = inner.pending_by_pod_uid.remove(pod_uid) else { + debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to activate"); + return Err(Status::not_found( + "pending supervisor pod registration not found", + )); + }; + inner + .activated_pod_uid + .insert(pod_uid.to_string(), Instant::now()); + pending + }; + + let pod_name = pending.identity.pod_name.clone(); + if pending.sender.try_send(Ok(activation)).is_err() { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + inner.activated_pod_uid.remove(pod_uid); + warn!( + pod = %pod_name, + pod_uid = %pod_uid, + "pending supervisor pod registration stream closed before activation" + ); + return Err(Status::unavailable( + "registered pod stream closed before activation", + )); + } + + info!( + pod = %pod_name, + pod_uid = %pod_uid, + pending_ms = pending.registered_at.elapsed().as_millis(), + "activated pending supervisor pod registration" + ); + Ok(()) + } + + #[allow(clippy::result_large_err)] + pub(crate) fn fail_pending(&self, pod_uid: &str, status: Status) -> Result<(), Status> { let pending = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_pod_uid.contains_key(pod_uid) { + return Err(Status::already_exists( + "supervisor pod has already been activated", + )); + } inner.pending_by_pod_uid.remove(pod_uid) }; let Some(pending) = pending else { - debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to activate"); - return Ok(false); + debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to fail"); + return Err(Status::not_found( + "pending supervisor pod registration not found", + )); }; let pod_name = pending.identity.pod_name.clone(); - pending.sender.try_send(Ok(activation)).map_err(|_| { + pending.sender.try_send(Err(status)).map_err(|_| { warn!( pod = %pod_name, pod_uid = %pod_uid, - "pending supervisor pod registration stream closed before activation" + "pending supervisor pod registration stream closed before failure notification" ); - Status::unavailable("registered pod stream closed before activation") + Status::unavailable("registered pod stream closed before failure notification") })?; info!( pod = %pod_name, pod_uid = %pod_uid, pending_ms = pending.registered_at.elapsed().as_millis(), - "activated pending supervisor pod registration" + "failed pending supervisor pod registration" ); - Ok(true) + Ok(()) } #[must_use] @@ -148,6 +219,16 @@ impl SupervisorPodRegistrationRegistry { .len() } + #[must_use] + #[allow(dead_code)] + pub fn activated_count(&self) -> usize { + self.inner + .lock() + .expect("pending pod registry poisoned") + .activated_pod_uid + .len() + } + fn remove_if_session(&self, pod_uid: &str, session_id: u64) { let removed = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); @@ -258,12 +339,11 @@ mod tests { token_expires_at_ms: 123, startup_metadata: HashMap::default(), }; - assert!( - registry - .activate("pod-uid-a", activation) - .expect("activate") - ); + registry + .activate("pod-uid-a", activation) + .expect("activate"); assert_eq!(registry.pending_count(), 0); + assert_eq!(registry.activated_count(), 1); let received = stream .next() @@ -273,4 +353,74 @@ mod tests { assert_eq!(received.sandbox_id, "sandbox-a"); assert!(stream.next().await.is_none()); } + + #[test] + fn activation_for_unknown_pod_uid_returns_not_found() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let activation = PodActivationMessage { + sandbox_id: "sandbox-a".to_string(), + sandbox_name: "sandbox-a".to_string(), + token: "token-a".to_string(), + token_expires_at_ms: 123, + startup_metadata: HashMap::default(), + }; + + let err = registry + .activate("pod-uid-a", activation) + .expect_err("unknown pod UID must fail"); + + assert_eq!(err.code(), tonic::Code::NotFound); + } + + #[tokio::test] + async fn activation_tombstone_rejects_duplicate_activation_and_registration() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let mut stream = registry + .register_pending(pod_identity("pod-uid-a")) + .expect("register pending"); + + let activation = PodActivationMessage { + sandbox_id: "sandbox-a".to_string(), + sandbox_name: "sandbox-a".to_string(), + token: "token-a".to_string(), + token_expires_at_ms: 123, + startup_metadata: HashMap::default(), + }; + registry + .activate("pod-uid-a", activation.clone()) + .expect("activate"); + let _ = stream.next().await.expect("activation message"); + + let err = registry + .activate("pod-uid-a", activation) + .expect_err("duplicate activation must fail"); + assert_eq!(err.code(), tonic::Code::AlreadyExists); + + let Err(err) = registry.register_pending(pod_identity("pod-uid-a")) else { + panic!("activated pod UID must not register again"); + }; + assert_eq!(err.code(), tonic::Code::AlreadyExists); + } + + #[test] + fn closed_pending_stream_cannot_be_activated() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let stream = registry + .register_pending(pod_identity("pod-uid-a")) + .expect("register pending"); + drop(stream); + + let activation = PodActivationMessage { + sandbox_id: "sandbox-a".to_string(), + sandbox_name: "sandbox-a".to_string(), + token: "token-a".to_string(), + token_expires_at_ms: 123, + startup_metadata: HashMap::default(), + }; + let err = registry + .activate("pod-uid-a", activation) + .expect_err("closed stream should remove pending registration"); + + assert_eq!(err.code(), tonic::Code::NotFound); + } } diff --git a/crates/openshell-server/src/warm_pod_activation.rs b/crates/openshell-server/src/warm_pod_activation.rs new file mode 100644 index 0000000000..881ff34651 --- /dev/null +++ b/crates/openshell-server/src/warm_pod_activation.rs @@ -0,0 +1,381 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Gateway-side warm supervisor pod activation. +//! +//! The Kubernetes claim controller will call this module after it observes and +//! revalidates a claim that binds a warm pod UID to an `OpenShell` sandbox. + +use crate::ServerState; +use crate::auth::principal::RegisteredPodIdentity; +use async_trait::async_trait; +use openshell_core::proto::{PodActivationMessage, Sandbox}; +use std::sync::Arc; +use tonic::Status; +use tracing::{info, warn}; + +#[derive(Debug, Clone, PartialEq, Eq)] +#[allow(dead_code)] +pub struct WarmPodActivationTarget { + pub pod_uid: String, + pub sandbox_id: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +#[allow(dead_code)] +pub struct ValidatedWarmPodActivation { + pub pod_uid: String, + pub pod_name: String, + pub sandbox_owner_uid: String, + pub sandbox_id: String, +} + +#[async_trait] +#[allow(dead_code)] +pub trait WarmPodActivationValidator: Send + Sync { + async fn validate( + &self, + target: &WarmPodActivationTarget, + pending: &RegisteredPodIdentity, + ) -> Result; +} + +#[allow(clippy::result_large_err)] +#[allow(dead_code)] +pub async fn activate_warm_pod( + state: &Arc, + validator: &V, + target: WarmPodActivationTarget, +) -> Result<(), Status> +where + V: WarmPodActivationValidator + ?Sized, +{ + let pending = state + .supervisor_pod_registrations + .pending_identity(&target.pod_uid)?; + let activation_result = async { + let validated = validator.validate(&target, &pending).await?; + ensure_validated_activation_matches_pending(&target, &pending, &validated)?; + mint_pod_activation(state, &validated.sandbox_id, "WarmPodActivation").await + } + .await; + let activation = match activation_result { + Ok(activation) => activation, + Err(status) => { + let stream_status = clone_status(&status); + state + .supervisor_pod_registrations + .fail_pending(&target.pod_uid, stream_status)?; + return Err(status); + } + }; + + state + .supervisor_pod_registrations + .activate(&target.pod_uid, activation)?; + Ok(()) +} + +fn clone_status(status: &Status) -> Status { + Status::new(status.code(), status.message().to_string()) +} + +#[allow(clippy::result_large_err)] +fn ensure_validated_activation_matches_pending( + target: &WarmPodActivationTarget, + pending: &RegisteredPodIdentity, + validated: &ValidatedWarmPodActivation, +) -> Result<(), Status> { + if validated.pod_uid != target.pod_uid || validated.pod_uid != pending.pod_uid { + return Err(Status::permission_denied( + "validated pod UID does not match pending registration", + )); + } + if validated.pod_name != pending.pod_name { + return Err(Status::permission_denied( + "validated pod name does not match pending registration", + )); + } + if validated.sandbox_owner_uid != pending.sandbox_owner_uid { + return Err(Status::permission_denied( + "validated Sandbox owner UID does not match pending registration", + )); + } + if validated.sandbox_id != target.sandbox_id { + return Err(Status::permission_denied( + "validated sandbox ID does not match activation target", + )); + } + Ok(()) +} + +#[allow(clippy::result_large_err)] +pub async fn mint_pod_activation( + state: &Arc, + sandbox_id: &str, + reason: &'static str, +) -> Result { + let issuer = state.sandbox_jwt_issuer.as_ref().ok_or_else(|| { + warn!( + sandbox_id = %sandbox_id, + reason, + "pod activation requested but sandbox JWT issuer is not configured" + ); + Status::unavailable("sandbox JWT minting is not configured on this gateway") + })?; + + let record = load_sandbox(state, sandbox_id).await?; + let minted = issuer.mint(sandbox_id)?; + let sandbox_name = record + .metadata + .as_ref() + .map_or_else(String::new, |m| m.name.clone()); + info!( + sandbox_id = %sandbox_id, + reason, + "issued gateway sandbox JWT for pod activation" + ); + + Ok(PodActivationMessage { + sandbox_id: sandbox_id.to_string(), + sandbox_name, + token: minted.token, + token_expires_at_ms: minted.expires_at_ms, + startup_metadata: std::collections::HashMap::default(), + }) +} + +pub async fn load_sandbox(state: &Arc, sandbox_id: &str) -> Result { + if sandbox_id.is_empty() { + return Err(Status::invalid_argument("sandbox_id is required")); + } + + state + .store + .get_message::(sandbox_id) + .await + .map_err(|e| Status::internal(format!("fetch sandbox failed: {e}")))? + .ok_or_else(|| Status::not_found("sandbox not found")) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::auth::sandbox_jwt::SandboxJwtIssuer; + use crate::compute::new_test_runtime; + use crate::persistence::Store; + use crate::sandbox_index::SandboxIndex; + use crate::sandbox_watch::SandboxWatchBus; + use crate::supervisor_session::SupervisorSessionRegistry; + use crate::tracing_bus::TracingLogBus; + use openshell_bootstrap::jwt::generate_jwt_key; + use openshell_core::Config; + use openshell_core::proto::datamodel::v1::ObjectMeta; + use openshell_core::proto::{Sandbox, SandboxPhase, SandboxSpec}; + use std::collections::HashMap; + use std::time::Duration; + use tokio_stream::StreamExt; + + struct StaticValidator { + validated: ValidatedWarmPodActivation, + } + + #[async_trait] + impl WarmPodActivationValidator for StaticValidator { + async fn validate( + &self, + _target: &WarmPodActivationTarget, + _pending: &RegisteredPodIdentity, + ) -> Result { + Ok(self.validated.clone()) + } + } + + async fn state_with_issuer() -> Arc { + let mat = generate_jwt_key().expect("jwt key"); + let store = Arc::new( + Store::connect("sqlite::memory:?cache=shared") + .await + .unwrap(), + ); + let compute = new_test_runtime(store.clone()).await; + let mut state = ServerState::new( + Config::new(None).with_database_url("sqlite::memory:?cache=shared"), + store, + compute, + SandboxIndex::new(), + SandboxWatchBus::new(), + TracingLogBus::new(), + Arc::new(SupervisorSessionRegistry::new()), + None, + ); + let issuer = SandboxJwtIssuer::from_pem( + mat.signing_key_pem.as_bytes(), + mat.kid, + "test-gateway", + Duration::from_secs(3600), + ) + .unwrap(); + state.sandbox_jwt_issuer = Some(Arc::new(issuer)); + let state = Arc::new(state); + insert_sandbox(&state, "sandbox-a").await; + state + } + + async fn insert_sandbox(state: &Arc, sandbox_id: &str) { + let mut sandbox = Sandbox { + metadata: Some(ObjectMeta { + id: sandbox_id.to_string(), + name: sandbox_id.to_string(), + created_at_ms: 1_000_000, + labels: HashMap::default(), + resource_version: 0, + }), + spec: Some(SandboxSpec { + policy: None, + ..Default::default() + }), + ..Default::default() + }; + sandbox.set_phase(SandboxPhase::Ready as i32); + state.store.put_message(&sandbox).await.unwrap(); + } + + fn pending_pod() -> RegisteredPodIdentity { + RegisteredPodIdentity { + pod_name: "warm-pod-a".to_string(), + pod_uid: "pod-uid-a".to_string(), + sandbox_id: None, + sandbox_owner_name: "sandbox-owner-a".to_string(), + sandbox_owner_uid: "owner-uid-a".to_string(), + } + } + + fn activation_target(sandbox_id: &str) -> WarmPodActivationTarget { + WarmPodActivationTarget { + pod_uid: "pod-uid-a".to_string(), + sandbox_id: sandbox_id.to_string(), + } + } + + fn validated_activation(sandbox_id: &str) -> ValidatedWarmPodActivation { + ValidatedWarmPodActivation { + pod_uid: "pod-uid-a".to_string(), + pod_name: "warm-pod-a".to_string(), + sandbox_owner_uid: "owner-uid-a".to_string(), + sandbox_id: sandbox_id.to_string(), + } + } + + #[tokio::test] + async fn pending_warm_pod_receives_activation_token() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_pod()) + .expect("register pending"); + let validator = StaticValidator { + validated: validated_activation("sandbox-a"), + }; + + activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect("activate"); + + let received = stream + .next() + .await + .expect("activation message") + .expect("activation OK"); + assert_eq!(received.sandbox_id, "sandbox-a"); + assert_eq!(received.sandbox_name, "sandbox-a"); + assert!(!received.token.is_empty()); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); + assert_eq!(state.supervisor_pod_registrations.activated_count(), 1); + } + + #[tokio::test] + async fn activation_for_unknown_pod_uid_fails_before_validation() { + let state = state_with_issuer().await; + let validator = StaticValidator { + validated: validated_activation("sandbox-a"), + }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("unknown pod UID must fail"); + + assert_eq!(err.code(), tonic::Code::NotFound); + } + + #[tokio::test] + async fn missing_target_sandbox_fails_without_token_emission() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_pod()) + .expect("register pending"); + let validator = StaticValidator { + validated: validated_activation("sandbox-deleted"), + }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-deleted")) + .await + .expect_err("missing sandbox must fail"); + + assert_eq!(err.code(), tonic::Code::NotFound); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); + let stream_err = stream + .next() + .await + .expect("failure message") + .expect_err("missing sandbox should fail stream"); + assert_eq!(stream_err.code(), tonic::Code::NotFound); + } + + #[tokio::test] + async fn duplicate_activation_for_same_pod_uid_is_rejected() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_pod()) + .expect("register pending"); + let validator = StaticValidator { + validated: validated_activation("sandbox-a"), + }; + + activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect("first activation"); + let _ = stream.next().await.expect("activation"); + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("duplicate activation must fail"); + + assert_eq!(err.code(), tonic::Code::AlreadyExists); + } + + #[tokio::test] + async fn revalidation_metadata_must_match_pending_registration() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_pod()) + .expect("register pending"); + let mut validated = validated_activation("sandbox-a"); + validated.sandbox_owner_uid = "other-owner".to_string(); + let validator = StaticValidator { validated }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("mismatched revalidation metadata must fail"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + let stream_err = stream + .next() + .await + .expect("failure message") + .expect_err("validation mismatch should fail stream"); + assert_eq!(stream_err.code(), tonic::Code::PermissionDenied); + } +} From 3162595d1c5e0454435168d2f3ce9f72c68106fd Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Tue, 21 Jul 2026 20:31:30 +0100 Subject: [PATCH 05/10] refactor(kubernetes): move supervisor bootstrap identity into driver Signed-off-by: Gordon Sim --- crates/openshell-core/src/lib.rs | 1 + .../src/supervisor_bootstrap.rs | 94 ++ .../src/bootstrap.rs | 907 ++++++++++++++++ crates/openshell-driver-kubernetes/src/lib.rs | 4 + .../src/auth/authenticator.rs | 4 +- crates/openshell-server/src/auth/guard.rs | 8 +- crates/openshell-server/src/auth/k8s_sa.rs | 995 ++---------------- crates/openshell-server/src/auth/principal.rs | 39 +- .../src/compute/driver_config.rs | 45 - crates/openshell-server/src/compute/mod.rs | 47 + crates/openshell-server/src/grpc/auth_rpc.rs | 91 +- crates/openshell-server/src/grpc/policy.rs | 2 +- crates/openshell-server/src/inference.rs | 2 +- crates/openshell-server/src/lib.rs | 42 - crates/openshell-server/src/multiplex.rs | 79 +- .../src/supervisor_pod_registration.rs | 182 ++-- .../src/warm_pod_activation.rs | 141 ++- 17 files changed, 1474 insertions(+), 1209 deletions(-) create mode 100644 crates/openshell-core/src/supervisor_bootstrap.rs create mode 100644 crates/openshell-driver-kubernetes/src/bootstrap.rs diff --git a/crates/openshell-core/src/lib.rs b/crates/openshell-core/src/lib.rs index 20c4dccb7d..f007fc178c 100644 --- a/crates/openshell-core/src/lib.rs +++ b/crates/openshell-core/src/lib.rs @@ -37,6 +37,7 @@ pub mod provider_credentials; pub mod sandbox_env; pub mod secrets; pub mod settings; +pub mod supervisor_bootstrap; pub mod telemetry; pub mod time; diff --git a/crates/openshell-core/src/supervisor_bootstrap.rs b/crates/openshell-core/src/supervisor_bootstrap.rs new file mode 100644 index 0000000000..3ffabb737c --- /dev/null +++ b/crates/openshell-core/src/supervisor_bootstrap.rs @@ -0,0 +1,94 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Driver-provided supervisor bootstrap capabilities. +//! +//! These types describe the narrow interface between the gateway and compute +//! drivers for supervisors that cannot start with a gateway-minted sandbox JWT. +//! Kubernetes is the initial implementation: the driver validates projected +//! `ServiceAccount` tokens and classifies pods as already bound or warm-pending. + +use tonic::Status; +use tonic::async_trait; + +/// Registration-only identity for a supervisor runtime instance. +/// +/// This is not a sandbox principal. The gateway may use it only on the +/// bootstrap registration path until a concrete sandbox token is minted. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SupervisorBootstrapIdentity { + /// Driver that produced the identity, for example `kubernetes`. + pub driver: String, + /// Driver-native stable instance ID, for example Kubernetes pod UID. + pub instance_id: String, + /// Driver-native human-readable instance name, for example Kubernetes pod + /// name. + pub instance_name: String, + /// Driver-native owner object name. + pub owner_name: String, + /// Driver-native owner object UID. + pub owner_uid: String, + /// Whether the instance is already bound to an `OpenShell` sandbox or is + /// waiting for a warm-pool claim. + pub binding: SupervisorBootstrapBinding, +} + +/// Binding state returned by a driver bootstrap identity provider. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum SupervisorBootstrapBinding { + /// The instance is already bound to a concrete `OpenShell` sandbox. + BoundSandbox { sandbox_id: String }, + /// The instance is valid but not yet claim-bound. + WarmPending, +} + +impl SupervisorBootstrapIdentity { + /// Return the bound sandbox ID when the identity is already activated by + /// driver state. + #[must_use] + pub fn bound_sandbox_id(&self) -> Option<&str> { + match &self.binding { + SupervisorBootstrapBinding::BoundSandbox { sandbox_id } => Some(sandbox_id.as_str()), + SupervisorBootstrapBinding::WarmPending => None, + } + } +} + +/// Driver-provided authentication for supervisor bootstrap registration. +#[async_trait] +pub trait SupervisorBootstrapIdentityProvider: Send + Sync { + /// Authenticate a driver-native bootstrap token. + /// + /// `Ok(None)` means the token did not authenticate and another + /// authenticator may try it. `Err` means authentication could not safely + /// complete or the token was authenticated but invalid for bootstrap. + async fn authenticate_registration( + &self, + token: &str, + ) -> Result, Status>; +} + +/// Request sent by a driver-side warm-pool controller to activate a pending +/// bootstrap stream. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SupervisorBootstrapActivationRequest { + /// Driver that produced the pending registration. + pub driver: String, + /// Driver-native stable instance ID to activate. + pub instance_id: String, + /// `OpenShell` sandbox ID the instance should become. + pub sandbox_id: String, + /// Driver-native owner UID observed during claim validation. + pub owner_uid: String, + /// Short reason for logs and audit messages. + pub reason: String, +} + +/// Gateway-owned activation callback passed to driver-side warm-pool logic. +#[async_trait] +pub trait SupervisorBootstrapActivator: Send + Sync { + async fn activate_registered_supervisor( + &self, + request: SupervisorBootstrapActivationRequest, + ) -> Result<(), Status>; +} diff --git a/crates/openshell-driver-kubernetes/src/bootstrap.rs b/crates/openshell-driver-kubernetes/src/bootstrap.rs new file mode 100644 index 0000000000..463339adc7 --- /dev/null +++ b/crates/openshell-driver-kubernetes/src/bootstrap.rs @@ -0,0 +1,907 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Kubernetes `ServiceAccount` supervisor bootstrap identity provider. +//! +//! Validates a projected SA token presented by a sandbox pod, reads the pod's +//! `openshell.io/sandbox-id` annotation, verifies the pod is controlled by the +//! corresponding Sandbox CR, and returns a registration-only driver bootstrap +//! identity. Warm pods may register before they are bound to a sandbox, so this +//! identity must not grant sandbox-scoped RPC access. +//! +//! This is the Kubernetes driver's apiserver-facing side of the supervisor +//! bootstrap boundary. The gateway owns the public registration stream and +//! token minting. + +use k8s_openapi::api::{ + authentication::v1::{TokenReview, TokenReviewSpec, TokenReviewStatus, UserInfo}, + core::v1::Pod, +}; +use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta; +use kube::Error as KubeError; +use kube::api::{Api, ApiResource, PostParams}; +use kube::core::{DynamicObject, gvk::GroupVersionKind}; +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapBinding, SupervisorBootstrapIdentity, SupervisorBootstrapIdentityProvider, +}; +use std::sync::Arc; +use tonic::Status; +use tonic::async_trait; +use tracing::{debug, info, warn}; + +/// Pod annotation that binds a sandbox pod to its UUID. Set by the +/// Kubernetes compute driver at pod-create time. +pub const SANDBOX_ID_ANNOTATION: &str = "openshell.io/sandbox-id"; +const SANDBOX_API_GROUP: &str = "agents.x-k8s.io"; +const SANDBOX_API_VERSION_V1BETA1: &str = "v1beta1"; +const SANDBOX_API_VERSION_V1ALPHA1: &str = "v1alpha1"; +const SANDBOX_API_VERSION_FULL_V1BETA1: &str = "agents.x-k8s.io/v1beta1"; +const SANDBOX_API_VERSION_FULL_V1ALPHA1: &str = "agents.x-k8s.io/v1alpha1"; +const SANDBOX_KIND: &str = "Sandbox"; +const SANDBOX_ID_LABEL: &str = "openshell.ai/sandbox-id"; +const POD_NAME_EXTRA: &str = "authentication.kubernetes.io/pod-name"; +const POD_UID_EXTRA: &str = "authentication.kubernetes.io/pod-uid"; + +/// Apiserver-facing operations the authenticator depends on. Split out so +/// tests can fake the apiserver without standing up a kube cluster. +#[async_trait] +pub trait K8sIdentityResolver: Send + Sync + 'static { + /// Validate `token` via `TokenReview` (`aud == openshell-gateway`), + /// extract the pod name/uid, then `GET` the pod and owning Sandbox CR. + /// Returns `Ok(None)` when the token is well-formed but does not + /// authenticate (e.g. wrong audience); returns `Err` for + /// transport/server errors. + async fn resolve(&self, token: &str) -> Result, Status>; +} + +/// Kubernetes implementation of the driver bootstrap identity provider. +pub struct KubernetesSupervisorBootstrapIdentityProvider { + resolver: Arc, +} + +impl std::fmt::Debug for KubernetesSupervisorBootstrapIdentityProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("KubernetesSupervisorBootstrapIdentityProvider") + .finish_non_exhaustive() + } +} + +impl KubernetesSupervisorBootstrapIdentityProvider { + pub fn new(resolver: Arc) -> Self { + Self { resolver } + } +} + +#[async_trait] +impl SupervisorBootstrapIdentityProvider for KubernetesSupervisorBootstrapIdentityProvider { + async fn authenticate_registration( + &self, + token: &str, + ) -> Result, Status> { + self.resolver.resolve(token).await + } +} + +#[derive(Debug)] +struct TokenReviewIdentity { + pod_name: String, + pod_uid: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct SandboxOwnerReference { + api_version: String, + name: String, + uid: String, +} + +/// Resolver backed by the apiserver's `TokenReview` API and `kube::Client` +/// for the per-pod annotation lookup. +pub struct LiveK8sResolver { + token_reviews_api: Api, + pods_api: Api, + sandboxes_api_v1beta1: Api, + sandboxes_api_v1alpha1: Api, + expected_audience: String, + sandbox_namespace: String, + expected_service_account: String, +} + +impl LiveK8sResolver { + pub fn new( + client: kube::Client, + namespace: &str, + expected_audience: String, + expected_service_account: String, + ) -> Self { + let token_reviews_api: Api = Api::all(client.clone()); + let pods_api: Api = Api::namespaced(client.clone(), namespace); + let sandbox_gvk_v1beta1 = + GroupVersionKind::gvk(SANDBOX_API_GROUP, SANDBOX_API_VERSION_V1BETA1, SANDBOX_KIND); + let sandbox_resource_v1beta1 = ApiResource::from_gvk(&sandbox_gvk_v1beta1); + let sandbox_gvk_v1alpha1 = GroupVersionKind::gvk( + SANDBOX_API_GROUP, + SANDBOX_API_VERSION_V1ALPHA1, + SANDBOX_KIND, + ); + let sandbox_resource_v1alpha1 = ApiResource::from_gvk(&sandbox_gvk_v1alpha1); + let sandboxes_api_v1beta1: Api = + Api::namespaced_with(client.clone(), namespace, &sandbox_resource_v1beta1); + let sandboxes_api_v1alpha1: Api = + Api::namespaced_with(client, namespace, &sandbox_resource_v1alpha1); + Self { + token_reviews_api, + pods_api, + sandboxes_api_v1beta1, + sandboxes_api_v1alpha1, + expected_audience, + sandbox_namespace: namespace.to_string(), + expected_service_account, + } + } + + async fn get_sandbox_cr_for_owner( + &self, + owner: &SandboxOwnerReference, + ) -> Result, KubeError> { + let apis = if owner.api_version == SANDBOX_API_VERSION_FULL_V1ALPHA1 { + [&self.sandboxes_api_v1alpha1, &self.sandboxes_api_v1beta1] + } else { + [&self.sandboxes_api_v1beta1, &self.sandboxes_api_v1alpha1] + }; + + for api in apis { + match api.get_opt(&owner.name).await { + Ok(Some(sandbox_cr)) => return Ok(Some(sandbox_cr)), + Ok(None) => {} + Err(err) if should_try_next_sandbox_api_version(&err) => {} + Err(err) => return Err(err), + } + } + + Ok(None) + } +} + +#[async_trait] +impl K8sIdentityResolver for LiveK8sResolver { + async fn resolve(&self, token: &str) -> Result, Status> { + let review = TokenReview { + metadata: ObjectMeta::default(), + spec: TokenReviewSpec { + audiences: Some(vec![self.expected_audience.clone()]), + token: Some(token.to_string()), + }, + status: None, + }; + + let review = self + .token_reviews_api + .create(&PostParams::default(), &review) + .await + .map_err(|e| { + warn!(error = %e, "K8s TokenReview failed"); + Status::internal(format!("tokenreview failed: {e}")) + })?; + let status = review + .status + .ok_or_else(|| Status::internal("TokenReview response missing status"))?; + let Some(identity) = token_review_identity( + &status, + &self.expected_audience, + &self.sandbox_namespace, + &self.expected_service_account, + )? + else { + return Ok(None); + }; + + info!( + pod_name = %identity.pod_name, + pod_uid = %identity.pod_uid, + service_account = %self.expected_service_account, + "validated K8s SA token via TokenReview" + ); + + // Look up the pod and read its sandbox-id annotation. + let pod = self + .pods_api + .get_opt(&identity.pod_name) + .await + .map_err(|e| { + warn!( + pod = %identity.pod_name, + error = %e, + "failed to fetch sandbox pod for annotation lookup" + ); + Status::internal(format!("pod GET failed: {e}")) + })?; + let Some(pod) = pod else { + warn!( + pod = %identity.pod_name, + "sandbox pod referenced by SA token not found in this namespace" + ); + return Err(Status::not_found("sandbox pod not found")); + }; + + // Defense-in-depth: confirm the pod UID matches the SA token's + // `kubernetes.io.pod.uid`. Prevents a replayed token from a + // recreated pod with the same name. + let actual_uid = pod.metadata.uid.as_deref().unwrap_or_default(); + if actual_uid != identity.pod_uid { + warn!( + pod = %identity.pod_name, + claimed_uid = %identity.pod_uid, + actual_uid = %actual_uid, + "SA token pod UID does not match live pod; rejecting" + ); + return Err(Status::permission_denied("SA token pod UID mismatch")); + } + + let sandbox_id = pod_sandbox_id(&pod); + + let owner = sandbox_owner_reference(&pod)?; + let sandbox_cr = self.get_sandbox_cr_for_owner(&owner).await.map_err(|e| { + warn!( + pod = %identity.pod_name, + sandbox_owner = %owner.name, + sandbox_owner_api_version = %owner.api_version, + error = %e, + "failed to fetch owning Sandbox CR for pod identity validation" + ); + Status::internal(format!("sandbox GET failed: {e}")) + })?; + let Some(sandbox_cr) = sandbox_cr else { + warn!( + pod = %identity.pod_name, + sandbox_owner = %owner.name, + sandbox_owner_api_version = %owner.api_version, + "pod ownerReference points to a Sandbox CR that does not exist" + ); + return Err(Status::permission_denied("sandbox owner not found")); + }; + validate_sandbox_owner_reference(&owner, &sandbox_cr)?; + if let Some(ref sandbox_id) = sandbox_id { + validate_sandbox_owner_binding(&owner, sandbox_id, &sandbox_cr)?; + } + + let binding = sandbox_id.map_or(SupervisorBootstrapBinding::WarmPending, |sandbox_id| { + SupervisorBootstrapBinding::BoundSandbox { sandbox_id } + }); + + Ok(Some(SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: identity.pod_name, + instance_id: identity.pod_uid, + owner_name: owner.name, + owner_uid: owner.uid, + binding, + })) + } +} + +#[allow(clippy::result_large_err)] +fn token_review_identity( + status: &TokenReviewStatus, + expected_audience: &str, + sandbox_namespace: &str, + expected_service_account: &str, +) -> Result, Status> { + if status.authenticated != Some(true) { + debug!( + error = status.error.as_deref().unwrap_or_default(), + "K8s TokenReview did not authenticate token" + ); + return Ok(None); + } + + let audiences = status.audiences.as_deref().unwrap_or_default(); + if !audiences.iter().any(|aud| aud == expected_audience) { + warn!( + expected_audience = %expected_audience, + audiences = ?audiences, + "K8s TokenReview authenticated token without expected audience" + ); + return Err(Status::unauthenticated("SA token audience not accepted")); + } + + let user = status + .user + .as_ref() + .ok_or_else(|| Status::permission_denied("TokenReview response missing user info"))?; + let username = user + .username + .as_deref() + .ok_or_else(|| Status::permission_denied("TokenReview response missing username"))?; + let expected_username = + format!("system:serviceaccount:{sandbox_namespace}:{expected_service_account}"); + if username != expected_username { + warn!( + username = %username, + sandbox_namespace = %sandbox_namespace, + service_account = %expected_service_account, + "K8s TokenReview principal is not the configured sandbox service account" + ); + return Err(Status::permission_denied( + "SA token is not from the configured sandbox service account", + )); + } + + let pod_name = user_extra_one(user, POD_NAME_EXTRA)?; + let pod_uid = user_extra_one(user, POD_UID_EXTRA)?; + Ok(Some(TokenReviewIdentity { pod_name, pod_uid })) +} + +#[allow(clippy::result_large_err)] +fn user_extra_one(user: &UserInfo, key: &str) -> Result { + let Some(values) = user.extra.as_ref().and_then(|extra| extra.get(key)) else { + return Err(Status::permission_denied("SA token is not pod-bound")); + }; + if values.len() != 1 || values[0].is_empty() { + return Err(Status::permission_denied( + "SA token has invalid pod binding", + )); + } + Ok(values[0].clone()) +} + +#[allow(clippy::result_large_err)] +fn pod_sandbox_id(pod: &Pod) -> Option { + pod.metadata + .annotations + .as_ref() + .and_then(|a| a.get(SANDBOX_ID_ANNOTATION)) + .cloned() + .filter(|sandbox_id| !sandbox_id.is_empty()) +} + +#[allow(clippy::result_large_err)] +fn sandbox_owner_reference(pod: &Pod) -> Result { + let owner_refs = pod.metadata.owner_references.as_deref().unwrap_or_default(); + let mut sandbox_refs = owner_refs + .iter() + .filter(|owner| is_supported_sandbox_owner_reference(owner)); + let Some(owner) = sandbox_refs.next() else { + let unsupported_sandbox_api_versions = owner_refs + .iter() + .filter(|owner| owner.kind == SANDBOX_KIND) + .map(|owner| owner.api_version.as_str()) + .collect::>(); + if !unsupported_sandbox_api_versions.is_empty() { + warn!( + api_versions = ?unsupported_sandbox_api_versions, + supported_api_versions = ?[ + SANDBOX_API_VERSION_FULL_V1BETA1, + SANDBOX_API_VERSION_FULL_V1ALPHA1, + ], + "pod Sandbox ownerReference uses unsupported apiVersion" + ); + } + return Err(Status::permission_denied( + "pod is not controlled by an OpenShell Sandbox", + )); + }; + if sandbox_refs.next().is_some() { + return Err(Status::permission_denied( + "pod has multiple OpenShell Sandbox owners", + )); + } + if owner.controller != Some(true) { + return Err(Status::permission_denied( + "pod Sandbox ownerReference is not controlling", + )); + } + if owner.name.is_empty() || owner.uid.is_empty() { + return Err(Status::permission_denied( + "pod Sandbox ownerReference is incomplete", + )); + } + Ok(SandboxOwnerReference { + api_version: owner.api_version.clone(), + name: owner.name.clone(), + uid: owner.uid.clone(), + }) +} + +fn is_supported_sandbox_owner_reference( + owner: &k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference, +) -> bool { + owner.kind == SANDBOX_KIND + && matches!( + owner.api_version.as_str(), + SANDBOX_API_VERSION_FULL_V1BETA1 | SANDBOX_API_VERSION_FULL_V1ALPHA1 + ) +} + +fn should_try_next_sandbox_api_version(err: &KubeError) -> bool { + // Kubernetes returns a structured 404 for some missing API resources and a + // raw "404 page not found" body for others. Both mean the probed + // group/version is unavailable and the next supported Sandbox API version + // should be tried. + matches!(err, KubeError::Api(api) if api.code == 404) +} + +#[allow(clippy::result_large_err)] +fn validate_sandbox_owner_reference( + owner: &SandboxOwnerReference, + sandbox_cr: &DynamicObject, +) -> Result<(), Status> { + let actual_uid = sandbox_cr.metadata.uid.as_deref().unwrap_or_default(); + if actual_uid != owner.uid { + warn!( + sandbox_owner = %owner.name, + owner_uid = %owner.uid, + actual_uid = %actual_uid, + "pod Sandbox ownerReference UID does not match live Sandbox CR" + ); + return Err(Status::permission_denied("sandbox owner UID mismatch")); + } + + Ok(()) +} + +#[allow(clippy::result_large_err)] +fn validate_sandbox_owner_binding( + owner: &SandboxOwnerReference, + sandbox_id: &str, + sandbox_cr: &DynamicObject, +) -> Result<(), Status> { + let actual_sandbox_id = sandbox_cr + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(SANDBOX_ID_LABEL)) + .map(String::as_str) + .unwrap_or_default(); + if actual_sandbox_id != sandbox_id { + warn!( + sandbox_owner = %owner.name, + owner_uid = %owner.uid, + pod_sandbox_id = %sandbox_id, + cr_sandbox_id = %actual_sandbox_id, + "pod sandbox annotation does not match owning Sandbox CR label" + ); + return Err(Status::permission_denied("sandbox owner ID mismatch")); + } + + Ok(()) +} + +#[cfg(test)] +pub mod test_support { + use super::*; + use std::sync::Mutex; + + /// Fake resolver for unit tests. Returns the configured outcome on + /// every call and records the tokens it observed. + pub struct FakeResolver { + pub outcome: Result, Status>, + pub seen_tokens: Mutex>, + } + + impl FakeResolver { + pub fn returning(outcome: Result, Status>) -> Self { + Self { + outcome, + seen_tokens: Mutex::new(Vec::new()), + } + } + } + + #[async_trait] + impl K8sIdentityResolver for FakeResolver { + async fn resolve( + &self, + token: &str, + ) -> Result, Status> { + self.seen_tokens.lock().unwrap().push(token.to_string()); + match &self.outcome { + Ok(opt) => Ok(opt.clone()), + Err(s) => Err(Status::new(s.code(), s.message())), + } + } + } +} + +#[cfg(test)] +mod tests { + use super::test_support::FakeResolver; + use super::*; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference; + use std::collections::BTreeMap; + use std::sync::Arc; + + fn kube_api_error(code: u16, message: &str) -> KubeError { + KubeError::Api(kube::core::ErrorResponse { + status: if code == 404 { + "404 Not Found".to_string() + } else { + "Failure".to_string() + }, + message: message.to_string(), + reason: "Failed to parse error data".to_string(), + code, + }) + } + + #[test] + fn sandbox_api_version_probe_retries_on_structured_and_raw_404() { + let structured = kube_api_error(404, "could not find the requested resource"); + assert!(should_try_next_sandbox_api_version(&structured)); + + let raw = kube_api_error(404, "404 page not found\n"); + assert!(should_try_next_sandbox_api_version(&raw)); + } + + #[test] + fn sandbox_api_version_probe_keeps_non_404_errors() { + let err = kube_api_error(403, "sandboxes.agents.x-k8s.io is forbidden"); + assert!(!should_try_next_sandbox_api_version(&err)); + } + + fn token_review_status( + authenticated: bool, + audiences: Vec<&str>, + username: &str, + extra: Vec<(&str, &str)>, + ) -> TokenReviewStatus { + TokenReviewStatus { + authenticated: Some(authenticated), + audiences: Some(audiences.into_iter().map(str::to_string).collect()), + error: None, + user: Some(UserInfo { + username: Some(username.to_string()), + uid: Some("sa-uid".to_string()), + groups: Some(vec![ + "system:serviceaccounts".to_string(), + "system:serviceaccounts:openshell".to_string(), + "system:authenticated".to_string(), + ]), + extra: Some( + extra + .into_iter() + .map(|(k, v)| (k.to_string(), vec![v.to_string()])) + .collect::>(), + ), + }), + } + } + + fn sandbox_owner(name: &str, uid: &str) -> OwnerReference { + sandbox_owner_with_api_version(SANDBOX_API_VERSION_FULL_V1BETA1, name, uid) + } + + fn sandbox_owner_with_api_version(api_version: &str, name: &str, uid: &str) -> OwnerReference { + OwnerReference { + api_version: api_version.to_string(), + block_owner_deletion: None, + controller: Some(true), + kind: SANDBOX_KIND.to_string(), + name: name.to_string(), + uid: uid.to_string(), + } + } + + fn pod_with_owner_refs(owner_references: Vec) -> Pod { + Pod { + metadata: ObjectMeta { + owner_references: Some(owner_references), + ..Default::default() + }, + ..Default::default() + } + } + + fn pod_with_sandbox_id(sandbox_id: Option<&str>) -> Pod { + Pod { + metadata: ObjectMeta { + annotations: sandbox_id.map(|id| { + BTreeMap::from([(SANDBOX_ID_ANNOTATION.to_string(), id.to_string())]) + }), + ..Default::default() + }, + ..Default::default() + } + } + + fn bootstrap_identity(binding: SupervisorBootstrapBinding) -> SupervisorBootstrapIdentity { + SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "openshell-sandbox-a".to_string(), + instance_id: "uid-a".to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "cr-uid-a".to_string(), + binding, + } + } + + fn sandbox_cr(name: &str, uid: &str, sandbox_id: &str) -> DynamicObject { + let sandbox_gvk = + GroupVersionKind::gvk(SANDBOX_API_GROUP, SANDBOX_API_VERSION_V1BETA1, SANDBOX_KIND); + let sandbox_resource = ApiResource::from_gvk(&sandbox_gvk); + let mut cr = DynamicObject::new(name, &sandbox_resource); + cr.metadata.uid = Some(uid.to_string()); + cr.metadata.labels = Some(BTreeMap::from([( + SANDBOX_ID_LABEL.to_string(), + sandbox_id.to_string(), + )])); + cr + } + + #[test] + fn token_review_identity_extracts_pod_binding() { + let status = token_review_status( + true, + vec!["openshell-gateway"], + "system:serviceaccount:openshell:default", + vec![ + (POD_NAME_EXTRA, "openshell-sandbox-a"), + (POD_UID_EXTRA, "uid-a"), + ], + ); + + let identity = token_review_identity(&status, "openshell-gateway", "openshell", "default") + .unwrap() + .expect("authenticated token should resolve"); + + assert_eq!(identity.pod_name, "openshell-sandbox-a"); + assert_eq!(identity.pod_uid, "uid-a"); + } + + #[test] + fn token_review_identity_returns_none_when_not_authenticated() { + let status = TokenReviewStatus { + authenticated: Some(false), + error: Some("invalid audience".to_string()), + ..Default::default() + }; + + assert!( + token_review_identity(&status, "openshell-gateway", "openshell", "default") + .unwrap() + .is_none() + ); + } + + #[test] + fn token_review_identity_requires_expected_audience() { + let status = token_review_status( + true, + vec!["kubernetes.default.svc"], + "system:serviceaccount:openshell:default", + vec![ + (POD_NAME_EXTRA, "openshell-sandbox-a"), + (POD_UID_EXTRA, "uid-a"), + ], + ); + + let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") + .expect_err("wrong audience must fail closed"); + assert_eq!(err.code(), tonic::Code::Unauthenticated); + } + + #[test] + fn token_review_identity_requires_sandbox_namespace() { + let status = token_review_status( + true, + vec!["openshell-gateway"], + "system:serviceaccount:other:default", + vec![ + (POD_NAME_EXTRA, "openshell-sandbox-a"), + (POD_UID_EXTRA, "uid-a"), + ], + ); + + let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") + .expect_err("other namespace must be rejected"); + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn token_review_identity_requires_configured_service_account() { + let status = token_review_status( + true, + vec!["openshell-gateway"], + "system:serviceaccount:openshell:other", + vec![ + (POD_NAME_EXTRA, "openshell-sandbox-a"), + (POD_UID_EXTRA, "uid-a"), + ], + ); + + let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") + .expect_err("other service account must be rejected"); + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn token_review_identity_requires_pod_bound_extras() { + let status = token_review_status( + true, + vec!["openshell-gateway"], + "system:serviceaccount:openshell:default", + vec![], + ); + + let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") + .expect_err("non pod-bound tokens must be rejected"); + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn pod_sandbox_id_is_optional_for_warm_registration() { + assert_eq!( + pod_sandbox_id(&pod_with_sandbox_id(Some("sandbox-id-a"))).as_deref(), + Some("sandbox-id-a") + ); + assert!(pod_sandbox_id(&pod_with_sandbox_id(None)).is_none()); + } + + #[test] + fn sandbox_owner_reference_extracts_controlling_sandbox_owner() { + let pod = pod_with_owner_refs(vec![sandbox_owner("sandbox-a", "cr-uid-a")]); + + let owner = sandbox_owner_reference(&pod).expect("expected Sandbox owner"); + + assert_eq!( + owner, + SandboxOwnerReference { + api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), + name: "sandbox-a".to_string(), + uid: "cr-uid-a".to_string(), + } + ); + } + + #[test] + fn sandbox_owner_reference_accepts_v1alpha1_owner() { + let pod = pod_with_owner_refs(vec![sandbox_owner_with_api_version( + SANDBOX_API_VERSION_FULL_V1ALPHA1, + "sandbox-a", + "cr-uid-a", + )]); + + let owner = sandbox_owner_reference(&pod).expect("expected v1alpha1 Sandbox owner"); + + assert_eq!( + owner, + SandboxOwnerReference { + api_version: SANDBOX_API_VERSION_FULL_V1ALPHA1.to_string(), + name: "sandbox-a".to_string(), + uid: "cr-uid-a".to_string(), + } + ); + } + + #[test] + fn sandbox_owner_reference_rejects_missing_owner() { + let pod = pod_with_owner_refs(vec![]); + + let err = sandbox_owner_reference(&pod).expect_err("missing owner must fail"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn sandbox_owner_reference_rejects_unsupported_sandbox_api_version() { + let pod = pod_with_owner_refs(vec![sandbox_owner_with_api_version( + "agents.x-k8s.io/v1", + "sandbox-a", + "cr-uid-a", + )]); + + let err = + sandbox_owner_reference(&pod).expect_err("unsupported apiVersion must fail closed"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn sandbox_owner_reference_requires_controlling_owner() { + let mut owner = sandbox_owner("sandbox-a", "cr-uid-a"); + owner.controller = Some(false); + let pod = pod_with_owner_refs(vec![owner]); + + let err = sandbox_owner_reference(&pod).expect_err("non-controller owner must fail"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn sandbox_owner_reference_rejects_ambiguous_sandbox_owners() { + let pod = pod_with_owner_refs(vec![ + sandbox_owner("sandbox-a", "cr-uid-a"), + sandbox_owner("sandbox-b", "cr-uid-b"), + ]); + + let err = sandbox_owner_reference(&pod).expect_err("multiple owners must fail"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn validate_sandbox_owner_reference_requires_matching_cr_uid() { + let owner = SandboxOwnerReference { + api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), + name: "sandbox-a".to_string(), + uid: "cr-uid-a".to_string(), + }; + let cr = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-a"); + validate_sandbox_owner_reference(&owner, &cr).expect("matching CR should be accepted"); + + let wrong_uid = sandbox_cr("sandbox-a", "cr-uid-b", "sandbox-id-a"); + let err = validate_sandbox_owner_reference(&owner, &wrong_uid) + .expect_err("wrong CR UID must fail"); + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn validate_sandbox_owner_binding_requires_matching_label() { + let owner = SandboxOwnerReference { + api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), + name: "sandbox-a".to_string(), + uid: "cr-uid-a".to_string(), + }; + let cr = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-a"); + validate_sandbox_owner_binding(&owner, "sandbox-id-a", &cr) + .expect("matching label should be accepted"); + let wrong_label = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-b"); + let err = validate_sandbox_owner_binding(&owner, "sandbox-id-a", &wrong_label) + .expect_err("wrong sandbox-id label must fail"); + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[tokio::test] + async fn provider_delegates_to_resolver() { + let resolved = bootstrap_identity(SupervisorBootstrapBinding::BoundSandbox { + sandbox_id: "sandbox-a".to_string(), + }); + let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved.clone())))); + let provider = KubernetesSupervisorBootstrapIdentityProvider::new(fake.clone()); + + let result = provider + .authenticate_registration("sa-jwt") + .await + .unwrap() + .expect("expected identity"); + + assert_eq!(result, resolved); + assert_eq!(fake.seen_tokens.lock().unwrap().as_slice(), ["sa-jwt"]); + } + + #[tokio::test] + async fn provider_allows_none_to_fall_through() { + let fake = Arc::new(FakeResolver::returning(Ok(None))); + let provider = KubernetesSupervisorBootstrapIdentityProvider::new(fake); + let result = provider.authenticate_registration("unknown").await.unwrap(); + assert!(result.is_none()); + } + + #[tokio::test] + async fn provider_accepts_warm_pending_identity() { + let resolved = bootstrap_identity(SupervisorBootstrapBinding::WarmPending); + let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved.clone())))); + let provider = KubernetesSupervisorBootstrapIdentityProvider::new(fake); + let identity = provider + .authenticate_registration("sa-jwt") + .await + .unwrap() + .expect("expected identity"); + + assert_eq!(identity.binding, SupervisorBootstrapBinding::WarmPending); + assert_eq!(identity.owner_uid, "cr-uid-a"); + } + + #[tokio::test] + async fn resolver_error_propagates() { + let fake = Arc::new(FakeResolver::returning(Err(Status::unavailable( + "apiserver down", + )))); + let provider = KubernetesSupervisorBootstrapIdentityProvider::new(fake); + let err = provider + .authenticate_registration("sa-jwt") + .await + .expect_err("resolver error must propagate"); + assert_eq!(err.code(), tonic::Code::Unavailable); + } +} diff --git a/crates/openshell-driver-kubernetes/src/lib.rs b/crates/openshell-driver-kubernetes/src/lib.rs index 7c56c8de5b..f967170e43 100644 --- a/crates/openshell-driver-kubernetes/src/lib.rs +++ b/crates/openshell-driver-kubernetes/src/lib.rs @@ -1,10 +1,14 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +pub mod bootstrap; pub mod config; pub mod driver; pub mod grpc; +pub use bootstrap::{ + K8sIdentityResolver, KubernetesSupervisorBootstrapIdentityProvider, LiveK8sResolver, +}; pub use config::{ AppArmorProfile, DEFAULT_PROXY_UID, DEFAULT_SANDBOX_SERVICE_ACCOUNT_NAME, DEFAULT_WORKSPACE_STORAGE_SIZE, KubernetesComputeConfig, KubernetesSidecarConfig, diff --git a/crates/openshell-server/src/auth/authenticator.rs b/crates/openshell-server/src/auth/authenticator.rs index bc52729224..a5f8c54bd2 100644 --- a/crates/openshell-server/src/auth/authenticator.rs +++ b/crates/openshell-server/src/auth/authenticator.rs @@ -14,8 +14,8 @@ //! //! Live authenticators slotting into the chain: //! - [`super::sandbox_jwt::SandboxJwtAuthenticator`] — gateway-minted JWTs -//! - [`super::k8s_sa::K8sServiceAccountAuthenticator`] — K8s projected SA -//! tokens (path-scoped to Kubernetes bootstrap RPCs) +//! - [`super::k8s_sa::SupervisorBootstrapAuthenticator`] — driver-provided +//! bootstrap tokens (path-scoped to supervisor bootstrap RPCs) //! - [`super::oidc::OidcAuthenticator`] — user OIDC Bearer tokens use super::principal::Principal; use async_trait::async_trait; diff --git a/crates/openshell-server/src/auth/guard.rs b/crates/openshell-server/src/auth/guard.rs index 481cfed2ff..c1e87b6e48 100644 --- a/crates/openshell-server/src/auth/guard.rs +++ b/crates/openshell-server/src/auth/guard.rs @@ -20,8 +20,8 @@ use tracing::info; /// scope at the router level). /// - [`Principal::Sandbox`] must reference the same canonical UUID it /// was authenticated with. -/// - [`Principal::K8sPod`] is rejected — pod registration identity is not a -/// sandbox-scoped credential. +/// - [`Principal::SupervisorBootstrap`] is rejected — bootstrap registration +/// identity is not a sandbox-scoped credential. /// - [`Principal::Anonymous`] is rejected — sandbox-class methods are /// never anonymously callable. /// @@ -46,7 +46,7 @@ pub fn ensure_sandbox_scope(principal: &Principal, claimed_sandbox_id: &str) -> )) } } - Principal::K8sPod(_) => Err(Status::permission_denied( + Principal::SupervisorBootstrap(_) => Err(Status::permission_denied( "sandbox-scoped methods require a sandbox principal", )), Principal::Anonymous => Err(Status::unauthenticated( @@ -89,7 +89,7 @@ pub fn ensure_sandbox_principal_scope( ensure_sandbox_scope(principal, claimed_sandbox_id)?; Ok(p.clone()) } - Principal::User(_) | Principal::K8sPod(_) => Err(Status::permission_denied( + Principal::User(_) | Principal::SupervisorBootstrap(_) => Err(Status::permission_denied( "supervisor RPCs require a sandbox principal", )), Principal::Anonymous => Err(Status::unauthenticated( diff --git a/crates/openshell-server/src/auth/k8s_sa.rs b/crates/openshell-server/src/auth/k8s_sa.rs index d288ebed91..fd547530ca 100644 --- a/crates/openshell-server/src/auth/k8s_sa.rs +++ b/crates/openshell-server/src/auth/k8s_sa.rs @@ -1,97 +1,58 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Kubernetes `ServiceAccount` bootstrap authenticator. +//! Supervisor bootstrap authenticator adapter. //! -//! Path-scoped to the Kubernetes supervisor bootstrap RPCs. Validates a projected SA token -//! presented by a sandbox pod, reads the pod's `openshell.io/sandbox-id` -//! annotation, verifies the pod is controlled by the corresponding Sandbox CR, -//! and returns either a sandbox principal for the legacy token RPC or a -//! registration-scoped pod principal for `RegisterSupervisorPod`. Warm pods may -//! register before they are bound to a sandbox, so registration identity must -//! not grant sandbox-scoped RPC access. -//! -//! This is the only authenticator that talks to the K8s apiserver. It is -//! optional — the gateway boots without it in singleplayer deployments. +//! The Kubernetes-specific `TokenReview` and pod lookup logic lives in the +//! Kubernetes driver. This module is intentionally gateway-small: it is +//! path-scoped to supervisor bootstrap RPCs, extracts a bearer token, delegates +//! validation to the active driver's bootstrap identity provider, and turns the +//! resulting registration-only identity into the principal shape expected by the +//! gateway auth router. use super::authenticator::Authenticator; -use super::principal::{Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal}; -use async_trait::async_trait; -use k8s_openapi::api::{ - authentication::v1::{TokenReview, TokenReviewSpec, TokenReviewStatus, UserInfo}, - core::v1::Pod, +use super::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapBinding, SupervisorBootstrapIdentityProvider, }; -use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta; -use kube::Error as KubeError; -use kube::api::{Api, ApiResource, PostParams}; -use kube::core::{DynamicObject, gvk::GroupVersionKind}; use std::sync::Arc; use tonic::Status; -use tracing::{debug, info, warn}; +use tonic::async_trait; +use tracing::{debug, warn}; /// Legacy gRPC method path accepted by this authenticator. All other -/// non-bootstrap paths fall through (return `Ok(None)`) so a gateway-minted JWT +/// non-bootstrap paths fall through so a gateway-minted JWT or user credential /// is required there. pub const ISSUE_SANDBOX_TOKEN_PATH: &str = "/openshell.v1.OpenShell/IssueSandboxToken"; -/// Supervisor pod registration path accepted by this authenticator for the -/// phase-1 cold Kubernetes activation flow. +/// Supervisor registration path accepted by this authenticator. pub const REGISTER_SUPERVISOR_POD_PATH: &str = "/openshell.v1.OpenShell/RegisterSupervisorPod"; -/// Pod annotation that binds a sandbox pod to its UUID. Set by the -/// Kubernetes compute driver at pod-create time. The gateway accepts this -/// annotation only after validating the pod's `TokenReview` binding, live UID, -/// and owning Sandbox CR. The K8s `Role` granted to the gateway must not -/// include `patch pods` (see plan §11.8). -pub const SANDBOX_ID_ANNOTATION: &str = "openshell.io/sandbox-id"; -const SANDBOX_API_GROUP: &str = "agents.x-k8s.io"; -const SANDBOX_API_VERSION_V1BETA1: &str = "v1beta1"; -const SANDBOX_API_VERSION_V1ALPHA1: &str = "v1alpha1"; -const SANDBOX_API_VERSION_FULL_V1BETA1: &str = "agents.x-k8s.io/v1beta1"; -const SANDBOX_API_VERSION_FULL_V1ALPHA1: &str = "agents.x-k8s.io/v1alpha1"; -const SANDBOX_KIND: &str = "Sandbox"; -const SANDBOX_ID_LABEL: &str = "openshell.ai/sandbox-id"; -const POD_NAME_EXTRA: &str = "authentication.kubernetes.io/pod-name"; -const POD_UID_EXTRA: &str = "authentication.kubernetes.io/pod-uid"; - -/// Apiserver-facing operations the authenticator depends on. Split out so -/// tests can fake the apiserver without standing up a kube cluster. -#[async_trait] -pub trait K8sIdentityResolver: Send + Sync + 'static { - /// Validate `token` via `TokenReview` (`aud == openshell-gateway`), - /// extract the pod name/uid, then `GET` the pod and owning Sandbox CR. - /// Returns `Ok(None)` when the token is well-formed but does not - /// authenticate (e.g. wrong audience); returns `Err` for - /// transport/server errors. - async fn resolve(&self, token: &str) -> Result, Status>; +/// Path-scoped authenticator backed by the active compute driver's bootstrap +/// identity provider. +pub struct SupervisorBootstrapAuthenticator { + provider: Arc, } -/// Authenticator wrapper around a [`K8sIdentityResolver`]. -pub struct K8sServiceAccountAuthenticator { - resolver: Arc, -} - -impl std::fmt::Debug for K8sServiceAccountAuthenticator { +impl std::fmt::Debug for SupervisorBootstrapAuthenticator { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("K8sServiceAccountAuthenticator") + f.debug_struct("SupervisorBootstrapAuthenticator") .finish_non_exhaustive() } } -impl K8sServiceAccountAuthenticator { - pub fn new(resolver: Arc) -> Self { - Self { resolver } +impl SupervisorBootstrapAuthenticator { + pub fn new(provider: Arc) -> Self { + Self { provider } } } #[async_trait] -impl Authenticator for K8sServiceAccountAuthenticator { +impl Authenticator for SupervisorBootstrapAuthenticator { async fn authenticate( &self, headers: &http::HeaderMap, path: &str, ) -> Result, Status> { - // Scope: only Kubernetes bootstrap RPCs. Other paths fall through so - // the SandboxJwtAuthenticator (or OIDC) handles them. if path != ISSUE_SANDBOX_TOKEN_PATH && path != REGISTER_SUPERVISOR_POD_PATH { return Ok(None); } @@ -104,428 +65,52 @@ impl Authenticator for K8sServiceAccountAuthenticator { return Ok(None); }; - let Some(resolved) = self.resolver.resolve(token).await? else { - debug!("K8s SA token did not authenticate; falling through"); + let Some(identity) = self.provider.authenticate_registration(token).await? else { + debug!("supervisor bootstrap token did not authenticate; falling through"); return Ok(None); }; if path == REGISTER_SUPERVISOR_POD_PATH { - return Ok(Some(Principal::K8sPod(resolved))); + return Ok(Some(Principal::SupervisorBootstrap(identity))); } - let sandbox_id = resolved.sandbox_id.clone().ok_or_else(|| { + let SupervisorBootstrapBinding::BoundSandbox { sandbox_id } = identity.binding else { warn!( - pod = %resolved.pod_name, - "pod missing openshell.io/sandbox-id annotation; rejecting legacy token issue" + driver = %identity.driver, + instance_name = %identity.instance_name, + instance_id = %identity.instance_id, + "bootstrap identity is not bound to a sandbox; rejecting legacy token issue" ); - Status::permission_denied("pod is not bound to a sandbox identity") - })?; + return Err(Status::permission_denied( + "supervisor instance is not bound to a sandbox identity", + )); + }; Ok(Some(Principal::Sandbox(SandboxPrincipal { sandbox_id, - source: SandboxIdentitySource::K8sServiceAccount { - pod_name: resolved.pod_name, - pod_uid: resolved.pod_uid, + source: SandboxIdentitySource::SupervisorBootstrap { + driver: identity.driver, + instance_name: identity.instance_name, + instance_id: identity.instance_id, }, trust_domain: Some("openshell".to_string()), }))) } } -#[derive(Debug)] -struct TokenReviewIdentity { - pod_name: String, - pod_uid: String, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct SandboxOwnerReference { - api_version: String, - name: String, - uid: String, -} - -/// Resolver backed by the apiserver's `TokenReview` API and `kube::Client` -/// for the per-pod annotation lookup. -pub struct LiveK8sResolver { - token_reviews_api: Api, - pods_api: Api, - sandboxes_api_v1beta1: Api, - sandboxes_api_v1alpha1: Api, - expected_audience: String, - sandbox_namespace: String, - expected_service_account: String, -} - -impl LiveK8sResolver { - pub fn new( - client: kube::Client, - namespace: &str, - expected_audience: String, - expected_service_account: String, - ) -> Self { - let token_reviews_api: Api = Api::all(client.clone()); - let pods_api: Api = Api::namespaced(client.clone(), namespace); - let sandbox_gvk_v1beta1 = - GroupVersionKind::gvk(SANDBOX_API_GROUP, SANDBOX_API_VERSION_V1BETA1, SANDBOX_KIND); - let sandbox_resource_v1beta1 = ApiResource::from_gvk(&sandbox_gvk_v1beta1); - let sandbox_gvk_v1alpha1 = GroupVersionKind::gvk( - SANDBOX_API_GROUP, - SANDBOX_API_VERSION_V1ALPHA1, - SANDBOX_KIND, - ); - let sandbox_resource_v1alpha1 = ApiResource::from_gvk(&sandbox_gvk_v1alpha1); - let sandboxes_api_v1beta1: Api = - Api::namespaced_with(client.clone(), namespace, &sandbox_resource_v1beta1); - let sandboxes_api_v1alpha1: Api = - Api::namespaced_with(client, namespace, &sandbox_resource_v1alpha1); - Self { - token_reviews_api, - pods_api, - sandboxes_api_v1beta1, - sandboxes_api_v1alpha1, - expected_audience, - sandbox_namespace: namespace.to_string(), - expected_service_account, - } - } - - async fn get_sandbox_cr_for_owner( - &self, - owner: &SandboxOwnerReference, - ) -> Result, KubeError> { - let apis = if owner.api_version == SANDBOX_API_VERSION_FULL_V1ALPHA1 { - [&self.sandboxes_api_v1alpha1, &self.sandboxes_api_v1beta1] - } else { - [&self.sandboxes_api_v1beta1, &self.sandboxes_api_v1alpha1] - }; - - for api in apis { - match api.get_opt(&owner.name).await { - Ok(Some(sandbox_cr)) => return Ok(Some(sandbox_cr)), - Ok(None) => {} - Err(err) if should_try_next_sandbox_api_version(&err) => {} - Err(err) => return Err(err), - } - } - - Ok(None) - } -} - -#[async_trait] -impl K8sIdentityResolver for LiveK8sResolver { - async fn resolve(&self, token: &str) -> Result, Status> { - let review = TokenReview { - metadata: ObjectMeta::default(), - spec: TokenReviewSpec { - audiences: Some(vec![self.expected_audience.clone()]), - token: Some(token.to_string()), - }, - status: None, - }; - - let review = self - .token_reviews_api - .create(&PostParams::default(), &review) - .await - .map_err(|e| { - warn!(error = %e, "K8s TokenReview failed"); - Status::internal(format!("tokenreview failed: {e}")) - })?; - let status = review - .status - .ok_or_else(|| Status::internal("TokenReview response missing status"))?; - let Some(identity) = token_review_identity( - &status, - &self.expected_audience, - &self.sandbox_namespace, - &self.expected_service_account, - )? - else { - return Ok(None); - }; - - info!( - pod_name = %identity.pod_name, - pod_uid = %identity.pod_uid, - service_account = %self.expected_service_account, - "validated K8s SA token via TokenReview" - ); - - // Look up the pod and read its sandbox-id annotation. - let pod = self - .pods_api - .get_opt(&identity.pod_name) - .await - .map_err(|e| { - warn!( - pod = %identity.pod_name, - error = %e, - "failed to fetch sandbox pod for annotation lookup" - ); - Status::internal(format!("pod GET failed: {e}")) - })?; - let Some(pod) = pod else { - warn!( - pod = %identity.pod_name, - "sandbox pod referenced by SA token not found in this namespace" - ); - return Err(Status::not_found("sandbox pod not found")); - }; - - // Defense-in-depth: confirm the pod UID matches the SA token's - // `kubernetes.io.pod.uid`. Prevents a replayed token from a - // recreated pod with the same name. - let actual_uid = pod.metadata.uid.as_deref().unwrap_or_default(); - if actual_uid != identity.pod_uid { - warn!( - pod = %identity.pod_name, - claimed_uid = %identity.pod_uid, - actual_uid = %actual_uid, - "SA token pod UID does not match live pod; rejecting" - ); - return Err(Status::permission_denied("SA token pod UID mismatch")); - } - - let sandbox_id = pod_sandbox_id(&pod); - - let owner = sandbox_owner_reference(&pod)?; - let sandbox_cr = self.get_sandbox_cr_for_owner(&owner).await.map_err(|e| { - warn!( - pod = %identity.pod_name, - sandbox_owner = %owner.name, - sandbox_owner_api_version = %owner.api_version, - error = %e, - "failed to fetch owning Sandbox CR for pod identity validation" - ); - Status::internal(format!("sandbox GET failed: {e}")) - })?; - let Some(sandbox_cr) = sandbox_cr else { - warn!( - pod = %identity.pod_name, - sandbox_owner = %owner.name, - sandbox_owner_api_version = %owner.api_version, - "pod ownerReference points to a Sandbox CR that does not exist" - ); - return Err(Status::permission_denied("sandbox owner not found")); - }; - validate_sandbox_owner_reference(&owner, &sandbox_cr)?; - if let Some(ref sandbox_id) = sandbox_id { - validate_sandbox_owner_binding(&owner, sandbox_id, &sandbox_cr)?; - } - - Ok(Some(RegisteredPodIdentity { - sandbox_id, - pod_name: identity.pod_name, - pod_uid: identity.pod_uid, - sandbox_owner_name: owner.name, - sandbox_owner_uid: owner.uid, - })) - } -} - -#[allow(clippy::result_large_err)] -fn token_review_identity( - status: &TokenReviewStatus, - expected_audience: &str, - sandbox_namespace: &str, - expected_service_account: &str, -) -> Result, Status> { - if status.authenticated != Some(true) { - debug!( - error = status.error.as_deref().unwrap_or_default(), - "K8s TokenReview did not authenticate token" - ); - return Ok(None); - } - - let audiences = status.audiences.as_deref().unwrap_or_default(); - if !audiences.iter().any(|aud| aud == expected_audience) { - warn!( - expected_audience = %expected_audience, - audiences = ?audiences, - "K8s TokenReview authenticated token without expected audience" - ); - return Err(Status::unauthenticated("SA token audience not accepted")); - } - - let user = status - .user - .as_ref() - .ok_or_else(|| Status::permission_denied("TokenReview response missing user info"))?; - let username = user - .username - .as_deref() - .ok_or_else(|| Status::permission_denied("TokenReview response missing username"))?; - let expected_username = - format!("system:serviceaccount:{sandbox_namespace}:{expected_service_account}"); - if username != expected_username { - warn!( - username = %username, - sandbox_namespace = %sandbox_namespace, - service_account = %expected_service_account, - "K8s TokenReview principal is not the configured sandbox service account" - ); - return Err(Status::permission_denied( - "SA token is not from the configured sandbox service account", - )); - } - - let pod_name = user_extra_one(user, POD_NAME_EXTRA)?; - let pod_uid = user_extra_one(user, POD_UID_EXTRA)?; - Ok(Some(TokenReviewIdentity { pod_name, pod_uid })) -} - -#[allow(clippy::result_large_err)] -fn user_extra_one(user: &UserInfo, key: &str) -> Result { - let Some(values) = user.extra.as_ref().and_then(|extra| extra.get(key)) else { - return Err(Status::permission_denied("SA token is not pod-bound")); - }; - if values.len() != 1 || values[0].is_empty() { - return Err(Status::permission_denied( - "SA token has invalid pod binding", - )); - } - Ok(values[0].clone()) -} - -#[allow(clippy::result_large_err)] -fn pod_sandbox_id(pod: &Pod) -> Option { - pod.metadata - .annotations - .as_ref() - .and_then(|a| a.get(SANDBOX_ID_ANNOTATION)) - .cloned() - .filter(|sandbox_id| !sandbox_id.is_empty()) -} - -#[allow(clippy::result_large_err)] -fn sandbox_owner_reference(pod: &Pod) -> Result { - let owner_refs = pod.metadata.owner_references.as_deref().unwrap_or_default(); - let mut sandbox_refs = owner_refs - .iter() - .filter(|owner| is_supported_sandbox_owner_reference(owner)); - let Some(owner) = sandbox_refs.next() else { - let unsupported_sandbox_api_versions = owner_refs - .iter() - .filter(|owner| owner.kind == SANDBOX_KIND) - .map(|owner| owner.api_version.as_str()) - .collect::>(); - if !unsupported_sandbox_api_versions.is_empty() { - warn!( - api_versions = ?unsupported_sandbox_api_versions, - supported_api_versions = ?[ - SANDBOX_API_VERSION_FULL_V1BETA1, - SANDBOX_API_VERSION_FULL_V1ALPHA1, - ], - "pod Sandbox ownerReference uses unsupported apiVersion" - ); - } - return Err(Status::permission_denied( - "pod is not controlled by an OpenShell Sandbox", - )); - }; - if sandbox_refs.next().is_some() { - return Err(Status::permission_denied( - "pod has multiple OpenShell Sandbox owners", - )); - } - if owner.controller != Some(true) { - return Err(Status::permission_denied( - "pod Sandbox ownerReference is not controlling", - )); - } - if owner.name.is_empty() || owner.uid.is_empty() { - return Err(Status::permission_denied( - "pod Sandbox ownerReference is incomplete", - )); - } - Ok(SandboxOwnerReference { - api_version: owner.api_version.clone(), - name: owner.name.clone(), - uid: owner.uid.clone(), - }) -} - -fn is_supported_sandbox_owner_reference( - owner: &k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference, -) -> bool { - owner.kind == SANDBOX_KIND - && matches!( - owner.api_version.as_str(), - SANDBOX_API_VERSION_FULL_V1BETA1 | SANDBOX_API_VERSION_FULL_V1ALPHA1 - ) -} - -fn should_try_next_sandbox_api_version(err: &KubeError) -> bool { - // Kubernetes returns a structured 404 for some missing API resources and a - // raw "404 page not found" body for others. Both mean the probed - // group/version is unavailable and the next supported Sandbox API version - // should be tried. - matches!(err, KubeError::Api(api) if api.code == 404) -} - -#[allow(clippy::result_large_err)] -fn validate_sandbox_owner_reference( - owner: &SandboxOwnerReference, - sandbox_cr: &DynamicObject, -) -> Result<(), Status> { - let actual_uid = sandbox_cr.metadata.uid.as_deref().unwrap_or_default(); - if actual_uid != owner.uid { - warn!( - sandbox_owner = %owner.name, - owner_uid = %owner.uid, - actual_uid = %actual_uid, - "pod Sandbox ownerReference UID does not match live Sandbox CR" - ); - return Err(Status::permission_denied("sandbox owner UID mismatch")); - } - - Ok(()) -} - -#[allow(clippy::result_large_err)] -fn validate_sandbox_owner_binding( - owner: &SandboxOwnerReference, - sandbox_id: &str, - sandbox_cr: &DynamicObject, -) -> Result<(), Status> { - let actual_sandbox_id = sandbox_cr - .metadata - .labels - .as_ref() - .and_then(|labels| labels.get(SANDBOX_ID_LABEL)) - .map(String::as_str) - .unwrap_or_default(); - if actual_sandbox_id != sandbox_id { - warn!( - sandbox_owner = %owner.name, - owner_uid = %owner.uid, - pod_sandbox_id = %sandbox_id, - cr_sandbox_id = %actual_sandbox_id, - "pod sandbox annotation does not match owning Sandbox CR label" - ); - return Err(Status::permission_denied("sandbox owner ID mismatch")); - } - - Ok(()) -} - #[cfg(test)] -pub mod test_support { +mod tests { use super::*; + use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentity; use std::sync::Mutex; - /// Fake resolver for unit tests. Returns the configured outcome on - /// every call and records the tokens it observed. - pub struct FakeResolver { - pub outcome: Result, Status>, - pub seen_tokens: Mutex>, + struct FakeProvider { + outcome: Result, Status>, + seen_tokens: Mutex>, } - impl FakeResolver { - pub fn returning(outcome: Result, Status>) -> Self { + impl FakeProvider { + fn returning(outcome: Result, Status>) -> Self { Self { outcome, seen_tokens: Mutex::new(Vec::new()), @@ -534,430 +119,89 @@ pub mod test_support { } #[async_trait] - impl K8sIdentityResolver for FakeResolver { - async fn resolve(&self, token: &str) -> Result, Status> { + impl SupervisorBootstrapIdentityProvider for FakeProvider { + async fn authenticate_registration( + &self, + token: &str, + ) -> Result, Status> { self.seen_tokens.lock().unwrap().push(token.to_string()); match &self.outcome { - Ok(opt) => Ok(opt.clone()), - Err(s) => Err(Status::new(s.code(), s.message())), + Ok(identity) => Ok(identity.clone()), + Err(status) => Err(Status::new(status.code(), status.message())), } } } -} - -#[cfg(test)] -mod tests { - use super::test_support::FakeResolver; - use super::*; - use k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference; - use std::collections::BTreeMap; fn bearer_headers(token: &str) -> http::HeaderMap { - let mut h = http::HeaderMap::new(); - h.insert( + let mut headers = http::HeaderMap::new(); + headers.insert( "authorization", http::HeaderValue::from_str(&format!("Bearer {token}")).unwrap(), ); - h - } - - fn kube_api_error(code: u16, message: &str) -> KubeError { - KubeError::Api(kube::core::ErrorResponse { - status: if code == 404 { - "404 Not Found".to_string() - } else { - "Failure".to_string() - }, - message: message.to_string(), - reason: "Failed to parse error data".to_string(), - code, - }) - } - - #[test] - fn sandbox_api_version_probe_retries_on_structured_and_raw_404() { - let structured = kube_api_error(404, "could not find the requested resource"); - assert!(should_try_next_sandbox_api_version(&structured)); - - let raw = kube_api_error(404, "404 page not found\n"); - assert!(should_try_next_sandbox_api_version(&raw)); - } - - #[test] - fn sandbox_api_version_probe_keeps_non_404_errors() { - let err = kube_api_error(403, "sandboxes.agents.x-k8s.io is forbidden"); - assert!(!should_try_next_sandbox_api_version(&err)); - } - - fn token_review_status( - authenticated: bool, - audiences: Vec<&str>, - username: &str, - extra: Vec<(&str, &str)>, - ) -> TokenReviewStatus { - TokenReviewStatus { - authenticated: Some(authenticated), - audiences: Some(audiences.into_iter().map(str::to_string).collect()), - error: None, - user: Some(UserInfo { - username: Some(username.to_string()), - uid: Some("sa-uid".to_string()), - groups: Some(vec![ - "system:serviceaccounts".to_string(), - "system:serviceaccounts:openshell".to_string(), - "system:authenticated".to_string(), - ]), - extra: Some( - extra - .into_iter() - .map(|(k, v)| (k.to_string(), vec![v.to_string()])) - .collect::>(), - ), - }), - } + headers } - fn sandbox_owner(name: &str, uid: &str) -> OwnerReference { - sandbox_owner_with_api_version(SANDBOX_API_VERSION_FULL_V1BETA1, name, uid) - } - - fn sandbox_owner_with_api_version(api_version: &str, name: &str, uid: &str) -> OwnerReference { - OwnerReference { - api_version: api_version.to_string(), - block_owner_deletion: None, - controller: Some(true), - kind: SANDBOX_KIND.to_string(), - name: name.to_string(), - uid: uid.to_string(), - } - } - - fn pod_with_owner_refs(owner_references: Vec) -> Pod { - Pod { - metadata: ObjectMeta { - owner_references: Some(owner_references), - ..Default::default() - }, - ..Default::default() - } - } - - fn pod_with_sandbox_id(sandbox_id: Option<&str>) -> Pod { - Pod { - metadata: ObjectMeta { - annotations: sandbox_id.map(|id| { - BTreeMap::from([(SANDBOX_ID_ANNOTATION.to_string(), id.to_string())]) - }), - ..Default::default() - }, - ..Default::default() + fn identity(binding: SupervisorBootstrapBinding) -> SupervisorBootstrapIdentity { + SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "openshell-sandbox-a".to_string(), + instance_id: "uid-a".to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "cr-uid-a".to_string(), + binding, } } - fn registered_pod(sandbox_id: Option<&str>) -> RegisteredPodIdentity { - RegisteredPodIdentity { - sandbox_id: sandbox_id.map(str::to_string), - pod_name: "openshell-sandbox-a".to_string(), - pod_uid: "uid-a".to_string(), - sandbox_owner_name: "sandbox-owner-a".to_string(), - sandbox_owner_uid: "cr-uid-a".to_string(), - } - } - - fn sandbox_cr(name: &str, uid: &str, sandbox_id: &str) -> DynamicObject { - let sandbox_gvk = - GroupVersionKind::gvk(SANDBOX_API_GROUP, SANDBOX_API_VERSION_V1BETA1, SANDBOX_KIND); - let sandbox_resource = ApiResource::from_gvk(&sandbox_gvk); - let mut cr = DynamicObject::new(name, &sandbox_resource); - cr.metadata.uid = Some(uid.to_string()); - cr.metadata.labels = Some(BTreeMap::from([( - SANDBOX_ID_LABEL.to_string(), - sandbox_id.to_string(), - )])); - cr - } - - #[test] - fn token_review_identity_extracts_pod_binding() { - let status = token_review_status( - true, - vec!["openshell-gateway"], - "system:serviceaccount:openshell:default", - vec![ - (POD_NAME_EXTRA, "openshell-sandbox-a"), - (POD_UID_EXTRA, "uid-a"), - ], - ); - - let identity = token_review_identity(&status, "openshell-gateway", "openshell", "default") - .unwrap() - .expect("authenticated token should resolve"); - - assert_eq!(identity.pod_name, "openshell-sandbox-a"); - assert_eq!(identity.pod_uid, "uid-a"); - } - - #[test] - fn token_review_identity_returns_none_when_not_authenticated() { - let status = TokenReviewStatus { - authenticated: Some(false), - error: Some("invalid audience".to_string()), - ..Default::default() - }; - - assert!( - token_review_identity(&status, "openshell-gateway", "openshell", "default") - .unwrap() - .is_none() - ); - } - - #[test] - fn token_review_identity_requires_expected_audience() { - let status = token_review_status( - true, - vec!["kubernetes.default.svc"], - "system:serviceaccount:openshell:default", - vec![ - (POD_NAME_EXTRA, "openshell-sandbox-a"), - (POD_UID_EXTRA, "uid-a"), - ], - ); - - let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") - .expect_err("wrong audience must fail closed"); - assert_eq!(err.code(), tonic::Code::Unauthenticated); - } - - #[test] - fn token_review_identity_requires_sandbox_namespace() { - let status = token_review_status( - true, - vec!["openshell-gateway"], - "system:serviceaccount:other:default", - vec![ - (POD_NAME_EXTRA, "openshell-sandbox-a"), - (POD_UID_EXTRA, "uid-a"), - ], - ); - - let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") - .expect_err("other namespace must be rejected"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn token_review_identity_requires_configured_service_account() { - let status = token_review_status( - true, - vec!["openshell-gateway"], - "system:serviceaccount:openshell:other", - vec![ - (POD_NAME_EXTRA, "openshell-sandbox-a"), - (POD_UID_EXTRA, "uid-a"), - ], - ); - - let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") - .expect_err("other service account must be rejected"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn token_review_identity_requires_pod_bound_extras() { - let status = token_review_status( - true, - vec!["openshell-gateway"], - "system:serviceaccount:openshell:default", - vec![], - ); - - let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") - .expect_err("non pod-bound tokens must be rejected"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn pod_sandbox_id_is_optional_for_warm_registration() { - assert_eq!( - pod_sandbox_id(&pod_with_sandbox_id(Some("sandbox-id-a"))).as_deref(), - Some("sandbox-id-a") - ); - assert!(pod_sandbox_id(&pod_with_sandbox_id(None)).is_none()); - } - - #[test] - fn sandbox_owner_reference_extracts_controlling_sandbox_owner() { - let pod = pod_with_owner_refs(vec![sandbox_owner("sandbox-a", "cr-uid-a")]); - - let owner = sandbox_owner_reference(&pod).expect("expected Sandbox owner"); - - assert_eq!( - owner, - SandboxOwnerReference { - api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), - name: "sandbox-a".to_string(), - uid: "cr-uid-a".to_string(), - } - ); - } - - #[test] - fn sandbox_owner_reference_accepts_v1alpha1_owner() { - let pod = pod_with_owner_refs(vec![sandbox_owner_with_api_version( - SANDBOX_API_VERSION_FULL_V1ALPHA1, - "sandbox-a", - "cr-uid-a", - )]); - - let owner = sandbox_owner_reference(&pod).expect("expected v1alpha1 Sandbox owner"); - - assert_eq!( - owner, - SandboxOwnerReference { - api_version: SANDBOX_API_VERSION_FULL_V1ALPHA1.to_string(), - name: "sandbox-a".to_string(), - uid: "cr-uid-a".to_string(), - } - ); - } - - #[test] - fn sandbox_owner_reference_rejects_missing_owner() { - let pod = pod_with_owner_refs(vec![]); - - let err = sandbox_owner_reference(&pod).expect_err("missing owner must fail"); - - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn sandbox_owner_reference_rejects_unsupported_sandbox_api_version() { - let pod = pod_with_owner_refs(vec![sandbox_owner_with_api_version( - "agents.x-k8s.io/v1", - "sandbox-a", - "cr-uid-a", - )]); - - let err = - sandbox_owner_reference(&pod).expect_err("unsupported apiVersion must fail closed"); - - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn sandbox_owner_reference_requires_controlling_owner() { - let mut owner = sandbox_owner("sandbox-a", "cr-uid-a"); - owner.controller = Some(false); - let pod = pod_with_owner_refs(vec![owner]); - - let err = sandbox_owner_reference(&pod).expect_err("non-controller owner must fail"); - - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn sandbox_owner_reference_rejects_ambiguous_sandbox_owners() { - let pod = pod_with_owner_refs(vec![ - sandbox_owner("sandbox-a", "cr-uid-a"), - sandbox_owner("sandbox-b", "cr-uid-b"), - ]); - - let err = sandbox_owner_reference(&pod).expect_err("multiple owners must fail"); - - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn validate_sandbox_owner_reference_requires_matching_cr_uid() { - let owner = SandboxOwnerReference { - api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), - name: "sandbox-a".to_string(), - uid: "cr-uid-a".to_string(), - }; - let cr = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-a"); - validate_sandbox_owner_reference(&owner, &cr).expect("matching CR should be accepted"); - - let wrong_uid = sandbox_cr("sandbox-a", "cr-uid-b", "sandbox-id-a"); - let err = validate_sandbox_owner_reference(&owner, &wrong_uid) - .expect_err("wrong CR UID must fail"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn validate_sandbox_owner_binding_requires_matching_label() { - let owner = SandboxOwnerReference { - api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), - name: "sandbox-a".to_string(), - uid: "cr-uid-a".to_string(), - }; - let cr = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-a"); - validate_sandbox_owner_binding(&owner, "sandbox-id-a", &cr) - .expect("matching label should be accepted"); - let wrong_label = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-b"); - let err = validate_sandbox_owner_binding(&owner, "sandbox-id-a", &wrong_label) - .expect_err("wrong sandbox-id label must fail"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - #[tokio::test] async fn authenticates_on_bootstrap_paths_only() { - let resolved = registered_pod(Some("sandbox-a")); - let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); - let auth = K8sServiceAccountAuthenticator::new(fake.clone()); + let provider = Arc::new(FakeProvider::returning(Ok(Some(identity( + SupervisorBootstrapBinding::BoundSandbox { + sandbox_id: "sandbox-a".to_string(), + }, + ))))); + let auth = SupervisorBootstrapAuthenticator::new(provider.clone()); - let on_issue = auth - .authenticate(&bearer_headers("sa-jwt"), ISSUE_SANDBOX_TOKEN_PATH) + let issue = auth + .authenticate(&bearer_headers("bootstrap-token"), ISSUE_SANDBOX_TOKEN_PATH) .await .unwrap() .expect("expected principal"); - match on_issue { + match issue { Principal::Sandbox(p) => { assert_eq!(p.sandbox_id, "sandbox-a"); assert!(matches!( p.source, - SandboxIdentitySource::K8sServiceAccount { .. } + SandboxIdentitySource::SupervisorBootstrap { .. } )); } _ => panic!("expected sandbox principal"), } - let on_register = auth - .authenticate(&bearer_headers("sa-jwt"), REGISTER_SUPERVISOR_POD_PATH) + let register = auth + .authenticate( + &bearer_headers("bootstrap-token"), + REGISTER_SUPERVISOR_POD_PATH, + ) .await .unwrap() .expect("expected principal"); - match on_register { - Principal::K8sPod(p) => { - assert_eq!(p.sandbox_id.as_deref(), Some("sandbox-a")); - assert_eq!(p.pod_name, "openshell-sandbox-a"); - assert_eq!(p.pod_uid, "uid-a"); - } - _ => panic!("expected pod registration principal"), - } + assert!(matches!(register, Principal::SupervisorBootstrap(_))); - let off_issue = auth + let off_path = auth .authenticate( - &bearer_headers("sa-jwt"), + &bearer_headers("bootstrap-token"), "/openshell.v1.OpenShell/GetSandboxConfig", ) .await .unwrap(); - assert!( - off_issue.is_none(), - "K8s SA authenticator must be scoped to bootstrap paths" - ); - assert_eq!( - fake.seen_tokens.lock().unwrap().len(), - 2, - "off-path call must not consult the apiserver" - ); + assert!(off_path.is_none()); + assert_eq!(provider.seen_tokens.lock().unwrap().len(), 2); } #[tokio::test] async fn missing_bearer_yields_none() { - let fake = Arc::new(FakeResolver::returning(Ok(None))); - let auth = K8sServiceAccountAuthenticator::new(fake); + let provider = Arc::new(FakeProvider::returning(Ok(None))); + let auth = SupervisorBootstrapAuthenticator::new(provider); let result = auth .authenticate(&http::HeaderMap::new(), ISSUE_SANDBOX_TOKEN_PATH) .await @@ -966,62 +210,39 @@ mod tests { } #[tokio::test] - async fn resolver_returning_none_falls_through() { - let fake = Arc::new(FakeResolver::returning(Ok(None))); - let auth = K8sServiceAccountAuthenticator::new(fake); - let result = auth - .authenticate( - &bearer_headers("not-a-real-sa-token"), - ISSUE_SANDBOX_TOKEN_PATH, - ) - .await - .unwrap(); - assert!(result.is_none(), "non-authenticating tokens fall through"); - } - - #[tokio::test] - async fn issue_token_rejects_unbound_registered_pod() { - let resolved = registered_pod(None); - let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); - let auth = K8sServiceAccountAuthenticator::new(fake); + async fn issue_token_rejects_unbound_identity() { + let provider = Arc::new(FakeProvider::returning(Ok(Some(identity( + SupervisorBootstrapBinding::WarmPending, + ))))); + let auth = SupervisorBootstrapAuthenticator::new(provider); let err = auth - .authenticate(&bearer_headers("sa-jwt"), ISSUE_SANDBOX_TOKEN_PATH) + .authenticate(&bearer_headers("bootstrap-token"), ISSUE_SANDBOX_TOKEN_PATH) .await - .expect_err("unbound pod must be rejected"); + .expect_err("unbound identity must be rejected"); assert_eq!(err.code(), tonic::Code::PermissionDenied); } #[tokio::test] - async fn register_supervisor_pod_accepts_unbound_registered_pod() { - let resolved = registered_pod(None); - let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); - let auth = K8sServiceAccountAuthenticator::new(fake); + async fn register_accepts_unbound_identity() { + let provider = Arc::new(FakeProvider::returning(Ok(Some(identity( + SupervisorBootstrapBinding::WarmPending, + ))))); + let auth = SupervisorBootstrapAuthenticator::new(provider); let principal = auth - .authenticate(&bearer_headers("sa-jwt"), REGISTER_SUPERVISOR_POD_PATH) + .authenticate( + &bearer_headers("bootstrap-token"), + REGISTER_SUPERVISOR_POD_PATH, + ) .await - .expect("unbound warm pod registration should authenticate") + .unwrap() .expect("expected principal"); match principal { - Principal::K8sPod(p) => { - assert!(p.sandbox_id.is_none()); - assert_eq!(p.pod_name, "openshell-sandbox-a"); - assert_eq!(p.sandbox_owner_uid, "cr-uid-a"); + Principal::SupervisorBootstrap(identity) => { + assert_eq!(identity.binding, SupervisorBootstrapBinding::WarmPending); + assert_eq!(identity.owner_uid, "cr-uid-a"); } - _ => panic!("expected pod registration principal"), + _ => panic!("expected bootstrap principal"), } } - - #[tokio::test] - async fn resolver_error_propagates() { - let fake = Arc::new(FakeResolver::returning(Err(Status::unavailable( - "apiserver down", - )))); - let auth = K8sServiceAccountAuthenticator::new(fake); - let err = auth - .authenticate(&bearer_headers("sa-jwt"), ISSUE_SANDBOX_TOKEN_PATH) - .await - .expect_err("resolver error must propagate"); - assert_eq!(err.code(), tonic::Code::Unavailable); - } } diff --git a/crates/openshell-server/src/auth/principal.rs b/crates/openshell-server/src/auth/principal.rs index d395805250..536bb08703 100644 --- a/crates/openshell-server/src/auth/principal.rs +++ b/crates/openshell-server/src/auth/principal.rs @@ -15,6 +15,7 @@ //! to prevent cross-sandbox access (see issue #1354). use super::identity::Identity; +use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentity; /// Who is calling. /// @@ -28,12 +29,13 @@ pub enum Principal { /// sandbox UUID. The wrapped `sandbox_id` MUST match any sandbox referenced /// in the request body for sandbox-class methods. Sandbox(#[allow(dead_code)] SandboxPrincipal), - /// Kubernetes pod authenticated for supervisor pod registration only. + /// Driver runtime instance authenticated for supervisor bootstrap + /// registration only. /// /// This is intentionally not a sandbox principal: warm pods can be valid - /// Kubernetes pods before they are claimed by a sandbox. The router only + /// driver instances before they are claimed by a sandbox. The router only /// allows this principal to call `RegisterSupervisorPod`. - K8sPod(RegisteredPodIdentity), + SupervisorBootstrap(SupervisorBootstrapIdentity), /// Truly unauthenticated caller (health probes, reflection). Sandbox-class /// and user-class methods reject this variant. #[allow(dead_code)] @@ -47,27 +49,6 @@ pub struct UserPrincipal { pub identity: Identity, } -/// Kubernetes pod identity validated by `TokenReview` and live pod lookup. -/// -/// `sandbox_id` is present only when the pod is already bound to a concrete -/// `OpenShell` sandbox. Unbound warm pods carry the owning Sandbox CR metadata -/// but must not receive a gateway sandbox JWT until a later claim binds the -/// exact pod UID to a sandbox. -#[derive(Debug, Clone)] -#[allow(dead_code)] -pub struct RegisteredPodIdentity { - /// Live pod name from the `TokenReview` binding. - pub pod_name: String, - /// Live pod UID from the `TokenReview` binding and pod lookup. - pub pod_uid: String, - /// `OpenShell` sandbox UUID from the pod annotation, when already bound. - pub sandbox_id: Option, - /// Owning Sandbox CR name from the pod ownerReference. - pub sandbox_owner_name: String, - /// Owning Sandbox CR UID from the pod ownerReference and live CR lookup. - pub sandbox_owner_uid: String, -} - /// Sandbox caller — bound to one specific sandbox UUID. /// /// `sandbox_id` and `source` are consumed by the router and handler guards. @@ -97,7 +78,11 @@ pub enum SandboxIdentitySource { /// Per-sandbox client certificate. Reserved for channel-bound sandbox /// identity. BootstrapCert { fingerprint: String }, - /// K8s `ServiceAccount` token used to bootstrap a gateway-minted JWT. - /// Populated only on Kubernetes bootstrap RPC paths. - K8sServiceAccount { pod_name: String, pod_uid: String }, + /// Driver-provided bootstrap token used to bootstrap a gateway-minted JWT. + /// Populated only on supervisor bootstrap RPC paths. + SupervisorBootstrap { + driver: String, + instance_name: String, + instance_id: String, + }, } diff --git a/crates/openshell-server/src/compute/driver_config.rs b/crates/openshell-server/src/compute/driver_config.rs index 59fed439bc..1d13277074 100644 --- a/crates/openshell-server/src/compute/driver_config.rs +++ b/crates/openshell-server/src/compute/driver_config.rs @@ -54,22 +54,6 @@ pub fn kubernetes_config_from_context( Ok(cfg) } -pub fn kubernetes_config_for_k8s_sa_bootstrap( - file: Option<&config_file::ConfigFile>, -) -> Result { - let Some(file) = file else { - return Err(Error::config( - "K8s ServiceAccount bootstrap requires [openshell.drivers.kubernetes] when sandbox JWT issuing is enabled in-cluster", - )); - }; - if !file.openshell.drivers.contains_key("kubernetes") { - return Err(Error::config( - "K8s ServiceAccount bootstrap requires [openshell.drivers.kubernetes] when sandbox JWT issuing is enabled in-cluster", - )); - } - driver_config_from_file(Some(file), ComputeDriverKind::Kubernetes.as_str()) -} - /// Build the selected Podman config from TOML plus runtime defaults. pub fn podman_config_from_context( context: DriverStartupContext<'_>, @@ -262,35 +246,6 @@ mod tests { } } - #[test] - fn k8s_sa_bootstrap_rejects_missing_kubernetes_driver_config() { - let err = kubernetes_config_for_k8s_sa_bootstrap(None).unwrap_err(); - assert!(err.to_string().contains("[openshell.drivers.kubernetes]")); - - let file: config_file::ConfigFile = - toml::from_str("[openshell.gateway]\n").expect("valid config"); - let err = kubernetes_config_for_k8s_sa_bootstrap(Some(&file)).unwrap_err(); - assert!(err.to_string().contains("[openshell.drivers.kubernetes]")); - } - - #[test] - fn k8s_sa_bootstrap_uses_configured_namespace_and_service_account() { - let file: config_file::ConfigFile = toml::from_str( - r#" -[openshell.gateway] - -[openshell.drivers.kubernetes] -namespace = "sandboxes" -service_account_name = "sandbox-sa" -"#, - ) - .expect("valid config"); - - let cfg = kubernetes_config_for_k8s_sa_bootstrap(Some(&file)).unwrap(); - assert_eq!(cfg.namespace, "sandboxes"); - assert_eq!(cfg.service_account_name, "sandbox-sa"); - } - #[test] fn podman_config_reads_bind_mount_opt_in_from_driver_table() { let file: config_file::ConfigFile = toml::from_str( diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index cf6b17c7e9..dd2f00bf6a 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -37,10 +37,12 @@ use openshell_core::proto::{ PlatformEvent, Sandbox, SandboxCondition, SandboxPhase, SandboxSpec, SandboxStatus, SandboxTemplate, ServiceEndpoint, SshSession, }; +use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentityProvider; use openshell_core::{ObjectLabels, ObjectWorkspace}; use openshell_driver_docker::DockerComputeDriver; use openshell_driver_kubernetes::{ ComputeDriverService as KubernetesDriverService, KubernetesComputeDriver, + KubernetesSupervisorBootstrapIdentityProvider, LiveK8sResolver, }; use openshell_driver_podman::{ComputeDriverService as PodmanDriverService, PodmanComputeDriver}; use prost::Message; @@ -231,6 +233,33 @@ impl Drop for ManagedDriverProcess { } } +async fn kubernetes_supervisor_bootstrap_identity_provider( + config: &KubernetesComputeConfig, +) -> Option> { + std::env::var_os("KUBERNETES_SERVICE_HOST")?; + + match kube::Client::try_default().await { + Ok(client) => { + let resolver = Arc::new(LiveK8sResolver::new( + client, + &config.namespace, + "openshell-gateway".to_string(), + config.service_account_name.clone(), + )); + Some(Arc::new( + KubernetesSupervisorBootstrapIdentityProvider::new(resolver), + )) + } + Err(err) => { + warn!( + error = %err, + "in-cluster K8s client construction failed; supervisor bootstrap is disabled" + ); + None + } + } +} + #[derive(Debug)] pub struct AcquiredRemoteDriverEndpoint { pub(crate) name: String, @@ -368,6 +397,7 @@ pub struct ComputeRuntime { sandbox_watch_bus: SandboxWatchBus, tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, + supervisor_bootstrap_identity: Option>, sync_lock: Arc>, delete_gates: Arc, gateway_bind_addresses: Vec, @@ -393,6 +423,7 @@ impl ComputeRuntime { sandbox_watch_bus: SandboxWatchBus, tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, + supervisor_bootstrap_identity: Option>, gateway_bind_addresses: Vec, ) -> Result { let capabilities = driver @@ -425,6 +456,7 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + supervisor_bootstrap_identity, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses, @@ -486,6 +518,7 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + None, gateway_bind_addresses, ) .await @@ -499,6 +532,8 @@ impl ComputeRuntime { tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, ) -> Result { + let supervisor_bootstrap_identity = + kubernetes_supervisor_bootstrap_identity_provider(&config).await; let driver = KubernetesComputeDriver::new(config) .await .map_err(|err| ComputeError::Message(err.to_string()))?; @@ -514,6 +549,7 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + supervisor_bootstrap_identity, Vec::new(), ) .await @@ -539,6 +575,7 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + None, Vec::new(), ) .await @@ -567,6 +604,7 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + None, Vec::new(), ) .await @@ -582,6 +620,13 @@ impl ComputeRuntime { std::slice::from_ref(&self.driver_info) } + #[must_use] + pub fn supervisor_bootstrap_identity_provider( + &self, + ) -> Option> { + self.supervisor_bootstrap_identity.clone() + } + #[must_use] pub fn driver_kind(&self) -> Option { self.driver_info.name.parse().ok() @@ -2759,6 +2804,7 @@ pub async fn new_test_runtime(store: Arc) -> ComputeRuntime { sandbox_watch_bus: SandboxWatchBus::new(), tracing_log_bus: TracingLogBus::new(), supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), + supervisor_bootstrap_identity: None, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses: Vec::new(), @@ -3228,6 +3274,7 @@ mod tests { sandbox_watch_bus: SandboxWatchBus::new(), tracing_log_bus: TracingLogBus::new(), supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), + supervisor_bootstrap_identity: None, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses: Vec::new(), diff --git a/crates/openshell-server/src/grpc/auth_rpc.rs b/crates/openshell-server/src/grpc/auth_rpc.rs index 7bfdcc820d..84c0a47a66 100644 --- a/crates/openshell-server/src/grpc/auth_rpc.rs +++ b/crates/openshell-server/src/grpc/auth_rpc.rs @@ -13,14 +13,13 @@ //! until their own `exp` and are bounded by the configured short TTL. use crate::ServerState; -use crate::auth::principal::{ - Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal, -}; +use crate::auth::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; use crate::warm_pod_activation::{load_sandbox, mint_pod_activation}; use openshell_core::proto::{ IssueSandboxTokenRequest, IssueSandboxTokenResponse, PodActivationMessage, RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RegisterSupervisorPodRequest, }; +use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentity; use std::sync::Arc; use std::{pin::Pin, result::Result as StdResult}; use tokio_stream::Stream; @@ -51,19 +50,22 @@ pub async fn handle_register_supervisor_pod( state: &Arc, request: Request, ) -> Result, Status> { - let pod = require_registered_pod(request.extensions())?; - if let Some(sandbox_id) = pod.sandbox_id.as_deref() { + let identity = require_bootstrap_identity(request.extensions())?; + if let Some(sandbox_id) = identity.bound_sandbox_id() { let activation = mint_pod_activation(state, sandbox_id, "RegisterSupervisorPod").await?; info!( sandbox_id = %activation.sandbox_id, - pod = %pod.pod_name, - pod_uid = %pod.pod_uid, - "activated bound supervisor pod" + driver = %identity.driver, + instance_name = %identity.instance_name, + instance_id = %identity.instance_id, + "activated bound supervisor instance" ); return Ok(Response::new(Box::pin(tokio_stream::once(Ok(activation))))); } - let stream = state.supervisor_pod_registrations.register_pending(pod)?; + let stream = state + .supervisor_pod_registrations + .register_pending(identity)?; Ok(Response::new(Box::pin(stream))) } @@ -135,12 +137,12 @@ fn require_k8s_bootstrap_sandbox( if !matches!( sandbox.source, - SandboxIdentitySource::K8sServiceAccount { .. } + SandboxIdentitySource::SupervisorBootstrap { .. } ) { debug!( sandbox_id = %sandbox.sandbox_id, rpc = rpc_name, - "bootstrap RPC rejected: non-K8s ServiceAccount principal source" + "bootstrap RPC rejected: non-bootstrap principal source" ); return Err(Status::permission_denied( "this principal cannot mint a sandbox token; use RefreshSandboxToken", @@ -150,9 +152,11 @@ fn require_k8s_bootstrap_sandbox( Ok(sandbox) } -fn require_registered_pod(extensions: &tonic::Extensions) -> Result { - if let Some(pod) = extensions.get::().cloned() { - return Ok(pod); +fn require_bootstrap_identity( + extensions: &tonic::Extensions, +) -> Result { + if let Some(identity) = extensions.get::().cloned() { + return Ok(identity); } let principal = extensions @@ -160,13 +164,13 @@ fn require_registered_pod(extensions: &tonic::Extensions) -> Result) -> RegisteredPodIdentity { - RegisteredPodIdentity { - pod_name: "pod-a".to_string(), - pod_uid: "uid-a".to_string(), - sandbox_id: sandbox_id.map(str::to_string), - sandbox_owner_name: "sandbox-owner-a".to_string(), - sandbox_owner_uid: "owner-uid-a".to_string(), + fn bootstrap_identity(binding: SupervisorBootstrapBinding) -> SupervisorBootstrapIdentity { + SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "pod-a".to_string(), + instance_id: "uid-a".to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "owner-uid-a".to_string(), + binding, } } @@ -298,9 +304,10 @@ mod tests { req.extensions_mut() .insert(Principal::Sandbox(SandboxPrincipal { sandbox_id: "sandbox-a".to_string(), - source: SandboxIdentitySource::K8sServiceAccount { - pod_name: "pod-a".to_string(), - pod_uid: "uid-a".to_string(), + source: SandboxIdentitySource::SupervisorBootstrap { + driver: "kubernetes".to_string(), + instance_name: "pod-a".to_string(), + instance_id: "uid-a".to_string(), }, trust_domain: Some("openshell".to_string()), })); @@ -317,7 +324,11 @@ mod tests { let state = state_with_issuer().await; let mut req = Request::new(RegisterSupervisorPodRequest {}); req.extensions_mut() - .insert(Principal::K8sPod(registered_pod(Some("sandbox-a")))); + .insert(Principal::SupervisorBootstrap(bootstrap_identity( + SupervisorBootstrapBinding::BoundSandbox { + sandbox_id: "sandbox-a".to_string(), + }, + ))); let mut stream = handle_register_supervisor_pod(&state, req) .await @@ -341,7 +352,9 @@ mod tests { let state = state_with_issuer().await; let mut req = Request::new(RegisterSupervisorPodRequest {}); req.extensions_mut() - .insert(Principal::K8sPod(registered_pod(None))); + .insert(Principal::SupervisorBootstrap(bootstrap_identity( + SupervisorBootstrapBinding::WarmPending, + ))); let mut stream = handle_register_supervisor_pod(&state, req) .await @@ -367,9 +380,10 @@ mod tests { req.extensions_mut() .insert(Principal::Sandbox(SandboxPrincipal { sandbox_id: "sandbox-deleted".to_string(), - source: SandboxIdentitySource::K8sServiceAccount { - pod_name: "pod-a".to_string(), - pod_uid: "uid-a".to_string(), + source: SandboxIdentitySource::SupervisorBootstrap { + driver: "kubernetes".to_string(), + instance_name: "pod-a".to_string(), + instance_id: "uid-a".to_string(), }, trust_domain: Some("openshell".to_string()), })); @@ -400,8 +414,8 @@ mod tests { } #[tokio::test] - async fn refresh_rejects_k8s_sa_principal() { - // K8s SA-bootstrap principals must use RegisterSupervisorPod, not + async fn refresh_rejects_bootstrap_principal() { + // Bootstrap principals must use RegisterSupervisorPod, not // RefreshSandboxToken. The refresh path assumes a still-valid // gateway-minted JWT already exists. use crate::auth::principal::SandboxIdentitySource; @@ -410,15 +424,16 @@ mod tests { req.extensions_mut() .insert(Principal::Sandbox(SandboxPrincipal { sandbox_id: "sandbox-a".to_string(), - source: SandboxIdentitySource::K8sServiceAccount { - pod_name: "pod-a".to_string(), - pod_uid: "uid-a".to_string(), + source: SandboxIdentitySource::SupervisorBootstrap { + driver: "kubernetes".to_string(), + instance_name: "pod-a".to_string(), + instance_id: "uid-a".to_string(), }, trust_domain: Some("openshell".to_string()), })); let err = handle_refresh_sandbox_token(&state, req) .await - .expect_err("K8s SA principal must not refresh"); + .expect_err("bootstrap principal must not refresh"); assert_eq!(err.code(), tonic::Code::PermissionDenied); } diff --git a/crates/openshell-server/src/grpc/policy.rs b/crates/openshell-server/src/grpc/policy.rs index 66261e3d5e..ee1cc51eba 100644 --- a/crates/openshell-server/src/grpc/policy.rs +++ b/crates/openshell-server/src/grpc/policy.rs @@ -1252,7 +1252,7 @@ async fn resolve_sandbox_by_name_for_principal( Ok(sandbox) } Principal::User(_) => sandbox.ok_or_else(|| Status::not_found("sandbox not found")), - Principal::K8sPod(_) => Err(Status::permission_denied( + Principal::SupervisorBootstrap(_) => Err(Status::permission_denied( "sandbox-scoped methods require a sandbox principal", )), Principal::Anonymous => Err(Status::unauthenticated( diff --git a/crates/openshell-server/src/inference.rs b/crates/openshell-server/src/inference.rs index 7fe80b36fa..26d40d52c3 100644 --- a/crates/openshell-server/src/inference.rs +++ b/crates/openshell-server/src/inference.rs @@ -932,7 +932,7 @@ fn authorize_inference_bundle( Some(crate::auth::principal::Principal::Sandbox(s)) => Ok(s.sandbox_id.clone()), Some( crate::auth::principal::Principal::User(_) - | crate::auth::principal::Principal::K8sPod(_), + | crate::auth::principal::Principal::SupervisorBootstrap(_), ) => Err(Status::permission_denied( "GetInferenceBundle requires a sandbox principal", )), diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index caecdb80d6..1aac8d8b55 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -153,11 +153,6 @@ pub struct ServerState { /// presenting a freshly minted token are recognized. pub sandbox_jwt_authenticator: Option>, - /// Optional K8s `ServiceAccount` authenticator that backs the - /// Kubernetes supervisor bootstrap paths. Only present when the gateway - /// runs in-cluster. - pub k8s_sa_authenticator: Option>, - /// Gateway-wide gRPC request rate limiter shared by every multiplex path. pub(crate) grpc_rate_limiter: Option, @@ -220,7 +215,6 @@ impl ServerState { oidc_cache, sandbox_jwt_issuer: None, sandbox_jwt_authenticator: None, - k8s_sa_authenticator: None, grpc_rate_limiter, gateway_interceptors: None, provider_profile_sources: @@ -395,42 +389,6 @@ pub(crate) async fn run_server( state.sandbox_jwt_authenticator = Some(Arc::new(authenticator)); } - // K8s ServiceAccount bootstrap authenticator. Only constructed when - // the gateway is running in-cluster (kubelet provides the API host - // env var) and has a sandbox JWT issuer to mint replacements against; - // outside the cluster we can't call the apiserver's TokenReview API, - // and without the issuer there's nothing to exchange the SA token for. - if state.sandbox_jwt_issuer.is_some() && std::env::var_os("KUBERNETES_SERVICE_HOST").is_some() { - // Pod lookups and TokenReview identity checks must match the sandbox - // namespace and service account used by the Kubernetes driver. - let kubernetes_config = - compute::driver_config::kubernetes_config_for_k8s_sa_bootstrap(config_file.as_ref())?; - let sandbox_namespace = kubernetes_config.namespace; - let sandbox_service_account = kubernetes_config.service_account_name; - match kube::Client::try_default().await { - Ok(client) => { - let resolver = Arc::new(auth::k8s_sa::LiveK8sResolver::new( - client, - &sandbox_namespace, - "openshell-gateway".to_string(), - sandbox_service_account.clone(), - )); - let authenticator = auth::k8s_sa::K8sServiceAccountAuthenticator::new(resolver); - state.k8s_sa_authenticator = Some(Arc::new(authenticator)); - info!( - namespace = %sandbox_namespace, - service_account = %sandbox_service_account, - "K8s ServiceAccount bootstrap authenticator enabled" - ); - } - Err(e) => warn!( - error = %e, - "in-cluster K8s client construction failed; \ - K8s ServiceAccount bootstrap is disabled" - ), - } - } - let state = Arc::new(state); let (shutdown_tx, shutdown_rx) = watch::channel(false); diff --git a/crates/openshell-server/src/multiplex.rs b/crates/openshell-server/src/multiplex.rs index e40a3f831a..65d5d349e3 100644 --- a/crates/openshell-server/src/multiplex.rs +++ b/crates/openshell-server/src/multiplex.rs @@ -463,7 +463,7 @@ fn gateway_principal_fields(principal: &Principal) -> BTreeMap { match &sandbox.source { SandboxIdentitySource::BootstrapJwt { .. } => "bootstrap_jwt", SandboxIdentitySource::BootstrapCert { .. } => "bootstrap_cert", - SandboxIdentitySource::K8sServiceAccount { .. } => "k8s_service_account", + SandboxIdentitySource::SupervisorBootstrap { .. } => "supervisor_bootstrap", } .to_string(), ); @@ -471,21 +471,13 @@ fn gateway_principal_fields(principal: &Principal) -> BTreeMap { fields.insert("trust_domain".to_string(), trust_domain.clone()); } } - Principal::K8sPod(pod) => { - fields.insert("kind".to_string(), "k8s_pod".to_string()); - fields.insert("pod_name".to_string(), pod.pod_name.clone()); - fields.insert("pod_uid".to_string(), pod.pod_uid.clone()); - fields.insert( - "sandbox_owner_name".to_string(), - pod.sandbox_owner_name.clone(), - ); - fields.insert( - "sandbox_owner_uid".to_string(), - pod.sandbox_owner_uid.clone(), - ); - if let Some(sandbox_id) = &pod.sandbox_id { - fields.insert("sandbox_id".to_string(), sandbox_id.clone()); - } + Principal::SupervisorBootstrap(identity) => { + fields.insert("kind".to_string(), "supervisor_bootstrap".to_string()); + fields.insert("driver".to_string(), identity.driver.clone()); + fields.insert("instance_id".to_string(), identity.instance_id.clone()); + fields.insert("instance_name".to_string(), identity.instance_name.clone()); + fields.insert("owner_name".to_string(), identity.owner_name.clone()); + fields.insert("owner_uid".to_string(), identity.owner_uid.clone()); } Principal::Anonymous => { fields.insert("kind".to_string(), "anonymous".to_string()); @@ -710,11 +702,12 @@ where /// Assemble the authenticator chain for the gateway. /// /// Chain order (first-match-wins): -/// 1. `K8sServiceAccountAuthenticator` (path-scoped to Kubernetes bootstrap -/// RPCs) — resolves a projected SA token to either a legacy -/// `Principal::Sandbox` for `IssueSandboxToken` or a registration-scoped -/// pod principal for `RegisterSupervisorPod`. No-op on every other path; -/// only present when the gateway runs in-cluster. +/// 1. `SupervisorBootstrapAuthenticator` (path-scoped to supervisor bootstrap +/// RPCs) — delegates driver-native bootstrap token validation to the active +/// compute driver and resolves it to either a legacy `Principal::Sandbox` +/// for `IssueSandboxToken` or a registration-scoped principal for +/// `RegisterSupervisorPod`. No-op on every other path; only present when +/// the active driver exposes a bootstrap identity provider. /// 2. `SandboxJwtAuthenticator` — validates gateway-minted JWTs. Recognized /// via a distinctive `kid` so non-matching Bearer tokens fall through. /// 3. `OidcAuthenticator` — validates user Bearer tokens against the @@ -732,8 +725,10 @@ where /// to pass-through unless mTLS or local unauthenticated users are enabled. fn build_authenticator_chain(state: &ServerState) -> Option { let mut authenticators: Vec> = Vec::new(); - if let Some(k8s) = state.k8s_sa_authenticator.clone() { - authenticators.push(k8s); + if let Some(provider) = state.compute.supervisor_bootstrap_identity_provider() { + authenticators.push(Arc::new( + crate::auth::k8s_sa::SupervisorBootstrapAuthenticator::new(provider), + )); } if let Some(jwt) = state.sandbox_jwt_authenticator.clone() { authenticators.push(jwt); @@ -760,7 +755,8 @@ fn build_authenticator_chain(state: &ServerState) -> Option /// `Principal::User` is gated by the RBAC `AuthzPolicy`. /// `Principal::Sandbox` is gated by a supervisor-method allowlist, then /// handlers enforce same-sandbox scope on request bodies. -/// `Principal::K8sPod` is gated by the pod-registration method only. +/// `Principal::SupervisorBootstrap` is gated by the pod-registration method +/// only. #[derive(Clone)] pub struct AuthGrpcRouter { inner: S, @@ -899,13 +895,13 @@ where ))); } } - Principal::K8sPod(ref pod) => { + Principal::SupervisorBootstrap(ref identity) => { if !crate::auth::method_authz::is_pod_registration_callable(&path) { return Ok(status_response(tonic::Status::permission_denied( - "Kubernetes pod principals may only register supervisor pods", + "supervisor bootstrap principals may only register supervisors", ))); } - req.extensions_mut().insert(pod.clone()); + req.extensions_mut().insert(identity.clone()); } Principal::Anonymous => { return Ok(status_response(tonic::Status::unauthenticated( @@ -1842,10 +1838,12 @@ mod tests { use crate::auth::authenticator::test_support::MockAuthenticator; use crate::auth::identity::{Identity, IdentityProvider}; use crate::auth::principal::{ - Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal, - UserPrincipal, + Principal, SandboxIdentitySource, SandboxPrincipal, UserPrincipal, }; use http_body_util::Full; + use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapBinding, SupervisorBootstrapIdentity, + }; use std::sync::Arc; use std::sync::Mutex; use tower::Service; @@ -1933,13 +1931,14 @@ mod tests { }) } - fn pod_registration_principal() -> Principal { - Principal::K8sPod(RegisteredPodIdentity { - pod_name: "pod-a".to_string(), - pod_uid: "pod-uid-a".to_string(), - sandbox_id: None, - sandbox_owner_name: "sandbox-owner-a".to_string(), - sandbox_owner_uid: "owner-uid-a".to_string(), + fn bootstrap_registration_principal() -> Principal { + Principal::SupervisorBootstrap(SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "pod-a".to_string(), + instance_id: "pod-uid-a".to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "owner-uid-a".to_string(), + binding: SupervisorBootstrapBinding::WarmPending, }) } @@ -2132,9 +2131,9 @@ mod tests { } #[tokio::test] - async fn pod_registration_principal_can_only_register_supervisor_pod() { + async fn bootstrap_registration_principal_can_only_register_supervisor_pod() { let mock = Arc::new(MockAuthenticator::returning(Ok(Some( - pod_registration_principal(), + bootstrap_registration_principal(), )))); let chain = AuthenticatorChain::new(vec![mock]); let (recorder, seen) = PrincipalRecorder::new(); @@ -2150,11 +2149,11 @@ mod tests { assert_eq!(res.status(), 200); assert!(matches!( seen.lock().unwrap().as_ref(), - Some(Principal::K8sPod(_)) + Some(Principal::SupervisorBootstrap(_)) )); let mock = Arc::new(MockAuthenticator::returning(Ok(Some( - pod_registration_principal(), + bootstrap_registration_principal(), )))); let chain = AuthenticatorChain::new(vec![mock]); let (recorder, seen) = PrincipalRecorder::new(); diff --git a/crates/openshell-server/src/supervisor_pod_registration.rs b/crates/openshell-server/src/supervisor_pod_registration.rs index a8eb9eafa0..805cfc27a2 100644 --- a/crates/openshell-server/src/supervisor_pod_registration.rs +++ b/crates/openshell-server/src/supervisor_pod_registration.rs @@ -8,8 +8,8 @@ //! assignment; the future claim controller will activate the stored stream once //! it binds the exact pod UID to a sandbox. -use crate::auth::principal::RegisteredPodIdentity; use openshell_core::proto::PodActivationMessage; +use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentity; use std::collections::HashMap; use std::pin::Pin; use std::result::Result as StdResult; @@ -31,13 +31,13 @@ pub struct SupervisorPodRegistrationRegistry { #[derive(Debug, Default)] struct Inner { - pending_by_pod_uid: HashMap, - activated_pod_uid: HashMap, + pending_by_instance_id: HashMap, + activated_instance_id: HashMap, } #[derive(Debug)] struct PendingRegistration { - identity: RegisteredPodIdentity, + identity: SupervisorBootstrapIdentity, sender: mpsc::Sender>, session_id: u64, registered_at: Instant, @@ -52,29 +52,32 @@ impl SupervisorPodRegistrationRegistry { #[allow(clippy::result_large_err)] pub fn register_pending( self: &Arc, - identity: RegisteredPodIdentity, + identity: SupervisorBootstrapIdentity, ) -> Result { - if identity.pod_uid.is_empty() { - return Err(Status::permission_denied("registered pod UID is required")); + if identity.instance_id.is_empty() { + return Err(Status::permission_denied( + "registered supervisor instance ID is required", + )); } let (sender, receiver) = mpsc::channel(1); let session_id = self.next_session_id.fetch_add(1, Ordering::Relaxed); - let pod_uid = identity.pod_uid.clone(); - let pod_name = identity.pod_name.clone(); - let sandbox_owner = identity.sandbox_owner_name.clone(); + let instance_id = identity.instance_id.clone(); + let instance_name = identity.instance_name.clone(); + let owner_name = identity.owner_name.clone(); + let driver = identity.driver.clone(); let replaced = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); - if inner.activated_pod_uid.contains_key(&pod_uid) { + if inner.activated_instance_id.contains_key(&instance_id) { return Err(Status::already_exists( - "supervisor pod has already been activated", + "supervisor instance has already been activated", )); } inner - .pending_by_pod_uid + .pending_by_instance_id .insert( - pod_uid.clone(), + instance_id.clone(), PendingRegistration { identity, sender, @@ -87,124 +90,132 @@ impl SupervisorPodRegistrationRegistry { if replaced { info!( - pod = %pod_name, - pod_uid = %pod_uid, - sandbox_owner = %sandbox_owner, - "replaced duplicate pending supervisor pod registration" + driver = %driver, + instance = %instance_name, + instance_id = %instance_id, + owner = %owner_name, + "replaced duplicate pending supervisor registration" ); } else { info!( - pod = %pod_name, - pod_uid = %pod_uid, - sandbox_owner = %sandbox_owner, - "registered warm supervisor pod pending activation" + driver = %driver, + instance = %instance_name, + instance_id = %instance_id, + owner = %owner_name, + "registered warm supervisor instance pending activation" ); } Ok(PendingRegistrationStream { registry: self.clone(), - pod_uid, + instance_id, session_id, inner: ReceiverStream::new(receiver), }) } #[allow(clippy::result_large_err)] - pub(crate) fn pending_identity(&self, pod_uid: &str) -> Result { + pub(crate) fn pending_identity( + &self, + instance_id: &str, + ) -> Result { let inner = self.inner.lock().expect("pending pod registry poisoned"); - if inner.activated_pod_uid.contains_key(pod_uid) { + if inner.activated_instance_id.contains_key(instance_id) { return Err(Status::already_exists( - "supervisor pod has already been activated", + "supervisor instance has already been activated", )); } inner - .pending_by_pod_uid - .get(pod_uid) + .pending_by_instance_id + .get(instance_id) .map(|pending| pending.identity.clone()) - .ok_or_else(|| Status::not_found("pending supervisor pod registration not found")) + .ok_or_else(|| Status::not_found("pending supervisor registration not found")) } #[allow(clippy::result_large_err)] pub(crate) fn activate( &self, - pod_uid: &str, + instance_id: &str, activation: PodActivationMessage, ) -> Result<(), Status> { let pending = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); - if inner.activated_pod_uid.contains_key(pod_uid) { + if inner.activated_instance_id.contains_key(instance_id) { return Err(Status::already_exists( - "supervisor pod has already been activated", + "supervisor instance has already been activated", )); } - let Some(pending) = inner.pending_by_pod_uid.remove(pod_uid) else { - debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to activate"); + let Some(pending) = inner.pending_by_instance_id.remove(instance_id) else { + debug!( + instance_id = %instance_id, + "no pending supervisor registration to activate" + ); return Err(Status::not_found( - "pending supervisor pod registration not found", + "pending supervisor registration not found", )); }; inner - .activated_pod_uid - .insert(pod_uid.to_string(), Instant::now()); + .activated_instance_id + .insert(instance_id.to_string(), Instant::now()); pending }; - let pod_name = pending.identity.pod_name.clone(); + let instance_name = pending.identity.instance_name.clone(); if pending.sender.try_send(Ok(activation)).is_err() { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); - inner.activated_pod_uid.remove(pod_uid); + inner.activated_instance_id.remove(instance_id); warn!( - pod = %pod_name, - pod_uid = %pod_uid, - "pending supervisor pod registration stream closed before activation" + instance = %instance_name, + instance_id = %instance_id, + "pending supervisor registration stream closed before activation" ); return Err(Status::unavailable( - "registered pod stream closed before activation", + "registered supervisor stream closed before activation", )); } info!( - pod = %pod_name, - pod_uid = %pod_uid, + instance = %instance_name, + instance_id = %instance_id, pending_ms = pending.registered_at.elapsed().as_millis(), - "activated pending supervisor pod registration" + "activated pending supervisor registration" ); Ok(()) } #[allow(clippy::result_large_err)] - pub(crate) fn fail_pending(&self, pod_uid: &str, status: Status) -> Result<(), Status> { + pub(crate) fn fail_pending(&self, instance_id: &str, status: Status) -> Result<(), Status> { let pending = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); - if inner.activated_pod_uid.contains_key(pod_uid) { + if inner.activated_instance_id.contains_key(instance_id) { return Err(Status::already_exists( - "supervisor pod has already been activated", + "supervisor instance has already been activated", )); } - inner.pending_by_pod_uid.remove(pod_uid) + inner.pending_by_instance_id.remove(instance_id) }; let Some(pending) = pending else { - debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to fail"); + debug!(instance_id = %instance_id, "no pending supervisor registration to fail"); return Err(Status::not_found( - "pending supervisor pod registration not found", + "pending supervisor registration not found", )); }; - let pod_name = pending.identity.pod_name.clone(); + let instance_name = pending.identity.instance_name.clone(); pending.sender.try_send(Err(status)).map_err(|_| { warn!( - pod = %pod_name, - pod_uid = %pod_uid, - "pending supervisor pod registration stream closed before failure notification" + instance = %instance_name, + instance_id = %instance_id, + "pending supervisor registration stream closed before failure notification" ); - Status::unavailable("registered pod stream closed before failure notification") + Status::unavailable("registered supervisor stream closed before failure notification") })?; info!( - pod = %pod_name, - pod_uid = %pod_uid, + instance = %instance_name, + instance_id = %instance_id, pending_ms = pending.registered_at.elapsed().as_millis(), - "failed pending supervisor pod registration" + "failed pending supervisor registration" ); Ok(()) } @@ -215,7 +226,7 @@ impl SupervisorPodRegistrationRegistry { self.inner .lock() .expect("pending pod registry poisoned") - .pending_by_pod_uid + .pending_by_instance_id .len() } @@ -225,19 +236,19 @@ impl SupervisorPodRegistrationRegistry { self.inner .lock() .expect("pending pod registry poisoned") - .activated_pod_uid + .activated_instance_id .len() } - fn remove_if_session(&self, pod_uid: &str, session_id: u64) { + fn remove_if_session(&self, instance_id: &str, session_id: u64) { let removed = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); if inner - .pending_by_pod_uid - .get(pod_uid) + .pending_by_instance_id + .get(instance_id) .is_some_and(|pending| pending.session_id == session_id) { - inner.pending_by_pod_uid.remove(pod_uid) + inner.pending_by_instance_id.remove(instance_id) } else { None } @@ -245,10 +256,10 @@ impl SupervisorPodRegistrationRegistry { if let Some(pending) = removed { debug!( - pod = %pending.identity.pod_name, - pod_uid = %pod_uid, + instance = %pending.identity.instance_name, + instance_id = %instance_id, pending_ms = pending.registered_at.elapsed().as_millis(), - "removed pending supervisor pod registration" + "removed pending supervisor registration" ); } } @@ -256,7 +267,7 @@ impl SupervisorPodRegistrationRegistry { pub struct PendingRegistrationStream { registry: Arc, - pod_uid: String, + instance_id: String, session_id: u64, inner: ReceiverStream>, } @@ -273,7 +284,7 @@ impl Stream for PendingRegistrationStream { impl Drop for PendingRegistrationStream { fn drop(&mut self) { self.registry - .remove_if_session(&self.pod_uid, self.session_id); + .remove_if_session(&self.instance_id, self.session_id); } } @@ -283,13 +294,14 @@ mod tests { use std::collections::HashMap; use tokio_stream::StreamExt; - fn pod_identity(pod_uid: &str) -> RegisteredPodIdentity { - RegisteredPodIdentity { - pod_name: "warm-pod-a".to_string(), - pod_uid: pod_uid.to_string(), - sandbox_id: None, - sandbox_owner_name: "sandbox-owner-a".to_string(), - sandbox_owner_uid: "owner-uid-a".to_string(), + fn bootstrap_identity(instance_id: &str) -> SupervisorBootstrapIdentity { + SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "warm-pod-a".to_string(), + instance_id: instance_id.to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "owner-uid-a".to_string(), + binding: openshell_core::supervisor_bootstrap::SupervisorBootstrapBinding::WarmPending, } } @@ -297,7 +309,7 @@ mod tests { fn dropping_stream_removes_pending_registration() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); let stream = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register pending"); assert_eq!(registry.pending_count(), 1); drop(stream); @@ -308,10 +320,10 @@ mod tests { fn duplicate_registration_replaces_prior_stream() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); let old = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register old"); let new = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register new"); assert_eq!(registry.pending_count(), 1); @@ -329,7 +341,7 @@ mod tests { async fn activation_sends_message_and_removes_pending_registration() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); let mut stream = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register pending"); let activation = PodActivationMessage { @@ -376,7 +388,7 @@ mod tests { async fn activation_tombstone_rejects_duplicate_activation_and_registration() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); let mut stream = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register pending"); let activation = PodActivationMessage { @@ -396,7 +408,7 @@ mod tests { .expect_err("duplicate activation must fail"); assert_eq!(err.code(), tonic::Code::AlreadyExists); - let Err(err) = registry.register_pending(pod_identity("pod-uid-a")) else { + let Err(err) = registry.register_pending(bootstrap_identity("pod-uid-a")) else { panic!("activated pod UID must not register again"); }; assert_eq!(err.code(), tonic::Code::AlreadyExists); @@ -406,7 +418,7 @@ mod tests { fn closed_pending_stream_cannot_be_activated() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); let stream = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register pending"); drop(stream); diff --git a/crates/openshell-server/src/warm_pod_activation.rs b/crates/openshell-server/src/warm_pod_activation.rs index 881ff34651..b93c9307ae 100644 --- a/crates/openshell-server/src/warm_pod_activation.rs +++ b/crates/openshell-server/src/warm_pod_activation.rs @@ -4,12 +4,15 @@ //! Gateway-side warm supervisor pod activation. //! //! The Kubernetes claim controller will call this module after it observes and -//! revalidates a claim that binds a warm pod UID to an `OpenShell` sandbox. +//! revalidates a claim that binds a warm supervisor instance to an `OpenShell` +//! sandbox. use crate::ServerState; -use crate::auth::principal::RegisteredPodIdentity; use async_trait::async_trait; use openshell_core::proto::{PodActivationMessage, Sandbox}; +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapActivationRequest, SupervisorBootstrapActivator, SupervisorBootstrapIdentity, +}; use std::sync::Arc; use tonic::Status; use tracing::{info, warn}; @@ -17,16 +20,19 @@ use tracing::{info, warn}; #[derive(Debug, Clone, PartialEq, Eq)] #[allow(dead_code)] pub struct WarmPodActivationTarget { - pub pod_uid: String, + pub driver: String, + pub instance_id: String, pub sandbox_id: String, + pub owner_uid: String, } #[derive(Debug, Clone, PartialEq, Eq)] #[allow(dead_code)] pub struct ValidatedWarmPodActivation { - pub pod_uid: String, - pub pod_name: String, - pub sandbox_owner_uid: String, + pub driver: String, + pub instance_id: String, + pub instance_name: String, + pub owner_uid: String, pub sandbox_id: String, } @@ -36,10 +42,59 @@ pub trait WarmPodActivationValidator: Send + Sync { async fn validate( &self, target: &WarmPodActivationTarget, - pending: &RegisteredPodIdentity, + pending: &SupervisorBootstrapIdentity, ) -> Result; } +#[derive(Debug)] +#[allow(dead_code)] +pub struct GatewaySupervisorBootstrapActivator { + state: Arc, +} + +impl GatewaySupervisorBootstrapActivator { + #[must_use] + #[allow(dead_code)] + pub fn new(state: Arc) -> Self { + Self { state } + } +} + +#[async_trait] +impl SupervisorBootstrapActivator for GatewaySupervisorBootstrapActivator { + async fn activate_registered_supervisor( + &self, + request: SupervisorBootstrapActivationRequest, + ) -> Result<(), Status> { + let target = WarmPodActivationTarget { + driver: request.driver, + instance_id: request.instance_id, + sandbox_id: request.sandbox_id, + owner_uid: request.owner_uid, + }; + activate_warm_pod(&self.state, &TrustDriverActivationRequest, target).await + } +} + +struct TrustDriverActivationRequest; + +#[async_trait] +impl WarmPodActivationValidator for TrustDriverActivationRequest { + async fn validate( + &self, + target: &WarmPodActivationTarget, + pending: &SupervisorBootstrapIdentity, + ) -> Result { + Ok(ValidatedWarmPodActivation { + driver: target.driver.clone(), + instance_id: target.instance_id.clone(), + instance_name: pending.instance_name.clone(), + owner_uid: target.owner_uid.clone(), + sandbox_id: target.sandbox_id.clone(), + }) + } +} + #[allow(clippy::result_large_err)] #[allow(dead_code)] pub async fn activate_warm_pod( @@ -52,7 +107,7 @@ where { let pending = state .supervisor_pod_registrations - .pending_identity(&target.pod_uid)?; + .pending_identity(&target.instance_id)?; let activation_result = async { let validated = validator.validate(&target, &pending).await?; ensure_validated_activation_matches_pending(&target, &pending, &validated)?; @@ -65,14 +120,14 @@ where let stream_status = clone_status(&status); state .supervisor_pod_registrations - .fail_pending(&target.pod_uid, stream_status)?; + .fail_pending(&target.instance_id, stream_status)?; return Err(status); } }; state .supervisor_pod_registrations - .activate(&target.pod_uid, activation)?; + .activate(&target.instance_id, activation)?; Ok(()) } @@ -83,22 +138,27 @@ fn clone_status(status: &Status) -> Status { #[allow(clippy::result_large_err)] fn ensure_validated_activation_matches_pending( target: &WarmPodActivationTarget, - pending: &RegisteredPodIdentity, + pending: &SupervisorBootstrapIdentity, validated: &ValidatedWarmPodActivation, ) -> Result<(), Status> { - if validated.pod_uid != target.pod_uid || validated.pod_uid != pending.pod_uid { + if validated.driver != target.driver || validated.driver != pending.driver { + return Err(Status::permission_denied( + "validated driver does not match pending registration", + )); + } + if validated.instance_id != target.instance_id || validated.instance_id != pending.instance_id { return Err(Status::permission_denied( - "validated pod UID does not match pending registration", + "validated instance ID does not match pending registration", )); } - if validated.pod_name != pending.pod_name { + if validated.instance_name != pending.instance_name { return Err(Status::permission_denied( - "validated pod name does not match pending registration", + "validated instance name does not match pending registration", )); } - if validated.sandbox_owner_uid != pending.sandbox_owner_uid { + if validated.owner_uid != target.owner_uid || validated.owner_uid != pending.owner_uid { return Err(Status::permission_denied( - "validated Sandbox owner UID does not match pending registration", + "validated owner UID does not match pending registration", )); } if validated.sandbox_id != target.sandbox_id { @@ -185,7 +245,7 @@ mod tests { async fn validate( &self, _target: &WarmPodActivationTarget, - _pending: &RegisteredPodIdentity, + _pending: &SupervisorBootstrapIdentity, ) -> Result { Ok(self.validated.clone()) } @@ -229,7 +289,10 @@ mod tests { name: sandbox_id.to_string(), created_at_ms: 1_000_000, labels: HashMap::default(), + annotations: HashMap::default(), resource_version: 0, + deletion_timestamp_ms: 0, + workspace: "default".to_string(), }), spec: Some(SandboxSpec { policy: None, @@ -241,28 +304,32 @@ mod tests { state.store.put_message(&sandbox).await.unwrap(); } - fn pending_pod() -> RegisteredPodIdentity { - RegisteredPodIdentity { - pod_name: "warm-pod-a".to_string(), - pod_uid: "pod-uid-a".to_string(), - sandbox_id: None, - sandbox_owner_name: "sandbox-owner-a".to_string(), - sandbox_owner_uid: "owner-uid-a".to_string(), + fn pending_identity() -> SupervisorBootstrapIdentity { + SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "warm-pod-a".to_string(), + instance_id: "pod-uid-a".to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "owner-uid-a".to_string(), + binding: openshell_core::supervisor_bootstrap::SupervisorBootstrapBinding::WarmPending, } } fn activation_target(sandbox_id: &str) -> WarmPodActivationTarget { WarmPodActivationTarget { - pod_uid: "pod-uid-a".to_string(), + driver: "kubernetes".to_string(), + instance_id: "pod-uid-a".to_string(), sandbox_id: sandbox_id.to_string(), + owner_uid: "owner-uid-a".to_string(), } } fn validated_activation(sandbox_id: &str) -> ValidatedWarmPodActivation { ValidatedWarmPodActivation { - pod_uid: "pod-uid-a".to_string(), - pod_name: "warm-pod-a".to_string(), - sandbox_owner_uid: "owner-uid-a".to_string(), + driver: "kubernetes".to_string(), + instance_id: "pod-uid-a".to_string(), + instance_name: "warm-pod-a".to_string(), + owner_uid: "owner-uid-a".to_string(), sandbox_id: sandbox_id.to_string(), } } @@ -272,7 +339,7 @@ mod tests { let state = state_with_issuer().await; let mut stream = state .supervisor_pod_registrations - .register_pending(pending_pod()) + .register_pending(pending_identity()) .expect("register pending"); let validator = StaticValidator { validated: validated_activation("sandbox-a"), @@ -295,7 +362,7 @@ mod tests { } #[tokio::test] - async fn activation_for_unknown_pod_uid_fails_before_validation() { + async fn activation_for_unknown_instance_id_fails_before_validation() { let state = state_with_issuer().await; let validator = StaticValidator { validated: validated_activation("sandbox-a"), @@ -303,7 +370,7 @@ mod tests { let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) .await - .expect_err("unknown pod UID must fail"); + .expect_err("unknown instance ID must fail"); assert_eq!(err.code(), tonic::Code::NotFound); } @@ -313,7 +380,7 @@ mod tests { let state = state_with_issuer().await; let mut stream = state .supervisor_pod_registrations - .register_pending(pending_pod()) + .register_pending(pending_identity()) .expect("register pending"); let validator = StaticValidator { validated: validated_activation("sandbox-deleted"), @@ -334,11 +401,11 @@ mod tests { } #[tokio::test] - async fn duplicate_activation_for_same_pod_uid_is_rejected() { + async fn duplicate_activation_for_same_instance_id_is_rejected() { let state = state_with_issuer().await; let mut stream = state .supervisor_pod_registrations - .register_pending(pending_pod()) + .register_pending(pending_identity()) .expect("register pending"); let validator = StaticValidator { validated: validated_activation("sandbox-a"), @@ -360,10 +427,10 @@ mod tests { let state = state_with_issuer().await; let mut stream = state .supervisor_pod_registrations - .register_pending(pending_pod()) + .register_pending(pending_identity()) .expect("register pending"); let mut validated = validated_activation("sandbox-a"); - validated.sandbox_owner_uid = "other-owner".to_string(); + validated.owner_uid = "other-owner".to_string(); let validator = StaticValidator { validated }; let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) From c9eb7fabc2a8d3a23f98be17dd15f03c514ef2ee Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Wed, 22 Jul 2026 09:31:12 +0100 Subject: [PATCH 06/10] feat(kubernetes): activate warm pods from SandboxClaims Signed-off-by: Gordon Sim --- crates/openshell-driver-kubernetes/src/lib.rs | 2 + .../src/sandboxclaim.rs | 747 ++++++++++++++++++ crates/openshell-server/src/compute/mod.rs | 47 +- crates/openshell-server/src/lib.rs | 4 + 4 files changed, 799 insertions(+), 1 deletion(-) create mode 100644 crates/openshell-driver-kubernetes/src/sandboxclaim.rs diff --git a/crates/openshell-driver-kubernetes/src/lib.rs b/crates/openshell-driver-kubernetes/src/lib.rs index f967170e43..8c077fc7e3 100644 --- a/crates/openshell-driver-kubernetes/src/lib.rs +++ b/crates/openshell-driver-kubernetes/src/lib.rs @@ -5,6 +5,7 @@ pub mod bootstrap; pub mod config; pub mod driver; pub mod grpc; +pub mod sandboxclaim; pub use bootstrap::{ K8sIdentityResolver, KubernetesSupervisorBootstrapIdentityProvider, LiveK8sResolver, @@ -16,3 +17,4 @@ pub use config::{ }; pub use driver::{KubernetesComputeDriver, KubernetesDriverError}; pub use grpc::ComputeDriverService; +pub use sandboxclaim::SandboxClaimActivationController; diff --git a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs new file mode 100644 index 0000000000..833dea9bab --- /dev/null +++ b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs @@ -0,0 +1,747 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Kubernetes `SandboxClaim` activation support. + +use futures::{StreamExt, TryStreamExt}; +use k8s_openapi::api::core::v1::Pod; +use k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference; +use kube::api::{Api, ApiResource, ListParams}; +use kube::core::DynamicObject; +use kube::core::gvk::GroupVersionKind; +use kube::runtime::watcher::{self, Event}; +use kube::{Client, Error as KubeError}; +use openshell_core::driver_utils::LABEL_SANDBOX_ID; +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapActivationRequest, SupervisorBootstrapActivator, +}; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::watch; +use tonic::Code; +use tracing::{debug, info, warn}; + +const KUBE_API_TIMEOUT: Duration = Duration::from_secs(30); +const DRIVER_NAME: &str = "kubernetes"; + +const SANDBOX_GROUP: &str = "agents.x-k8s.io"; +const SANDBOX_KIND: &str = "Sandbox"; +const SANDBOX_VERSION_V1BETA1: &str = "v1beta1"; +const SANDBOX_VERSION_V1ALPHA1: &str = "v1alpha1"; +const SANDBOX_VERSIONS: &[&str] = &[SANDBOX_VERSION_V1BETA1, SANDBOX_VERSION_V1ALPHA1]; +const SANDBOX_API_VERSION_FULL_V1BETA1: &str = "agents.x-k8s.io/v1beta1"; +const SANDBOX_API_VERSION_FULL_V1ALPHA1: &str = "agents.x-k8s.io/v1alpha1"; + +const SANDBOX_CLAIM_GROUP: &str = "extensions.agents.x-k8s.io"; +const SANDBOX_CLAIM_KIND: &str = "SandboxClaim"; +const SANDBOX_CLAIM_VERSION_V1BETA1: &str = "v1beta1"; +const SANDBOX_CLAIM_VERSION_V1ALPHA1: &str = "v1alpha1"; +const SANDBOX_CLAIM_VERSIONS: &[&str] = &[ + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_VERSION_V1ALPHA1, +]; + +/// Watches Kubernetes `SandboxClaim` resources and activates pending +/// supervisor bootstrap streams after revalidating live driver state. +#[derive(Clone)] +pub struct SandboxClaimActivationController { + client: Client, + watch_client: Client, + namespace: String, +} + +impl SandboxClaimActivationController { + #[must_use] + pub fn new(client: Client, watch_client: Client, namespace: impl Into) -> Self { + Self { + client, + watch_client, + namespace: namespace.into(), + } + } + + pub fn spawn( + &self, + activator: Arc, + shutdown_rx: watch::Receiver, + ) { + let controller = self.clone(); + tokio::spawn(async move { + controller.run(activator, shutdown_rx).await; + }); + } + + async fn run( + self, + activator: Arc, + mut shutdown_rx: watch::Receiver, + ) { + let claim_api = match self + .supported_sandbox_claim_api(self.watch_client.clone()) + .await + { + Ok(api) => api, + Err(err) => { + debug!( + namespace = %self.namespace, + error = %err, + "SandboxClaim API is not available; warm-pool claim activation disabled" + ); + return; + } + }; + + let mut stream = watcher::watcher(claim_api.api, watcher::Config::default()).boxed(); + info!( + namespace = %self.namespace, + sandbox_claim_api_version = %claim_api.version, + "Watching Kubernetes SandboxClaims for warm-pool activation" + ); + + loop { + tokio::select! { + changed = shutdown_rx.changed() => { + if changed.is_err() || *shutdown_rx.borrow() { + break; + } + } + event = stream.try_next() => match event { + Ok(Some(Event::Applied(claim))) => { + self.handle_claim(claim, activator.as_ref()).await; + } + Ok(Some(Event::Restarted(claims))) => { + for claim in claims { + self.handle_claim(claim, activator.as_ref()).await; + } + } + Ok(Some(Event::Deleted(_))) => {} + Ok(None) => break, + Err(err) => { + warn!( + namespace = %self.namespace, + error = %err, + "SandboxClaim watch failed; warm-pool claim activation stopped" + ); + break; + } + } + } + } + } + + async fn handle_claim( + &self, + claim: DynamicObject, + activator: &(dyn SupervisorBootstrapActivator + '_), + ) { + let claim = match parse_sandbox_claim(&claim) { + Ok(claim) => claim, + Err(err) => { + warn!( + namespace = %self.namespace, + error = %err, + "Ignoring invalid SandboxClaim" + ); + return; + } + }; + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox_claim_uid = %claim.uid, + warm_pool = ?claim.warm_pool_name, + sandbox_template = ?claim.sandbox_template_name, + sandbox = ?claim.sandbox_name, + pod_ips = ?claim.pod_ips, + "Processing SandboxClaim for warm-pool activation" + ); + let Some(sandbox_name) = claim.sandbox_name.as_deref() else { + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + "SandboxClaim has not selected a Sandbox yet" + ); + return; + }; + + let sandbox_cr = match self.get_sandbox_cr(sandbox_name).await { + Ok(Some(sandbox_cr)) => sandbox_cr, + Ok(None) => { + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + "SandboxClaim selected Sandbox is not available yet" + ); + return; + } + Err(err) => { + warn!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + error = %err, + "Failed to read Sandbox selected by SandboxClaim" + ); + return; + } + }; + + let pod = match self.resolve_controlled_pod(&sandbox_cr).await { + Ok(Some(pod)) => pod, + Ok(None) => return, + Err(err) => { + warn!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + error = %err, + "Failed to resolve Sandbox pod for SandboxClaim activation" + ); + return; + } + }; + + let request = match activation_request_from_claim_state(&claim, &sandbox_cr, &pod) { + Ok(Some(request)) => request, + Ok(None) => return, + Err(err) => { + warn!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + error = %err, + "SandboxClaim activation validation failed" + ); + return; + } + }; + + match activator + .activate_registered_supervisor(request.clone()) + .await + { + Ok(()) => { + info!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + "Activated warm-pool supervisor from SandboxClaim" + ); + } + Err(status) if status.code() == Code::AlreadyExists => { + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + "Warm-pool supervisor was already activated" + ); + } + Err(status) if status.code() == Code::NotFound => { + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + "Warm-pool supervisor registration is not pending yet" + ); + } + Err(status) => { + warn!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + code = ?status.code(), + message = %status.message(), + "Failed to activate warm-pool supervisor from SandboxClaim" + ); + } + } + } + + fn sandbox_claim_api(&self, client: Client, version: &'static str) -> SandboxClaimApi { + let gvk = GroupVersionKind::gvk(SANDBOX_CLAIM_GROUP, version, SANDBOX_CLAIM_KIND); + let resource = ApiResource::from_gvk(&gvk); + let api = Api::namespaced_with(client, &self.namespace, &resource); + SandboxClaimApi { api, version } + } + + async fn supported_sandbox_claim_api(&self, client: Client) -> Result { + for version in SANDBOX_CLAIM_VERSIONS { + let claim_api = self.sandbox_claim_api(client.clone(), version); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + claim_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => return Ok(claim_api), + Ok(Err(err)) if should_try_next_api_version(&err) => {} + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Err(format!( + "no supported SandboxClaim API version is available; tried {}", + SANDBOX_CLAIM_VERSIONS.join(", ") + )) + } + + fn sandbox_api(&self, client: Client, version: &'static str) -> Api { + let gvk = GroupVersionKind::gvk(SANDBOX_GROUP, version, SANDBOX_KIND); + let resource = ApiResource::from_gvk(&gvk); + Api::namespaced_with(client, &self.namespace, &resource) + } + + async fn get_sandbox_cr(&self, name: &str) -> Result, String> { + for version in SANDBOX_VERSIONS { + let sandbox_api = self.sandbox_api(self.client.clone(), version); + match tokio::time::timeout(KUBE_API_TIMEOUT, sandbox_api.get(name)).await { + Ok(Ok(sandbox_cr)) => return Ok(Some(sandbox_cr)), + Ok(Err(KubeError::Api(err))) if err.code == 404 => {} + Ok(Err(err)) if should_try_next_api_version(&err) => {} + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(None) + } + + async fn resolve_controlled_pod( + &self, + sandbox_cr: &DynamicObject, + ) -> Result, String> { + let Some(owner_uid) = sandbox_cr.metadata.uid.as_deref() else { + return Err("Sandbox CR is missing uid".to_string()); + }; + + let pods_api: Api = Api::namespaced(self.client.clone(), &self.namespace); + let pods = match sandbox_selector(sandbox_cr) { + Some(selector) => { + let params = ListParams::default().labels(&selector); + list_pods(&pods_api, ¶ms).await? + } + None => list_pods(&pods_api, &ListParams::default()).await?, + }; + + let controlled = pods + .into_iter() + .filter(|pod| pod_has_sandbox_owner(pod, owner_uid)) + .collect::>(); + + match controlled.as_slice() { + [pod] => Ok(Some(pod.clone())), + [] => { + debug!( + namespace = %self.namespace, + sandbox = %sandbox_cr.metadata.name.as_deref().unwrap_or_default(), + sandbox_uid = %owner_uid, + "No controlled pod found for SandboxClaim activation" + ); + Ok(None) + } + _ => Err(format!( + "expected one controlled Sandbox pod, found {}", + controlled.len() + )), + } + } +} + +struct SandboxClaimApi { + api: Api, + version: &'static str, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +struct ParsedSandboxClaim { + name: String, + uid: String, + warm_pool_name: Option, + sandbox_template_name: Option, + sandbox_name: Option, + pod_ips: Vec, +} + +fn parse_sandbox_claim(obj: &DynamicObject) -> Result { + Ok(ParsedSandboxClaim { + name: required_metadata_field(&obj.metadata.name, "name")?, + uid: required_metadata_field(&obj.metadata.uid, "uid")?, + warm_pool_name: string_at(&obj.data, &["spec", "warmPoolRef", "name"]) + .or_else(|| string_at(&obj.data, &["spec", "warmpool"])), + sandbox_template_name: string_at(&obj.data, &["spec", "sandboxTemplateRef", "name"]), + sandbox_name: string_at(&obj.data, &["status", "sandbox", "name"]), + pod_ips: strings_at(&obj.data, &["status", "sandbox", "podIPs"]), + }) +} + +fn required_metadata_field(value: &Option, field: &str) -> Result { + value + .as_ref() + .filter(|value| !value.is_empty()) + .cloned() + .ok_or_else(|| format!("SandboxClaim is missing metadata.{field}")) +} + +fn activation_request_from_claim_state( + claim: &ParsedSandboxClaim, + sandbox_cr: &DynamicObject, + pod: &Pod, +) -> Result, String> { + let Some(claim_sandbox_name) = claim.sandbox_name.as_deref() else { + return Ok(None); + }; + let sandbox_name = sandbox_cr.metadata.name.as_deref().unwrap_or_default(); + if sandbox_name != claim_sandbox_name { + return Err(format!( + "SandboxClaim selected Sandbox {claim_sandbox_name}, but live Sandbox is {sandbox_name}" + )); + } + let owner_uid = sandbox_cr + .metadata + .uid + .as_ref() + .filter(|uid| !uid.is_empty()) + .cloned() + .ok_or_else(|| "Sandbox CR is missing uid".to_string())?; + let sandbox_id = sandbox_cr + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(LABEL_SANDBOX_ID)) + .filter(|id| !id.is_empty()) + .cloned() + .ok_or_else(|| "Sandbox CR is missing OpenShell sandbox id label".to_string())?; + let instance_id = pod + .metadata + .uid + .as_ref() + .filter(|uid| !uid.is_empty()) + .cloned() + .ok_or_else(|| "Sandbox pod is missing uid".to_string())?; + if !pod_has_sandbox_owner(pod, &owner_uid) { + return Err("Sandbox pod is not controlled by the selected Sandbox".to_string()); + } + + Ok(Some(SupervisorBootstrapActivationRequest { + driver: DRIVER_NAME.to_string(), + instance_id, + sandbox_id, + owner_uid, + reason: format!("SandboxClaim/{}", claim.name), + })) +} + +async fn list_pods(api: &Api, params: &ListParams) -> Result, String> { + match tokio::time::timeout(KUBE_API_TIMEOUT, api.list(params)).await { + Ok(Ok(list)) => Ok(list.items), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } +} + +fn sandbox_selector(obj: &DynamicObject) -> Option { + string_at(&obj.data, &["status", "selector"]) +} + +fn pod_has_sandbox_owner(pod: &Pod, owner_uid: &str) -> bool { + pod.metadata + .owner_references + .as_deref() + .unwrap_or_default() + .iter() + .any(|owner| is_supported_sandbox_owner_reference(owner) && owner.uid == owner_uid) +} + +fn is_supported_sandbox_owner_reference(owner: &OwnerReference) -> bool { + owner.kind == SANDBOX_KIND + && owner.controller == Some(true) + && matches!( + owner.api_version.as_str(), + SANDBOX_API_VERSION_FULL_V1BETA1 | SANDBOX_API_VERSION_FULL_V1ALPHA1 + ) +} + +fn should_try_next_api_version(err: &KubeError) -> bool { + matches!(err, KubeError::Api(api) if api.code == 404) +} + +fn string_at(value: &serde_json::Value, path: &[&str]) -> Option { + path.iter() + .try_fold(value, |current, segment| current.get(segment)) + .and_then(|value| value.as_str()) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn strings_at(value: &serde_json::Value, path: &[&str]) -> Vec { + path.iter() + .try_fold(value, |current, segment| current.get(segment)) + .and_then(|value| value.as_array()) + .into_iter() + .flatten() + .filter_map(|value| value.as_str()) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use kube::core::ObjectMeta; + use serde_json::json; + use std::collections::BTreeMap; + + fn dynamic_object( + group: &'static str, + version: &'static str, + kind: &'static str, + name: &str, + uid: &str, + data: serde_json::Value, + ) -> DynamicObject { + let gvk = GroupVersionKind::gvk(group, version, kind); + let resource = ApiResource::from_gvk(&gvk); + let mut obj = DynamicObject::new(name, &resource); + obj.metadata.uid = Some(uid.to_string()); + obj.data = data; + obj + } + + fn sandbox_cr(name: &str, uid: &str, sandbox_id: &str, version: &'static str) -> DynamicObject { + let mut obj = dynamic_object( + SANDBOX_GROUP, + version, + SANDBOX_KIND, + name, + uid, + json!({"status": {"selector": "agents.x-k8s.io/sandbox=sandbox-a"}}), + ); + obj.metadata + .labels + .get_or_insert_with(BTreeMap::new) + .insert(LABEL_SANDBOX_ID.to_string(), sandbox_id.to_string()); + obj + } + + fn sandbox_pod(name: &str, uid: &str, owner_uid: &str, api_version: &str) -> Pod { + Pod { + metadata: ObjectMeta { + name: Some(name.to_string()), + uid: Some(uid.to_string()), + owner_references: Some(vec![OwnerReference { + api_version: api_version.to_string(), + kind: SANDBOX_KIND.to_string(), + name: "sandbox-a".to_string(), + uid: owner_uid.to_string(), + controller: Some(true), + block_owner_deletion: None, + }]), + ..ObjectMeta::default() + }, + ..Pod::default() + } + } + + #[test] + fn parses_v1beta1_sandbox_claim_status() { + let obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({ + "spec": {"warmPoolRef": {"name": "pool-a"}}, + "status": { + "sandbox": { + "name": "sandbox-a", + "podIPs": ["10.0.0.10", "fd00::10"] + } + } + }), + ); + + let claim = parse_sandbox_claim(&obj).unwrap(); + + assert_eq!(claim.name, "claim-a"); + assert_eq!(claim.uid, "claim-uid"); + assert_eq!(claim.warm_pool_name.as_deref(), Some("pool-a")); + assert_eq!(claim.sandbox_name.as_deref(), Some("sandbox-a")); + assert_eq!(claim.pod_ips, vec!["10.0.0.10", "fd00::10"]); + } + + #[test] + fn parses_v1alpha1_sandbox_claim_status() { + let obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1ALPHA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({ + "spec": { + "sandboxTemplateRef": {"name": "template-a"}, + "warmpool": "pool-a" + }, + "status": {"sandbox": {"name": "sandbox-a"}} + }), + ); + + let claim = parse_sandbox_claim(&obj).unwrap(); + + assert_eq!(claim.warm_pool_name.as_deref(), Some("pool-a")); + assert_eq!(claim.sandbox_template_name.as_deref(), Some("template-a")); + assert_eq!(claim.sandbox_name.as_deref(), Some("sandbox-a")); + } + + #[test] + fn claim_without_selected_sandbox_does_not_activate() { + let obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({"spec": {"warmPoolRef": {"name": "pool-a"}}}), + ); + let claim = parse_sandbox_claim(&obj).unwrap(); + let sandbox = sandbox_cr( + "sandbox-a", + "sandbox-uid", + "openshell-sandbox", + SANDBOX_VERSION_V1BETA1, + ); + let pod = sandbox_pod( + "pod-a", + "pod-uid", + "sandbox-uid", + SANDBOX_API_VERSION_FULL_V1BETA1, + ); + + assert_eq!( + activation_request_from_claim_state(&claim, &sandbox, &pod).unwrap(), + None + ); + } + + #[test] + fn activation_request_uses_live_sandbox_and_pod_identity() { + let obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({ + "spec": {"warmPoolRef": {"name": "pool-a"}}, + "status": {"sandbox": {"name": "sandbox-a"}} + }), + ); + let claim = parse_sandbox_claim(&obj).unwrap(); + let sandbox = sandbox_cr( + "sandbox-a", + "sandbox-uid", + "openshell-sandbox", + SANDBOX_VERSION_V1ALPHA1, + ); + let pod = sandbox_pod( + "pod-a", + "pod-uid", + "sandbox-uid", + SANDBOX_API_VERSION_FULL_V1ALPHA1, + ); + + let request = activation_request_from_claim_state(&claim, &sandbox, &pod) + .unwrap() + .unwrap(); + + assert_eq!(request.driver, "kubernetes"); + assert_eq!(request.instance_id, "pod-uid"); + assert_eq!(request.sandbox_id, "openshell-sandbox"); + assert_eq!(request.owner_uid, "sandbox-uid"); + assert_eq!(request.reason, "SandboxClaim/claim-a"); + } + + #[test] + fn activation_request_rejects_uncontrolled_pod() { + let obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({"status": {"sandbox": {"name": "sandbox-a"}}}), + ); + let claim = parse_sandbox_claim(&obj).unwrap(); + let sandbox = sandbox_cr( + "sandbox-a", + "sandbox-uid", + "openshell-sandbox", + SANDBOX_VERSION_V1BETA1, + ); + let pod = sandbox_pod( + "pod-a", + "pod-uid", + "other-sandbox-uid", + SANDBOX_API_VERSION_FULL_V1BETA1, + ); + + assert!(activation_request_from_claim_state(&claim, &sandbox, &pod).is_err()); + } + + #[test] + fn sandbox_selector_reads_status_selector() { + let sandbox = sandbox_cr( + "sandbox-a", + "sandbox-uid", + "openshell-sandbox", + SANDBOX_VERSION_V1BETA1, + ); + + assert_eq!( + sandbox_selector(&sandbox).as_deref(), + Some("agents.x-k8s.io/sandbox=sandbox-a") + ); + } + + #[test] + fn api_version_probe_retries_only_404_errors() { + let unavailable = KubeError::Api(kube::core::ErrorResponse { + status: "404 Not Found".to_string(), + message: "could not find the requested resource".to_string(), + reason: "NotFound".to_string(), + code: 404, + }); + let forbidden = KubeError::Api(kube::core::ErrorResponse { + status: "Failure".to_string(), + message: "forbidden".to_string(), + reason: "Forbidden".to_string(), + code: 403, + }); + + assert!(should_try_next_api_version(&unavailable)); + assert!(!should_try_next_api_version(&forbidden)); + } +} diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index dd2f00bf6a..08d9b1f03e 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -37,12 +37,15 @@ use openshell_core::proto::{ PlatformEvent, Sandbox, SandboxCondition, SandboxPhase, SandboxSpec, SandboxStatus, SandboxTemplate, ServiceEndpoint, SshSession, }; -use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentityProvider; +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapActivator, SupervisorBootstrapIdentityProvider, +}; use openshell_core::{ObjectLabels, ObjectWorkspace}; use openshell_driver_docker::DockerComputeDriver; use openshell_driver_kubernetes::{ ComputeDriverService as KubernetesDriverService, KubernetesComputeDriver, KubernetesSupervisorBootstrapIdentityProvider, LiveK8sResolver, + SandboxClaimActivationController, }; use openshell_driver_podman::{ComputeDriverService as PodmanDriverService, PodmanComputeDriver}; use prost::Message; @@ -260,6 +263,27 @@ async fn kubernetes_supervisor_bootstrap_identity_provider( } } +async fn kubernetes_sandbox_claim_activation_controller( + config: &KubernetesComputeConfig, +) -> Option> { + std::env::var_os("KUBERNETES_SERVICE_HOST")?; + + match kube::Client::try_default().await { + Ok(client) => Some(Arc::new(SandboxClaimActivationController::new( + client.clone(), + client, + config.namespace.clone(), + ))), + Err(err) => { + warn!( + error = %err, + "in-cluster K8s client construction failed; SandboxClaim activation is disabled" + ); + None + } + } +} + #[derive(Debug)] pub struct AcquiredRemoteDriverEndpoint { pub(crate) name: String, @@ -398,6 +422,7 @@ pub struct ComputeRuntime { tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, supervisor_bootstrap_identity: Option>, + sandbox_claim_activation: Option>, sync_lock: Arc>, delete_gates: Arc, gateway_bind_addresses: Vec, @@ -424,6 +449,7 @@ impl ComputeRuntime { tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, supervisor_bootstrap_identity: Option>, + sandbox_claim_activation: Option>, gateway_bind_addresses: Vec, ) -> Result { let capabilities = driver @@ -457,6 +483,7 @@ impl ComputeRuntime { tracing_log_bus, supervisor_sessions, supervisor_bootstrap_identity, + sandbox_claim_activation, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses, @@ -519,6 +546,7 @@ impl ComputeRuntime { tracing_log_bus, supervisor_sessions, None, + None, gateway_bind_addresses, ) .await @@ -534,6 +562,8 @@ impl ComputeRuntime { ) -> Result { let supervisor_bootstrap_identity = kubernetes_supervisor_bootstrap_identity_provider(&config).await; + let sandbox_claim_activation = + kubernetes_sandbox_claim_activation_controller(&config).await; let driver = KubernetesComputeDriver::new(config) .await .map_err(|err| ComputeError::Message(err.to_string()))?; @@ -550,6 +580,7 @@ impl ComputeRuntime { tracing_log_bus, supervisor_sessions, supervisor_bootstrap_identity, + sandbox_claim_activation, Vec::new(), ) .await @@ -576,6 +607,7 @@ impl ComputeRuntime { tracing_log_bus, supervisor_sessions, None, + None, Vec::new(), ) .await @@ -605,6 +637,7 @@ impl ComputeRuntime { tracing_log_bus, supervisor_sessions, None, + None, Vec::new(), ) .await @@ -627,6 +660,16 @@ impl ComputeRuntime { self.supervisor_bootstrap_identity.clone() } + pub(crate) fn spawn_sandbox_claim_activation( + &self, + activator: Arc, + shutdown_rx: watch::Receiver, + ) { + if let Some(controller) = self.sandbox_claim_activation.clone() { + controller.spawn(activator, shutdown_rx); + } + } + #[must_use] pub fn driver_kind(&self) -> Option { self.driver_info.name.parse().ok() @@ -2805,6 +2848,7 @@ pub async fn new_test_runtime(store: Arc) -> ComputeRuntime { tracing_log_bus: TracingLogBus::new(), supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), supervisor_bootstrap_identity: None, + sandbox_claim_activation: None, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses: Vec::new(), @@ -3275,6 +3319,7 @@ mod tests { tracing_log_bus: TracingLogBus::new(), supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), supervisor_bootstrap_identity: None, + sandbox_claim_activation: None, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses: Vec::new(), diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index 1aac8d8b55..24896d3c98 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -404,6 +404,10 @@ pub(crate) async fn run_server( } state.compute.spawn_watchers(shutdown_rx.clone()); + state.compute.spawn_sandbox_claim_activation( + Arc::new(warm_pod_activation::GatewaySupervisorBootstrapActivator::new(state.clone())), + shutdown_rx.clone(), + ); ssh_sessions::spawn_session_reaper(store.clone(), Duration::from_secs(3600)); supervisor_session::spawn_relay_reaper(state.clone(), Duration::from_secs(30)); provider_refresh::spawn_refresh_worker(state.clone(), Duration::from_secs(60)); From 6cdab13e47b9af68e7ae86bc014e172334545bfb Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Thu, 23 Jul 2026 20:28:51 +0100 Subject: [PATCH 07/10] feat(kubernetes): create SandboxClaims for matching warm pools Signed-off-by: Gordon Sim --- Cargo.lock | 1 + architecture/compute-runtimes.md | 9 +- crates/openshell-core/src/grpc_client.rs | 74 +- crates/openshell-driver-kubernetes/Cargo.toml | 1 + crates/openshell-driver-kubernetes/README.md | 20 + .../openshell-driver-kubernetes/src/config.rs | 45 + .../openshell-driver-kubernetes/src/driver.rs | 1837 ++++++++++++++++- crates/openshell-driver-kubernetes/src/lib.rs | 2 +- .../openshell-driver-kubernetes/src/main.rs | 3 +- .../src/sandboxclaim.rs | 425 +++- crates/openshell-sandbox/src/lib.rs | 18 +- .../openshell-supervisor-process/src/run.rs | 4 + deploy/helm/openshell/README.md | 1 + .../helm/openshell/templates/clusterrole.yaml | 22 + .../openshell/templates/gateway-config.yaml | 3 + .../templates/warm-pool-profile-role.yaml | 27 + .../warm-pool-profile-rolebinding.yaml | 19 + .../openshell/tests/gateway_config_test.yaml | 9 + .../tests/warm_pool_profiles_rbac_test.yaml | 58 + deploy/helm/openshell/values.yaml | 5 + docs/kubernetes/setup.mdx | 6 +- docs/reference/gateway-config.mdx | 5 + docs/reference/sandbox-compute-drivers.mdx | 10 + 23 files changed, 2430 insertions(+), 174 deletions(-) create mode 100644 deploy/helm/openshell/templates/warm-pool-profile-role.yaml create mode 100644 deploy/helm/openshell/templates/warm-pool-profile-rolebinding.yaml create mode 100644 deploy/helm/openshell/tests/warm_pool_profiles_rbac_test.yaml diff --git a/Cargo.lock b/Cargo.lock index 21c81cc945..56e0d443e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3943,6 +3943,7 @@ dependencies = [ "prost-types", "serde", "serde_json", + "sha2 0.10.9", "temp-env", "thiserror 2.0.18", "tokio", diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index 7145e90ad3..43539e1322 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -64,7 +64,7 @@ delete, reconciliation removes the row; otherwise it can remain `Deleting`. |---|---|---|---| | Docker | Local development with Docker available. | Container plus nested sandbox namespace. | Uses host networking so loopback gateway endpoints work from the supervisor. | | Podman | Rootless or single-machine deployments. | Container plus nested sandbox namespace. | Uses the Podman REST API, OCI image volumes, and CDI GPU devices when available. | -| Kubernetes | Cluster deployment through Helm. | Pod plus nested sandbox namespace. | Uses Kubernetes API objects, service accounts, secrets, PVC-backed workspace storage, and GPU resources. | +| Kubernetes | Cluster deployment through Helm. | Pod plus nested sandbox namespace. | Uses Kubernetes API objects, service accounts, secrets, PVC-backed workspace storage, GPU resources, and optionally Agent Sandbox v1beta1 warm-pool claims. | | VM | Experimental microVM isolation. | Per-sandbox libkrun VM. | Managed endpoint-backed driver. The gateway spawns `openshell-driver-vm`, waits for its Unix socket, and then consumes it through the same remote `compute_driver.proto` path used by unmanaged endpoint drivers. The VM driver boots a cached bootstrap `rootfs.ext4`, prepares requested OCI images inside a bootstrap VM with `umoci`, attaches the prepared image disk read-only, and gives each sandbox a writable `overlay.ext4` for merged-root changes and runtime material. The driver persists each accepted launch request beside the overlay and restarts those VMs on driver startup without recreating the overlay. | | Extension | Out-of-tree drivers operated alongside the gateway. | Whatever boundary the driver implements. | Selected by a non-reserved custom `compute_drivers = [""]` entry with `[openshell.drivers.].socket_path`, or at launch time by pairing `--drivers ` with `--compute-driver-socket=`. Reserved built-in names such as `vm`, `docker`, `podman`, and `kubernetes` cannot be used as unmanaged socket endpoints. The gateway connects to a UDS the operator already provisioned, runs `GetCapabilities`, logs the advertised `driver_name`, and dispatches all sandbox lifecycle calls through `compute_driver.proto`. The driver process and socket lifecycle are operator-owned; the gateway does not spawn, supervise, or remove unmanaged extension drivers. The trust boundary is the socket's filesystem permissions: the operator must ensure only the gateway uid can read/write it. | @@ -134,6 +134,13 @@ JWT. The supervisor installs that JWT in memory and starts the normal pod IPs or require an inbound activation port; activation is supervisor initiated over the existing outbound gRPC connection. +When compatible OpenShell-enabled Agent Sandbox warm pools exist, the +Kubernetes driver can create a `SandboxClaim` instead of a direct `Sandbox`. +The driver watches warm pools and templates into a local cache so create +requests do not list or fetch extension objects on the hot path. This remains +behind the compute-driver boundary: other drivers and the driver-agnostic +gateway lifecycle continue to operate on OpenShell sandbox identity and status. + Kubernetes can run the supervisor in the default combined topology or in a sidecar topology. Combined mode keeps network and process supervision in the agent container. Sidecar mode runs network enforcement, the proxy, and gateway diff --git a/crates/openshell-core/src/grpc_client.rs b/crates/openshell-core/src/grpc_client.rs index c27ff4dbb7..c261f3eee7 100644 --- a/crates/openshell-core/src/grpc_client.rs +++ b/crates/openshell-core/src/grpc_client.rs @@ -25,7 +25,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use crate::proto::{ DenialSummary, GetDraftPolicyRequest, GetInferenceBundleRequest, GetInferenceBundleResponse, GetSandboxConfigRequest, GetSandboxProviderEnvironmentRequest, NetworkActivitySummary, - PolicyChunk, PolicySource, PolicyStatus, RefreshSandboxTokenRequest, + PodActivationMessage, PolicyChunk, PolicySource, PolicyStatus, RefreshSandboxTokenRequest, RegisterSupervisorPodRequest, ReportPolicyStatusRequest, SandboxPolicy as ProtoSandboxPolicy, SubmitPolicyAnalysisRequest, SubmitPolicyAnalysisResponse, UpdateConfigRequest, inference_client::InferenceClient, open_shell_client::OpenShellClient, @@ -272,7 +272,9 @@ async fn acquire_sandbox_token(endpoint: &str, plain_channel: &Channel) -> Resul && !sa_path.is_empty() { return Ok(AcquiredToken { - token: acquire_k8s_sandbox_token(endpoint, plain_channel, &sa_path).await?, + token: acquire_k8s_supervisor_activation(endpoint, plain_channel, &sa_path) + .await? + .token, refresh_mode: RefreshMode::GatewayJwt(TokenSource::K8sServiceAccount), }); } @@ -290,6 +292,18 @@ async fn acquire_k8s_sandbox_token( plain_channel: &Channel, sa_path: &str, ) -> Result { + Ok( + acquire_k8s_supervisor_activation(endpoint, plain_channel, sa_path) + .await? + .token, + ) +} + +async fn acquire_k8s_supervisor_activation( + endpoint: &str, + plain_channel: &Channel, + sa_path: &str, +) -> Result { let sa_token = std::fs::read_to_string(sa_path) .into_diagnostic() .wrap_err_with(|| format!("failed to read K8s SA token from {sa_path}"))? @@ -324,7 +338,61 @@ async fn acquire_k8s_sandbox_token( sandbox_name = %activation.sandbox_name, "received supervisor pod activation" ); - Ok(activation.token) + Ok(activation) +} + +/// Register a Kubernetes supervisor pod and install the activated gateway JWT. +/// +/// Warm pods call this before policy loading. The registration stream remains +/// pending until the Kubernetes driver observes a claim binding and the gateway +/// sends activation. +pub async fn register_k8s_supervisor_pod(endpoint: &str) -> Result { + let sa_path = std::env::var(sandbox_env::K8S_SA_TOKEN_FILE) + .ok() + .filter(|path| !path.is_empty()) + .ok_or_else(|| { + miette::miette!( + "{} must be set to register Kubernetes supervisor pod", + sandbox_env::K8S_SA_TOKEN_FILE + ) + })?; + + let mut backoff = Duration::from_secs(1); + loop { + let activation = { + let guard = TOKEN_INIT_LOCK.lock().await; + if TOKEN_SLOT.get().is_some() { + return Err(miette::miette!( + "Kubernetes supervisor pod registration was requested after sandbox token initialization" + )); + } + + match async { + let plain_channel = build_plain_channel(endpoint).await?; + acquire_k8s_supervisor_activation(endpoint, &plain_channel, &sa_path).await + } + .await + { + Ok(activation) => activation, + Err(err) => { + warn!( + endpoint = %endpoint, + error = %err, + retry_after_secs = backoff.as_secs(), + "Kubernetes supervisor pod registration failed; retrying: {err:#}" + ); + drop(guard); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(Duration::from_secs(30)); + continue; + } + } + }; + + let _slot = install_token_slot(&activation.token)?; + let _ = TOKEN_REFRESH_MODE.set(RefreshMode::GatewayJwt(TokenSource::K8sServiceAccount)); + return Ok(activation); + } } /// Build an authenticated channel for direct external use (e.g. the diff --git a/crates/openshell-driver-kubernetes/Cargo.toml b/crates/openshell-driver-kubernetes/Cargo.toml index 2c02f864ab..7ba2794b8a 100644 --- a/crates/openshell-driver-kubernetes/Cargo.toml +++ b/crates/openshell-driver-kubernetes/Cargo.toml @@ -29,6 +29,7 @@ kube-runtime = { workspace = true } k8s-openapi = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +sha2 = { workspace = true } clap = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/crates/openshell-driver-kubernetes/README.md b/crates/openshell-driver-kubernetes/README.md index ce8a88d55e..b0a3847f8d 100644 --- a/crates/openshell-driver-kubernetes/README.md +++ b/crates/openshell-driver-kubernetes/README.md @@ -26,6 +26,26 @@ by the gateway. Kubernetes API calls use explicit timeouts so gRPC handlers do not block indefinitely when the API server is slow or unavailable. +## Warm-Pool Allocation + +The driver can use Agent Sandbox extension CRDs for transparent warm-pool +allocation. When `warm_pooling.enabled` is true, the driver watches and caches +`extensions.agents.x-k8s.io/v1beta1` `SandboxWarmPool` resources labelled +`openshell.ai/enabled=true`, resolves each pool's referenced +`SandboxTemplate`, and computes an in-memory fingerprint of the template spec. + +On create, the driver renders the Kubernetes spec that direct `Sandbox` +creation would use, strips per-sandbox identity values from the fingerprint, +and matches it against the warm-pool cache for the target Kubernetes namespace. +If exactly one pool matches, the driver creates a v1beta1 `SandboxClaim` with +`spec.warmPoolRef.name` set to that pool. If no pool matches, multiple pools +match, RBAC prevents cache maintenance, or the v1beta1 extension APIs are +unavailable, the driver falls back to direct `Sandbox` creation. + +The extension CRDs are intentionally v1beta1-only. The core +`agents.x-k8s.io/Sandbox` CRD still supports the existing v1beta1-to-v1alpha1 +fallback. + ## Workspace Persistence Sandbox pods use a PVC-backed `/sandbox` workspace. An init container seeds the diff --git a/crates/openshell-driver-kubernetes/src/config.rs b/crates/openshell-driver-kubernetes/src/config.rs index b9b43de6fe..35c12cfc5d 100644 --- a/crates/openshell-driver-kubernetes/src/config.rs +++ b/crates/openshell-driver-kubernetes/src/config.rs @@ -125,6 +125,20 @@ impl KubernetesSidecarConfig { } } +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct KubernetesWarmPoolingConfig { + /// Allow the Kubernetes driver to satisfy compatible create requests by + /// creating v1beta1 Agent Sandbox `SandboxClaim` resources. + pub enabled: bool, +} + +impl Default for KubernetesWarmPoolingConfig { + fn default() -> Self { + Self { enabled: true } + } +} + /// Kubernetes `AppArmor` profile requested for the sandbox agent container. #[derive(Debug, Clone, PartialEq, Eq)] pub enum AppArmorProfile { @@ -253,6 +267,8 @@ pub struct KubernetesComputeConfig { pub topology: SupervisorTopology, /// Sidecar-only settings used when `topology = "sidecar"`. pub sidecar: KubernetesSidecarConfig, + /// Warm-pool allocation settings. + pub warm_pooling: KubernetesWarmPoolingConfig, pub grpc_endpoint: String, pub ssh_socket_path: String, pub client_tls_secret_name: String, @@ -340,6 +356,7 @@ impl Default for KubernetesComputeConfig { supervisor_sideload_method: SupervisorSideloadMethod::default(), topology: SupervisorTopology::default(), sidecar: KubernetesSidecarConfig::default(), + warm_pooling: KubernetesWarmPoolingConfig::default(), grpc_endpoint: String::new(), ssh_socket_path: "/run/openshell/ssh.sock".to_string(), client_tls_secret_name: String::new(), @@ -533,6 +550,34 @@ mod tests { assert!(cfg.sidecar.process_binary_aware_network_policy); } + #[test] + fn default_warm_pooling_is_enabled() { + let cfg = KubernetesComputeConfig::default(); + assert!(cfg.warm_pooling.enabled); + } + + #[test] + fn serde_override_warm_pooling_enabled() { + let json = serde_json::json!({ + "warm_pooling": { + "enabled": false + } + }); + let cfg: KubernetesComputeConfig = serde_json::from_value(json).unwrap(); + assert!(!cfg.warm_pooling.enabled); + } + + #[test] + fn serde_rejects_unknown_warm_pooling_field() { + let json = serde_json::json!({ + "warm_pooling": { + "mode": "always" + } + }); + let err = serde_json::from_value::(json).unwrap_err(); + assert!(err.to_string().contains("unknown field")); + } + #[test] fn serde_override_topology_sidecar() { let json = serde_json::json!({ diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index 3c69b80a83..cee6939cc6 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -9,7 +9,7 @@ use crate::config::{ DEFAULT_WORKSPACE_STORAGE_SIZE, KubernetesComputeConfig, SupervisorSideloadMethod, SupervisorTopology, }; -use futures::{Stream, StreamExt, TryStreamExt}; +use futures::{Stream, StreamExt, TryStreamExt, stream}; use k8s_openapi::api::core::v1::{ Event as KubeEventObj, Namespace, Node, PersistentVolumeClaimVolumeSource, Volume, VolumeMount, }; @@ -38,12 +38,13 @@ use openshell_core::proto::compute::v1::{ }; use openshell_core::proto_struct::{struct_to_json_object, value_to_json}; use serde::Deserialize; +use sha2::{Digest, Sha256}; use std::collections::{BTreeMap, HashSet}; use std::path::Path; use std::pin::Pin; use std::sync::Arc; use std::time::Duration; -use tokio::sync::{OnceCell, mpsc}; +use tokio::sync::{OnceCell, RwLock, mpsc}; use tokio_stream::wrappers::ReceiverStream; use tracing::{debug, info, warn}; @@ -86,6 +87,7 @@ impl From for openshell_core::ComputeDriverError { /// This prevents gRPC handlers from blocking indefinitely when the k8s /// API server is unreachable or slow. const KUBE_API_TIMEOUT: Duration = Duration::from_secs(30); +const WARM_POOL_CACHE_RETRY_DELAY: Duration = Duration::from_secs(10); const SANDBOX_GROUP: &str = "agents.x-k8s.io"; const SANDBOX_VERSION_V1BETA1: &str = "v1beta1"; @@ -93,6 +95,16 @@ const SANDBOX_VERSION_V1ALPHA1: &str = "v1alpha1"; const SANDBOX_VERSIONS: &[&str] = &[SANDBOX_VERSION_V1BETA1, SANDBOX_VERSION_V1ALPHA1]; pub const SANDBOX_KIND: &str = "Sandbox"; +const EXTENSIONS_GROUP: &str = "extensions.agents.x-k8s.io"; +const EXTENSIONS_VERSION_V1BETA1: &str = "v1beta1"; +const SANDBOX_CLAIM_KIND: &str = "SandboxClaim"; +const SANDBOX_WARM_POOL_KIND: &str = "SandboxWarmPool"; +const SANDBOX_TEMPLATE_KIND: &str = "SandboxTemplate"; +const LABEL_WARM_POOL_ENABLED: &str = "openshell.ai/enabled"; +const LABEL_ALLOCATION: &str = "openshell.ai/allocation"; +const LABEL_ALLOCATION_SANDBOX_CLAIM: &str = "sandbox-claim"; +const POD_ANNOTATION_SANDBOX_ID: &str = "openshell.io/sandbox-id"; + const GPU_RESOURCE_NAME: &str = "nvidia.com/gpu"; const SPIFFE_WORKLOAD_API_VOLUME_NAME: &str = "spiffe-workload-api"; @@ -101,6 +113,100 @@ struct AgentSandboxApi { resource: ApiResource, } +struct ExtensionApi { + api: Api, + resource: ApiResource, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +struct EnabledWarmPool { + namespace: String, + name: String, + template_namespace: String, + template_name: String, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +struct WarmPoolCacheEntry { + namespace: String, + name: String, + template_namespace: String, + template_name: String, + template_fingerprint: String, +} + +#[derive(Debug, Clone, Default)] +struct WarmPoolCache { + state: Arc>, +} + +#[derive(Debug, Default)] +struct WarmPoolCacheState { + ready: bool, + entries: BTreeMap<(String, String), WarmPoolCacheEntry>, +} + +#[derive(Debug, Eq, PartialEq)] +enum WarmPoolCacheLookup { + NotReady, + NoMatch, + Match(WarmPoolCacheEntry), + Ambiguous(usize), +} + +impl WarmPoolCache { + async fn replace_entries(&self, entries: Vec) { + let mut state = self.state.write().await; + state.ready = true; + state.entries = entries + .into_iter() + .map(|entry| ((entry.namespace.clone(), entry.name.clone()), entry)) + .collect(); + } + + async fn mark_not_ready(&self) { + let mut state = self.state.write().await; + state.ready = false; + state.entries.clear(); + } + + async fn matching_pool( + &self, + namespace: &str, + template_fingerprint: &str, + ) -> WarmPoolCacheLookup { + let state = self.state.read().await; + if !state.ready { + return WarmPoolCacheLookup::NotReady; + } + + let mut matches = state + .entries + .values() + .filter(|entry| { + entry.namespace == namespace && entry.template_fingerprint == template_fingerprint + }) + .cloned() + .collect::>(); + + match matches.len() { + 0 => WarmPoolCacheLookup::NoMatch, + 1 => WarmPoolCacheLookup::Match(matches.pop().expect("one match exists")), + count => WarmPoolCacheLookup::Ambiguous(count), + } + } + + async fn entries_for_namespace(&self, namespace: &str) -> Vec { + let state = self.state.read().await; + state + .entries + .values() + .filter(|entry| entry.namespace == namespace) + .cloned() + .collect() + } +} + // This POC treats the selected Struct as a driver-local typed schema. Once the // Kubernetes shape stabilizes, these serde structs may move to driver-local // protobuf definitions, but the typed decode should stay inside this driver. @@ -431,6 +537,7 @@ pub struct KubernetesComputeDriver { client: Client, watch_client: Client, sandbox_api_version: Arc>, + warm_pool_cache: WarmPoolCache, config: KubernetesComputeConfig, } @@ -476,12 +583,17 @@ impl KubernetesComputeDriver { let watch_client = Client::try_from(watch_kube_config).map_err(KubernetesDriverError::from_kube)?; - Ok(Self { + let driver = Self { client, watch_client, sandbox_api_version: Arc::new(OnceCell::new()), + warm_pool_cache: WarmPoolCache::default(), config, - }) + }; + if driver.config.warm_pooling.enabled { + driver.spawn_warm_pool_cache_controller(); + } + Ok(driver) } pub fn capabilities(&self) -> Result { @@ -525,6 +637,54 @@ impl KubernetesComputeDriver { AgentSandboxApi { api, resource } } + fn extension_resource(kind: &str) -> ApiResource { + let gvk = GroupVersionKind::gvk(EXTENSIONS_GROUP, EXTENSIONS_VERSION_V1BETA1, kind); + ApiResource::from_gvk(&gvk) + } + + fn namespaced_extension_api(client: Client, namespace: &str, kind: &str) -> ExtensionApi { + let resource = Self::extension_resource(kind); + let api = Api::namespaced_with(client, namespace, &resource); + ExtensionApi { api, resource } + } + + fn all_extension_api(client: Client, kind: &str) -> ExtensionApi { + let resource = Self::extension_resource(kind); + let api = Api::all_with(client, &resource); + ExtensionApi { api, resource } + } + + fn sandbox_claim_api(client: Client, namespace: &str) -> ExtensionApi { + Self::namespaced_extension_api(client, namespace, SANDBOX_CLAIM_KIND) + } + + fn sandbox_claim_api_all(client: Client) -> ExtensionApi { + Self::all_extension_api(client, SANDBOX_CLAIM_KIND) + } + + fn sandbox_warm_pool_api_all(client: Client) -> ExtensionApi { + Self::all_extension_api(client, SANDBOX_WARM_POOL_KIND) + } + + fn sandbox_warm_pool_api(client: Client, namespace: &str) -> ExtensionApi { + Self::namespaced_extension_api(client, namespace, SANDBOX_WARM_POOL_KIND) + } + + fn sandbox_template_api(client: Client, namespace: &str) -> ExtensionApi { + Self::namespaced_extension_api(client, namespace, SANDBOX_TEMPLATE_KIND) + } + + fn spawn_warm_pool_cache_controller(&self) { + let cache = self.warm_pool_cache.clone(); + let client = self.client.clone(); + let watch_client = self.watch_client.clone(); + let fallback_namespace = self.config.namespace.clone(); + + tokio::spawn(async move { + run_warm_pool_cache_controller(cache, client, watch_client, fallback_namespace).await; + }); + } + async fn supported_agent_sandbox_api(&self, client: Client) -> Result { let sandbox_api_version = self.supported_sandbox_api_version(client.clone()).await?; Ok(self.agent_sandbox_api(client, sandbox_api_version)) @@ -700,17 +860,17 @@ impl KubernetesComputeDriver { format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE},{LABEL_SANDBOX_ID}={sandbox_id}"); let lp = ListParams::default().labels(&selector); match tokio::time::timeout(KUBE_API_TIMEOUT, agent_sandbox_api.api.list(&lp)).await { - Ok(Ok(list)) => list.items.into_iter().next().map_or_else( - || { - debug!(sandbox_id = %sandbox_id, "Sandbox not found in Kubernetes"); - Ok(None) - }, - |obj| { - Ok(sandbox_from_object(&self.config.namespace, obj) + Ok(Ok(list)) => { + if let Some(obj) = list.items.into_iter().next() { + return Ok(sandbox_from_object(&self.config.namespace, obj) .ok() - .map(|(_, s)| s)) - }, - ), + .map(|(_, s)| s)); + } + { + debug!(sandbox_id = %sandbox_id, "Sandbox not found in Kubernetes"); + self.get_claim_backed_sandbox(sandbox_id).await + } + } Ok(Err(err)) => { warn!( sandbox_id = %sandbox_id, @@ -765,6 +925,13 @@ impl KubernetesComputeDriver { } }) .collect(); + let mut claim_sandboxes = self.list_claim_backed_sandboxes().await?; + let mut claim_ids = claim_sandboxes + .iter() + .map(|sandbox| sandbox.id.clone()) + .collect::>(); + sandboxes.retain(|sandbox| !claim_ids.remove(&sandbox.id)); + sandboxes.append(&mut claim_sandboxes); sandboxes.sort_by(|left, right| { left.name .cmp(&right.name) @@ -794,6 +961,381 @@ impl KubernetesComputeDriver { } } + async fn list_claim_backed_sandboxes(&self) -> Result, String> { + let selector = format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE}"); + let lp = ListParams::default().labels(&selector); + let claim_api = Self::sandbox_claim_api_all(self.client.clone()); + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(Vec::new()), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + let claim_api = + Self::sandbox_claim_api(self.client.clone(), &self.config.namespace); + match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(Vec::new()), + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + + Ok(list + .items + .into_iter() + .filter_map(|obj| sandbox_from_claim_object(&self.config.namespace, obj).ok()) + .collect()) + } + + async fn get_claim_backed_sandbox(&self, sandbox_id: &str) -> Result, String> { + let selector = + format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE},{LABEL_SANDBOX_ID}={sandbox_id}"); + let lp = ListParams::default().labels(&selector); + let claim_api = Self::sandbox_claim_api_all(self.client.clone()); + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(None), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + let claim_api = + Self::sandbox_claim_api(self.client.clone(), &self.config.namespace); + match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(None), + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + + Ok(list + .items + .into_iter() + .find_map(|obj| sandbox_from_claim_object(&self.config.namespace, obj).ok())) + } + + async fn list_enabled_warm_pools_with_client( + client: Client, + fallback_namespace: &str, + ) -> Result, String> { + let lp = ListParams::default().labels(&format!("{LABEL_WARM_POOL_ENABLED}=true")); + let warm_pool_api = Self::sandbox_warm_pool_api_all(client.clone()); + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, warm_pool_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(Vec::new()), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + let warm_pool_api = Self::sandbox_warm_pool_api(client, fallback_namespace); + match tokio::time::timeout(KUBE_API_TIMEOUT, warm_pool_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(Vec::new()), + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + + Ok(list + .items + .into_iter() + .filter_map(enabled_warm_pool_from_object) + .collect()) + } + + async fn template_fingerprint_for_warm_pool_with_client( + client: Client, + pool: &EnabledWarmPool, + ) -> Result, String> { + let template_api = Self::sandbox_template_api(client, &pool.template_namespace); + match tokio::time::timeout(KUBE_API_TIMEOUT, template_api.api.get(&pool.template_name)) + .await + { + Ok(Ok(template)) => sandbox_template_fingerprint(&template).map(Some), + Ok(Err(KubeError::Api(err))) if err.code == 404 => Ok(None), + Ok(Err(err)) if extension_api_unavailable(&err) => Ok(None), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } + } + + async fn extension_api_for_watch( + client: Client, + fallback_namespace: &str, + kind: &str, + ) -> Result, String> { + let all_api = Self::all_extension_api(client.clone(), kind); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + all_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => Ok(all_api.api), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + let namespaced_api = + Self::namespaced_extension_api(client, fallback_namespace, kind); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + namespaced_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => Ok(namespaced_api.api), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } + } + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } + } + + async fn matching_warm_pool( + &self, + target_namespace: &str, + request_fingerprint: &str, + ) -> Option { + match self + .warm_pool_cache + .matching_pool(target_namespace, request_fingerprint) + .await + { + WarmPoolCacheLookup::NotReady => { + debug!( + namespace = %target_namespace, + "Warm-pool cache is not ready; falling back to direct Sandbox" + ); + None + } + WarmPoolCacheLookup::NoMatch => { + let cached_warm_pool_templates = self + .warm_pool_cache + .entries_for_namespace(target_namespace) + .await + .into_iter() + .map(|entry| { + format!( + "{}/{}:{}", + entry.template_namespace, + entry.template_name, + entry.template_fingerprint + ) + }) + .collect::>(); + debug!( + namespace = %target_namespace, + request_fingerprint, + cached_warm_pool_template_count = cached_warm_pool_templates.len(), + cached_warm_pool_templates = ?cached_warm_pool_templates, + "No OpenShell-enabled warm pool matches sandbox request; falling back to direct Sandbox" + ); + None + } + WarmPoolCacheLookup::Match(warm_pool) => { + debug!( + namespace = %target_namespace, + request_fingerprint, + warm_pool = %warm_pool.name, + template_namespace = %warm_pool.template_namespace, + template = %warm_pool.template_name, + "OpenShell-enabled warm pool matches sandbox request" + ); + Some(warm_pool) + } + WarmPoolCacheLookup::Ambiguous(count) => { + warn!( + namespace = %target_namespace, + count, + "Multiple OpenShell-enabled warm pools match sandbox request; falling back to direct Sandbox" + ); + None + } + } + } + + async fn try_create_sandbox_claim( + &self, + sandbox: &Sandbox, + target_namespace: &str, + request_fingerprint: &str, + ) -> Result { + if !self.config.warm_pooling.enabled { + return Ok(false); + } + + let Some(warm_pool) = self + .matching_warm_pool(target_namespace, request_fingerprint) + .await + else { + return Ok(false); + }; + + let claim_api = Self::sandbox_claim_api(self.client.clone(), target_namespace); + let claim = sandbox_claim_to_k8s_object(sandbox, &warm_pool.name, &claim_api.resource); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + claim_api.api.create(&PostParams::default(), &claim), + ) + .await + { + Ok(Ok(_result)) => { + info!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + warm_pool = %warm_pool.name, + "SandboxClaim created in Kubernetes successfully" + ); + Ok(true) + } + Ok(Err(err)) if extension_api_unavailable(&err) => { + debug!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + error = %err, + "SandboxClaim API is unavailable; falling back to direct Sandbox" + ); + Ok(false) + } + Ok(Err(err)) if matches!(&err, KubeError::Api(api) if api.code == 409) => { + Err(KubernetesDriverError::from_kube(err)) + } + Ok(Err(err)) => { + warn!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + warm_pool = %warm_pool.name, + error = %err, + "Failed to create SandboxClaim; falling back to direct Sandbox" + ); + Ok(false) + } + Err(_elapsed) => { + warn!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + timeout_secs = KUBE_API_TIMEOUT.as_secs(), + "Timed out creating SandboxClaim; falling back to direct Sandbox" + ); + Ok(false) + } + } + } + + async fn delete_sandbox_claim(&self, sandbox_id: &str) -> Result { + let selector = + format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE},{LABEL_SANDBOX_ID}={sandbox_id}"); + let lp = ListParams::default().labels(&selector); + let mut claim_api = Self::sandbox_claim_api_all(self.client.clone()); + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(false), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + claim_api = Self::sandbox_claim_api(self.client.clone(), &self.config.namespace); + match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(false), + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + + let Some(obj) = list.items.into_iter().next() else { + return Ok(false); + }; + let Some(claim_name) = obj.metadata.name.clone() else { + return Ok(false); + }; + let claim_namespace = obj + .metadata + .namespace + .clone() + .unwrap_or_else(|| self.config.namespace.clone()); + let preconditions = Preconditions { + uid: obj.metadata.uid, + resource_version: obj.metadata.resource_version, + }; + + let claim_api = Self::sandbox_claim_api(self.client.clone(), &claim_namespace); + let dp = DeleteParams::default().preconditions(preconditions); + match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.delete(&claim_name, &dp)).await { + Ok(Ok(_response)) => { + info!( + sandbox_id = %sandbox_id, + namespace = %claim_namespace, + sandbox_claim = %claim_name, + "SandboxClaim deleted from Kubernetes" + ); + Ok(true) + } + Ok(Err(KubeError::Api(err))) if err.code == 404 || err.code == 409 => Ok(false), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } + } + #[allow(clippy::similar_names)] pub async fn create_sandbox(&self, sandbox: &Sandbox) -> Result<(), KubernetesDriverError> { let gpu_requirements = sandbox @@ -861,6 +1403,36 @@ impl KubernetesComputeDriver { let data = sandbox_to_k8s_spec(sandbox.spec.as_ref(), ¶ms) .map_err(KubernetesDriverError::InvalidArgument)?; + let target_namespace = self.config.namespace.clone(); + if self.config.warm_pooling.enabled { + match sandbox_spec_fingerprint(&data) { + Ok(fingerprint) => { + if self + .try_create_sandbox_claim(sandbox, &target_namespace, &fingerprint) + .await? + { + return Ok(()); + } + debug!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + request_fingerprint = %fingerprint, + normalized_request_spec = %sandbox_spec_fingerprint_debug_json(&data) + .unwrap_or_else(|err| format!("failed to render normalized request spec: {err}")), + "Warm-pool request fingerprint input" + ); + } + Err(err) => { + debug!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + error = %err, + "Could not fingerprint sandbox request for warm-pool matching; falling back to direct Sandbox" + ); + } + } + } let kube_name = kube_resource_name(&sandbox.workspace, name); let mut obj = DynamicObject::new(&kube_name, &agent_sandbox_api.resource); // Copy only the SCC-related annotations onto the Sandbox CR for @@ -877,7 +1449,7 @@ impl KubernetesComputeDriver { } obj.metadata = ObjectMeta { name: Some(kube_name), - namespace: Some(self.config.namespace.clone()), + namespace: Some(target_namespace), labels: Some(sandbox_labels(sandbox)), annotations: Some(annotations), ..Default::default() @@ -929,6 +1501,10 @@ impl KubernetesComputeDriver { "Deleting sandbox from Kubernetes" ); + if self.delete_sandbox_claim(sandbox_id).await? { + return Ok(true); + } + let agent_sandbox_api = self .supported_agent_sandbox_api(self.client.clone()) .await?; @@ -1024,7 +1600,12 @@ impl KubernetesComputeDriver { format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE},{LABEL_SANDBOX_ID}={sandbox_id}"); let lp = ListParams::default().labels(&selector); match tokio::time::timeout(KUBE_API_TIMEOUT, agent_sandbox_api.api.list(&lp)).await { - Ok(Ok(list)) => Ok(!list.items.is_empty()), + Ok(Ok(list)) => { + if !list.items.is_empty() { + return Ok(true); + } + Ok(self.get_claim_backed_sandbox(sandbox_id).await?.is_some()) + } Ok(Err(err)) => Err(err.to_string()), Err(_elapsed) => Err(format!( "timed out after {}s waiting for Kubernetes API", @@ -1043,8 +1624,55 @@ impl KubernetesComputeDriver { let event_api: Api = Api::namespaced(self.watch_client.clone(), &namespace); let watcher_config = watcher::Config::default().labels(&openshell_sandbox_label_selector()); let mut sandbox_stream = watcher::watcher(agent_sandbox_api.api, watcher_config).boxed(); - let mut event_stream = watcher::watcher(event_api, watcher::Config::default()).boxed(); - let (tx, rx) = mpsc::channel(256); + let claim_watcher_config = watcher::Config::default() + .labels(&format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE}")); + let claim_api = if self.config.warm_pooling.enabled { + let all_api = Self::sandbox_claim_api_all(self.watch_client.clone()); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + all_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => Some(all_api.api), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + let namespaced_api = + Self::sandbox_claim_api(self.watch_client.clone(), &namespace); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + namespaced_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => Some(namespaced_api.api), + Ok(Err(err)) if extension_api_unavailable(&err) => None, + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(Err(err)) if extension_api_unavailable(&err) => None, + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } else { + None + }; + let mut claim_stream = claim_api.map_or_else( + || stream::pending().boxed(), + |api| watcher::watcher(api, claim_watcher_config).boxed(), + ); + let mut event_stream = watcher::watcher(event_api, watcher::Config::default()).boxed(); + let (tx, rx) = mpsc::channel(256); tokio::spawn(async move { let mut sandbox_name_to_id = std::collections::HashMap::::new(); @@ -1107,6 +1735,59 @@ impl KubernetesComputeDriver { break; } }, + result = claim_stream.try_next() => match result { + Ok(Some(Event::Applied(obj))) => { + if let Ok(sandbox) = sandbox_from_claim_object(&namespace, obj) { + let event = WatchSandboxesEvent { + payload: Some(watch_sandboxes_event::Payload::Sandbox( + WatchSandboxesSandboxEvent { sandbox: Some(sandbox) } + )), + }; + if tx.send(Ok(event)).await.is_err() { + break; + } + } + } + Ok(Some(Event::Deleted(obj))) => { + if is_openshell_managed(&obj) + && let Ok(sandbox_id) = sandbox_id_from_object(&obj) + { + remove_indexes(&mut sandbox_name_to_id, &mut agent_pod_to_id, &sandbox_id); + let event = WatchSandboxesEvent { + payload: Some(watch_sandboxes_event::Payload::Deleted( + WatchSandboxesDeletedEvent { sandbox_id } + )), + }; + if tx.send(Ok(event)).await.is_err() { + break; + } + } + } + Ok(Some(Event::Restarted(objs))) => { + for obj in objs { + if let Ok(sandbox) = sandbox_from_claim_object(&namespace, obj) { + let event = WatchSandboxesEvent { + payload: Some(watch_sandboxes_event::Payload::Sandbox( + WatchSandboxesSandboxEvent { sandbox: Some(sandbox) } + )), + }; + if tx.send(Ok(event)).await.is_err() { + return; + } + } + } + } + Ok(None) => { + let _ = tx.send(Err(KubernetesDriverError::Message( + "sandbox claim watcher stream ended unexpectedly".to_string() + ))).await; + break; + } + Err(err) => { + let _ = tx.send(Err(KubernetesDriverError::Message(err.to_string()))).await; + break; + } + }, result = event_stream.try_next() => match result { Ok(Some(Event::Applied(obj))) => { if let Some((sandbox_id, event)) = map_kube_event_to_platform( @@ -1148,63 +1829,603 @@ impl KubernetesComputeDriver { } } -fn should_try_next_sandbox_api_version(err: &KubeError) -> bool { - // Kubernetes returns a structured 404 for some missing API resources and a - // raw "404 page not found" body for others. Both mean the probed - // group/version is unavailable and the next supported Sandbox API version - // should be tried. - matches!(err, KubeError::Api(api) if api.code == 404) +async fn run_warm_pool_cache_controller( + cache: WarmPoolCache, + client: Client, + watch_client: Client, + fallback_namespace: String, +) { + loop { + refresh_warm_pool_cache(&cache, client.clone(), &fallback_namespace).await; + + let warm_pool_api = match KubernetesComputeDriver::extension_api_for_watch( + watch_client.clone(), + &fallback_namespace, + SANDBOX_WARM_POOL_KIND, + ) + .await + { + Ok(api) => api, + Err(err) => { + cache.mark_not_ready().await; + warn!( + namespace = %fallback_namespace, + error = %err, + "Warm-pool watch is unavailable; warm-pool allocation disabled until retry" + ); + tokio::time::sleep(WARM_POOL_CACHE_RETRY_DELAY).await; + continue; + } + }; + let template_api = match KubernetesComputeDriver::extension_api_for_watch( + watch_client.clone(), + &fallback_namespace, + SANDBOX_TEMPLATE_KIND, + ) + .await + { + Ok(api) => api, + Err(err) => { + cache.mark_not_ready().await; + warn!( + namespace = %fallback_namespace, + error = %err, + "SandboxTemplate watch is unavailable; warm-pool allocation disabled until retry" + ); + tokio::time::sleep(WARM_POOL_CACHE_RETRY_DELAY).await; + continue; + } + }; + + let warm_pool_config = + watcher::Config::default().labels(&format!("{LABEL_WARM_POOL_ENABLED}=true")); + let template_config = watcher::Config::default(); + let mut warm_pool_stream = watcher::watcher(warm_pool_api, warm_pool_config).boxed(); + let mut template_stream = watcher::watcher(template_api, template_config).boxed(); + + loop { + tokio::select! { + event = warm_pool_stream.try_next() => { + if !handle_warm_pool_cache_watch_event( + &cache, + client.clone(), + &fallback_namespace, + event, + SANDBOX_WARM_POOL_KIND, + ) + .await + { + break; + } + } + event = template_stream.try_next() => { + if !handle_warm_pool_cache_watch_event( + &cache, + client.clone(), + &fallback_namespace, + event, + SANDBOX_TEMPLATE_KIND, + ) + .await + { + break; + } + } + } + } + + cache.mark_not_ready().await; + tokio::time::sleep(WARM_POOL_CACHE_RETRY_DELAY).await; + } +} + +async fn handle_warm_pool_cache_watch_event( + cache: &WarmPoolCache, + client: Client, + fallback_namespace: &str, + event: Result>, watcher::Error>, + kind: &str, +) -> bool { + match event { + Ok(Some(Event::Applied(_) | Event::Deleted(_) | Event::Restarted(_))) => { + refresh_warm_pool_cache(cache, client, fallback_namespace).await; + true + } + Ok(None) => { + warn!( + kind, + "Warm-pool cache watch stream ended; warm-pool allocation disabled until retry" + ); + false + } + Err(err) => { + warn!( + kind, + error = %err, + "Warm-pool cache watch failed; warm-pool allocation disabled until retry" + ); + false + } + } +} + +async fn refresh_warm_pool_cache(cache: &WarmPoolCache, client: Client, fallback_namespace: &str) { + let pools = match KubernetesComputeDriver::list_enabled_warm_pools_with_client( + client.clone(), + fallback_namespace, + ) + .await + { + Ok(pools) => pools, + Err(err) => { + cache.mark_not_ready().await; + warn!( + namespace = %fallback_namespace, + error = %err, + "Failed to refresh warm-pool cache; warm-pool allocation disabled until retry" + ); + return; + } + }; + + let mut entries = Vec::new(); + for pool in pools { + let template_fingerprint = + match KubernetesComputeDriver::template_fingerprint_for_warm_pool_with_client( + client.clone(), + &pool, + ) + .await + { + Ok(Some(fingerprint)) => fingerprint, + Ok(None) => { + debug!( + namespace = %pool.namespace, + warm_pool = %pool.name, + template_namespace = %pool.template_namespace, + template = %pool.template_name, + "Skipping warm pool with missing or unreadable SandboxTemplate" + ); + continue; + } + Err(err) => { + warn!( + namespace = %pool.namespace, + warm_pool = %pool.name, + template_namespace = %pool.template_namespace, + template = %pool.template_name, + error = %err, + "Skipping warm pool after SandboxTemplate fingerprint failed" + ); + continue; + } + }; + + debug!( + namespace = %pool.namespace, + warm_pool = %pool.name, + template_namespace = %pool.template_namespace, + template = %pool.template_name, + template_fingerprint = %template_fingerprint, + "Cached OpenShell-enabled warm-pool template fingerprint" + ); + + entries.push(WarmPoolCacheEntry { + namespace: pool.namespace, + name: pool.name, + template_namespace: pool.template_namespace, + template_name: pool.template_name, + template_fingerprint, + }); + } + + let count = entries.len(); + cache.replace_entries(entries).await; + debug!( + namespace = %fallback_namespace, + count, + "Warm-pool cache refreshed" + ); +} + +fn should_try_next_sandbox_api_version(err: &KubeError) -> bool { + // Kubernetes returns a structured 404 for some missing API resources and a + // raw "404 page not found" body for others. Both mean the probed + // group/version is unavailable and the next supported Sandbox API version + // should be tried. + matches!(err, KubeError::Api(api) if api.code == 404) +} + +fn validate_gpu_request( + gpu_requirements: Option<&GpuResourceRequirements>, +) -> Result<(), tonic::Status> { + let _ = + effective_driver_gpu_count(gpu_requirements).map_err(tonic::Status::invalid_argument)?; + Ok(()) +} + +fn kube_resource_name(workspace: &str, name: &str) -> String { + format!("{workspace}--{name}") +} + +const MAX_KUBE_NAME_LEN: usize = 63; + +fn validate_kube_resource_name_length(workspace: &str, name: &str) -> Result<(), tonic::Status> { + let combined = workspace.len() + 2 + name.len(); // "--" separator + if combined > MAX_KUBE_NAME_LEN { + return Err(tonic::Status::invalid_argument(format!( + "combined Kubernetes resource name '{workspace}--{name}' is {combined} characters, \ + exceeding the DNS-1123 limit of {MAX_KUBE_NAME_LEN}" + ))); + } + Ok(()) +} + +fn sandbox_labels(sandbox: &Sandbox) -> BTreeMap { + let mut labels = BTreeMap::new(); + labels.insert(LABEL_SANDBOX_ID.to_string(), sandbox.id.clone()); + labels.insert(LABEL_SANDBOX_NAME.to_string(), sandbox.name.clone()); + labels.insert( + LABEL_SANDBOX_WORKSPACE.to_string(), + sandbox.workspace.clone(), + ); + labels.insert( + LABEL_MANAGED_BY.to_string(), + LABEL_MANAGED_BY_VALUE.to_string(), + ); + labels +} + +fn sandbox_annotations(sandbox: &Sandbox) -> BTreeMap { + let mut annotations = BTreeMap::new(); + annotations.insert(LABEL_SANDBOX_ID.to_string(), sandbox.id.clone()); + annotations.insert(LABEL_SANDBOX_NAME.to_string(), sandbox.name.clone()); + annotations.insert( + LABEL_SANDBOX_WORKSPACE.to_string(), + sandbox.workspace.clone(), + ); + annotations +} + +fn sandbox_claim_to_k8s_object( + sandbox: &Sandbox, + warm_pool_name: &str, + resource: &ApiResource, +) -> DynamicObject { + let kube_name = kube_resource_name(&sandbox.workspace, &sandbox.name); + let mut obj = DynamicObject::new(&kube_name, resource); + let mut labels = sandbox_labels(sandbox); + labels.insert( + LABEL_ALLOCATION.to_string(), + LABEL_ALLOCATION_SANDBOX_CLAIM.to_string(), + ); + obj.metadata = ObjectMeta { + name: Some(kube_name), + labels: Some(labels), + annotations: Some(sandbox_annotations(sandbox)), + ..Default::default() + }; + obj.data = serde_json::json!({ + "spec": { + "warmPoolRef": { + "name": warm_pool_name + } + } + }); + obj +} + +fn enabled_warm_pool_from_object(obj: DynamicObject) -> Option { + let namespace = obj.metadata.namespace.clone()?; + let name = obj.metadata.name.clone()?; + let template_name = string_at(&obj.data, &["spec", "sandboxTemplateRef", "name"])?; + let template_namespace = string_at(&obj.data, &["spec", "sandboxTemplateRef", "namespace"]) + .unwrap_or_else(|| namespace.clone()); + Some(EnabledWarmPool { + namespace, + name, + template_namespace, + template_name, + }) +} + +fn sandbox_template_fingerprint(obj: &DynamicObject) -> Result { + sandbox_spec_fingerprint(&obj.data) +} + +fn sandbox_spec_fingerprint(data: &serde_json::Value) -> Result { + let spec = data + .get("spec") + .ok_or_else(|| "object is missing spec".to_string())?; + let mut normalized = spec.clone(); + normalize_template_spec_for_fingerprint(&mut normalized); + stable_json_fingerprint(&normalized) +} + +fn sandbox_spec_fingerprint_debug_json(data: &serde_json::Value) -> Result { + let spec = data + .get("spec") + .ok_or_else(|| "object is missing spec".to_string())?; + let mut normalized = spec.clone(); + normalize_template_spec_for_fingerprint(&mut normalized); + serde_json::to_string(&canonical_json_value(&normalized)) + .map_err(|err| format!("failed to serialize normalized template: {err}")) +} + +fn normalize_template_spec_for_fingerprint(value: &mut serde_json::Value) { + remove_path_if_value( + value, + &["envVarsInjectionPolicy"], + &serde_json::json!("Disallowed"), + ); + remove_path_if_value( + value, + &["networkPolicyManagement"], + &serde_json::json!("Managed"), + ); + remove_path_if_value(value, &["operatingMode"], &serde_json::json!("Running")); + remove_path_if_value(value, &["replicas"], &serde_json::json!(1)); + remove_path_if_value(value, &["shutdownPolicy"], &serde_json::json!("Retain")); + remove_path_if_value( + value, + &["podTemplate", "spec", "dnsPolicy"], + &serde_json::json!("ClusterFirst"), + ); + remove_path_if_value( + value, + &["volumeClaimTemplatesPolicy"], + &serde_json::json!("Disallowed"), + ); + remove_path( + value, + &[ + "podTemplate", + "metadata", + "annotations", + POD_ANNOTATION_SANDBOX_ID, + ], + ); + remove_path( + value, + &["podTemplate", "metadata", "labels", LABEL_SANDBOX_ID], + ); + remove_empty_object_path(value, &["podTemplate", "metadata", "annotations"]); + remove_empty_object_path(value, &["podTemplate", "metadata", "labels"]); + remove_empty_object_path(value, &["podTemplate", "metadata"]); + if let Some(containers) = value + .pointer_mut("/podTemplate/spec/containers") + .and_then(serde_json::Value::as_array_mut) + { + for container in containers { + remove_sandbox_identity_env(container); + remove_empty_container_resources(container); + remove_default_volume_mount_read_only(container); + } + } + if let Some(init_containers) = value + .pointer_mut("/podTemplate/spec/initContainers") + .and_then(serde_json::Value::as_array_mut) + { + for container in init_containers { + remove_empty_container_resources(container); + remove_default_volume_mount_read_only(container); + } + } +} + +fn remove_sandbox_identity_env(container: &mut serde_json::Value) { + let Some(env) = container + .get_mut("env") + .and_then(serde_json::Value::as_array_mut) + else { + return; + }; + env.retain(|entry| { + !matches!( + entry.get("name").and_then(serde_json::Value::as_str), + Some(name) + if name == openshell_core::sandbox_env::SANDBOX_ID + || name == openshell_core::sandbox_env::SANDBOX + ) + }); +} + +fn remove_path(value: &mut serde_json::Value, path: &[&str]) { + let Some((last, parents)) = path.split_last() else { + return; + }; + let mut current = value; + for key in parents { + let Some(next) = current.get_mut(*key) else { + return; + }; + current = next; + } + if let Some(object) = current.as_object_mut() { + object.remove(*last); + } +} + +fn remove_path_if_value( + value: &mut serde_json::Value, + path: &[&str], + expected: &serde_json::Value, +) { + let Some((last, parents)) = path.split_last() else { + return; + }; + let mut current = value; + for key in parents { + let Some(next) = current.get_mut(*key) else { + return; + }; + current = next; + } + if current.get(*last) == Some(expected) + && let Some(object) = current.as_object_mut() + { + object.remove(*last); + } +} + +fn remove_empty_container_resources(container: &mut serde_json::Value) { + if container + .get("resources") + .is_some_and(|resources| resources.as_object().is_some_and(serde_json::Map::is_empty)) + && let Some(object) = container.as_object_mut() + { + object.remove("resources"); + } +} + +fn remove_default_volume_mount_read_only(container: &mut serde_json::Value) { + let Some(volume_mounts) = container + .get_mut("volumeMounts") + .and_then(serde_json::Value::as_array_mut) + else { + return; + }; + for mount in volume_mounts { + remove_path_if_value(mount, &["readOnly"], &serde_json::json!(false)); + } +} + +fn remove_empty_object_path(value: &mut serde_json::Value, path: &[&str]) { + let Some((last, parents)) = path.split_last() else { + return; + }; + let mut current = value; + for key in parents { + let Some(next) = current.get_mut(*key) else { + return; + }; + current = next; + } + if current + .get(*last) + .is_some_and(|entry| entry.as_object().is_some_and(serde_json::Map::is_empty)) + && let Some(object) = current.as_object_mut() + { + object.remove(*last); + } +} + +fn stable_json_fingerprint(value: &serde_json::Value) -> Result { + let canonical = canonical_json_value(value); + let bytes = serde_json::to_vec(&canonical) + .map_err(|err| format!("failed to serialize canonical template: {err}"))?; + let digest = Sha256::digest(bytes); + Ok(hex_encode(&digest)) +} + +fn canonical_json_value(value: &serde_json::Value) -> serde_json::Value { + match value { + serde_json::Value::Array(items) => { + serde_json::Value::Array(items.iter().map(canonical_json_value).collect()) + } + serde_json::Value::Object(object) => { + let mut sorted = serde_json::Map::new(); + let mut keys = object.keys().collect::>(); + keys.sort(); + for key in keys { + if let Some(value) = object.get(key) { + sorted.insert(key.clone(), canonical_json_value(value)); + } + } + serde_json::Value::Object(sorted) + } + other => other.clone(), + } +} + +fn hex_encode(bytes: &[u8]) -> String { + const HEX: &[u8; 16] = b"0123456789abcdef"; + let mut out = String::with_capacity(bytes.len() * 2); + for byte in bytes { + out.push(HEX[(byte >> 4) as usize] as char); + out.push(HEX[(byte & 0x0f) as usize] as char); + } + out } -fn validate_gpu_request( - gpu_requirements: Option<&GpuResourceRequirements>, -) -> Result<(), tonic::Status> { - let _ = - effective_driver_gpu_count(gpu_requirements).map_err(tonic::Status::invalid_argument)?; - Ok(()) +fn string_at(data: &serde_json::Value, path: &[&str]) -> Option { + let mut current = data; + for key in path { + current = current.get(*key)?; + } + current + .as_str() + .filter(|value| !value.is_empty()) + .map(ToString::to_string) } -fn kube_resource_name(workspace: &str, name: &str) -> String { - format!("{workspace}--{name}") +fn sandbox_from_claim_object( + default_namespace: &str, + obj: DynamicObject, +) -> Result { + let kube_name = obj.metadata.name.clone().unwrap_or_default(); + if !is_openshell_managed(&obj) { + debug!(object = %kube_name, "skipping SandboxClaim not managed by openshell"); + return Err(format!("SandboxClaim {kube_name} not managed by openshell")); + } + let id = sandbox_id_from_object(&obj)?; + let Some(name) = annotation_or_label(&obj, LABEL_SANDBOX_NAME) else { + return Err(format!("SandboxClaim {kube_name} missing sandbox name")); + }; + let Some(workspace) = annotation_or_label(&obj, LABEL_SANDBOX_WORKSPACE) else { + return Err(format!( + "SandboxClaim {kube_name} missing sandbox workspace" + )); + }; + let namespace = obj + .metadata + .namespace + .clone() + .unwrap_or_else(|| default_namespace.to_string()); + Ok(Sandbox { + id, + name, + namespace, + spec: None, + status: Some(claim_status_from_object(&obj)), + workspace, + }) } -const MAX_KUBE_NAME_LEN: usize = 63; +fn claim_status_from_object(obj: &DynamicObject) -> SandboxStatus { + let status_obj = obj + .data + .get("status") + .and_then(serde_json::Value::as_object); + let sandbox_name = status_obj + .and_then(|status| status.get("sandbox")) + .and_then(|value| value.get("name")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string(); + let conditions = status_obj + .and_then(|status| status.get("conditions")) + .and_then(serde_json::Value::as_array) + .map(|items| { + items + .iter() + .filter_map(condition_from_value) + .collect::>() + }) + .unwrap_or_default(); -fn validate_kube_resource_name_length(workspace: &str, name: &str) -> Result<(), tonic::Status> { - let combined = workspace.len() + 2 + name.len(); // "--" separator - if combined > MAX_KUBE_NAME_LEN { - return Err(tonic::Status::invalid_argument(format!( - "combined Kubernetes resource name '{workspace}--{name}' is {combined} characters, \ - exceeding the DNS-1123 limit of {MAX_KUBE_NAME_LEN}" - ))); + SandboxStatus { + sandbox_name, + instance_id: String::new(), + agent_fd: String::new(), + sandbox_fd: String::new(), + conditions, + deleting: obj.metadata.deletion_timestamp.is_some(), } - Ok(()) -} - -fn sandbox_labels(sandbox: &Sandbox) -> BTreeMap { - let mut labels = BTreeMap::new(); - labels.insert(LABEL_SANDBOX_ID.to_string(), sandbox.id.clone()); - labels.insert(LABEL_SANDBOX_NAME.to_string(), sandbox.name.clone()); - labels.insert( - LABEL_SANDBOX_WORKSPACE.to_string(), - sandbox.workspace.clone(), - ); - labels.insert( - LABEL_MANAGED_BY.to_string(), - LABEL_MANAGED_BY_VALUE.to_string(), - ); - labels } -fn sandbox_annotations(sandbox: &Sandbox) -> BTreeMap { - let mut annotations = BTreeMap::new(); - annotations.insert(LABEL_SANDBOX_ID.to_string(), sandbox.id.clone()); - annotations.insert(LABEL_SANDBOX_NAME.to_string(), sandbox.name.clone()); - annotations.insert( - LABEL_SANDBOX_WORKSPACE.to_string(), - sandbox.workspace.clone(), - ); - annotations +fn extension_api_unavailable(err: &KubeError) -> bool { + matches!(err, KubeError::Api(api) if api.code == 404) } fn sandbox_id_from_object(obj: &DynamicObject) -> Result { @@ -3180,6 +4401,492 @@ mod tests { assert!(!should_try_next_sandbox_api_version(&err)); } + #[test] + fn sandbox_claim_object_renders_v1beta1_warm_pool_ref() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_CLAIM_KIND); + let sandbox = Sandbox { + id: "sandbox-id".to_string(), + name: "dev".to_string(), + namespace: String::new(), + workspace: "workspace-a".to_string(), + spec: None, + status: None, + }; + + let claim = sandbox_claim_to_k8s_object(&sandbox, "pool-a", &resource); + + assert_eq!(claim.metadata.name.as_deref(), Some("workspace-a--dev")); + assert_eq!( + claim + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(LABEL_ALLOCATION)) + .map(String::as_str), + Some(LABEL_ALLOCATION_SANDBOX_CLAIM) + ); + assert_eq!( + string_at(&claim.data, &["spec", "warmPoolRef", "name"]).as_deref(), + Some("pool-a") + ); + assert_eq!( + string_at(&claim.data, &["spec", "sandboxTemplateRef", "name"]), + None + ); + } + + #[test] + fn enabled_warm_pool_parse_defaults_template_namespace_to_pool_namespace() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_WARM_POOL_KIND); + let mut warm_pool = DynamicObject::new("pool-a", &resource); + warm_pool.metadata.namespace = Some("team-a".to_string()); + warm_pool.metadata.labels = Some(BTreeMap::from([( + LABEL_WARM_POOL_ENABLED.to_string(), + "true".to_string(), + )])); + warm_pool.data = serde_json::json!({ + "spec": { + "sandboxTemplateRef": { + "name": "template-a" + } + } + }); + + let parsed = enabled_warm_pool_from_object(warm_pool).unwrap(); + + assert_eq!(parsed.namespace, "team-a"); + assert_eq!(parsed.name, "pool-a"); + assert_eq!(parsed.template_namespace, "team-a"); + assert_eq!(parsed.template_name, "template-a"); + } + + fn warm_pool_cache_entry( + namespace: &str, + name: &str, + template_fingerprint: &str, + ) -> WarmPoolCacheEntry { + WarmPoolCacheEntry { + namespace: namespace.to_string(), + name: name.to_string(), + template_namespace: namespace.to_string(), + template_name: format!("{name}-template"), + template_fingerprint: template_fingerprint.to_string(), + } + } + + #[test] + fn warm_pool_cache_requires_successful_refresh_before_matching() { + futures::executor::block_on(async { + let cache = WarmPoolCache::default(); + + assert_eq!( + cache.matching_pool("team-a", "fingerprint-a").await, + WarmPoolCacheLookup::NotReady + ); + + cache.replace_entries(Vec::new()).await; + + assert_eq!( + cache.matching_pool("team-a", "fingerprint-a").await, + WarmPoolCacheLookup::NoMatch + ); + }); + } + + #[test] + fn warm_pool_cache_matches_by_namespace_and_fingerprint() { + futures::executor::block_on(async { + let cache = WarmPoolCache::default(); + cache + .replace_entries(vec![ + warm_pool_cache_entry("team-a", "pool-a", "fingerprint-a"), + warm_pool_cache_entry("team-b", "pool-b", "fingerprint-a"), + ]) + .await; + + assert_eq!( + cache.matching_pool("team-a", "fingerprint-a").await, + WarmPoolCacheLookup::Match(warm_pool_cache_entry( + "team-a", + "pool-a", + "fingerprint-a" + )) + ); + assert_eq!( + cache.matching_pool("team-a", "fingerprint-b").await, + WarmPoolCacheLookup::NoMatch + ); + }); + } + + #[test] + fn warm_pool_cache_reports_ambiguous_matches() { + futures::executor::block_on(async { + let cache = WarmPoolCache::default(); + cache + .replace_entries(vec![ + warm_pool_cache_entry("team-a", "pool-a", "fingerprint-a"), + warm_pool_cache_entry("team-a", "pool-b", "fingerprint-a"), + ]) + .await; + + assert_eq!( + cache.matching_pool("team-a", "fingerprint-a").await, + WarmPoolCacheLookup::Ambiguous(2) + ); + }); + } + + #[test] + fn template_fingerprint_ignores_sandbox_identity_values() { + let base = serde_json::json!({ + "spec": { + "podTemplate": { + "metadata": { + "annotations": { + POD_ANNOTATION_SANDBOX_ID: "sandbox-a", + "stable": "value" + }, + "labels": { + LABEL_SANDBOX_ID: "sandbox-a", + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE + } + }, + "spec": { + "containers": [{ + "name": "agent", + "image": "example.test/sandbox:latest", + "env": [ + {"name": openshell_core::sandbox_env::SANDBOX_ID, "value": "sandbox-a"}, + {"name": openshell_core::sandbox_env::SANDBOX, "value": "dev-a"}, + {"name": "ENDPOINT", "value": "https://gateway"} + ] + }] + } + } + } + }); + let changed = serde_json::json!({ + "spec": { + "podTemplate": { + "metadata": { + "labels": { + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE, + LABEL_SANDBOX_ID: "sandbox-b" + }, + "annotations": { + "stable": "value", + POD_ANNOTATION_SANDBOX_ID: "sandbox-b" + } + }, + "spec": { + "containers": [{ + "env": [ + {"name": openshell_core::sandbox_env::SANDBOX_ID, "value": "sandbox-b"}, + {"name": openshell_core::sandbox_env::SANDBOX, "value": "dev-b"}, + {"name": "ENDPOINT", "value": "https://gateway"} + ], + "image": "example.test/sandbox:latest", + "name": "agent" + }] + } + } + } + }); + + assert_eq!( + sandbox_spec_fingerprint(&base).unwrap(), + sandbox_spec_fingerprint(&changed).unwrap() + ); + } + + #[test] + fn template_fingerprint_ignores_agent_sandbox_defaulted_values() { + let pre_create = serde_json::json!({ + "spec": { + "podTemplate": { + "metadata": { + "labels": { + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE + } + }, + "spec": { + "containers": [{ + "name": "agent", + "image": "example.test/sandbox:latest", + "env": [ + {"name": "OPENSHELL_ENDPOINT", "value": "https://gateway"}, + {"name": openshell_core::sandbox_env::SANDBOX_UID, "value": "1000"} + ] + }], + "initContainers": [{ + "name": "openshell-supervisor-install", + "image": "example.test/supervisor:latest" + }] + } + }, + "volumeClaimTemplates": [{ + "metadata": {"name": "workspace"}, + "spec": {"resources": {"requests": {"storage": "2Gi"}}} + }] + } + }); + let defaulted = serde_json::json!({ + "spec": { + "envVarsInjectionPolicy": "Disallowed", + "networkPolicyManagement": "Managed", + "operatingMode": "Running", + "replicas": 1, + "shutdownPolicy": "Retain", + "volumeClaimTemplatesPolicy": "Disallowed", + "podTemplate": { + "metadata": { + "labels": { + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE + } + }, + "spec": { + "containers": [{ + "name": "agent", + "image": "example.test/sandbox:latest", + "env": [ + {"name": "OPENSHELL_ENDPOINT", "value": "https://gateway"}, + {"name": openshell_core::sandbox_env::SANDBOX_UID, "value": "1000"} + ], + "resources": {} + }], + "initContainers": [{ + "name": "openshell-supervisor-install", + "image": "example.test/supervisor:latest", + "resources": {} + }] + } + }, + "volumeClaimTemplates": [{ + "metadata": {"name": "workspace"}, + "spec": {"resources": {"requests": {"storage": "2Gi"}}} + }] + } + }); + + assert_eq!( + sandbox_spec_fingerprint(&pre_create).unwrap(), + sandbox_spec_fingerprint(&defaulted).unwrap() + ); + } + + #[test] + fn template_fingerprint_prunes_empty_identity_metadata() { + let sandbox = serde_json::json!({ + "spec": { + "podTemplate": { + "metadata": { + "annotations": { + POD_ANNOTATION_SANDBOX_ID: "sandbox-a" + }, + "labels": { + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE, + LABEL_SANDBOX_ID: "sandbox-a" + } + }, + "spec": { + "containers": [{ + "name": "agent", + "image": "example.test/sandbox:latest" + }] + } + } + } + }); + let template = serde_json::json!({ + "spec": { + "podTemplate": { + "metadata": { + "labels": { + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE + } + }, + "spec": { + "containers": [{ + "name": "agent", + "image": "example.test/sandbox:latest" + }] + } + } + } + }); + + assert_eq!( + sandbox_spec_fingerprint(&sandbox).unwrap(), + sandbox_spec_fingerprint(&template).unwrap() + ); + } + + #[test] + fn template_fingerprint_ignores_default_volume_mount_read_only_false() { + let explicit_default = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "containers": [{ + "name": "agent", + "image": "image", + "volumeMounts": [{ + "name": "workspace", + "mountPath": "/sandbox", + "readOnly": false + }] + }], + "initContainers": [{ + "name": "openshell-supervisor-install", + "image": "supervisor", + "volumeMounts": [{ + "name": "openshell-supervisor-bin", + "mountPath": "/opt/openshell/bin", + "readOnly": false + }] + }] + } + } + } + }); + let implicit_default = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "containers": [{ + "name": "agent", + "image": "image", + "volumeMounts": [{ + "name": "workspace", + "mountPath": "/sandbox" + }] + }], + "initContainers": [{ + "name": "openshell-supervisor-install", + "image": "supervisor", + "volumeMounts": [{ + "name": "openshell-supervisor-bin", + "mountPath": "/opt/openshell/bin" + }] + }] + } + } + } + }); + + assert_eq!( + sandbox_spec_fingerprint(&explicit_default).unwrap(), + sandbox_spec_fingerprint(&implicit_default).unwrap() + ); + } + + #[test] + fn template_fingerprint_ignores_default_cluster_first_dns_policy() { + let explicit_default = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "dnsPolicy": "ClusterFirst", + "containers": [{ + "name": "agent", + "image": "image" + }] + } + } + } + }); + let implicit_default = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "containers": [{ + "name": "agent", + "image": "image" + }] + } + } + } + }); + + assert_eq!( + sandbox_spec_fingerprint(&explicit_default).unwrap(), + sandbox_spec_fingerprint(&implicit_default).unwrap() + ); + } + + #[test] + fn sandbox_from_claim_object_preserves_claim_identity() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_CLAIM_KIND); + let mut claim = DynamicObject::new("workspace-a--dev", &resource); + claim.metadata.namespace = Some("team-a".to_string()); + claim.metadata.labels = Some(BTreeMap::from([ + ( + LABEL_MANAGED_BY.to_string(), + LABEL_MANAGED_BY_VALUE.to_string(), + ), + (LABEL_SANDBOX_ID.to_string(), "sandbox-id".to_string()), + (LABEL_SANDBOX_NAME.to_string(), "dev".to_string()), + ( + LABEL_SANDBOX_WORKSPACE.to_string(), + "workspace-a".to_string(), + ), + ])); + claim.data = serde_json::json!({ + "status": { + "sandbox": { + "name": "pool-a-7x9sk" + }, + "conditions": [{ + "type": "Ready", + "status": "False", + "reason": "Binding", + "message": "waiting for warm pool" + }] + } + }); + + let sandbox = sandbox_from_claim_object("default", claim).unwrap(); + + assert_eq!(sandbox.id, "sandbox-id"); + assert_eq!(sandbox.name, "dev"); + assert_eq!(sandbox.workspace, "workspace-a"); + assert_eq!(sandbox.namespace, "team-a"); + assert_eq!( + sandbox + .status + .as_ref() + .map(|status| status.sandbox_name.as_str()), + Some("pool-a-7x9sk") + ); + } + + #[test] + fn sandbox_from_pending_claim_object_has_empty_status() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_CLAIM_KIND); + let mut claim = DynamicObject::new("workspace-a--dev", &resource); + claim.metadata.namespace = Some("team-a".to_string()); + claim.metadata.labels = Some(BTreeMap::from([ + ( + LABEL_MANAGED_BY.to_string(), + LABEL_MANAGED_BY_VALUE.to_string(), + ), + (LABEL_SANDBOX_ID.to_string(), "sandbox-id".to_string()), + (LABEL_SANDBOX_NAME.to_string(), "dev".to_string()), + ( + LABEL_SANDBOX_WORKSPACE.to_string(), + "workspace-a".to_string(), + ), + ])); + + let sandbox = sandbox_from_claim_object("default", claim).unwrap(); + + let status = sandbox.status.expect("pending claim should have status"); + assert!(status.sandbox_name.is_empty()); + assert!(status.instance_id.is_empty()); + assert!(!status.deleting); + } + fn rendered_env<'a>(container: &'a serde_json::Value, name: &str) -> Option<&'a str> { container["env"] .as_array()? diff --git a/crates/openshell-driver-kubernetes/src/lib.rs b/crates/openshell-driver-kubernetes/src/lib.rs index 8c077fc7e3..bad24c8759 100644 --- a/crates/openshell-driver-kubernetes/src/lib.rs +++ b/crates/openshell-driver-kubernetes/src/lib.rs @@ -13,7 +13,7 @@ pub use bootstrap::{ pub use config::{ AppArmorProfile, DEFAULT_PROXY_UID, DEFAULT_SANDBOX_SERVICE_ACCOUNT_NAME, DEFAULT_WORKSPACE_STORAGE_SIZE, KubernetesComputeConfig, KubernetesSidecarConfig, - SupervisorSideloadMethod, SupervisorTopology, + KubernetesWarmPoolingConfig, SupervisorSideloadMethod, SupervisorTopology, }; pub use driver::{KubernetesComputeDriver, KubernetesDriverError}; pub use grpc::ComputeDriverService; diff --git a/crates/openshell-driver-kubernetes/src/main.rs b/crates/openshell-driver-kubernetes/src/main.rs index e6e9aac33e..14e23dc0b0 100644 --- a/crates/openshell-driver-kubernetes/src/main.rs +++ b/crates/openshell-driver-kubernetes/src/main.rs @@ -12,7 +12,7 @@ use openshell_core::proto::compute::v1::compute_driver_server::ComputeDriverServ use openshell_driver_kubernetes::{ AppArmorProfile, ComputeDriverService, DEFAULT_PROXY_UID, DEFAULT_SANDBOX_SERVICE_ACCOUNT_NAME, KubernetesComputeConfig, KubernetesComputeDriver, KubernetesSidecarConfig, - SupervisorSideloadMethod, SupervisorTopology, + KubernetesWarmPoolingConfig, SupervisorSideloadMethod, SupervisorTopology, }; #[derive(Parser, Debug)] @@ -148,6 +148,7 @@ async fn main() -> Result<()> { proxy_uid: args.sidecar_proxy_uid, process_binary_aware_network_policy: args.sidecar_process_binary_aware_network_policy, }, + warm_pooling: KubernetesWarmPoolingConfig::default(), grpc_endpoint: args.grpc_endpoint.unwrap_or_default(), ssh_socket_path: args.sandbox_ssh_socket_path, client_tls_secret_name: args.client_tls_secret_name.unwrap_or_default(), diff --git a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs index 833dea9bab..2d26d95e6e 100644 --- a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs +++ b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs @@ -22,6 +22,9 @@ use tonic::Code; use tracing::{debug, info, warn}; const KUBE_API_TIMEOUT: Duration = Duration::from_secs(30); +const ACTIVATION_RETRY_ATTEMPTS: usize = 8; +const ACTIVATION_RETRY_INITIAL_DELAY: Duration = Duration::from_millis(250); +const ACTIVATION_RETRY_MAX_DELAY: Duration = Duration::from_secs(5); const DRIVER_NAME: &str = "kubernetes"; const SANDBOX_GROUP: &str = "agents.x-k8s.io"; @@ -31,15 +34,11 @@ const SANDBOX_VERSION_V1ALPHA1: &str = "v1alpha1"; const SANDBOX_VERSIONS: &[&str] = &[SANDBOX_VERSION_V1BETA1, SANDBOX_VERSION_V1ALPHA1]; const SANDBOX_API_VERSION_FULL_V1BETA1: &str = "agents.x-k8s.io/v1beta1"; const SANDBOX_API_VERSION_FULL_V1ALPHA1: &str = "agents.x-k8s.io/v1alpha1"; +const SANDBOX_POD_NAME_ANNOTATION: &str = "agents.x-k8s.io/pod-name"; const SANDBOX_CLAIM_GROUP: &str = "extensions.agents.x-k8s.io"; const SANDBOX_CLAIM_KIND: &str = "SandboxClaim"; const SANDBOX_CLAIM_VERSION_V1BETA1: &str = "v1beta1"; -const SANDBOX_CLAIM_VERSION_V1ALPHA1: &str = "v1alpha1"; -const SANDBOX_CLAIM_VERSIONS: &[&str] = &[ - SANDBOX_CLAIM_VERSION_V1BETA1, - SANDBOX_CLAIM_VERSION_V1ALPHA1, -]; /// Watches Kubernetes `SandboxClaim` resources and activates pending /// supervisor bootstrap streams after revalidating live driver state. @@ -217,53 +216,30 @@ impl SandboxClaimActivationController { } }; - match activator - .activate_registered_supervisor(request.clone()) - .await - { - Ok(()) => { - info!( - namespace = %self.namespace, - sandbox_claim = %claim.name, - sandbox = %sandbox_name, - sandbox_id = %request.sandbox_id, - instance_id = %request.instance_id, - "Activated warm-pool supervisor from SandboxClaim" - ); - } - Err(status) if status.code() == Code::AlreadyExists => { - debug!( - namespace = %self.namespace, - sandbox_claim = %claim.name, - sandbox = %sandbox_name, - sandbox_id = %request.sandbox_id, - instance_id = %request.instance_id, - "Warm-pool supervisor was already activated" - ); - } - Err(status) if status.code() == Code::NotFound => { - debug!( - namespace = %self.namespace, - sandbox_claim = %claim.name, - sandbox = %sandbox_name, - sandbox_id = %request.sandbox_id, - instance_id = %request.instance_id, - "Warm-pool supervisor registration is not pending yet" - ); - } - Err(status) => { - warn!( - namespace = %self.namespace, - sandbox_claim = %claim.name, - sandbox = %sandbox_name, - sandbox_id = %request.sandbox_id, - instance_id = %request.instance_id, - code = ?status.code(), - message = %status.message(), - "Failed to activate warm-pool supervisor from SandboxClaim" - ); - } - } + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + owner_uid = %request.owner_uid, + pod = %pod.metadata.name.as_deref().unwrap_or_default(), + "Resolved SandboxClaim activation target" + ); + + activate_registered_supervisor_with_retry( + activator, + request, + ActivationLogContext { + namespace: self.namespace.as_str(), + sandbox_claim: claim.name.as_str(), + sandbox: sandbox_name, + }, + ACTIVATION_RETRY_ATTEMPTS, + ACTIVATION_RETRY_INITIAL_DELAY, + ACTIVATION_RETRY_MAX_DELAY, + ) + .await; } fn sandbox_claim_api(&self, client: Client, version: &'static str) -> SandboxClaimApi { @@ -274,29 +250,23 @@ impl SandboxClaimActivationController { } async fn supported_sandbox_claim_api(&self, client: Client) -> Result { - for version in SANDBOX_CLAIM_VERSIONS { - let claim_api = self.sandbox_claim_api(client.clone(), version); - match tokio::time::timeout( - KUBE_API_TIMEOUT, - claim_api.api.list(&ListParams::default().limit(1)), - ) - .await - { - Ok(Ok(_)) => return Ok(claim_api), - Ok(Err(err)) if should_try_next_api_version(&err) => {} - Ok(Err(err)) => return Err(err.to_string()), - Err(_) => { - return Err(format!( - "timed out after {}s waiting for Kubernetes API", - KUBE_API_TIMEOUT.as_secs() - )); - } - } + let claim_api = self.sandbox_claim_api(client, SANDBOX_CLAIM_VERSION_V1BETA1); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + claim_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => Ok(claim_api), + Ok(Err(err)) if should_try_next_api_version(&err) => Err( + "SandboxClaim API extensions.agents.x-k8s.io/v1beta1 is not available".to_string(), + ), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), } - Err(format!( - "no supported SandboxClaim API version is available; tried {}", - SANDBOX_CLAIM_VERSIONS.join(", ") - )) } fn sandbox_api(&self, client: Client, version: &'static str) -> Api { @@ -333,6 +303,35 @@ impl SandboxClaimActivationController { }; let pods_api: Api = Api::namespaced(self.client.clone(), &self.namespace); + if let Some(pod_name) = sandbox_pod_name(sandbox_cr) { + let pod = match tokio::time::timeout(KUBE_API_TIMEOUT, pods_api.get(&pod_name)).await { + Ok(Ok(pod)) => pod, + Ok(Err(KubeError::Api(err))) if err.code == 404 => { + debug!( + namespace = %self.namespace, + sandbox = %sandbox_cr.metadata.name.as_deref().unwrap_or_default(), + sandbox_uid = %owner_uid, + pod = %pod_name, + "Annotated Sandbox pod was not found for SandboxClaim activation" + ); + return Ok(None); + } + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + if !pod_has_sandbox_owner(&pod, owner_uid) { + return Err(format!( + "annotated Sandbox pod {pod_name} is not controlled by the selected Sandbox" + )); + } + return Ok(Some(pod)); + } + let pods = match sandbox_selector(sandbox_cr) { Some(selector) => { let params = ListParams::default().labels(&selector); @@ -374,18 +373,115 @@ struct SandboxClaimApi { struct ParsedSandboxClaim { name: String, uid: String, + sandbox_id: Option, warm_pool_name: Option, sandbox_template_name: Option, sandbox_name: Option, pod_ips: Vec, } +struct ActivationLogContext<'a> { + namespace: &'a str, + sandbox_claim: &'a str, + sandbox: &'a str, +} + +async fn activate_registered_supervisor_with_retry( + activator: &(dyn SupervisorBootstrapActivator + '_), + request: SupervisorBootstrapActivationRequest, + context: ActivationLogContext<'_>, + attempts: usize, + initial_delay: Duration, + max_delay: Duration, +) { + let attempts = attempts.max(1); + let mut delay = initial_delay; + + for attempt in 1..=attempts { + match activator + .activate_registered_supervisor(request.clone()) + .await + { + Ok(()) => { + info!( + namespace = %context.namespace, + sandbox_claim = %context.sandbox_claim, + sandbox = %context.sandbox, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + attempt, + "Activated warm-pool supervisor from SandboxClaim" + ); + return; + } + Err(status) if status.code() == Code::AlreadyExists => { + debug!( + namespace = %context.namespace, + sandbox_claim = %context.sandbox_claim, + sandbox = %context.sandbox, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + attempt, + "Warm-pool supervisor was already activated" + ); + return; + } + Err(status) if status.code() == Code::NotFound && attempt < attempts => { + debug!( + namespace = %context.namespace, + sandbox_claim = %context.sandbox_claim, + sandbox = %context.sandbox, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + attempt, + retry_after_ms = delay.as_millis(), + "Warm-pool supervisor registration is not pending yet; retrying activation" + ); + tokio::time::sleep(delay).await; + delay = (delay * 2).min(max_delay); + } + Err(status) if status.code() == Code::NotFound => { + debug!( + namespace = %context.namespace, + sandbox_claim = %context.sandbox_claim, + sandbox = %context.sandbox, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + attempts, + "Warm-pool supervisor registration is not pending after retries" + ); + return; + } + Err(status) => { + warn!( + namespace = %context.namespace, + sandbox_claim = %context.sandbox_claim, + sandbox = %context.sandbox, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + attempt, + code = ?status.code(), + message = %status.message(), + "Failed to activate warm-pool supervisor from SandboxClaim" + ); + return; + } + } + } +} + fn parse_sandbox_claim(obj: &DynamicObject) -> Result { Ok(ParsedSandboxClaim { name: required_metadata_field(&obj.metadata.name, "name")?, uid: required_metadata_field(&obj.metadata.uid, "uid")?, - warm_pool_name: string_at(&obj.data, &["spec", "warmPoolRef", "name"]) - .or_else(|| string_at(&obj.data, &["spec", "warmpool"])), + sandbox_id: obj + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(LABEL_SANDBOX_ID)) + .filter(|id| !id.is_empty()) + .cloned(), + warm_pool_name: string_at(&obj.data, &["spec", "warmPoolRef", "name"]), sandbox_template_name: string_at(&obj.data, &["spec", "sandboxTemplateRef", "name"]), sandbox_name: string_at(&obj.data, &["status", "sandbox", "name"]), pod_ips: strings_at(&obj.data, &["status", "sandbox", "podIPs"]), @@ -428,7 +524,10 @@ fn activation_request_from_claim_state( .and_then(|labels| labels.get(LABEL_SANDBOX_ID)) .filter(|id| !id.is_empty()) .cloned() - .ok_or_else(|| "Sandbox CR is missing OpenShell sandbox id label".to_string())?; + .or_else(|| claim.sandbox_id.clone()) + .ok_or_else(|| { + "SandboxClaim and selected Sandbox are missing OpenShell sandbox id label".to_string() + })?; let instance_id = pod .metadata .uid @@ -464,6 +563,15 @@ fn sandbox_selector(obj: &DynamicObject) -> Option { string_at(&obj.data, &["status", "selector"]) } +fn sandbox_pod_name(obj: &DynamicObject) -> Option { + obj.metadata + .annotations + .as_ref() + .and_then(|annotations| annotations.get(SANDBOX_POD_NAME_ANNOTATION)) + .filter(|pod_name| !pod_name.is_empty()) + .cloned() +} + fn pod_has_sandbox_owner(pod: &Pod, owner_uid: &str) -> bool { pod.metadata .owner_references @@ -512,6 +620,44 @@ mod tests { use kube::core::ObjectMeta; use serde_json::json; use std::collections::BTreeMap; + use std::sync::Mutex; + use tonic::Status; + use tonic::async_trait; + + struct FakeActivator { + outcomes: Mutex>>, + requests: Mutex>, + } + + impl FakeActivator { + fn new(outcomes: Vec>) -> Self { + Self { + outcomes: Mutex::new(outcomes), + requests: Mutex::new(Vec::new()), + } + } + + fn request_count(&self) -> usize { + self.requests.lock().expect("requests mutex poisoned").len() + } + } + + #[async_trait] + impl SupervisorBootstrapActivator for FakeActivator { + async fn activate_registered_supervisor( + &self, + request: SupervisorBootstrapActivationRequest, + ) -> Result<(), Status> { + self.requests + .lock() + .expect("requests mutex poisoned") + .push(request); + self.outcomes + .lock() + .expect("outcomes mutex poisoned") + .remove(0) + } + } fn dynamic_object( group: &'static str, @@ -587,35 +733,12 @@ mod tests { assert_eq!(claim.name, "claim-a"); assert_eq!(claim.uid, "claim-uid"); + assert_eq!(claim.sandbox_id, None); assert_eq!(claim.warm_pool_name.as_deref(), Some("pool-a")); assert_eq!(claim.sandbox_name.as_deref(), Some("sandbox-a")); assert_eq!(claim.pod_ips, vec!["10.0.0.10", "fd00::10"]); } - #[test] - fn parses_v1alpha1_sandbox_claim_status() { - let obj = dynamic_object( - SANDBOX_CLAIM_GROUP, - SANDBOX_CLAIM_VERSION_V1ALPHA1, - SANDBOX_CLAIM_KIND, - "claim-a", - "claim-uid", - json!({ - "spec": { - "sandboxTemplateRef": {"name": "template-a"}, - "warmpool": "pool-a" - }, - "status": {"sandbox": {"name": "sandbox-a"}} - }), - ); - - let claim = parse_sandbox_claim(&obj).unwrap(); - - assert_eq!(claim.warm_pool_name.as_deref(), Some("pool-a")); - assert_eq!(claim.sandbox_template_name.as_deref(), Some("template-a")); - assert_eq!(claim.sandbox_name.as_deref(), Some("sandbox-a")); - } - #[test] fn claim_without_selected_sandbox_does_not_activate() { let obj = dynamic_object( @@ -684,6 +807,90 @@ mod tests { assert_eq!(request.reason, "SandboxClaim/claim-a"); } + #[test] + fn activation_request_uses_claim_sandbox_id_when_sandbox_lacks_label() { + let mut obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({"status": {"sandbox": {"name": "sandbox-a"}}}), + ); + obj.metadata.labels = Some(BTreeMap::from([( + LABEL_SANDBOX_ID.to_string(), + "openshell-sandbox".to_string(), + )])); + let claim = parse_sandbox_claim(&obj).unwrap(); + let mut sandbox = sandbox_cr("sandbox-a", "sandbox-uid", "", SANDBOX_VERSION_V1BETA1); + sandbox.metadata.labels = None; + let pod = sandbox_pod( + "pod-a", + "pod-uid", + "sandbox-uid", + SANDBOX_API_VERSION_FULL_V1BETA1, + ); + + let request = activation_request_from_claim_state(&claim, &sandbox, &pod) + .unwrap() + .unwrap(); + + assert_eq!(request.sandbox_id, "openshell-sandbox"); + } + + #[tokio::test] + async fn activation_retry_handles_registration_after_claim_binding() { + let activator = FakeActivator::new(vec![Err(Status::not_found("not pending")), Ok(())]); + let request = SupervisorBootstrapActivationRequest { + driver: DRIVER_NAME.to_string(), + instance_id: "pod-uid".to_string(), + sandbox_id: "sandbox-id".to_string(), + owner_uid: "sandbox-uid".to_string(), + reason: "SandboxClaim/claim-a".to_string(), + }; + + activate_registered_supervisor_with_retry( + &activator, + request, + ActivationLogContext { + namespace: "openshell", + sandbox_claim: "claim-a", + sandbox: "sandbox-a", + }, + 2, + Duration::ZERO, + Duration::ZERO, + ) + .await; + + assert_eq!(activator.request_count(), 2); + } + + #[test] + fn activation_request_rejects_missing_sandbox_id_labels() { + let obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({"status": {"sandbox": {"name": "sandbox-a"}}}), + ); + let claim = parse_sandbox_claim(&obj).unwrap(); + let mut sandbox = sandbox_cr("sandbox-a", "sandbox-uid", "", SANDBOX_VERSION_V1BETA1); + sandbox.metadata.labels = None; + let pod = sandbox_pod( + "pod-a", + "pod-uid", + "sandbox-uid", + SANDBOX_API_VERSION_FULL_V1BETA1, + ); + + let err = activation_request_from_claim_state(&claim, &sandbox, &pod).unwrap_err(); + + assert!(err.contains("missing OpenShell sandbox id label")); + } + #[test] fn activation_request_rejects_uncontrolled_pod() { let obj = dynamic_object( @@ -726,6 +933,22 @@ mod tests { ); } + #[test] + fn sandbox_pod_name_reads_pod_name_annotation() { + let mut sandbox = sandbox_cr( + "sandbox-a", + "sandbox-uid", + "openshell-sandbox", + SANDBOX_VERSION_V1BETA1, + ); + sandbox.metadata.annotations = Some(BTreeMap::from([( + SANDBOX_POD_NAME_ANNOTATION.to_string(), + "pod-a".to_string(), + )])); + + assert_eq!(sandbox_pod_name(&sandbox).as_deref(), Some("pod-a")); + } + #[test] fn api_version_probe_retries_only_404_errors() { let unavailable = KubeError::Api(kube::core::ErrorResponse { diff --git a/crates/openshell-sandbox/src/lib.rs b/crates/openshell-sandbox/src/lib.rs index 271f82710c..0b824d73c9 100644 --- a/crates/openshell-sandbox/src/lib.rs +++ b/crates/openshell-sandbox/src/lib.rs @@ -102,6 +102,8 @@ pub async fn run_sandbox( let (program, args) = command .split_first() .ok_or_else(|| miette::miette!("No command specified"))?; + let mut sandbox_id = sandbox_id; + let mut sandbox = sandbox; // Initialize the process-wide OCSF context early so that events emitted // during policy loading (filesystem config, validation) have a context. @@ -151,6 +153,19 @@ pub async fn run_sandbox( // Load policy and initialize OPA engine let openshell_endpoint_for_proxy = openshell_endpoint.clone(); + if sidecar_bootstrap.is_none() + && sandbox_id.is_none() + && let Some(endpoint) = openshell_endpoint.as_deref() + && std::env::var(openshell_core::sandbox_env::K8S_SA_TOKEN_FILE) + .ok() + .is_some_and(|path| !path.is_empty()) + { + let activation = openshell_core::grpc_client::register_k8s_supervisor_pod(endpoint).await?; + if sandbox.is_none() && !activation.sandbox_name.is_empty() { + sandbox = Some(activation.sandbox_name); + } + sandbox_id = Some(activation.sandbox_id); + } let sandbox_name_for_agg = sandbox.clone(); let ( mut policy, @@ -2085,7 +2100,8 @@ async fn load_policy( Err(miette::miette!( "Sandbox policy required. Provide one of:\n\ - --policy-rules and --policy-data (or OPENSHELL_POLICY_RULES and OPENSHELL_POLICY_DATA env vars)\n\ - - --sandbox-id and --openshell-endpoint (or OPENSHELL_SANDBOX_ID and OPENSHELL_ENDPOINT env vars)" + - --sandbox-id and --openshell-endpoint (or OPENSHELL_SANDBOX_ID and OPENSHELL_ENDPOINT env vars)\n\ + - OPENSHELL_ENDPOINT and OPENSHELL_K8S_SA_TOKEN_FILE for Kubernetes supervisor registration" )) } diff --git a/crates/openshell-supervisor-process/src/run.rs b/crates/openshell-supervisor-process/src/run.rs index 3cb29d1e6b..3c84fbadbc 100644 --- a/crates/openshell-supervisor-process/src/run.rs +++ b/crates/openshell-supervisor-process/src/run.rs @@ -107,6 +107,10 @@ pub async fn run_process( // tasks. By this point the orchestrator has finished privileged startup // helpers (network namespace setup, nftables probes via run_networking), // and the SSH listener and entrypoint child have not been exposed yet. + #[cfg(target_os = "linux")] + if enforcement_mode.uses_privileged_process_setup() { + crate::process::prepare_supervisor_identity_mount_namespace_from_env()?; + } crate::sandbox::apply_supervisor_startup_hardening()?; // Spawn the bypass detection monitor. It tails dmesg for nftables LOG diff --git a/deploy/helm/openshell/README.md b/deploy/helm/openshell/README.md index 8723535d1a..c12760443e 100644 --- a/deploy/helm/openshell/README.md +++ b/deploy/helm/openshell/README.md @@ -228,6 +228,7 @@ add `ci/values-spire.yaml` to the OpenShell release values files. | server.tls.certSecretName | string | `"openshell-server-tls"` | K8s secret (type kubernetes.io/tls) with tls.crt and tls.key for the server. | | server.tls.clientCaSecretName | string | `"openshell-server-client-ca"` | K8s secret with ca.crt for client certificate verification (mTLS). Set to "" to disable mTLS and run HTTPS-only (use OIDC for auth instead). | | server.tls.clientTlsSecretName | string | `"openshell-client-tls"` | K8s secret mounted into sandbox pods for mTLS to the server. | +| server.warmPooling | object | `{"enabled":true}` | Enable transparent Kubernetes warm-pool allocation through Agent Sandbox v1beta1 SandboxClaim resources when a compatible OpenShell-enabled SandboxWarmPool exists in the target namespace. | | server.workspaceDefaultStorageSize | string | `""` | Default storage size for the workspace PVC in sandbox pods. Uses Kubernetes quantity syntax (e.g. "2Gi", "10Gi", "500Mi"). Empty = built-in default (2Gi). | | service.healthPort | int | `8081` | Gateway health service port. | | service.metricsPort | int | `9090` | Gateway metrics service port. | diff --git a/deploy/helm/openshell/templates/clusterrole.yaml b/deploy/helm/openshell/templates/clusterrole.yaml index 073c8835ec..b8401ac850 100644 --- a/deploy/helm/openshell/templates/clusterrole.yaml +++ b/deploy/helm/openshell/templates/clusterrole.yaml @@ -31,3 +31,25 @@ rules: - namespaces verbs: - get + # Discover OpenShell-enabled Agent Sandbox warm pools across namespaces. + # SandboxClaim create/delete is namespaced, but the ClusterRoleBinding lets + # workspace-to-namespace mappings use warm pools outside the release namespace. + - apiGroups: + - extensions.agents.x-k8s.io + resources: + - sandboxclaims + verbs: + - create + - delete + - get + - list + - watch + - apiGroups: + - extensions.agents.x-k8s.io + resources: + - sandboxtemplates + - sandboxwarmpools + verbs: + - get + - list + - watch diff --git a/deploy/helm/openshell/templates/gateway-config.yaml b/deploy/helm/openshell/templates/gateway-config.yaml index 36a579250b..0314391459 100644 --- a/deploy/helm/openshell/templates/gateway-config.yaml +++ b/deploy/helm/openshell/templates/gateway-config.yaml @@ -145,6 +145,9 @@ data: supervisor_image_pull_policy = {{ .Values.supervisor.image.pullPolicy | quote }} {{- end }} + [openshell.drivers.kubernetes.warm_pooling] + enabled = {{ .Values.server.warmPooling.enabled }} + [openshell.drivers.kubernetes.sidecar] proxy_uid = {{ .Values.supervisor.sidecar.proxyUid | default 1337 }} process_binary_aware_network_policy = {{ .Values.supervisor.sidecar.processBinaryAwareNetworkPolicy }} diff --git a/deploy/helm/openshell/templates/warm-pool-profile-role.yaml b/deploy/helm/openshell/templates/warm-pool-profile-role.yaml new file mode 100644 index 0000000000..a01b53e5f0 --- /dev/null +++ b/deploy/helm/openshell/templates/warm-pool-profile-role.yaml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +{{- if .Values.server.warmPooling.profiles.enabled }} +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: {{ include "openshell.fullname" . }}-warm-pool-profiles + namespace: {{ default .Release.Namespace .Values.server.warmPooling.profiles.namespace }} + labels: + {{- include "openshell.labels" . | nindent 4 }} +rules: + - apiGroups: + - "" + resources: + - configmaps + verbs: + - get + - list + - watch + - patch + - apiGroups: + - "" + resources: + - events + verbs: + - create +{{- end }} diff --git a/deploy/helm/openshell/templates/warm-pool-profile-rolebinding.yaml b/deploy/helm/openshell/templates/warm-pool-profile-rolebinding.yaml new file mode 100644 index 0000000000..a3c82c98be --- /dev/null +++ b/deploy/helm/openshell/templates/warm-pool-profile-rolebinding.yaml @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +{{- if .Values.server.warmPooling.profiles.enabled }} +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: {{ include "openshell.fullname" . }}-warm-pool-profiles + namespace: {{ default .Release.Namespace .Values.server.warmPooling.profiles.namespace }} + labels: + {{- include "openshell.labels" . | nindent 4 }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: {{ include "openshell.fullname" . }}-warm-pool-profiles +subjects: + - kind: ServiceAccount + name: {{ include "openshell.serviceAccountName" . }} + namespace: {{ .Release.Namespace }} +{{- end }} diff --git a/deploy/helm/openshell/tests/gateway_config_test.yaml b/deploy/helm/openshell/tests/gateway_config_test.yaml index aee396c38f..215ed1f67a 100644 --- a/deploy/helm/openshell/tests/gateway_config_test.yaml +++ b/deploy/helm/openshell/tests/gateway_config_test.yaml @@ -159,6 +159,15 @@ tests: path: data["gateway.toml"] pattern: '(?ms)\[openshell\.drivers\.kubernetes\.sidecar\].*?process_binary_aware_network_policy\s*=\s*false' + - it: renders warm pooling config under [openshell.drivers.kubernetes.warm_pooling] + template: templates/gateway-config.yaml + set: + server.warmPooling.enabled: false + asserts: + - matchRegex: + path: data["gateway.toml"] + pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\].*?enabled\s*=\s*false' + - it: renders sandbox image pull secrets under [openshell.drivers.kubernetes] template: templates/gateway-config.yaml set: diff --git a/deploy/helm/openshell/tests/warm_pool_profiles_rbac_test.yaml b/deploy/helm/openshell/tests/warm_pool_profiles_rbac_test.yaml new file mode 100644 index 0000000000..6815a3e664 --- /dev/null +++ b/deploy/helm/openshell/tests/warm_pool_profiles_rbac_test.yaml @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +suite: warm pool profile RBAC +templates: + - templates/clusterrole.yaml + - templates/warm-pool-profile-role.yaml + - templates/warm-pool-profile-rolebinding.yaml +release: + name: openshell + namespace: gateway-ns + +tests: + - it: does not render profile namespace RBAC when profile reconciliation is explicitly disabled + template: templates/warm-pool-profile-role.yaml + set: + server.warmPooling.profiles.enabled: false + asserts: + - hasDocuments: + count: 0 + + - it: renders profile ConfigMap RBAC in the release namespace by default + template: templates/warm-pool-profile-role.yaml + asserts: + - equal: + path: metadata.namespace + value: gateway-ns + - contains: + path: rules[0].resources + content: configmaps + - contains: + path: rules[0].verbs + content: patch + + - it: renders profile ConfigMap RBAC in the configured namespace + template: templates/warm-pool-profile-rolebinding.yaml + set: + server.warmPooling.profiles.namespace: profiles-ns + asserts: + - equal: + path: metadata.namespace + value: profiles-ns + - equal: + path: subjects[0].namespace + value: gateway-ns + + - it: adds generated warm-pool write verbs to the ClusterRole when enabled + template: templates/clusterrole.yaml + asserts: + - contains: + path: rules[4].verbs + content: create + - contains: + path: rules[4].verbs + content: patch + - contains: + path: rules[4].verbs + content: delete diff --git a/deploy/helm/openshell/values.yaml b/deploy/helm/openshell/values.yaml index e89a234912..fa1c2923a6 100644 --- a/deploy/helm/openshell/values.yaml +++ b/deploy/helm/openshell/values.yaml @@ -193,6 +193,11 @@ server: # Set to a RuntimeClass name (e.g. "kata-containers", "nvidia") to apply it # to all sandboxes that don't explicitly override it. defaultRuntimeClassName: "" + # -- Enable transparent Kubernetes warm-pool allocation through Agent Sandbox + # v1beta1 SandboxClaim resources when a compatible OpenShell-enabled + # SandboxWarmPool exists in the target namespace. + warmPooling: + enabled: true # -- gRPC endpoint sandboxes call back into the gateway. Leave empty to derive # it from the chart fullname, release namespace, service port, and # disableTls flag, for example https://openshell.openshell.svc.cluster.local:8080. diff --git a/docs/kubernetes/setup.mdx b/docs/kubernetes/setup.mdx index c2fca827f1..a0bc78141a 100644 --- a/docs/kubernetes/setup.mdx +++ b/docs/kubernetes/setup.mdx @@ -216,12 +216,16 @@ The namespaced Role covers sandbox lifecycle and identity: | `""` | `events` | get, list, watch | | `""` | `pods` | get | -The ClusterRole grants node inspection and token validation: +The ClusterRole grants node inspection, token validation, namespace reads, and +warm-pool extension access: | API Group | Resource | Verbs | |---|---|---| | `authentication.k8s.io` | `tokenreviews` | create | | `""` | `nodes` | get, list, watch | +| `""` | `namespaces` | get | +| `extensions.agents.x-k8s.io` | `sandboxclaims` | create, delete, get, list, watch | +| `extensions.agents.x-k8s.io` | `sandboxtemplates`, `sandboxwarmpools` | get, list, watch | To use an existing ServiceAccount instead of creating one, set `serviceAccount.create=false` and supply its name: diff --git a/docs/reference/gateway-config.mdx b/docs/reference/gateway-config.mdx index 56f5d15348..7c69e6043d 100644 --- a/docs/reference/gateway-config.mdx +++ b/docs/reference/gateway-config.mdx @@ -294,6 +294,11 @@ provider_spiffe_workload_api_socket_path = "/spiffe-workload-api/spire-agent.soc # sandbox_uid = 1500 # sandbox_gid = 1500 +[openshell.drivers.kubernetes.warm_pooling] +# When enabled, compatible creates may use v1beta1 SandboxClaim resources +# backed by the driver's cache of OpenShell-enabled SandboxWarmPool resources. +enabled = true + [openshell.drivers.kubernetes.sidecar] # UID used by relaxed long-running network sidecars. Strict process/binary-aware # sidecars run as UID 0 so Kubernetes grants the required /proc inspection diff --git a/docs/reference/sandbox-compute-drivers.mdx b/docs/reference/sandbox-compute-drivers.mdx index 280849b5f6..5790ea3210 100644 --- a/docs/reference/sandbox-compute-drivers.mdx +++ b/docs/reference/sandbox-compute-drivers.mdx @@ -340,6 +340,16 @@ process/binary identity through `/proc/`. The Kubernetes driver creates namespaced `agents.x-k8s.io` `Sandbox` resources from the Kubernetes SIG Apps [agent-sandbox](https://github.com/kubernetes-sigs/agent-sandbox) project. It detects the served Sandbox API at runtime, caches the selected API version for the gateway process, and uses `v1beta1` when available before falling back to `v1alpha1`, so supported Agent Sandbox installations work without version-specific operator configuration. The Agent Sandbox controller turns those resources into sandbox pods and related storage. +When Kubernetes warm pooling is enabled, the driver can satisfy compatible +create requests by creating `extensions.agents.x-k8s.io/v1beta1` +`SandboxClaim` resources instead of direct `Sandbox` resources. The driver +discovers `SandboxWarmPool` resources labelled `openshell.ai/enabled=true` +across namespaces it can observe, resolves each pool's `SandboxTemplate`, and +matches only against pools in the Kubernetes namespace selected for the +OpenShell workspace. Set +`openshell.drivers.kubernetes.warm_pooling.enabled = false` to force direct +`Sandbox` creation. + If Agent Sandbox is upgraded in place, restart the OpenShell gateway after the controller and CRD rollout completes so the gateway can detect the served API versions again. From 1a6a4b1809a2abbe03b46df5f6a35dba7f5cb408 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Fri, 24 Jul 2026 19:00:30 +0100 Subject: [PATCH 08/10] feat(kubernetes): add support for automatic warm pool and template creation Signed-off-by: Gordon Sim --- Cargo.lock | 1 + architecture/compute-runtimes.md | 6 + crates/openshell-driver-kubernetes/Cargo.toml | 1 + crates/openshell-driver-kubernetes/README.md | 9 + .../openshell-driver-kubernetes/src/config.rs | 93 +- .../openshell-driver-kubernetes/src/driver.rs | 1291 ++++++++++++++++- deploy/helm/openshell/README.md | 6 +- .../helm/openshell/templates/clusterrole.yaml | 6 + .../openshell/templates/gateway-config.yaml | 6 + .../openshell/tests/gateway_config_test.yaml | 21 + deploy/helm/openshell/values.yaml | 12 + docs/reference/gateway-config.mdx | 7 + docs/reference/sandbox-compute-drivers.mdx | 25 + .../cpu-0-2-configmap.yaml | 18 + .../default-configmap.yaml | 15 + .../env-foo-configmap.yaml | 18 + 16 files changed, 1467 insertions(+), 68 deletions(-) create mode 100644 examples/kubernetes-warm-pool-config/cpu-0-2-configmap.yaml create mode 100644 examples/kubernetes-warm-pool-config/default-configmap.yaml create mode 100644 examples/kubernetes-warm-pool-config/env-foo-configmap.yaml diff --git a/Cargo.lock b/Cargo.lock index 56e0d443e6..c60f93d1f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3948,6 +3948,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-stream", + "toml", "tonic", "tracing", "tracing-subscriber", diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index 43539e1322..cd426d7822 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -140,6 +140,12 @@ The driver watches warm pools and templates into a local cache so create requests do not list or fetch extension objects on the hot path. This remains behind the compute-driver boundary: other drivers and the driver-agnostic gateway lifecycle continue to operate on OpenShell sandbox identity and status. +The Kubernetes driver can optionally run a separate ConfigMap profile +reconciler that turns admin-authored TOML profiles into generated +`SandboxTemplate` and `SandboxWarmPool` resources. That reconciler is +independent of the create-path matcher: admins can use it, Helm, kubectl, or +another controller to create warm pools, and the matcher only considers the +resulting enabled warm-pool resources and their fingerprints. Kubernetes can run the supervisor in the default combined topology or in a sidecar topology. Combined mode keeps network and process supervision in the diff --git a/crates/openshell-driver-kubernetes/Cargo.toml b/crates/openshell-driver-kubernetes/Cargo.toml index 7ba2794b8a..6a1a0b2bcc 100644 --- a/crates/openshell-driver-kubernetes/Cargo.toml +++ b/crates/openshell-driver-kubernetes/Cargo.toml @@ -29,6 +29,7 @@ kube-runtime = { workspace = true } k8s-openapi = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +toml = { workspace = true } sha2 = { workspace = true } clap = { workspace = true } tracing = { workspace = true } diff --git a/crates/openshell-driver-kubernetes/README.md b/crates/openshell-driver-kubernetes/README.md index b0a3847f8d..9af04155d0 100644 --- a/crates/openshell-driver-kubernetes/README.md +++ b/crates/openshell-driver-kubernetes/README.md @@ -46,6 +46,15 @@ The extension CRDs are intentionally v1beta1-only. The core `agents.x-k8s.io/Sandbox` CRD still supports the existing v1beta1-to-v1alpha1 fallback. +The driver can also reconcile admin-authored warm-pool profiles from labelled +ConfigMaps when `warm_pooling.profiles.enabled` is true. Profiles live in the +configured profile namespace, use the `warm-pool.toml` data key, and expose a +narrow CLI-shaped schema: `workspace`, `replicas`, `image`, +`runtime_class_name`, `[environment]`, and `[resources]` with `cpu`, `memory`, +and `gpu_count`. The reconciler generates ordinary `SandboxTemplate` and +`SandboxWarmPool` resources labelled `openshell.ai/enabled=true`; the existing +matching cache then handles allocation unchanged. + ## Workspace Persistence Sandbox pods use a PVC-backed `/sandbox` workspace. An init container seeds the diff --git a/crates/openshell-driver-kubernetes/src/config.rs b/crates/openshell-driver-kubernetes/src/config.rs index 35c12cfc5d..32bccd940f 100644 --- a/crates/openshell-driver-kubernetes/src/config.rs +++ b/crates/openshell-driver-kubernetes/src/config.rs @@ -18,6 +18,12 @@ pub const DEFAULT_WORKSPACE_STORAGE_SIZE: &str = "2Gi"; /// Default non-root UID for relaxed Kubernetes network supervisor sidecars. pub const DEFAULT_PROXY_UID: u32 = 1337; +/// Label selector for ConfigMap-backed warm-pool profiles. +pub const DEFAULT_WARM_POOL_PROFILE_LABEL_SELECTOR: &str = "openshell.ai/warm-pool-profile=true"; + +/// Conservative upper bound for profile-declared warm-pool replicas. +pub const DEFAULT_WARM_POOL_PROFILE_MAX_REPLICAS: u32 = 100; + /// How the supervisor binary is delivered into sandbox pods. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] @@ -131,11 +137,73 @@ pub struct KubernetesWarmPoolingConfig { /// Allow the Kubernetes driver to satisfy compatible create requests by /// creating v1beta1 Agent Sandbox `SandboxClaim` resources. pub enabled: bool, + /// Optional ConfigMap-backed warm-pool profile reconciler. + pub profiles: KubernetesWarmPoolProfilesConfig, } impl Default for KubernetesWarmPoolingConfig { fn default() -> Self { - Self { enabled: true } + Self { + enabled: true, + profiles: KubernetesWarmPoolProfilesConfig::default(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct KubernetesWarmPoolProfilesConfig { + /// Watch labelled `ConfigMap` resources and reconcile generated + /// `SandboxTemplate` and `SandboxWarmPool` resources from their TOML + /// profiles. + pub enabled: bool, + /// Namespace containing admin-authored warm-pool profile `ConfigMap` + /// resources. When empty, the driver uses its configured Kubernetes + /// namespace. + pub namespace: String, + /// Label selector used to select warm-pool profile `ConfigMap` resources. + pub label_selector: String, + /// Maximum allowed `replicas` value in a profile. + pub max_replicas: u32, +} + +impl Default for KubernetesWarmPoolProfilesConfig { + fn default() -> Self { + Self { + enabled: true, + namespace: String::new(), + label_selector: DEFAULT_WARM_POOL_PROFILE_LABEL_SELECTOR.to_string(), + max_replicas: DEFAULT_WARM_POOL_PROFILE_MAX_REPLICAS, + } + } +} + +impl KubernetesWarmPoolProfilesConfig { + #[must_use] + pub fn effective_namespace<'a>(&'a self, fallback_namespace: &'a str) -> &'a str { + if self.namespace.trim().is_empty() { + fallback_namespace + } else { + self.namespace.as_str() + } + } + + #[must_use] + pub fn effective_label_selector(&self) -> &str { + if self.label_selector.trim().is_empty() { + DEFAULT_WARM_POOL_PROFILE_LABEL_SELECTOR + } else { + self.label_selector.as_str() + } + } + + #[must_use] + pub fn effective_max_replicas(&self) -> u32 { + if self.max_replicas == 0 { + DEFAULT_WARM_POOL_PROFILE_MAX_REPLICAS + } else { + self.max_replicas + } } } @@ -554,6 +622,7 @@ mod tests { fn default_warm_pooling_is_enabled() { let cfg = KubernetesComputeConfig::default(); assert!(cfg.warm_pooling.enabled); + assert!(cfg.warm_pooling.profiles.enabled); } #[test] @@ -567,6 +636,28 @@ mod tests { assert!(!cfg.warm_pooling.enabled); } + #[test] + fn serde_override_warm_pool_profile_reconciliation() { + let json = serde_json::json!({ + "warm_pooling": { + "profiles": { + "enabled": false, + "namespace": "profiles", + "label_selector": "openshell.ai/warm-pool-profile=true,tier=gpu", + "max_replicas": 7 + } + } + }); + let cfg: KubernetesComputeConfig = serde_json::from_value(json).unwrap(); + assert!(!cfg.warm_pooling.profiles.enabled); + assert_eq!(cfg.warm_pooling.profiles.namespace, "profiles"); + assert_eq!( + cfg.warm_pooling.profiles.label_selector, + "openshell.ai/warm-pool-profile=true,tier=gpu" + ); + assert_eq!(cfg.warm_pooling.profiles.max_replicas, 7); + } + #[test] fn serde_rejects_unknown_warm_pooling_field() { let json = serde_json::json!({ diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index cee6939cc6..6329ee9489 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -11,9 +11,12 @@ use crate::config::{ }; use futures::{Stream, StreamExt, TryStreamExt, stream}; use k8s_openapi::api::core::v1::{ - Event as KubeEventObj, Namespace, Node, PersistentVolumeClaimVolumeSource, Volume, VolumeMount, + ConfigMap, Event as KubeEventObj, Namespace, Node, PersistentVolumeClaimVolumeSource, Volume, + VolumeMount, +}; +use kube::api::{ + Api, ApiResource, DeleteParams, ListParams, Patch, PatchParams, PostParams, Preconditions, }; -use kube::api::{Api, ApiResource, DeleteParams, ListParams, PostParams, Preconditions}; use kube::core::gvk::GroupVersionKind; use kube::core::{DynamicObject, ObjectMeta}; use kube::runtime::watcher::{self, Event}; @@ -30,11 +33,11 @@ use openshell_core::progress::{ }; use openshell_core::proto::compute::v1::{ DriverCondition as SandboxCondition, DriverPlatformEvent as PlatformEvent, - DriverSandbox as Sandbox, DriverSandboxSpec as SandboxSpec, + DriverResourceRequirements, DriverSandbox as Sandbox, DriverSandboxSpec as SandboxSpec, DriverSandboxStatus as SandboxStatus, DriverSandboxTemplate as SandboxTemplate, - GetCapabilitiesResponse, GpuResourceRequirements, WatchSandboxesDeletedEvent, - WatchSandboxesEvent, WatchSandboxesPlatformEvent, WatchSandboxesSandboxEvent, - watch_sandboxes_event, + GetCapabilitiesResponse, GpuResourceRequirements, ResourceRequirements, + WatchSandboxesDeletedEvent, WatchSandboxesEvent, WatchSandboxesPlatformEvent, + WatchSandboxesSandboxEvent, watch_sandboxes_event, }; use openshell_core::proto_struct::{struct_to_json_object, value_to_json}; use serde::Deserialize; @@ -103,7 +106,28 @@ const SANDBOX_TEMPLATE_KIND: &str = "SandboxTemplate"; const LABEL_WARM_POOL_ENABLED: &str = "openshell.ai/enabled"; const LABEL_ALLOCATION: &str = "openshell.ai/allocation"; const LABEL_ALLOCATION_SANDBOX_CLAIM: &str = "sandbox-claim"; +const LABEL_WARM_POOL_PROFILE: &str = "openshell.ai/warm-pool-profile"; +const LABEL_WARM_POOL_PROFILE_UID: &str = "openshell.ai/warm-pool-profile-uid"; +const LABEL_WARM_POOL_MANAGED_BY: &str = "openshell.ai/managed-by"; +const LABEL_WARM_POOL_MANAGED_BY_VALUE: &str = "openshell-kubernetes-driver"; +const ANNOTATION_WARM_POOL_PROFILE_NAMESPACE: &str = "openshell.ai/warm-pool-profile-namespace"; +const ANNOTATION_WARM_POOL_PROFILE_NAME: &str = "openshell.ai/warm-pool-profile-name"; +const ANNOTATION_WARM_POOL_PROFILE_UID: &str = "openshell.ai/warm-pool-profile-uid"; +const ANNOTATION_WARM_POOL_SOURCE_RESOURCE_VERSION: &str = "openshell.ai/source-resource-version"; +const ANNOTATION_WARM_POOL_TEMPLATE_FINGERPRINT: &str = "openshell.ai/template-fingerprint"; +const ANNOTATION_WARM_POOL_STATUS: &str = "openshell.ai/warm-pool-status"; +const ANNOTATION_WARM_POOL_OBSERVED_RESOURCE_VERSION: &str = + "openshell.ai/warm-pool-observed-resource-version"; +const ANNOTATION_WARM_POOL_TARGET_NAMESPACE: &str = "openshell.ai/warm-pool-target-namespace"; +const ANNOTATION_WARM_POOL_TEMPLATE: &str = "openshell.ai/warm-pool-template"; +const ANNOTATION_WARM_POOL_FINGERPRINT: &str = "openshell.ai/warm-pool-fingerprint"; +const ANNOTATION_WARM_POOL_LAST_ERROR: &str = "openshell.ai/warm-pool-last-error"; const POD_ANNOTATION_SANDBOX_ID: &str = "openshell.io/sandbox-id"; +const WARM_POOL_PROFILE_DATA_KEY: &str = "warm-pool.toml"; +const WARM_POOL_PROFILE_VERSION: u32 = 1; +const WARM_POOL_PROFILE_REQUEUE_DELAY: Duration = Duration::from_secs(10); +const WARM_POOL_PROFILE_RESYNC_DELAY: Duration = Duration::from_secs(300); +const WARM_POOL_PROFILE_NAME_PREFIX: &str = "openshell-wp"; const GPU_RESOURCE_NAME: &str = "nvidia.com/gpu"; const SPIFFE_WORKLOAD_API_VOLUME_NAME: &str = "spiffe-workload-api"; @@ -135,6 +159,56 @@ struct WarmPoolCacheEntry { template_fingerprint: String, } +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct WarmPoolProfile { + version: u32, + workspace: String, + #[serde(default = "default_warm_pool_profile_replicas")] + replicas: u32, + #[serde(default)] + image: String, + #[serde(default)] + runtime_class_name: String, + #[serde(default)] + environment: BTreeMap, + #[serde(default)] + resources: WarmPoolProfileResources, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +struct WarmPoolProfileResources { + #[serde(default)] + cpu: String, + #[serde(default)] + memory: String, + gpu_count: Option, +} + +fn default_warm_pool_profile_replicas() -> u32 { + 1 +} + +#[derive(Debug, Clone)] +struct WarmPoolProfileSource { + namespace: String, + name: String, + uid: String, + resource_version: String, +} + +#[derive(Debug, Clone)] +struct RenderedWarmPoolProfile { + source: WarmPoolProfileSource, + workspace: String, + target_namespace: String, + generated_name: String, + replicas: u32, + template_spec: serde_json::Value, + fingerprint: String, +} + #[derive(Debug, Clone, Default)] struct WarmPoolCache { state: Arc>, @@ -592,6 +666,9 @@ impl KubernetesComputeDriver { }; if driver.config.warm_pooling.enabled { driver.spawn_warm_pool_cache_controller(); + if driver.config.warm_pooling.profiles.enabled { + driver.spawn_warm_pool_profile_reconciler(); + } } Ok(driver) } @@ -685,6 +762,16 @@ impl KubernetesComputeDriver { }); } + fn spawn_warm_pool_profile_reconciler(&self) { + let client = self.client.clone(); + let watch_client = self.watch_client.clone(); + let config = self.config.clone(); + + tokio::spawn(async move { + run_warm_pool_profile_reconciler(client, watch_client, config).await; + }); + } + async fn supported_agent_sandbox_api(&self, client: Client) -> Result { let sandbox_api_version = self.supported_sandbox_api_version(client.clone()).await?; Ok(self.agent_sandbox_api(client, sandbox_api_version)) @@ -751,66 +838,12 @@ impl KubernetesComputeDriver { /// annotations. /// - If neither config nor `OpenShift` is found, returns `(1000, 1000, {})` as defaults. async fn resolve_sandbox_identity(&self) -> (u32, u32, BTreeMap) { - // Explicit config takes priority — skip namespace lookup entirely. - if self.config.sandbox_uid.is_some() { - let uid = self.config.resolve_sandbox_uid(None); - let gid = self.config.resolve_sandbox_gid(uid, None); - return (uid, gid, BTreeMap::new()); - } - - // Try to read namespace annotations for OpenShift SCC. - // Namespace is namespaced so Api::all works (it's cluster-scoped but - // can list all namespaces) and we filter by name, or use Api::namespaced. - let ns_api: Api = Api::all(self.client.clone()); - match tokio::time::timeout(KUBE_API_TIMEOUT, ns_api.get(self.config.namespace.as_str())) - .await - { - Ok(Ok(ns)) => { - let anns = ns.metadata.annotations.unwrap_or_default(); - tracing::info!( - namespace = %self.config.namespace, - uid_range = ?anns.get(crate::config::ANNOTATION_SCC_UID_RANGE), - sup_groups = ?anns.get(crate::config::ANNOTATION_SCC_SUPPLEMENTAL_GROUPS), - "Resolved namespace annotations for sandbox identity" - ); - let uid = self.config.resolve_sandbox_uid(Some(&anns)); - // Explicit sandbox_gid config wins; SCC annotation only applies when not set. - let baseline_gid = self.config.resolve_sandbox_gid(uid, None); - let gid = self.config.sandbox_gid.map_or_else( - || { - anns.get(crate::config::ANNOTATION_SCC_SUPPLEMENTAL_GROUPS) - .and_then(|sup_range| { - KubernetesComputeConfig::from_open_shift_supplemental_groups( - sup_range, - ) - }) - .unwrap_or(baseline_gid) - }, - |_| baseline_gid, - ); - tracing::info!(uid, gid, "Resolved sandbox identity"); - (uid, gid, anns) - } - Ok(Err(e)) => { - tracing::warn!( - namespace = %self.config.namespace, - error = %e, - "Failed to fetch namespace for SCC annotations, falling back to defaults" - ); - let uid = DEFAULT_SANDBOX_UID; - let gid = self.config.resolve_sandbox_gid(uid, None); - (uid, gid, BTreeMap::new()) - } - Err(_) => { - tracing::warn!( - namespace = %self.config.namespace, - "Namespace fetch timed out, falling back to defaults" - ); - let uid = DEFAULT_SANDBOX_UID; - let gid = self.config.resolve_sandbox_gid(uid, None); - (uid, gid, BTreeMap::new()) - } - } + resolve_sandbox_identity_for_config( + self.client.clone(), + &self.config, + self.config.namespace.as_str(), + ) + .await } async fn has_gpu_capacity(&self) -> Result { @@ -2028,6 +2061,937 @@ async fn refresh_warm_pool_cache(cache: &WarmPoolCache, client: Client, fallback ); } +async fn run_warm_pool_profile_reconciler( + client: Client, + watch_client: Client, + config: KubernetesComputeConfig, +) { + let profile_namespace = config + .warm_pooling + .profiles + .effective_namespace(&config.namespace) + .to_string(); + let label_selector = config + .warm_pooling + .profiles + .effective_label_selector() + .to_string(); + + loop { + reconcile_existing_warm_pool_profiles( + client.clone(), + &config, + &profile_namespace, + &label_selector, + ) + .await; + + let api: Api = Api::namespaced(watch_client.clone(), &profile_namespace); + let watcher_config = watcher::Config::default().labels(&label_selector); + let mut stream = watcher::watcher(api, watcher_config).boxed(); + + loop { + tokio::select! { + event = stream.try_next() => { + match event { + Ok(Some(Event::Applied(config_map))) => { + reconcile_warm_pool_profile_config_map(client.clone(), &config, config_map).await; + } + Ok(Some(Event::Deleted(config_map))) => { + garbage_collect_warm_pool_profile(client.clone(), &config, &config_map).await; + } + Ok(Some(Event::Restarted(config_maps))) => { + for config_map in config_maps { + reconcile_warm_pool_profile_config_map(client.clone(), &config, config_map).await; + } + } + Ok(None) => { + warn!( + namespace = %profile_namespace, + "Warm-pool profile watch stream ended; retrying" + ); + break; + } + Err(err) => { + warn!( + namespace = %profile_namespace, + error = %err, + "Warm-pool profile watch failed; retrying" + ); + break; + } + } + } + () = tokio::time::sleep(WARM_POOL_PROFILE_RESYNC_DELAY) => { + reconcile_existing_warm_pool_profiles( + client.clone(), + &config, + &profile_namespace, + &label_selector, + ) + .await; + } + } + } + + tokio::time::sleep(WARM_POOL_PROFILE_REQUEUE_DELAY).await; + } +} + +async fn resolve_sandbox_identity_for_config( + client: Client, + config: &KubernetesComputeConfig, + namespace: &str, +) -> (u32, u32, BTreeMap) { + // Explicit config takes priority — skip namespace lookup entirely. + if config.sandbox_uid.is_some() { + let uid = config.resolve_sandbox_uid(None); + let gid = config.resolve_sandbox_gid(uid, None); + return (uid, gid, BTreeMap::new()); + } + + let ns_api: Api = Api::all(client); + match tokio::time::timeout(KUBE_API_TIMEOUT, ns_api.get(namespace)).await { + Ok(Ok(ns)) => { + let anns = ns.metadata.annotations.unwrap_or_default(); + tracing::info!( + namespace, + uid_range = ?anns.get(crate::config::ANNOTATION_SCC_UID_RANGE), + sup_groups = ?anns.get(crate::config::ANNOTATION_SCC_SUPPLEMENTAL_GROUPS), + "Resolved namespace annotations for sandbox identity" + ); + let uid = config.resolve_sandbox_uid(Some(&anns)); + let baseline_gid = config.resolve_sandbox_gid(uid, None); + let gid = config.sandbox_gid.map_or_else( + || { + anns.get(crate::config::ANNOTATION_SCC_SUPPLEMENTAL_GROUPS) + .and_then(|sup_range| { + KubernetesComputeConfig::from_open_shift_supplemental_groups(sup_range) + }) + .unwrap_or(baseline_gid) + }, + |_| baseline_gid, + ); + tracing::info!(uid, gid, "Resolved sandbox identity"); + (uid, gid, anns) + } + Ok(Err(e)) => { + tracing::warn!( + namespace, + error = %e, + "Failed to fetch namespace for SCC annotations, falling back to defaults" + ); + let uid = DEFAULT_SANDBOX_UID; + let gid = config.resolve_sandbox_gid(uid, None); + (uid, gid, BTreeMap::new()) + } + Err(_) => { + tracing::warn!( + namespace, + "Namespace fetch timed out, falling back to defaults" + ); + let uid = DEFAULT_SANDBOX_UID; + let gid = config.resolve_sandbox_gid(uid, None); + (uid, gid, BTreeMap::new()) + } + } +} + +async fn reconcile_existing_warm_pool_profiles( + client: Client, + config: &KubernetesComputeConfig, + profile_namespace: &str, + label_selector: &str, +) { + let api: Api = Api::namespaced(client.clone(), profile_namespace); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + api.list(&ListParams::default().labels(label_selector)), + ) + .await + { + Ok(Ok(list)) => { + let active_source_uids = list + .items + .iter() + .filter_map(|config_map| config_map.metadata.uid.clone()) + .collect::>(); + for config_map in list.items { + reconcile_warm_pool_profile_config_map(client.clone(), config, config_map).await; + } + if let Err(err) = garbage_collect_orphaned_warm_pool_profile_resources( + client, + &config.namespace, + &active_source_uids, + ) + .await + { + warn!( + namespace = %config.namespace, + error = %err, + "Failed to garbage-collect orphaned warm-pool profile resources" + ); + } + } + Ok(Err(err)) => { + warn!( + namespace = %profile_namespace, + error = %err, + "Failed to list warm-pool profile ConfigMaps" + ); + } + Err(_) => { + warn!( + namespace = %profile_namespace, + "Timed out listing warm-pool profile ConfigMaps" + ); + } + } +} + +async fn reconcile_warm_pool_profile_config_map( + client: Client, + config: &KubernetesComputeConfig, + config_map: ConfigMap, +) { + match render_warm_pool_profile(client.clone(), config, &config_map).await { + Ok(rendered) => { + if let Err(err) = apply_rendered_warm_pool_profile(client.clone(), &rendered).await { + warn!( + namespace = %rendered.source.namespace, + config_map = %rendered.source.name, + error = %err, + "Failed to apply warm-pool profile" + ); + patch_warm_pool_profile_status( + client, + &rendered.source.namespace, + &rendered.source.name, + "Error", + Some(&rendered), + &err, + ) + .await; + return; + } + if let Err(err) = + garbage_collect_superseded_warm_pool_profile_resources(client.clone(), &rendered) + .await + { + warn!( + namespace = %rendered.source.namespace, + config_map = %rendered.source.name, + error = %err, + "Failed to garbage-collect superseded warm-pool profile resources" + ); + } + patch_warm_pool_profile_status( + client, + &rendered.source.namespace, + &rendered.source.name, + "Ready", + Some(&rendered), + "", + ) + .await; + } + Err(err) => { + let namespace = config_map + .metadata + .namespace + .as_deref() + .unwrap_or_else(|| { + config + .warm_pooling + .profiles + .effective_namespace(&config.namespace) + }) + .to_string(); + let Some(name) = config_map.metadata.name.as_deref() else { + warn!(error = %err, "Warm-pool profile ConfigMap is missing a name"); + return; + }; + warn!( + namespace = %namespace, + config_map = %name, + error = %err, + "Warm-pool profile is invalid" + ); + patch_warm_pool_profile_status(client, &namespace, name, "Error", None, &err).await; + } + } +} + +async fn render_warm_pool_profile( + client: Client, + config: &KubernetesComputeConfig, + config_map: &ConfigMap, +) -> Result { + let source = warm_pool_profile_source(config, config_map)?; + let data = config_map + .data + .as_ref() + .and_then(|data| data.get(WARM_POOL_PROFILE_DATA_KEY)) + .ok_or_else(|| format!("missing data key '{WARM_POOL_PROFILE_DATA_KEY}'"))?; + let profile = parse_warm_pool_profile(data)?; + validate_warm_pool_profile( + &profile, + config.warm_pooling.profiles.effective_max_replicas(), + )?; + + let target_namespace = config.namespace.clone(); + let spec = warm_pool_profile_to_sandbox_spec(&profile); + let (profile_user_id, profile_group_id, _) = + resolve_sandbox_identity_for_config(client, config, &target_namespace).await; + let params = SandboxPodParams { + default_image: &config.default_image, + image_pull_policy: &config.image_pull_policy, + image_pull_secrets: &config.image_pull_secrets, + supervisor_image: &config.supervisor_image, + supervisor_image_pull_policy: &config.supervisor_image_pull_policy, + supervisor_sideload_method: config.supervisor_sideload_method, + topology: config.topology, + proxy_uid: config.sidecar.proxy_uid, + process_binary_aware_network_policy: config.sidecar.process_binary_aware_network_policy, + service_account_name: &config.service_account_name, + sandbox_id: "", + sandbox_name: "", + grpc_endpoint: &config.grpc_endpoint, + ssh_socket_path: &config.ssh_socket_path, + client_tls_secret_name: &config.client_tls_secret_name, + host_gateway_ip: &config.host_gateway_ip, + enable_user_namespaces: config.enable_user_namespaces, + app_armor_profile: config.app_armor_profile.as_ref(), + workspace_default_storage_size: &config.workspace_default_storage_size, + default_runtime_class_name: &config.default_runtime_class_name, + sa_token_ttl_secs: config.effective_sa_token_ttl_secs(), + provider_spiffe_enabled: config.provider_spiffe_enabled(), + provider_spiffe_workload_api_socket_path: &config.provider_spiffe_workload_api_socket_path, + sandbox_uid: profile_user_id, + sandbox_gid: profile_group_id, + }; + validate_sidecar_proxy_identity(¶ms).map_err(|err| err.to_string())?; + + let mut rendered = sandbox_to_k8s_spec(Some(&spec), ¶ms)?; + remove_warm_pool_template_identity_env(&mut rendered); + ensure_pod_dns_policy(&mut rendered); + let fingerprint = sandbox_spec_fingerprint(&rendered)?; + let generated_name = generated_warm_pool_profile_name(&source.name, &fingerprint); + validate_kubernetes_dns1123_label(&generated_name, "generated warm-pool resource name")?; + let template_spec = rendered + .get("spec") + .cloned() + .ok_or_else(|| "rendered sandbox spec is missing spec".to_string())?; + + Ok(RenderedWarmPoolProfile { + source, + workspace: profile.workspace, + target_namespace, + generated_name, + replicas: profile.replicas, + template_spec, + fingerprint, + }) +} + +fn warm_pool_profile_source( + config: &KubernetesComputeConfig, + config_map: &ConfigMap, +) -> Result { + let namespace = config_map.metadata.namespace.clone().unwrap_or_else(|| { + config + .warm_pooling + .profiles + .effective_namespace(&config.namespace) + .to_string() + }); + let name = config_map + .metadata + .name + .clone() + .ok_or_else(|| "ConfigMap metadata.name is required".to_string())?; + let uid = config_map + .metadata + .uid + .clone() + .ok_or_else(|| "ConfigMap metadata.uid is required".to_string())?; + let resource_version = config_map + .metadata + .resource_version + .clone() + .unwrap_or_default(); + Ok(WarmPoolProfileSource { + namespace, + name, + uid, + resource_version, + }) +} + +fn parse_warm_pool_profile(data: &str) -> Result { + toml::from_str::(data) + .map_err(|err| format!("failed to parse {WARM_POOL_PROFILE_DATA_KEY}: {err}")) +} + +fn validate_warm_pool_profile(profile: &WarmPoolProfile, max_replicas: u32) -> Result<(), String> { + if profile.version != WARM_POOL_PROFILE_VERSION { + return Err(format!( + "unsupported warm-pool profile version {}; expected {WARM_POOL_PROFILE_VERSION}", + profile.version + )); + } + if profile.workspace.trim().is_empty() { + return Err("workspace must not be empty".to_string()); + } + if profile.workspace.chars().any(char::is_control) { + return Err("workspace must not contain control characters".to_string()); + } + if profile.replicas == 0 { + return Err("replicas must be greater than 0".to_string()); + } + if profile.replicas > max_replicas { + return Err(format!( + "replicas exceeds maximum ({} > {max_replicas})", + profile.replicas + )); + } + validate_warm_pool_profile_environment(&profile.environment)?; + if let Some(count) = profile.resources.gpu_count + && count == 0 + { + return Err("resources.gpu_count must be greater than 0".to_string()); + } + Ok(()) +} + +fn validate_warm_pool_profile_environment(env: &BTreeMap) -> Result<(), String> { + for (key, value) in env { + if !is_valid_env_key(key) { + return Err(format!( + "environment keys must match ^[A-Za-z_][A-Za-z0-9_]*$; got '{key}'" + )); + } + if key.starts_with("OPENSHELL_") { + return Err(format!( + "environment keys starting with OPENSHELL_ are reserved; got '{key}'" + )); + } + if value.chars().any(char::is_control) { + return Err(format!( + "environment value for '{key}' must not contain control characters" + )); + } + } + Ok(()) +} + +fn is_valid_env_key(key: &str) -> bool { + let mut chars = key.chars(); + let Some(first) = chars.next() else { + return false; + }; + (first == '_' || first.is_ascii_alphabetic()) + && chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric()) +} + +fn warm_pool_profile_to_sandbox_spec(profile: &WarmPoolProfile) -> SandboxSpec { + let resources = (!profile.resources.cpu.is_empty() || !profile.resources.memory.is_empty()) + .then(|| DriverResourceRequirements { + cpu_limit: profile.resources.cpu.clone(), + memory_limit: profile.resources.memory.clone(), + ..Default::default() + }); + + SandboxSpec { + environment: profile.environment.clone().into_iter().collect(), + template: Some(SandboxTemplate { + image: profile.image.clone(), + platform_config: platform_config_for_warm_pool_profile(profile), + resources, + ..Default::default() + }), + resource_requirements: profile + .resources + .gpu_count + .map(|count| ResourceRequirements { + gpu: Some(GpuResourceRequirements { count: Some(count) }), + }), + ..Default::default() + } +} + +fn platform_config_for_warm_pool_profile(profile: &WarmPoolProfile) -> Option { + use prost_types::{Struct, Value, value::Kind}; + + if profile.runtime_class_name.is_empty() { + return None; + } + + Some(Struct { + fields: BTreeMap::from([( + "runtime_class_name".to_string(), + Value { + kind: Some(Kind::StringValue(profile.runtime_class_name.clone())), + }, + )]), + }) +} + +fn ensure_pod_dns_policy(rendered: &mut serde_json::Value) { + if let Some(spec) = rendered + .pointer_mut("/spec/podTemplate/spec") + .and_then(serde_json::Value::as_object_mut) + { + spec.entry("dnsPolicy".to_string()) + .or_insert_with(|| serde_json::json!("ClusterFirst")); + } +} + +fn remove_warm_pool_template_identity_env(rendered: &mut serde_json::Value) { + if let Some(containers) = rendered + .pointer_mut("/spec/podTemplate/spec/containers") + .and_then(serde_json::Value::as_array_mut) + { + for container in containers { + remove_sandbox_identity_env(container); + } + } +} + +fn generated_warm_pool_profile_name(profile_name: &str, fingerprint: &str) -> String { + let hash = fingerprint + .strip_prefix("sha256:") + .unwrap_or(fingerprint) + .chars() + .filter(char::is_ascii_hexdigit) + .take(8) + .collect::() + .to_ascii_lowercase(); + let suffix = if hash.is_empty() { + "00000000".to_string() + } else { + hash + }; + let prefix = sanitize_dns_label_segment(profile_name); + let reserved = WARM_POOL_PROFILE_NAME_PREFIX.len() + 1 + 1 + suffix.len(); + let max_prefix_len = MAX_KUBE_NAME_LEN.saturating_sub(reserved); + let mut trimmed = prefix.chars().take(max_prefix_len).collect::(); + trimmed = trimmed.trim_matches('-').to_string(); + if trimmed.is_empty() { + trimmed = "profile".to_string(); + } + format!("{WARM_POOL_PROFILE_NAME_PREFIX}-{trimmed}-{suffix}") +} + +fn sanitize_dns_label_segment(value: &str) -> String { + let mut out = String::new(); + let mut last_dash = false; + for byte in value.bytes() { + let ch = if byte.is_ascii_lowercase() || byte.is_ascii_digit() { + byte as char + } else if byte.is_ascii_uppercase() { + (byte as char).to_ascii_lowercase() + } else { + '-' + }; + if ch == '-' { + if !last_dash { + out.push(ch); + } + last_dash = true; + } else { + out.push(ch); + last_dash = false; + } + } + let trimmed = out.trim_matches('-').to_string(); + if trimmed.is_empty() { + "profile".to_string() + } else { + trimmed + } +} + +async fn apply_rendered_warm_pool_profile( + client: Client, + rendered: &RenderedWarmPoolProfile, +) -> Result<(), String> { + let template_api = + KubernetesComputeDriver::sandbox_template_api(client.clone(), &rendered.target_namespace); + let warm_pool_api = + KubernetesComputeDriver::sandbox_warm_pool_api(client.clone(), &rendered.target_namespace); + let template = rendered_warm_pool_template_object(rendered, &template_api.resource); + let warm_pool = rendered_warm_pool_object(rendered, &warm_pool_api.resource); + apply_dynamic_object(&template_api.api, &rendered.generated_name, &template).await?; + apply_dynamic_object(&warm_pool_api.api, &rendered.generated_name, &warm_pool).await?; + Ok(()) +} + +async fn apply_dynamic_object( + api: &Api, + name: &str, + obj: &DynamicObject, +) -> Result<(), String> { + match tokio::time::timeout(KUBE_API_TIMEOUT, api.create(&PostParams::default(), obj)).await { + Ok(Ok(_)) => Ok(()), + Ok(Err(KubeError::Api(err))) if err.code == 409 => { + let patch = dynamic_object_merge_patch(obj); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + api.patch(name, &PatchParams::default(), &Patch::Merge(&patch)), + ) + .await + { + Ok(Ok(_)) => Ok(()), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } + } + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } +} + +fn dynamic_object_merge_patch(obj: &DynamicObject) -> serde_json::Value { + serde_json::json!({ + "metadata": { + "labels": obj.metadata.labels, + "annotations": obj.metadata.annotations, + }, + "spec": obj.data.get("spec").cloned().unwrap_or_else(|| serde_json::json!({})), + }) +} + +fn rendered_warm_pool_template_object( + rendered: &RenderedWarmPoolProfile, + resource: &ApiResource, +) -> DynamicObject { + let mut obj = DynamicObject::new(&rendered.generated_name, resource); + obj.metadata = ObjectMeta { + name: Some(rendered.generated_name.clone()), + namespace: Some(rendered.target_namespace.clone()), + labels: Some(warm_pool_profile_generated_labels(rendered)), + annotations: Some(warm_pool_profile_generated_annotations(rendered)), + ..Default::default() + }; + obj.data = serde_json::json!({ + "spec": rendered.template_spec, + }); + obj +} + +fn rendered_warm_pool_object( + rendered: &RenderedWarmPoolProfile, + resource: &ApiResource, +) -> DynamicObject { + let mut obj = DynamicObject::new(&rendered.generated_name, resource); + obj.metadata = ObjectMeta { + name: Some(rendered.generated_name.clone()), + namespace: Some(rendered.target_namespace.clone()), + labels: Some(warm_pool_profile_generated_labels(rendered)), + annotations: Some(warm_pool_profile_generated_annotations(rendered)), + ..Default::default() + }; + obj.data = serde_json::json!({ + "spec": { + "replicas": rendered.replicas, + "sandboxTemplateRef": { + "name": rendered.generated_name + } + } + }); + obj +} + +fn warm_pool_profile_generated_labels( + rendered: &RenderedWarmPoolProfile, +) -> BTreeMap { + BTreeMap::from([ + (LABEL_WARM_POOL_ENABLED.to_string(), "true".to_string()), + ( + LABEL_WARM_POOL_MANAGED_BY.to_string(), + LABEL_WARM_POOL_MANAGED_BY_VALUE.to_string(), + ), + ( + LABEL_WARM_POOL_PROFILE.to_string(), + label_value_for_profile_name(&rendered.source.name), + ), + ( + LABEL_WARM_POOL_PROFILE_UID.to_string(), + rendered.source.uid.clone(), + ), + ]) +} + +fn label_value_for_profile_name(name: &str) -> String { + let sanitized = sanitize_dns_label_segment(name); + let mut value = sanitized.chars().take(63).collect::(); + value = value.trim_matches('-').to_string(); + if value.is_empty() { + "profile".to_string() + } else { + value + } +} + +fn warm_pool_profile_generated_annotations( + rendered: &RenderedWarmPoolProfile, +) -> BTreeMap { + BTreeMap::from([ + ( + ANNOTATION_WARM_POOL_PROFILE_NAMESPACE.to_string(), + rendered.source.namespace.clone(), + ), + ( + ANNOTATION_WARM_POOL_PROFILE_NAME.to_string(), + rendered.source.name.clone(), + ), + ( + ANNOTATION_WARM_POOL_PROFILE_UID.to_string(), + rendered.source.uid.clone(), + ), + ( + ANNOTATION_WARM_POOL_SOURCE_RESOURCE_VERSION.to_string(), + rendered.source.resource_version.clone(), + ), + ( + ANNOTATION_WARM_POOL_TEMPLATE_FINGERPRINT.to_string(), + rendered.fingerprint.clone(), + ), + ( + LABEL_SANDBOX_WORKSPACE.to_string(), + rendered.workspace.clone(), + ), + ]) +} + +async fn garbage_collect_superseded_warm_pool_profile_resources( + client: Client, + rendered: &RenderedWarmPoolProfile, +) -> Result<(), String> { + garbage_collect_warm_pool_profile_resources( + client, + &rendered.target_namespace, + &rendered.source.uid, + Some(&rendered.generated_name), + ) + .await +} + +async fn garbage_collect_warm_pool_profile( + client: Client, + config: &KubernetesComputeConfig, + cm: &ConfigMap, +) { + let Some(uid) = cm.metadata.uid.as_deref() else { + return; + }; + let namespace = config.namespace.clone(); + if let Err(err) = + garbage_collect_warm_pool_profile_resources(client, &namespace, uid, None).await + { + warn!( + namespace = %namespace, + config_map_uid = %uid, + error = %err, + "Failed to garbage-collect deleted warm-pool profile resources" + ); + } +} + +async fn garbage_collect_warm_pool_profile_resources( + client: Client, + target_namespace: &str, + source_uid: &str, + keep_name: Option<&str>, +) -> Result<(), String> { + let selector = format!( + "{LABEL_WARM_POOL_MANAGED_BY}={LABEL_WARM_POOL_MANAGED_BY_VALUE},{LABEL_WARM_POOL_PROFILE_UID}={source_uid}" + ); + let lp = ListParams::default().labels(&selector); + let warm_pool_api = + KubernetesComputeDriver::sandbox_warm_pool_api(client.clone(), target_namespace); + delete_matching_dynamic_objects(&warm_pool_api.api, &lp, keep_name).await?; + let template_api = KubernetesComputeDriver::sandbox_template_api(client, target_namespace); + delete_matching_dynamic_objects(&template_api.api, &lp, keep_name).await?; + Ok(()) +} + +async fn garbage_collect_orphaned_warm_pool_profile_resources( + client: Client, + target_namespace: &str, + active_source_uids: &HashSet, +) -> Result<(), String> { + let selector = format!("{LABEL_WARM_POOL_MANAGED_BY}={LABEL_WARM_POOL_MANAGED_BY_VALUE}"); + let lp = ListParams::default().labels(&selector); + let warm_pool_api = + KubernetesComputeDriver::sandbox_warm_pool_api(client.clone(), target_namespace); + delete_orphaned_dynamic_objects(&warm_pool_api.api, &lp, active_source_uids).await?; + let template_api = KubernetesComputeDriver::sandbox_template_api(client, target_namespace); + delete_orphaned_dynamic_objects(&template_api.api, &lp, active_source_uids).await?; + Ok(()) +} + +async fn delete_orphaned_dynamic_objects( + api: &Api, + lp: &ListParams, + active_source_uids: &HashSet, +) -> Result<(), String> { + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, api.list(lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + for obj in list.items { + if !warm_pool_profile_generated_resource_is_orphaned(&obj, active_source_uids) { + continue; + } + let Some(name) = obj.metadata.name.as_deref() else { + continue; + }; + match tokio::time::timeout(KUBE_API_TIMEOUT, api.delete(name, &DeleteParams::default())) + .await + { + Ok(Ok(_)) => {} + Ok(Err(KubeError::Api(err))) if err.code == 404 => {} + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(()) +} + +fn warm_pool_profile_generated_resource_is_orphaned( + obj: &DynamicObject, + active_source_uids: &HashSet, +) -> bool { + let labels = obj.metadata.labels.as_ref(); + let managed_by_driver = labels + .and_then(|labels| labels.get(LABEL_WARM_POOL_MANAGED_BY)) + .is_some_and(|value| value == LABEL_WARM_POOL_MANAGED_BY_VALUE); + let Some(source_uid) = labels.and_then(|labels| labels.get(LABEL_WARM_POOL_PROFILE_UID)) else { + return false; + }; + managed_by_driver && !active_source_uids.contains(source_uid) +} + +async fn delete_matching_dynamic_objects( + api: &Api, + lp: &ListParams, + keep_name: Option<&str>, +) -> Result<(), String> { + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, api.list(lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + for obj in list.items { + let Some(name) = obj.metadata.name.as_deref() else { + continue; + }; + if keep_name == Some(name) { + continue; + } + match tokio::time::timeout(KUBE_API_TIMEOUT, api.delete(name, &DeleteParams::default())) + .await + { + Ok(Ok(_)) => {} + Ok(Err(KubeError::Api(err))) if err.code == 404 => {} + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(()) +} + +async fn patch_warm_pool_profile_status( + client: Client, + namespace: &str, + name: &str, + status: &str, + rendered: Option<&RenderedWarmPoolProfile>, + error: &str, +) { + let api: Api = Api::namespaced(client, namespace); + let mut annotations = serde_json::Map::new(); + annotations.insert( + ANNOTATION_WARM_POOL_STATUS.to_string(), + serde_json::json!(status), + ); + annotations.insert( + ANNOTATION_WARM_POOL_LAST_ERROR.to_string(), + serde_json::json!(error), + ); + if let Some(rendered) = rendered { + annotations.insert( + ANNOTATION_WARM_POOL_OBSERVED_RESOURCE_VERSION.to_string(), + serde_json::json!(rendered.source.resource_version), + ); + annotations.insert( + ANNOTATION_WARM_POOL_TARGET_NAMESPACE.to_string(), + serde_json::json!(rendered.target_namespace), + ); + annotations.insert( + ANNOTATION_WARM_POOL_TEMPLATE.to_string(), + serde_json::json!(rendered.generated_name), + ); + annotations.insert( + ANNOTATION_WARM_POOL_FINGERPRINT.to_string(), + serde_json::json!(rendered.fingerprint), + ); + } + let patch = serde_json::json!({ + "metadata": { + "annotations": annotations + } + }); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + api.patch(name, &PatchParams::default(), &Patch::Merge(&patch)), + ) + .await + { + Ok(Ok(_)) => {} + Ok(Err(err)) => warn!( + namespace, + config_map = name, + error = %err, + "Failed to patch warm-pool profile status annotations" + ), + Err(_) => warn!( + namespace, + config_map = name, + "Timed out patching warm-pool profile status annotations" + ), + } +} + fn should_try_next_sandbox_api_version(err: &KubeError) -> bool { // Kubernetes returns a structured 404 for some missing API resources and a // raw "404 page not found" body for others. Both mean the probed @@ -7629,4 +8593,199 @@ mod tests { }; assert!(sandbox_id_from_object(&obj).is_err()); } + + #[test] + fn warm_pool_profile_parses_cli_shaped_schema() { + let profile = parse_warm_pool_profile( + r#" +version = 1 +workspace = "research" +replicas = 3 +image = "ghcr.io/acme/agent-python:latest" +runtime_class_name = "nvidia" + +[environment] +FOO = "bar" + +[resources] +cpu = "4" +memory = "8Gi" +gpu_count = 1 +"#, + ) + .expect("profile should parse"); + + assert_eq!(profile.version, 1); + assert_eq!(profile.workspace, "research"); + assert_eq!(profile.replicas, 3); + assert_eq!(profile.image, "ghcr.io/acme/agent-python:latest"); + assert_eq!(profile.runtime_class_name, "nvidia"); + assert_eq!(profile.environment.get("FOO").unwrap(), "bar"); + assert_eq!(profile.resources.cpu, "4"); + assert_eq!(profile.resources.memory, "8Gi"); + assert_eq!(profile.resources.gpu_count, Some(1)); + } + + #[test] + fn warm_pool_profile_rejects_unknown_fields() { + let err = parse_warm_pool_profile( + r#" +version = 1 +workspace = "research" +replicas = 1 +agent_socket = "/tmp/agent.sock" +"#, + ) + .unwrap_err(); + + assert!(err.contains("unknown field")); + } + + #[test] + fn warm_pool_profile_converts_resources_to_driver_spec() { + let profile = parse_warm_pool_profile( + r#" +version = 1 +workspace = "research" +image = "ghcr.io/acme/agent-python:latest" +runtime_class_name = "kata" + +[environment] +FOO = "bar" + +[resources] +cpu = "500m" +memory = "2Gi" +gpu_count = 2 +"#, + ) + .expect("profile should parse"); + + let spec = warm_pool_profile_to_sandbox_spec(&profile); + assert_eq!(spec.environment.get("FOO").unwrap(), "bar"); + let template = spec.template.as_ref().unwrap(); + assert_eq!(template.image, "ghcr.io/acme/agent-python:latest"); + let resources = template.resources.as_ref().unwrap(); + assert_eq!(resources.cpu_limit, "500m"); + assert_eq!(resources.memory_limit, "2Gi"); + assert_eq!( + platform_config_string(template, "runtime_class_name").as_deref(), + Some("kata") + ); + assert_eq!( + spec.resource_requirements + .as_ref() + .and_then(|requirements| requirements.gpu.as_ref()) + .and_then(|gpu| gpu.count), + Some(2) + ); + } + + #[test] + fn ensure_pod_dns_policy_sets_cluster_first() { + let mut rendered = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "containers": [] + } + } + } + }); + + ensure_pod_dns_policy(&mut rendered); + + assert_eq!( + rendered.pointer("/spec/podTemplate/spec/dnsPolicy"), + Some(&serde_json::json!("ClusterFirst")) + ); + } + + #[test] + fn warm_pool_template_identity_env_is_removed() { + let mut rendered = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "containers": [{ + "name": "agent", + "env": [ + {"name": openshell_core::sandbox_env::SANDBOX_ID, "value": ""}, + {"name": openshell_core::sandbox_env::SANDBOX, "value": ""}, + {"name": openshell_core::sandbox_env::ENDPOINT, "value": "http://gateway"} + ] + }] + } + } + } + }); + + remove_warm_pool_template_identity_env(&mut rendered); + + let env = rendered + .pointer("/spec/podTemplate/spec/containers/0/env") + .and_then(serde_json::Value::as_array) + .expect("env should exist"); + let names = env + .iter() + .filter_map(|entry| entry.get("name").and_then(serde_json::Value::as_str)) + .collect::>(); + assert_eq!(names, vec![openshell_core::sandbox_env::ENDPOINT]); + } + + #[test] + fn orphan_gc_only_selects_driver_managed_inactive_profile_resources() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_WARM_POOL_KIND); + let mut active = HashSet::new(); + active.insert("active-uid".to_string()); + + let mut orphan = DynamicObject::new("orphan", &resource); + orphan.metadata.labels = Some(BTreeMap::from([ + ( + LABEL_WARM_POOL_MANAGED_BY.to_string(), + LABEL_WARM_POOL_MANAGED_BY_VALUE.to_string(), + ), + ( + LABEL_WARM_POOL_PROFILE_UID.to_string(), + "deleted-uid".to_string(), + ), + ])); + assert!(warm_pool_profile_generated_resource_is_orphaned( + &orphan, &active + )); + + let mut still_active = DynamicObject::new("active", &resource); + still_active.metadata.labels = Some(BTreeMap::from([ + ( + LABEL_WARM_POOL_MANAGED_BY.to_string(), + LABEL_WARM_POOL_MANAGED_BY_VALUE.to_string(), + ), + ( + LABEL_WARM_POOL_PROFILE_UID.to_string(), + "active-uid".to_string(), + ), + ])); + assert!(!warm_pool_profile_generated_resource_is_orphaned( + &still_active, + &active + )); + + let mut external = DynamicObject::new("external", &resource); + external.metadata.labels = Some(BTreeMap::from([( + LABEL_WARM_POOL_PROFILE_UID.to_string(), + "deleted-uid".to_string(), + )])); + assert!(!warm_pool_profile_generated_resource_is_orphaned( + &external, &active + )); + } + + #[test] + fn generated_warm_pool_profile_name_is_stable_dns_label() { + let name = + generated_warm_pool_profile_name("GPU_Python.Profile", "sha256:abcdef1234567890"); + + assert_eq!(name, "openshell-wp-gpu-python-profile-abcdef12"); + assert!(is_dns_label(&name)); + } } diff --git a/deploy/helm/openshell/README.md b/deploy/helm/openshell/README.md index c12760443e..13f7686f34 100644 --- a/deploy/helm/openshell/README.md +++ b/deploy/helm/openshell/README.md @@ -228,7 +228,11 @@ add `ci/values-spire.yaml` to the OpenShell release values files. | server.tls.certSecretName | string | `"openshell-server-tls"` | K8s secret (type kubernetes.io/tls) with tls.crt and tls.key for the server. | | server.tls.clientCaSecretName | string | `"openshell-server-client-ca"` | K8s secret with ca.crt for client certificate verification (mTLS). Set to "" to disable mTLS and run HTTPS-only (use OIDC for auth instead). | | server.tls.clientTlsSecretName | string | `"openshell-client-tls"` | K8s secret mounted into sandbox pods for mTLS to the server. | -| server.warmPooling | object | `{"enabled":true}` | Enable transparent Kubernetes warm-pool allocation through Agent Sandbox v1beta1 SandboxClaim resources when a compatible OpenShell-enabled SandboxWarmPool exists in the target namespace. | +| server.warmPooling | object | `{"enabled":true,"profiles":{"enabled":true,"labelSelector":"openshell.ai/warm-pool-profile=true","maxReplicas":100,"namespace":""}}` | Enable transparent Kubernetes warm-pool allocation through Agent Sandbox v1beta1 SandboxClaim resources when a compatible OpenShell-enabled SandboxWarmPool exists in the target namespace. | +| server.warmPooling.profiles | object | `{"enabled":true,"labelSelector":"openshell.ai/warm-pool-profile=true","maxReplicas":100,"namespace":""}` | Reconcile labelled ConfigMap warm-pool profiles into generated SandboxTemplate and SandboxWarmPool resources. Admins can still create warm pools by any other mechanism. | +| server.warmPooling.profiles.labelSelector | string | `"openshell.ai/warm-pool-profile=true"` | Label selector used to select warm-pool profile ConfigMaps. | +| server.warmPooling.profiles.maxReplicas | int | `100` | Maximum replicas accepted from a single warm-pool profile. | +| server.warmPooling.profiles.namespace | string | `""` | Namespace containing warm-pool profile ConfigMaps. Empty defaults to the Helm release namespace. | | server.workspaceDefaultStorageSize | string | `""` | Default storage size for the workspace PVC in sandbox pods. Uses Kubernetes quantity syntax (e.g. "2Gi", "10Gi", "500Mi"). Empty = built-in default (2Gi). | | service.healthPort | int | `8081` | Gateway health service port. | | service.metricsPort | int | `9090` | Gateway metrics service port. | diff --git a/deploy/helm/openshell/templates/clusterrole.yaml b/deploy/helm/openshell/templates/clusterrole.yaml index b8401ac850..eea70f26ad 100644 --- a/deploy/helm/openshell/templates/clusterrole.yaml +++ b/deploy/helm/openshell/templates/clusterrole.yaml @@ -53,3 +53,9 @@ rules: - get - list - watch + {{- if .Values.server.warmPooling.profiles.enabled }} + - create + - patch + - update + - delete + {{- end }} diff --git a/deploy/helm/openshell/templates/gateway-config.yaml b/deploy/helm/openshell/templates/gateway-config.yaml index 0314391459..486dc7bda6 100644 --- a/deploy/helm/openshell/templates/gateway-config.yaml +++ b/deploy/helm/openshell/templates/gateway-config.yaml @@ -148,6 +148,12 @@ data: [openshell.drivers.kubernetes.warm_pooling] enabled = {{ .Values.server.warmPooling.enabled }} + [openshell.drivers.kubernetes.warm_pooling.profiles] + enabled = {{ .Values.server.warmPooling.profiles.enabled }} + namespace = {{ default .Release.Namespace .Values.server.warmPooling.profiles.namespace | quote }} + label_selector = {{ .Values.server.warmPooling.profiles.labelSelector | quote }} + max_replicas = {{ .Values.server.warmPooling.profiles.maxReplicas | default 100 }} + [openshell.drivers.kubernetes.sidecar] proxy_uid = {{ .Values.supervisor.sidecar.proxyUid | default 1337 }} process_binary_aware_network_policy = {{ .Values.supervisor.sidecar.processBinaryAwareNetworkPolicy }} diff --git a/deploy/helm/openshell/tests/gateway_config_test.yaml b/deploy/helm/openshell/tests/gateway_config_test.yaml index 215ed1f67a..3c4d0e7f30 100644 --- a/deploy/helm/openshell/tests/gateway_config_test.yaml +++ b/deploy/helm/openshell/tests/gateway_config_test.yaml @@ -168,6 +168,27 @@ tests: path: data["gateway.toml"] pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\].*?enabled\s*=\s*false' + - it: renders warm pool profile config under [openshell.drivers.kubernetes.warm_pooling.profiles] + template: templates/gateway-config.yaml + set: + server.warmPooling.profiles.enabled: true + server.warmPooling.profiles.namespace: warm-profile-ns + server.warmPooling.profiles.labelSelector: openshell.ai/warm-pool-profile=true,tier=gpu + server.warmPooling.profiles.maxReplicas: 7 + asserts: + - matchRegex: + path: data["gateway.toml"] + pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\.profiles\].*?enabled\s*=\s*true' + - matchRegex: + path: data["gateway.toml"] + pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\.profiles\].*?namespace\s*=\s*"warm-profile-ns"' + - matchRegex: + path: data["gateway.toml"] + pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\.profiles\].*?label_selector\s*=\s*"openshell.ai/warm-pool-profile=true,tier=gpu"' + - matchRegex: + path: data["gateway.toml"] + pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\.profiles\].*?max_replicas\s*=\s*7' + - it: renders sandbox image pull secrets under [openshell.drivers.kubernetes] template: templates/gateway-config.yaml set: diff --git a/deploy/helm/openshell/values.yaml b/deploy/helm/openshell/values.yaml index fa1c2923a6..ceaf7e14e8 100644 --- a/deploy/helm/openshell/values.yaml +++ b/deploy/helm/openshell/values.yaml @@ -198,6 +198,18 @@ server: # SandboxWarmPool exists in the target namespace. warmPooling: enabled: true + # -- Reconcile labelled ConfigMap warm-pool profiles into generated + # SandboxTemplate and SandboxWarmPool resources. Admins can still create + # warm pools by any other mechanism. + profiles: + enabled: true + # -- Namespace containing warm-pool profile ConfigMaps. Empty defaults to + # the Helm release namespace. + namespace: "" + # -- Label selector used to select warm-pool profile ConfigMaps. + labelSelector: "openshell.ai/warm-pool-profile=true" + # -- Maximum replicas accepted from a single warm-pool profile. + maxReplicas: 100 # -- gRPC endpoint sandboxes call back into the gateway. Leave empty to derive # it from the chart fullname, release namespace, service port, and # disableTls flag, for example https://openshell.openshell.svc.cluster.local:8080. diff --git a/docs/reference/gateway-config.mdx b/docs/reference/gateway-config.mdx index 7c69e6043d..fc90a0eb3e 100644 --- a/docs/reference/gateway-config.mdx +++ b/docs/reference/gateway-config.mdx @@ -299,6 +299,13 @@ provider_spiffe_workload_api_socket_path = "/spiffe-workload-api/spire-agent.soc # backed by the driver's cache of OpenShell-enabled SandboxWarmPool resources. enabled = true +[openshell.drivers.kubernetes.warm_pooling.profiles] +# Optional ConfigMap-backed reconciler for admin-declared warm pools. +enabled = true +namespace = "openshell" +label_selector = "openshell.ai/warm-pool-profile=true" +max_replicas = 100 + [openshell.drivers.kubernetes.sidecar] # UID used by relaxed long-running network sidecars. Strict process/binary-aware # sidecars run as UID 0 so Kubernetes grants the required /proc inspection diff --git a/docs/reference/sandbox-compute-drivers.mdx b/docs/reference/sandbox-compute-drivers.mdx index 5790ea3210..c44a536f2c 100644 --- a/docs/reference/sandbox-compute-drivers.mdx +++ b/docs/reference/sandbox-compute-drivers.mdx @@ -350,6 +350,31 @@ OpenShell workspace. Set `openshell.drivers.kubernetes.warm_pooling.enabled = false` to force direct `Sandbox` creation. +When `openshell.drivers.kubernetes.warm_pooling.profiles.enabled = true`, the +driver reconciles labelled ConfigMaps into generated `SandboxTemplate` and +`SandboxWarmPool` resources. This is enabled by default with the Helm chart. +Each profile ConfigMap must contain +`data["warm-pool.toml"]` with a narrow CLI-shaped schema: + +```toml +version = 1 +workspace = "default" +replicas = 3 +image = "ghcr.io/acme/agent-python:latest" +runtime_class_name = "nvidia" + +[environment] +FOO = "bar" + +[resources] +cpu = "4" +memory = "8Gi" +gpu_count = 1 +``` + +The generated warm pools carry `openshell.ai/enabled=true`, so the existing +matching path discovers them like any externally managed warm pool. + If Agent Sandbox is upgraded in place, restart the OpenShell gateway after the controller and CRD rollout completes so the gateway can detect the served API versions again. diff --git a/examples/kubernetes-warm-pool-config/cpu-0-2-configmap.yaml b/examples/kubernetes-warm-pool-config/cpu-0-2-configmap.yaml new file mode 100644 index 0000000000..264d94e8ed --- /dev/null +++ b/examples/kubernetes-warm-pool-config/cpu-0-2-configmap.yaml @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +apiVersion: v1 +kind: ConfigMap +metadata: + name: openshell-warm-pool-cpu-0-2 + namespace: openshell + labels: + openshell.ai/warm-pool-profile: "true" +data: + warm-pool.toml: | + version = 1 + workspace = "openshell" + replicas = 1 + + [resources] + cpu = "0.2" diff --git a/examples/kubernetes-warm-pool-config/default-configmap.yaml b/examples/kubernetes-warm-pool-config/default-configmap.yaml new file mode 100644 index 0000000000..46050ba93f --- /dev/null +++ b/examples/kubernetes-warm-pool-config/default-configmap.yaml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +apiVersion: v1 +kind: ConfigMap +metadata: + name: openshell-warm-pool-default + namespace: openshell + labels: + openshell.ai/warm-pool-profile: "true" +data: + warm-pool.toml: | + version = 1 + workspace = "openshell" + replicas = 1 diff --git a/examples/kubernetes-warm-pool-config/env-foo-configmap.yaml b/examples/kubernetes-warm-pool-config/env-foo-configmap.yaml new file mode 100644 index 0000000000..608f0d5d85 --- /dev/null +++ b/examples/kubernetes-warm-pool-config/env-foo-configmap.yaml @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +apiVersion: v1 +kind: ConfigMap +metadata: + name: openshell-warm-pool-env-foo + namespace: openshell + labels: + openshell.ai/warm-pool-profile: "true" +data: + warm-pool.toml: | + version = 1 + workspace = "openshell" + replicas = 1 + + [environment] + FOO = "bar" From 950ff61aeda414017b3142dd5203998bc8a54184 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Fri, 24 Jul 2026 21:15:07 +0100 Subject: [PATCH 09/10] fix(kubernetes): prevent warm pod activation spoofing Signed-off-by: Gordon Sim --- Cargo.lock | 1 + architecture/compute-runtimes.md | 15 +- architecture/gateway.md | 31 ++- .../src/supervisor_bootstrap.rs | 2 + crates/openshell-driver-docker/src/lib.rs | 4 +- crates/openshell-driver-kubernetes/Cargo.toml | 1 + .../openshell-driver-kubernetes/src/driver.rs | 46 ++-- .../openshell-driver-kubernetes/src/grpc.rs | 7 +- .../src/sandboxclaim.rs | 114 ++++++++-- crates/openshell-driver-podman/src/grpc.rs | 4 +- crates/openshell-driver-vm/src/driver.rs | 4 +- crates/openshell-server/src/compute/mod.rs | 70 +++++- .../openshell-server/src/persistence/mod.rs | 89 ++++++++ .../openshell-server/src/persistence/tests.rs | 95 ++++++++- crates/openshell-server/src/test_support.rs | 4 +- .../src/warm_pod_activation.rs | 199 +++++++++++++++++- proto/compute_driver.proto | 7 +- 17 files changed, 623 insertions(+), 70 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c60f93d1f1..1b93758bdc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3931,6 +3931,7 @@ dependencies = [ name = "openshell-driver-kubernetes" version = "0.0.0" dependencies = [ + "base64 0.22.1", "clap", "futures", "k8s-openapi", diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index cd426d7822..91c20b0b22 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -136,10 +136,17 @@ initiated over the existing outbound gRPC connection. When compatible OpenShell-enabled Agent Sandbox warm pools exist, the Kubernetes driver can create a `SandboxClaim` instead of a direct `Sandbox`. -The driver watches warm pools and templates into a local cache so create -requests do not list or fetch extension objects on the hot path. This remains -behind the compute-driver boundary: other drivers and the driver-agnostic -gateway lifecycle continue to operate on OpenShell sandbox identity and status. +The driver returns an opaque activation token with the create response for +those delayed warm allocations. The gateway stores only a hash of that token +against the sandbox ID and later requires the claim activation controller to +present the same token before minting a sandbox-scoped supervisor JWT. The +token is derived by the driver from immutable live claim identity, including +the claim UID; Kubernetes labels remain useful for diagnostics and cross-checks +but are not sufficient activation proof. The driver watches warm pools and +templates into a local cache so create requests do not list or fetch extension +objects on the hot path. This remains behind the compute-driver boundary: +other drivers and the driver-agnostic gateway lifecycle continue to operate on +OpenShell sandbox identity and status. The Kubernetes driver can optionally run a separate ConfigMap profile reconciler that turns admin-authored TOML profiles into generated `SandboxTemplate` and `SandboxWarmPool` resources. That reconciler is diff --git a/architecture/gateway.md b/architecture/gateway.md index defb0019c2..5b58798401 100644 --- a/architecture/gateway.md +++ b/architecture/gateway.md @@ -166,9 +166,17 @@ Podman, and VM drivers deliver the initial token through supervisor-only runtime material; Kubernetes supervisors exchange a projected ServiceAccount token through `RegisterSupervisorPod`. The gateway validates that projected token with Kubernetes `TokenReview`, requires the configured sandbox service -account, checks the returned pod binding against the live pod UID, and verifies -the pod's controlling `Sandbox` ownerReference against the live Sandbox CR UID -and sandbox-id label before sending an activation message with the gateway JWT. +account, and checks the returned pod binding against the live pod UID. Warm +pods that are not yet bound receive only a pending bootstrap identity and keep +the registration stream open. Later warm activation must match that pending +pod identity, the live Sandbox owner UID, and a gateway-stored opaque +activation token hash returned by the driver during sandbox creation. A missing +token mapping is retryable and leaves the pending stream intact; an empty or +mismatched token is fatal. Kubernetes sandbox-id labels are cross-checks and +diagnostics only, not authority to choose which sandbox receives the gateway +JWT. Already-bound pods still verify the pod's controlling `Sandbox` +ownerReference against the live Sandbox CR UID and sandbox-id label before +sending an activation message with the gateway JWT. The bootstrap path accepts both `agents.x-k8s.io/v1beta1` ownerReferences from newer Agent Sandbox controllers and `agents.x-k8s.io/v1alpha1` ownerReferences from existing deployments. `IssueSandboxToken` remains as a compatibility shim @@ -275,13 +283,16 @@ default WAL journal mode), which mirror the same sensitive contents. Persisted state includes sandboxes, providers, provider credential refresh state, SSH sessions, policy revisions, settings, inference configuration, and -deployment records. Provider refresh material is stored as a separate object -scoped to the provider instance through `objects.scope`; the provider record -keeps only the current injectable credential values and optional per-credential -expiry timestamps. A refresh normally mints one credential, but a strategy may -co-mint several (AWS STS mints the access key, secret key, and session token in -one call); the refresh state pins the resolved set of env keys it owns so -collision checks reserve all of them before the first mint. +deployment records. Warm-pool activation token mappings are stored as +generic objects keyed by sandbox ID with only the token hash in the payload; +sandbox deletion removes the mapping opportunistically after the sandbox row is +gone. Provider refresh material is stored as a separate object scoped to the +provider instance through `objects.scope`; the provider record keeps only the +current injectable credential values and optional per-credential expiry +timestamps. A refresh normally mints one credential, but a strategy may co-mint +several (AWS STS mints the access key, secret key, and session token in one +call); the refresh state pins the resolved set of env keys it owns so collision +checks reserve all of them before the first mint. ### Optimistic Concurrency (CAS) diff --git a/crates/openshell-core/src/supervisor_bootstrap.rs b/crates/openshell-core/src/supervisor_bootstrap.rs index 3ffabb737c..bda04d12e9 100644 --- a/crates/openshell-core/src/supervisor_bootstrap.rs +++ b/crates/openshell-core/src/supervisor_bootstrap.rs @@ -80,6 +80,8 @@ pub struct SupervisorBootstrapActivationRequest { pub sandbox_id: String, /// Driver-native owner UID observed during claim validation. pub owner_uid: String, + /// Opaque driver-issued activation token returned during sandbox create. + pub activation_token: String, /// Short reason for logs and audit messages. pub reason: String, } diff --git a/crates/openshell-driver-docker/src/lib.rs b/crates/openshell-driver-docker/src/lib.rs index 2f89c2229e..f71957e703 100644 --- a/crates/openshell-driver-docker/src/lib.rs +++ b/crates/openshell-driver-docker/src/lib.rs @@ -1430,7 +1430,9 @@ impl ComputeDriver for DockerComputeDriver { .sandbox .ok_or_else(|| Status::invalid_argument("sandbox is required"))?; self.create_sandbox_inner(&sandbox).await?; - Ok(Response::new(CreateSandboxResponse {})) + Ok(Response::new(CreateSandboxResponse { + activation_token: String::new(), + })) } async fn stop_sandbox( diff --git a/crates/openshell-driver-kubernetes/Cargo.toml b/crates/openshell-driver-kubernetes/Cargo.toml index 6a1a0b2bcc..983e366923 100644 --- a/crates/openshell-driver-kubernetes/Cargo.toml +++ b/crates/openshell-driver-kubernetes/Cargo.toml @@ -31,6 +31,7 @@ serde = { workspace = true } serde_json = { workspace = true } toml = { workspace = true } sha2 = { workspace = true } +base64 = { workspace = true } clap = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index 6329ee9489..8786a513e6 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -9,6 +9,7 @@ use crate::config::{ DEFAULT_WORKSPACE_STORAGE_SIZE, KubernetesComputeConfig, SupervisorSideloadMethod, SupervisorTopology, }; +use crate::sandboxclaim::sandbox_claim_activation_token; use futures::{Stream, StreamExt, TryStreamExt, stream}; use k8s_openapi::api::core::v1::{ ConfigMap, Event as KubeEventObj, Namespace, Node, PersistentVolumeClaimVolumeSource, Volume, @@ -1234,16 +1235,16 @@ impl KubernetesComputeDriver { sandbox: &Sandbox, target_namespace: &str, request_fingerprint: &str, - ) -> Result { + ) -> Result, KubernetesDriverError> { if !self.config.warm_pooling.enabled { - return Ok(false); + return Ok(None); } let Some(warm_pool) = self .matching_warm_pool(target_namespace, request_fingerprint) .await else { - return Ok(false); + return Ok(None); }; let claim_api = Self::sandbox_claim_api(self.client.clone(), target_namespace); @@ -1254,7 +1255,25 @@ impl KubernetesComputeDriver { ) .await { - Ok(Ok(_result)) => { + Ok(Ok(result)) => { + let claim_name = result + .metadata + .name + .as_deref() + .filter(|name| !name.is_empty()) + .unwrap_or(sandbox.name.as_str()); + let claim_uid = result + .metadata + .uid + .as_deref() + .filter(|uid| !uid.is_empty()) + .ok_or_else(|| { + KubernetesDriverError::Precondition( + "created SandboxClaim is missing Kubernetes UID".to_string(), + ) + })?; + let activation_token = + sandbox_claim_activation_token(target_namespace, claim_name, claim_uid); info!( sandbox_id = %sandbox.id, sandbox_name = %sandbox.name, @@ -1262,7 +1281,7 @@ impl KubernetesComputeDriver { warm_pool = %warm_pool.name, "SandboxClaim created in Kubernetes successfully" ); - Ok(true) + Ok(Some(activation_token)) } Ok(Err(err)) if extension_api_unavailable(&err) => { debug!( @@ -1272,7 +1291,7 @@ impl KubernetesComputeDriver { error = %err, "SandboxClaim API is unavailable; falling back to direct Sandbox" ); - Ok(false) + Ok(None) } Ok(Err(err)) if matches!(&err, KubeError::Api(api) if api.code == 409) => { Err(KubernetesDriverError::from_kube(err)) @@ -1286,7 +1305,7 @@ impl KubernetesComputeDriver { error = %err, "Failed to create SandboxClaim; falling back to direct Sandbox" ); - Ok(false) + Ok(None) } Err(_elapsed) => { warn!( @@ -1296,7 +1315,7 @@ impl KubernetesComputeDriver { timeout_secs = KUBE_API_TIMEOUT.as_secs(), "Timed out creating SandboxClaim; falling back to direct Sandbox" ); - Ok(false) + Ok(None) } } } @@ -1370,7 +1389,10 @@ impl KubernetesComputeDriver { } #[allow(clippy::similar_names)] - pub async fn create_sandbox(&self, sandbox: &Sandbox) -> Result<(), KubernetesDriverError> { + pub async fn create_sandbox( + &self, + sandbox: &Sandbox, + ) -> Result, KubernetesDriverError> { let gpu_requirements = sandbox .spec .as_ref() @@ -1440,11 +1462,11 @@ impl KubernetesComputeDriver { if self.config.warm_pooling.enabled { match sandbox_spec_fingerprint(&data) { Ok(fingerprint) => { - if self + if let Some(activation_token) = self .try_create_sandbox_claim(sandbox, &target_namespace, &fingerprint) .await? { - return Ok(()); + return Ok(Some(activation_token)); } debug!( sandbox_id = %sandbox.id, @@ -1501,7 +1523,7 @@ impl KubernetesComputeDriver { sandbox_name = %name, "Sandbox created in Kubernetes successfully" ); - Ok(()) + Ok(None) } Ok(Err(err)) => { warn!( diff --git a/crates/openshell-driver-kubernetes/src/grpc.rs b/crates/openshell-driver-kubernetes/src/grpc.rs index fccfa9464b..5cf8fa6416 100644 --- a/crates/openshell-driver-kubernetes/src/grpc.rs +++ b/crates/openshell-driver-kubernetes/src/grpc.rs @@ -93,11 +93,14 @@ impl ComputeDriver for ComputeDriverService { .into_inner() .sandbox .ok_or_else(|| Status::invalid_argument("sandbox is required"))?; - self.driver + let activation_token = self + .driver .create_sandbox(&sandbox) .await .map_err(|e| Status::from(openshell_core::ComputeDriverError::from(e)))?; - Ok(Response::new(CreateSandboxResponse {})) + Ok(Response::new(CreateSandboxResponse { + activation_token: activation_token.unwrap_or_default(), + })) } async fn stop_sandbox( diff --git a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs index 2d26d95e6e..d1f361f5d8 100644 --- a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs +++ b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs @@ -3,6 +3,7 @@ //! Kubernetes `SandboxClaim` activation support. +use base64::Engine as _; use futures::{StreamExt, TryStreamExt}; use k8s_openapi::api::core::v1::Pod; use k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference; @@ -11,10 +12,13 @@ use kube::core::DynamicObject; use kube::core::gvk::GroupVersionKind; use kube::runtime::watcher::{self, Event}; use kube::{Client, Error as KubeError}; -use openshell_core::driver_utils::LABEL_SANDBOX_ID; +use openshell_core::driver_utils::{ + LABEL_MANAGED_BY, LABEL_MANAGED_BY_VALUE, LABEL_SANDBOX_ID, LABEL_SANDBOX_WORKSPACE, +}; use openshell_core::supervisor_bootstrap::{ SupervisorBootstrapActivationRequest, SupervisorBootstrapActivator, }; +use sha2::{Digest, Sha256}; use std::sync::Arc; use std::time::Duration; use tokio::sync::watch; @@ -39,6 +43,8 @@ const SANDBOX_POD_NAME_ANNOTATION: &str = "agents.x-k8s.io/pod-name"; const SANDBOX_CLAIM_GROUP: &str = "extensions.agents.x-k8s.io"; const SANDBOX_CLAIM_KIND: &str = "SandboxClaim"; const SANDBOX_CLAIM_VERSION_V1BETA1: &str = "v1beta1"; +const LABEL_ALLOCATION: &str = "openshell.ai/allocation"; +const LABEL_ALLOCATION_SANDBOX_CLAIM: &str = "sandbox-claim"; /// Watches Kubernetes `SandboxClaim` resources and activates pending /// supervisor bootstrap streams after revalidating live driver state. @@ -372,8 +378,12 @@ struct SandboxClaimApi { #[derive(Debug, Clone, Eq, PartialEq)] struct ParsedSandboxClaim { name: String, + namespace: String, uid: String, sandbox_id: Option, + workspace: Option, + managed_by: Option, + allocation: Option, warm_pool_name: Option, sandbox_template_name: Option, sandbox_name: Option, @@ -473,14 +483,12 @@ async fn activate_registered_supervisor_with_retry( fn parse_sandbox_claim(obj: &DynamicObject) -> Result { Ok(ParsedSandboxClaim { name: required_metadata_field(&obj.metadata.name, "name")?, + namespace: obj.metadata.namespace.clone().unwrap_or_default(), uid: required_metadata_field(&obj.metadata.uid, "uid")?, - sandbox_id: obj - .metadata - .labels - .as_ref() - .and_then(|labels| labels.get(LABEL_SANDBOX_ID)) - .filter(|id| !id.is_empty()) - .cloned(), + sandbox_id: label_value(obj, LABEL_SANDBOX_ID), + workspace: label_value(obj, LABEL_SANDBOX_WORKSPACE), + managed_by: label_value(obj, LABEL_MANAGED_BY), + allocation: label_value(obj, LABEL_ALLOCATION), warm_pool_name: string_at(&obj.data, &["spec", "warmPoolRef", "name"]), sandbox_template_name: string_at(&obj.data, &["spec", "sandboxTemplateRef", "name"]), sandbox_name: string_at(&obj.data, &["status", "sandbox", "name"]), @@ -488,6 +496,15 @@ fn parse_sandbox_claim(obj: &DynamicObject) -> Result Option { + obj.metadata + .labels + .as_ref() + .and_then(|labels| labels.get(key)) + .filter(|value| !value.is_empty()) + .cloned() +} + fn required_metadata_field(value: &Option, field: &str) -> Result { value .as_ref() @@ -496,6 +513,13 @@ fn required_metadata_field(value: &Option, field: &str) -> Result String { + let input = + format!("openshell-kubernetes-sandbox-claim-activation-v1\0{namespace}\0{name}\0{uid}"); + let digest = Sha256::digest(input.as_bytes()); + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest) +} + fn activation_request_from_claim_state( claim: &ParsedSandboxClaim, sandbox_cr: &DynamicObject, @@ -510,6 +534,12 @@ fn activation_request_from_claim_state( "SandboxClaim selected Sandbox {claim_sandbox_name}, but live Sandbox is {sandbox_name}" )); } + if claim.managed_by.as_deref() != Some(LABEL_MANAGED_BY_VALUE) { + return Err("SandboxClaim is not OpenShell-managed".to_string()); + } + if claim.allocation.as_deref() != Some(LABEL_ALLOCATION_SANDBOX_CLAIM) { + return Err("SandboxClaim is missing OpenShell sandbox-claim allocation label".to_string()); + } let owner_uid = sandbox_cr .metadata .uid @@ -517,17 +547,31 @@ fn activation_request_from_claim_state( .filter(|uid| !uid.is_empty()) .cloned() .ok_or_else(|| "Sandbox CR is missing uid".to_string())?; - let sandbox_id = sandbox_cr + let sandbox_id = claim + .sandbox_id + .clone() + .ok_or_else(|| "SandboxClaim is missing OpenShell sandbox id label".to_string())?; + if let Some(live_sandbox_id) = sandbox_cr .metadata .labels .as_ref() .and_then(|labels| labels.get(LABEL_SANDBOX_ID)) .filter(|id| !id.is_empty()) - .cloned() - .or_else(|| claim.sandbox_id.clone()) - .ok_or_else(|| { - "SandboxClaim and selected Sandbox are missing OpenShell sandbox id label".to_string() - })?; + && live_sandbox_id != &sandbox_id + { + return Err("SandboxClaim sandbox id does not match selected Sandbox label".to_string()); + } + if let Some(claim_workspace) = claim.workspace.as_deref() + && let Some(sandbox_workspace) = sandbox_cr + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(LABEL_SANDBOX_WORKSPACE)) + .filter(|workspace| !workspace.is_empty()) + && sandbox_workspace != claim_workspace + { + return Err("SandboxClaim workspace does not match selected Sandbox label".to_string()); + } let instance_id = pod .metadata .uid @@ -544,6 +588,7 @@ fn activation_request_from_claim_state( instance_id, sandbox_id, owner_uid, + activation_token: sandbox_claim_activation_token(&claim.namespace, &claim.name, &claim.uid), reason: format!("SandboxClaim/{}", claim.name), })) } @@ -670,11 +715,27 @@ mod tests { let gvk = GroupVersionKind::gvk(group, version, kind); let resource = ApiResource::from_gvk(&gvk); let mut obj = DynamicObject::new(name, &resource); + obj.metadata.namespace = Some("openshell".to_string()); obj.metadata.uid = Some(uid.to_string()); obj.data = data; obj } + fn claim_labels(sandbox_id: &str) -> BTreeMap { + BTreeMap::from([ + ( + LABEL_MANAGED_BY.to_string(), + LABEL_MANAGED_BY_VALUE.to_string(), + ), + ( + LABEL_ALLOCATION.to_string(), + LABEL_ALLOCATION_SANDBOX_CLAIM.to_string(), + ), + (LABEL_SANDBOX_ID.to_string(), sandbox_id.to_string()), + (LABEL_SANDBOX_WORKSPACE.to_string(), "default".to_string()), + ]) + } + fn sandbox_cr(name: &str, uid: &str, sandbox_id: &str, version: &'static str) -> DynamicObject { let mut obj = dynamic_object( SANDBOX_GROUP, @@ -688,6 +749,10 @@ mod tests { .labels .get_or_insert_with(BTreeMap::new) .insert(LABEL_SANDBOX_ID.to_string(), sandbox_id.to_string()); + obj.metadata + .labels + .get_or_insert_with(BTreeMap::new) + .insert(LABEL_SANDBOX_WORKSPACE.to_string(), "default".to_string()); obj } @@ -771,7 +836,7 @@ mod tests { #[test] fn activation_request_uses_live_sandbox_and_pod_identity() { - let obj = dynamic_object( + let mut obj = dynamic_object( SANDBOX_CLAIM_GROUP, SANDBOX_CLAIM_VERSION_V1BETA1, SANDBOX_CLAIM_KIND, @@ -782,6 +847,7 @@ mod tests { "status": {"sandbox": {"name": "sandbox-a"}} }), ); + obj.metadata.labels = Some(claim_labels("openshell-sandbox")); let claim = parse_sandbox_claim(&obj).unwrap(); let sandbox = sandbox_cr( "sandbox-a", @@ -804,6 +870,10 @@ mod tests { assert_eq!(request.instance_id, "pod-uid"); assert_eq!(request.sandbox_id, "openshell-sandbox"); assert_eq!(request.owner_uid, "sandbox-uid"); + assert_eq!( + request.activation_token, + sandbox_claim_activation_token("openshell", "claim-a", "claim-uid") + ); assert_eq!(request.reason, "SandboxClaim/claim-a"); } @@ -817,10 +887,7 @@ mod tests { "claim-uid", json!({"status": {"sandbox": {"name": "sandbox-a"}}}), ); - obj.metadata.labels = Some(BTreeMap::from([( - LABEL_SANDBOX_ID.to_string(), - "openshell-sandbox".to_string(), - )])); + obj.metadata.labels = Some(claim_labels("openshell-sandbox")); let claim = parse_sandbox_claim(&obj).unwrap(); let mut sandbox = sandbox_cr("sandbox-a", "sandbox-uid", "", SANDBOX_VERSION_V1BETA1); sandbox.metadata.labels = None; @@ -846,6 +913,7 @@ mod tests { instance_id: "pod-uid".to_string(), sandbox_id: "sandbox-id".to_string(), owner_uid: "sandbox-uid".to_string(), + activation_token: "activation-token".to_string(), reason: "SandboxClaim/claim-a".to_string(), }; @@ -868,7 +936,7 @@ mod tests { #[test] fn activation_request_rejects_missing_sandbox_id_labels() { - let obj = dynamic_object( + let mut obj = dynamic_object( SANDBOX_CLAIM_GROUP, SANDBOX_CLAIM_VERSION_V1BETA1, SANDBOX_CLAIM_KIND, @@ -876,6 +944,9 @@ mod tests { "claim-uid", json!({"status": {"sandbox": {"name": "sandbox-a"}}}), ); + let mut labels = claim_labels(""); + labels.remove(LABEL_SANDBOX_ID); + obj.metadata.labels = Some(labels); let claim = parse_sandbox_claim(&obj).unwrap(); let mut sandbox = sandbox_cr("sandbox-a", "sandbox-uid", "", SANDBOX_VERSION_V1BETA1); sandbox.metadata.labels = None; @@ -893,7 +964,7 @@ mod tests { #[test] fn activation_request_rejects_uncontrolled_pod() { - let obj = dynamic_object( + let mut obj = dynamic_object( SANDBOX_CLAIM_GROUP, SANDBOX_CLAIM_VERSION_V1BETA1, SANDBOX_CLAIM_KIND, @@ -901,6 +972,7 @@ mod tests { "claim-uid", json!({"status": {"sandbox": {"name": "sandbox-a"}}}), ); + obj.metadata.labels = Some(claim_labels("openshell-sandbox")); let claim = parse_sandbox_claim(&obj).unwrap(); let sandbox = sandbox_cr( "sandbox-a", diff --git a/crates/openshell-driver-podman/src/grpc.rs b/crates/openshell-driver-podman/src/grpc.rs index 8e68a91e72..10cd85712b 100644 --- a/crates/openshell-driver-podman/src/grpc.rs +++ b/crates/openshell-driver-podman/src/grpc.rs @@ -96,7 +96,9 @@ impl ComputeDriver for ComputeDriverService { .create_sandbox(&sandbox) .await .map_err(Status::from)?; - Ok(Response::new(CreateSandboxResponse {})) + Ok(Response::new(CreateSandboxResponse { + activation_token: String::new(), + })) } async fn stop_sandbox( diff --git a/crates/openshell-driver-vm/src/driver.rs b/crates/openshell-driver-vm/src/driver.rs index 7af0ddc389..ab6e9c99cd 100644 --- a/crates/openshell-driver-vm/src/driver.rs +++ b/crates/openshell-driver-vm/src/driver.rs @@ -622,7 +622,9 @@ impl VmDriver { task.abort(); } - Ok(CreateSandboxResponse {}) + Ok(CreateSandboxResponse { + activation_token: String::new(), + }) } async fn provision_sandbox( diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index 08d9b1f03e..b945d44fb7 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -750,7 +750,43 @@ impl ComputeRuntime { })) .await { - Ok(_) => { + Ok(response) => { + let activation_token = response.into_inner().activation_token; + if !activation_token.is_empty() + && let Err(err) = self + .store + .put_activation_token(sandbox.object_id(), &activation_token) + .await + { + warn!( + sandbox_id = %sandbox.object_id(), + error = %err, + "Failed to persist warm-pool activation token mapping after driver create" + ); + match self + .driver + .delete_sandbox(Request::new(DeleteSandboxRequest { + sandbox_id: sandbox.object_id().to_string(), + sandbox_name: sandbox.object_name().to_string(), + })) + .await + { + Ok(_) => {} + Err(cleanup_err) => warn!( + sandbox_id = %sandbox.object_id(), + error = %cleanup_err.message(), + "Failed to clean up driver resource after activation token persistence failure" + ), + } + let _ = self + .store + .delete(Sandbox::object_type(), sandbox.object_id()) + .await; + self.sandbox_index.remove_sandbox(sandbox.object_id()); + return Err(Status::internal( + "persist activation token mapping failed after driver create", + )); + } self.sandbox_watch_bus.notify(sandbox.object_id()); if let Some(metadata) = sandbox.metadata.as_mut() { metadata.resource_version = result.resource_version; @@ -886,6 +922,8 @@ impl ComputeRuntime { if deleted { self.cleanup_local_state_if_sandbox_absent(&delete_guard, &target.sandbox_id) .await?; + self.cleanup_activation_token_mapping(&target.sandbox_id) + .await; } else if !self .remove_deleting_sandbox_record(&delete_guard, &target.sandbox_id) .await @@ -1059,6 +1097,7 @@ impl ComputeRuntime { .map_err(|err| err.to_string())? else { self.cleanup_removed_sandbox_state(sandbox_id); + self.cleanup_activation_token_mapping(sandbox_id).await; return Ok(false); }; @@ -1086,10 +1125,12 @@ impl ComputeRuntime { { Ok(true) => { self.cleanup_removed_sandbox_state(sandbox_id); + self.cleanup_activation_token_mapping(sandbox_id).await; Ok(true) } Ok(false) => { self.cleanup_removed_sandbox_state(sandbox_id); + self.cleanup_activation_token_mapping(sandbox_id).await; Ok(false) } Err(crate::persistence::PersistenceError::Conflict { @@ -2005,10 +2046,25 @@ impl ComputeRuntime { .map_err(|err| Status::internal(format!("fetch sandbox failed: {err}")))?; if record.is_none() { self.cleanup_removed_sandbox_state(sandbox_id); + self.cleanup_activation_token_mapping(sandbox_id).await; } Ok(()) } + async fn cleanup_activation_token_mapping(&self, sandbox_id: &str) { + if let Err(err) = self + .store + .delete_activation_token_for_sandbox(sandbox_id) + .await + { + warn!( + sandbox_id, + error = %err, + "Failed to delete warm-pool activation token mapping" + ); + } + } + fn cleanup_removed_sandbox_state(&self, sandbox_id: &str) { self.sandbox_index.remove_sandbox(sandbox_id); self.sandbox_watch_bus.notify(sandbox_id); @@ -2797,7 +2853,9 @@ impl ComputeDriver for NoopTestDriver { ) -> Result, Status> { Ok(tonic::Response::new( - openshell_core::proto::compute::v1::CreateSandboxResponse {}, + openshell_core::proto::compute::v1::CreateSandboxResponse { + activation_token: String::new(), + }, )) } @@ -3060,7 +3118,9 @@ mod tests { &self, _request: Request, ) -> Result, Status> { - Ok(tonic::Response::new(CreateSandboxResponse {})) + Ok(tonic::Response::new(CreateSandboxResponse { + activation_token: String::new(), + })) } async fn stop_sandbox( @@ -3240,7 +3300,9 @@ mod tests { &self, _request: Request, ) -> Result, Status> { - Ok(tonic::Response::new(CreateSandboxResponse {})) + Ok(tonic::Response::new(CreateSandboxResponse { + activation_token: String::new(), + })) } async fn stop_sandbox( diff --git a/crates/openshell-server/src/persistence/mod.rs b/crates/openshell-server/src/persistence/mod.rs index 291e2eafd7..d8e6267c25 100644 --- a/crates/openshell-server/src/persistence/mod.rs +++ b/crates/openshell-server/src/persistence/mod.rs @@ -13,6 +13,7 @@ pub use openshell_core::proto::{ use openshell_core::{Error as CoreError, Result as CoreResult}; use prost::Message; use rand::Rng; +use sha2::{Digest, Sha256}; use std::collections::HashMap; use thiserror::Error; @@ -23,6 +24,12 @@ pub use sqlite::SqliteStore; pub const POLICY_OBJECT_TYPE: &str = "sandbox_policy"; /// Object type string for draft policy chunk records. pub const DRAFT_CHUNK_OBJECT_TYPE: &str = "draft_policy_chunk"; +/// Object type string for warm-pool activation token hashes. +pub const ACTIVATION_TOKEN_OBJECT_TYPE: &str = "activation_token"; + +fn activation_token_record_id(sandbox_id: &str) -> String { + format!("{ACTIVATION_TOKEN_OBJECT_TYPE}/{sandbox_id}") +} pub type PersistenceResult = Result; @@ -210,6 +217,88 @@ impl Store { store_dispatch!(self.ping()) } + /// Hash an opaque warm-pool activation token for durable storage. + #[must_use] + pub fn activation_token_hash(token: &str) -> String { + let digest = Sha256::digest(token.as_bytes()); + hex::encode(digest) + } + + /// Store the hash of a driver-issued activation token for a sandbox. + /// + /// Re-inserting the same token for the same sandbox is idempotent. Changing + /// the token for a sandbox, or binding an existing token hash to another + /// sandbox, is rejected. + pub async fn put_activation_token( + &self, + sandbox_id: &str, + token: &str, + ) -> PersistenceResult<()> { + if sandbox_id.is_empty() { + return Err(PersistenceError::Encode( + "sandbox_id is required for activation token mapping".to_string(), + )); + } + if token.is_empty() { + return Err(PersistenceError::Encode( + "activation token cannot be empty".to_string(), + )); + } + + let token_hash = Self::activation_token_hash(token); + if let Some(existing) = self.get_activation_token_hash(sandbox_id).await? { + if existing == token_hash { + return Ok(()); + } + return Err(PersistenceError::Conflict { + current_resource_version: None, + }); + } + + let payload = token_hash.as_bytes(); + self.put_if( + ACTIVATION_TOKEN_OBJECT_TYPE, + &activation_token_record_id(sandbox_id), + &token_hash, + "", + payload, + None, + WriteCondition::MustCreate, + ) + .await + .map(|_| ()) + } + + /// Fetch a stored activation token hash for a sandbox id. + pub async fn get_activation_token_hash( + &self, + sandbox_id: &str, + ) -> PersistenceResult> { + self.get( + ACTIVATION_TOKEN_OBJECT_TYPE, + &activation_token_record_id(sandbox_id), + ) + .await? + .map(|record| { + String::from_utf8(record.payload).map_err(|e| { + PersistenceError::Decode(format!("activation token hash is not valid UTF-8: {e}")) + }) + }) + .transpose() + } + + /// Delete a sandbox's activation token mapping. + pub async fn delete_activation_token_for_sandbox( + &self, + sandbox_id: &str, + ) -> PersistenceResult { + self.delete( + ACTIVATION_TOKEN_OBJECT_TYPE, + &activation_token_record_id(sandbox_id), + ) + .await + } + /// Test support only: close the underlying connection pool. /// /// There is no runtime shutdown path yet. If we add graceful shutdown, diff --git a/crates/openshell-server/src/persistence/tests.rs b/crates/openshell-server/src/persistence/tests.rs index 9539eab49d..de7716655b 100644 --- a/crates/openshell-server/src/persistence/tests.rs +++ b/crates/openshell-server/src/persistence/tests.rs @@ -1,7 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use super::{ObjectType, PersistenceError, Store, generate_name, test_store}; +use super::{ + ACTIVATION_TOKEN_OBJECT_TYPE, ObjectType, PersistenceError, Store, generate_name, test_store, +}; use crate::policy_store::{AtomicPolicyRevisionWrite, PolicyStoreExt}; use openshell_core::proto::datamodel::v1::ObjectMeta as ProtoObjectMeta; use openshell_core::proto::{ObjectForTest, Sandbox, SandboxPolicy, SandboxSpec}; @@ -266,6 +268,97 @@ async fn sqlite_delete_behavior() { assert!(!deleted_again); } +#[tokio::test] +async fn activation_token_mapping_stores_hash_and_deletes() { + let store = test_store().await; + + store + .put_activation_token("sandbox-a", "raw-activation-token") + .await + .unwrap(); + let hash = store + .get_activation_token_hash("sandbox-a") + .await + .unwrap() + .unwrap(); + + assert_eq!(hash, Store::activation_token_hash("raw-activation-token")); + assert_ne!(hash, "raw-activation-token"); + + let record = store + .get(ACTIVATION_TOKEN_OBJECT_TYPE, "activation_token/sandbox-a") + .await + .unwrap() + .unwrap(); + assert_eq!(record.payload, hash.as_bytes()); + assert_ne!(record.payload, b"raw-activation-token"); + + assert!( + store + .delete_activation_token_for_sandbox("sandbox-a") + .await + .unwrap() + ); + assert!( + store + .get_activation_token_hash("sandbox-a") + .await + .unwrap() + .is_none() + ); +} + +#[tokio::test] +async fn activation_token_mapping_is_idempotent_for_same_token() { + let store = test_store().await; + + store + .put_activation_token("sandbox-a", "raw-activation-token") + .await + .unwrap(); + store + .put_activation_token("sandbox-a", "raw-activation-token") + .await + .unwrap(); + + assert_eq!( + store.get_activation_token_hash("sandbox-a").await.unwrap(), + Some(Store::activation_token_hash("raw-activation-token")) + ); +} + +#[tokio::test] +async fn activation_token_mapping_rejects_token_change_for_sandbox() { + let store = test_store().await; + + store + .put_activation_token("sandbox-a", "raw-activation-token") + .await + .unwrap(); + let err = store + .put_activation_token("sandbox-a", "other-token") + .await + .unwrap_err(); + + assert!(matches!(err, PersistenceError::Conflict { .. })); +} + +#[tokio::test] +async fn activation_token_mapping_rejects_token_reuse_for_other_sandbox() { + let store = test_store().await; + + store + .put_activation_token("sandbox-a", "raw-activation-token") + .await + .unwrap(); + let err = store + .put_activation_token("sandbox-b", "raw-activation-token") + .await + .unwrap_err(); + + assert!(matches!(err, PersistenceError::UniqueViolation { .. })); +} + #[tokio::test] async fn sqlite_protobuf_round_trip() { let store = test_store().await; diff --git a/crates/openshell-server/src/test_support.rs b/crates/openshell-server/src/test_support.rs index 9cd80d6ed8..589d55105c 100644 --- a/crates/openshell-server/src/test_support.rs +++ b/crates/openshell-server/src/test_support.rs @@ -257,7 +257,9 @@ impl ComputeDriver for FakeComputeDriver { .calls .push(FakeComputeDriverCall::CreateSandbox { sandbox }); }); - Ok(Response::new(CreateSandboxResponse {})) + Ok(Response::new(CreateSandboxResponse { + activation_token: String::new(), + })) } async fn stop_sandbox( diff --git a/crates/openshell-server/src/warm_pod_activation.rs b/crates/openshell-server/src/warm_pod_activation.rs index b93c9307ae..e19b6a135a 100644 --- a/crates/openshell-server/src/warm_pod_activation.rs +++ b/crates/openshell-server/src/warm_pod_activation.rs @@ -8,6 +8,7 @@ //! sandbox. use crate::ServerState; +use crate::persistence::Store; use async_trait::async_trait; use openshell_core::proto::{PodActivationMessage, Sandbox}; use openshell_core::supervisor_bootstrap::{ @@ -36,6 +37,22 @@ pub struct ValidatedWarmPodActivation { pub sandbox_id: String, } +#[derive(Debug)] +pub enum WarmPodActivationError { + RetryLater(Status), + Fatal(Status), +} + +impl WarmPodActivationError { + fn fatal(status: Status) -> Self { + Self::Fatal(status) + } + + fn retry_later(status: Status) -> Self { + Self::RetryLater(status) + } +} + #[async_trait] #[allow(dead_code)] pub trait WarmPodActivationValidator: Send + Sync { @@ -43,7 +60,7 @@ pub trait WarmPodActivationValidator: Send + Sync { &self, target: &WarmPodActivationTarget, pending: &SupervisorBootstrapIdentity, - ) -> Result; + ) -> Result; } #[derive(Debug)] @@ -72,19 +89,54 @@ impl SupervisorBootstrapActivator for GatewaySupervisorBootstrapActivator { sandbox_id: request.sandbox_id, owner_uid: request.owner_uid, }; - activate_warm_pod(&self.state, &TrustDriverActivationRequest, target).await + let validator = TokenBackedActivationRequest { + state: self.state.clone(), + activation_token: request.activation_token, + }; + activate_warm_pod(&self.state, &validator, target).await } } -struct TrustDriverActivationRequest; +struct TokenBackedActivationRequest { + state: Arc, + activation_token: String, +} #[async_trait] -impl WarmPodActivationValidator for TrustDriverActivationRequest { +impl WarmPodActivationValidator for TokenBackedActivationRequest { async fn validate( &self, target: &WarmPodActivationTarget, pending: &SupervisorBootstrapIdentity, - ) -> Result { + ) -> Result { + if self.activation_token.is_empty() { + return Err(WarmPodActivationError::fatal(Status::invalid_argument( + "activation token is required", + ))); + } + + let expected_hash = self + .state + .store + .get_activation_token_hash(&target.sandbox_id) + .await + .map_err(|err| { + WarmPodActivationError::retry_later(Status::unavailable(format!( + "activation token mapping unavailable: {err}" + ))) + })? + .ok_or_else(|| { + WarmPodActivationError::retry_later(Status::unavailable( + "activation token mapping is not available yet", + )) + })?; + let request_hash = Store::activation_token_hash(&self.activation_token); + if request_hash != expected_hash { + return Err(WarmPodActivationError::fatal(Status::permission_denied( + "activation token does not match sandbox", + ))); + } + Ok(ValidatedWarmPodActivation { driver: target.driver.clone(), instance_id: target.instance_id.clone(), @@ -110,13 +162,25 @@ where .pending_identity(&target.instance_id)?; let activation_result = async { let validated = validator.validate(&target, &pending).await?; - ensure_validated_activation_matches_pending(&target, &pending, &validated)?; - mint_pod_activation(state, &validated.sandbox_id, "WarmPodActivation").await + ensure_validated_activation_matches_pending(&target, &pending, &validated) + .map_err(WarmPodActivationError::fatal)?; + mint_pod_activation(state, &validated.sandbox_id, "WarmPodActivation") + .await + .map_err(|status| { + if status.code() == tonic::Code::NotFound { + WarmPodActivationError::retry_later(status) + } else { + WarmPodActivationError::fatal(status) + } + }) } .await; let activation = match activation_result { Ok(activation) => activation, - Err(status) => { + Err(WarmPodActivationError::RetryLater(status)) => { + return Err(status); + } + Err(WarmPodActivationError::Fatal(status)) => { let stream_status = clone_status(&status); state .supervisor_pod_registrations @@ -246,7 +310,7 @@ mod tests { &self, _target: &WarmPodActivationTarget, _pending: &SupervisorBootstrapIdentity, - ) -> Result { + ) -> Result { Ok(self.validated.clone()) } } @@ -391,13 +455,126 @@ mod tests { .expect_err("missing sandbox must fail"); assert_eq!(err.code(), tonic::Code::NotFound); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 1); + assert!( + tokio::time::timeout(Duration::from_millis(10), stream.next()) + .await + .is_err(), + "retryable missing sandbox must leave stream pending" + ); + } + + #[tokio::test] + async fn token_backed_activation_accepts_matching_token() { + let state = state_with_issuer().await; + state + .store + .put_activation_token("sandbox-a", "activation-token-a") + .await + .unwrap(); + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = TokenBackedActivationRequest { + state: state.clone(), + activation_token: "activation-token-a".to_string(), + }; + + activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect("activate"); + + let received = stream + .next() + .await + .expect("activation message") + .expect("activation OK"); + assert_eq!(received.sandbox_id, "sandbox-a"); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); + } + + #[tokio::test] + async fn token_backed_activation_missing_mapping_is_retryable() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = TokenBackedActivationRequest { + state: state.clone(), + activation_token: "activation-token-a".to_string(), + }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("missing mapping must fail"); + + assert_eq!(err.code(), tonic::Code::Unavailable); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 1); + assert!( + tokio::time::timeout(Duration::from_millis(10), stream.next()) + .await + .is_err(), + "retryable token mapping miss must leave stream pending" + ); + } + + #[tokio::test] + async fn token_backed_activation_rejects_token_mismatch_as_fatal() { + let state = state_with_issuer().await; + state + .store + .put_activation_token("sandbox-a", "activation-token-a") + .await + .unwrap(); + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = TokenBackedActivationRequest { + state: state.clone(), + activation_token: "wrong-token".to_string(), + }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("mismatched token must fail"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); + let stream_err = stream + .next() + .await + .expect("failure message") + .expect_err("mismatch should fail stream"); + assert_eq!(stream_err.code(), tonic::Code::PermissionDenied); + } + + #[tokio::test] + async fn token_backed_activation_rejects_empty_token_as_fatal() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = TokenBackedActivationRequest { + state: state.clone(), + activation_token: String::new(), + }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("empty token must fail"); + + assert_eq!(err.code(), tonic::Code::InvalidArgument); assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); let stream_err = stream .next() .await .expect("failure message") - .expect_err("missing sandbox should fail stream"); - assert_eq!(stream_err.code(), tonic::Code::NotFound); + .expect_err("empty token should fail stream"); + assert_eq!(stream_err.code(), tonic::Code::InvalidArgument); } #[tokio::test] diff --git a/proto/compute_driver.proto b/proto/compute_driver.proto index c99b7756b0..55fa888c6f 100644 --- a/proto/compute_driver.proto +++ b/proto/compute_driver.proto @@ -242,7 +242,12 @@ message CreateSandboxRequest { DriverSandbox sandbox = 1; } -message CreateSandboxResponse {} +message CreateSandboxResponse { + // Opaque driver-issued token that must be presented during delayed warm + // activation. Empty for cold creates and drivers that do not use delayed + // activation. + string activation_token = 1 [(openshell.options.v1.secret) = true]; +} message StopSandboxRequest { // Stable sandbox ID stored by the gateway. From cdfe4be0875e6c9cf6d639381ef65baeebf658c5 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Fri, 24 Jul 2026 22:52:47 +0100 Subject: [PATCH 10/10] test(kubernetes): add warm pool e2e coverage Signed-off-by: Gordon Sim --- e2e/rust/Cargo.toml | 5 + e2e/rust/src/harness/kubernetes.rs | 219 ++++++++ e2e/rust/src/harness/mod.rs | 1 + e2e/rust/tests/kubernetes_warm_pool.rs | 671 +++++++++++++++++++++++++ e2e/with-kube-gateway.sh | 4 + 5 files changed, 900 insertions(+) create mode 100644 e2e/rust/src/harness/kubernetes.rs create mode 100644 e2e/rust/tests/kubernetes_warm_pool.rs diff --git a/e2e/rust/Cargo.toml b/e2e/rust/Cargo.toml index e4d99dc45c..d5ec1364c1 100644 --- a/e2e/rust/Cargo.toml +++ b/e2e/rust/Cargo.toml @@ -87,6 +87,11 @@ name = "host_gateway_alias" path = "tests/host_gateway_alias.rs" required-features = ["e2e-host-gateway"] +[[test]] +name = "kubernetes_warm_pool" +path = "tests/kubernetes_warm_pool.rs" +required-features = ["e2e-kubernetes"] + [[test]] name = "forward_proxy_l7_bypass" path = "tests/forward_proxy_l7_bypass.rs" diff --git a/e2e/rust/src/harness/kubernetes.rs b/e2e/rust/src/harness/kubernetes.rs new file mode 100644 index 0000000000..5be3a6dbae --- /dev/null +++ b/e2e/rust/src/harness/kubernetes.rs @@ -0,0 +1,219 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Lightweight Kubernetes helpers for e2e tests. +//! +//! These helpers intentionally shell out to `kubectl` so the e2e crate does not +//! need a Kubernetes client dependency. They inherit `KUBECONFIG` and, when +//! present, `OPENSHELL_E2E_KUBE_CONTEXT` from the wrapper. + +use std::process::Stdio; +use std::time::{Duration, Instant}; + +use serde_json::Value; +use tokio::process::Command; +use tokio::time::sleep; + +pub async fn kubectl(args: &[&str]) -> Result { + let mut cmd = Command::new("kubectl"); + if let Ok(context) = std::env::var("OPENSHELL_E2E_KUBE_CONTEXT") + && !context.trim().is_empty() + { + cmd.arg("--context").arg(context); + } + cmd.args(args).stdout(Stdio::piped()).stderr(Stdio::piped()); + + let output = cmd + .output() + .await + .map_err(|err| format!("failed to run kubectl {args:?}: {err}"))?; + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + let combined = format!("{stdout}{stderr}"); + + if !output.status.success() { + return Err(format!( + "kubectl {args:?} failed (exit {:?}):\n{combined}", + output.status.code() + )); + } + + Ok(stdout) +} + +pub async fn kubectl_json(args: &[&str]) -> Result { + let output = kubectl(args).await?; + serde_json::from_str(&output) + .map_err(|err| format!("kubectl {args:?} did not return valid JSON: {err}\n{output}")) +} + +pub async fn has_crd(plural: &str, group: &str) -> bool { + let name = format!("{plural}.{group}"); + kubectl(&["get", "crd", &name]).await.is_ok() +} + +pub async fn wait_for_jsonpath( + namespace: &str, + kind: &str, + name: &str, + jsonpath: &str, + expected: &str, + timeout: Duration, +) -> Result<(), String> { + let deadline = Instant::now() + timeout; + let mut last_output: Option; + let output_arg = format!("jsonpath={jsonpath}"); + + loop { + match kubectl(&["get", kind, name, "-n", namespace, "-o", &output_arg]).await { + Ok(output) => { + let trimmed = output.trim().to_string(); + if trimmed == expected { + return Ok(()); + } + last_output = Some(trimmed); + } + Err(err) => last_output = Some(err), + } + + if Instant::now() >= deadline { + let last_output = last_output.unwrap_or_else(|| "".to_string()); + return Err(format!( + "timed out after {}s waiting for {kind}/{name} jsonpath {jsonpath} to equal {expected:?}. Last output:\n{last_output}", + timeout.as_secs() + )); + } + sleep(Duration::from_secs(2)).await; + } +} + +pub async fn wait_for_resource_by_label( + namespace: &str, + kind: &str, + selector: &str, + timeout: Duration, +) -> Result { + let deadline = Instant::now() + timeout; + let mut last_output: Option; + + loop { + match kubectl_json(&["get", kind, "-n", namespace, "-l", selector, "-o", "json"]).await { + Ok(value) if !items(&value).is_empty() => return Ok(value), + Ok(value) => last_output = Some(value.to_string()), + Err(err) => last_output = Some(err), + } + + if Instant::now() >= deadline { + let last_output = last_output.unwrap_or_else(|| "".to_string()); + return Err(format!( + "timed out after {}s waiting for {kind} in namespace {namespace} with selector {selector}. Last output:\n{last_output}", + timeout.as_secs() + )); + } + sleep(Duration::from_secs(2)).await; + } +} + +pub async fn wait_for_resource_absent_by_label( + namespace: &str, + kind: &str, + selector: &str, + timeout: Duration, +) -> Result<(), String> { + let deadline = Instant::now() + timeout; + let mut last_output: Option; + + loop { + match kubectl_json(&["get", kind, "-n", namespace, "-l", selector, "-o", "json"]).await { + Ok(value) if items(&value).is_empty() => return Ok(()), + Ok(value) => last_output = Some(value.to_string()), + Err(err) if err.contains("the server doesn't have a resource type") => return Ok(()), + Err(err) => last_output = Some(err), + } + + if Instant::now() >= deadline { + let last_output = last_output.unwrap_or_else(|| "".to_string()); + return Err(format!( + "timed out after {}s waiting for {kind} in namespace {namespace} with selector {selector} to be absent. Last output:\n{last_output}", + timeout.as_secs() + )); + } + sleep(Duration::from_secs(2)).await; + } +} + +pub async fn delete_resource(namespace: &str, kind: &str, name: &str) -> Result { + kubectl(&[ + "delete", + kind, + name, + "-n", + namespace, + "--ignore-not-found=true", + ]) + .await +} + +pub async fn dump_namespace_diagnostics(namespace: &str) { + eprintln!("=== Kubernetes diagnostics for namespace {namespace} ==="); + for args in [ + vec![ + "get", + "sandboxclaims.extensions.agents.x-k8s.io", + "-n", + namespace, + "-o", + "wide", + ], + vec![ + "get", + "sandboxwarmpools.extensions.agents.x-k8s.io", + "-n", + namespace, + "-o", + "wide", + ], + vec![ + "get", + "sandboxtemplates.extensions.agents.x-k8s.io", + "-n", + namespace, + "-o", + "wide", + ], + vec![ + "get", + "sandboxes.agents.x-k8s.io", + "-n", + namespace, + "-o", + "wide", + ], + vec!["get", "pods", "-n", namespace, "-o", "wide"], + vec!["get", "events", "-n", namespace, "--sort-by=.lastTimestamp"], + vec![ + "logs", + "-n", + namespace, + "-l", + "app.kubernetes.io/instance=openshell", + "--tail=200", + "--all-containers", + "--prefix", + ], + ] { + eprintln!("--- kubectl {args:?} ---"); + match kubectl(&args).await { + Ok(output) => eprintln!("{output}"), + Err(err) => eprintln!("{err}"), + } + } + eprintln!("=== end Kubernetes diagnostics ==="); +} + +pub fn items(value: &Value) -> &[Value] { + value + .get("items") + .and_then(Value::as_array) + .map_or(&[], Vec::as_slice) +} diff --git a/e2e/rust/src/harness/mod.rs b/e2e/rust/src/harness/mod.rs index f2dfd5ec9c..272e1e303b 100644 --- a/e2e/rust/src/harness/mod.rs +++ b/e2e/rust/src/harness/mod.rs @@ -7,6 +7,7 @@ pub mod binary; pub mod cli; pub mod container; pub mod gateway; +pub mod kubernetes; pub mod output; pub mod port; pub mod sandbox; diff --git a/e2e/rust/tests/kubernetes_warm_pool.rs b/e2e/rust/tests/kubernetes_warm_pool.rs new file mode 100644 index 0000000000..290bedb918 --- /dev/null +++ b/e2e/rust/tests/kubernetes_warm_pool.rs @@ -0,0 +1,671 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![cfg(feature = "e2e-kubernetes")] + +use std::io::Write; +use std::time::Instant; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use openshell_e2e::harness::cli::{run_cli, sandbox_names, wait_for_sandbox_exec_contains}; +use openshell_e2e::harness::kubernetes::{ + delete_resource, dump_namespace_diagnostics, has_crd, items, kubectl, kubectl_json, + wait_for_jsonpath, wait_for_resource_absent_by_label, wait_for_resource_by_label, +}; +use serde_json::Value; +use tempfile::NamedTempFile; + +const EXTENSIONS_GROUP: &str = "extensions.agents.x-k8s.io"; +const AGENTS_GROUP: &str = "agents.x-k8s.io"; +const CLAIM_KIND: &str = "sandboxclaims.extensions.agents.x-k8s.io"; +const WARM_POOL_KIND: &str = "sandboxwarmpools.extensions.agents.x-k8s.io"; +const TEMPLATE_KIND: &str = "sandboxtemplates.extensions.agents.x-k8s.io"; +const SANDBOX_KIND: &str = "sandboxes.agents.x-k8s.io"; + +const LABEL_ENABLED: &str = "openshell.ai/enabled"; +const LABEL_MANAGED_BY: &str = "openshell.ai/managed-by"; +const LABEL_PROFILE_UID: &str = "openshell.ai/warm-pool-profile-uid"; +const LABEL_ALLOCATION: &str = "openshell.ai/allocation"; +const LABEL_SANDBOX_NAME: &str = "openshell.ai/sandbox-name"; +const LABEL_SANDBOX_WORKSPACE: &str = "openshell.ai/sandbox-workspace"; + +const ANNOTATION_STATUS: &str = "openshell.ai/warm-pool-status"; +const ANNOTATION_TEMPLATE: &str = "openshell.ai/warm-pool-template"; +const ANNOTATION_FINGERPRINT: &str = "openshell.ai/warm-pool-fingerprint"; + +#[derive(Clone, Copy)] +enum ProfileCase { + Default, + Env, + Cpu, +} + +struct GeneratedProfile { + source_name: String, + selector: String, + template_name: String, + warm_pool_name: String, +} + +struct TestContext { + namespace: String, + profiles: Vec, + sandboxes: Vec, +} + +impl TestContext { + fn new(namespace: String) -> Self { + Self { + namespace, + profiles: Vec::new(), + sandboxes: Vec::new(), + } + } + + async fn cleanup(&mut self) { + for sandbox in self.sandboxes.iter().rev() { + let _ = run_cli(&["sandbox", "delete", sandbox]).await; + } + for profile in self.profiles.iter().rev() { + let _ = delete_resource(&self.namespace, "configmap", profile).await; + } + } +} + +#[tokio::test] +async fn kubernetes_warm_pool_profiles_claim_and_fallback() { + if !env_flag("OPENSHELL_E2E_KUBE_WARM_POOL") { + eprintln!( + "Skipping Kubernetes warm-pool e2e test: set OPENSHELL_E2E_KUBE_WARM_POOL=1 to enable" + ); + return; + } + + let namespace = + std::env::var("OPENSHELL_E2E_SANDBOX_NAMESPACE").unwrap_or_else(|_| "openshell".into()); + let strict = env_flag("OPENSHELL_E2E_KUBE_WARM_POOL_STRICT"); + if let Err(err) = ensure_required_crds(strict).await { + if strict { + panic!("{err}"); + } + eprintln!("Skipping Kubernetes warm-pool e2e test: {err}"); + return; + } + + let mut ctx = TestContext::new(namespace); + let result = run_warm_pool_e2e(&mut ctx).await; + if let Err(err) = result { + dump_namespace_diagnostics(&ctx.namespace).await; + ctx.cleanup().await; + panic!("{err}"); + } + ctx.cleanup().await; +} + +async fn run_warm_pool_e2e(ctx: &mut TestContext) -> Result<(), String> { + let profile_suffix = unique_profile_suffix(); + let sandbox_suffix = unique_sandbox_suffix(); + + let default_profile = apply_and_assert_profile( + ctx, + &format!("openshell-wp-e2e-default-{profile_suffix}"), + ProfileCase::Default, + ) + .await?; + let env_profile = apply_and_assert_profile( + ctx, + &format!("openshell-wp-e2e-env-{profile_suffix}"), + ProfileCase::Env, + ) + .await?; + let cpu_profile = apply_and_assert_profile( + ctx, + &format!("openshell-wp-e2e-cpu-{profile_suffix}"), + ProfileCase::Cpu, + ) + .await?; + + create_and_assert_claimed( + ctx, + &format!("wpd-{sandbox_suffix}"), + &[], + "warm-pool-ok", + &default_profile, + ) + .await?; + create_and_assert_claimed( + ctx, + &format!("wpe-{sandbox_suffix}"), + &["--env", "FOO=bar"], + "env-ok", + &env_profile, + ) + .await?; + create_and_assert_claimed( + ctx, + &format!("wpc-{sandbox_suffix}"), + &["--cpu", "0.2"], + "cpu-ok", + &cpu_profile, + ) + .await?; + + create_and_assert_direct_fallback(ctx, &format!("wpf-{sandbox_suffix}")).await?; + delete_profile_and_assert_gc(ctx, &env_profile, &[&default_profile, &cpu_profile]).await?; + + Ok(()) +} + +async fn ensure_required_crds(strict: bool) -> Result<(), String> { + let required = [ + ("sandboxes", AGENTS_GROUP), + ("sandboxclaims", EXTENSIONS_GROUP), + ("sandboxtemplates", EXTENSIONS_GROUP), + ("sandboxwarmpools", EXTENSIONS_GROUP), + ]; + let mut missing = Vec::new(); + for (plural, group) in required { + if !has_crd(plural, group).await { + missing.push(format!("{plural}.{group}")); + } + } + if missing.is_empty() { + Ok(()) + } else if strict { + Err(format!( + "required Agent Sandbox CRDs are missing in strict mode: {}", + missing.join(", ") + )) + } else { + Err(format!( + "required Agent Sandbox CRDs are missing: {}", + missing.join(", ") + )) + } +} + +async fn apply_and_assert_profile( + ctx: &mut TestContext, + name: &str, + case: ProfileCase, +) -> Result { + ctx.profiles.push(name.to_string()); + let manifest = profile_configmap_yaml(&ctx.namespace, name, case); + let mut file = + NamedTempFile::new().map_err(|err| format!("create profile manifest tempfile: {err}"))?; + file.write_all(manifest.as_bytes()) + .map_err(|err| format!("write profile manifest tempfile: {err}"))?; + let path = file + .path() + .to_str() + .ok_or_else(|| "profile manifest tempfile path is not UTF-8".to_string())?; + kubectl(&["apply", "-f", path]).await?; + + wait_for_jsonpath( + &ctx.namespace, + "configmap", + name, + "{.metadata.annotations.openshell\\.ai/warm-pool-status}", + "Ready", + Duration::from_secs(120), + ) + .await?; + + let source = + kubectl_json(&["get", "configmap", name, "-n", &ctx.namespace, "-o", "json"]).await?; + let source_uid = source["metadata"]["uid"] + .as_str() + .ok_or_else(|| format!("ConfigMap/{name} is missing metadata.uid"))? + .to_string(); + let annotations = source["metadata"]["annotations"] + .as_object() + .ok_or_else(|| format!("ConfigMap/{name} is missing status annotations"))?; + let template_annotation = annotations + .get(ANNOTATION_TEMPLATE) + .and_then(Value::as_str) + .ok_or_else(|| format!("ConfigMap/{name} is missing {ANNOTATION_TEMPLATE} annotation"))?; + let fingerprint = annotations + .get(ANNOTATION_FINGERPRINT) + .and_then(Value::as_str) + .ok_or_else(|| format!("ConfigMap/{name} is missing {ANNOTATION_FINGERPRINT} annotation"))? + .to_string(); + if fingerprint.is_empty() { + return Err(format!( + "ConfigMap/{name} has empty {ANNOTATION_FINGERPRINT} annotation" + )); + } + if annotations.get(ANNOTATION_STATUS).and_then(Value::as_str) != Some("Ready") { + return Err(format!( + "ConfigMap/{name} did not reconcile Ready: {source}" + )); + } + + let selector = + format!("{LABEL_MANAGED_BY}=openshell-kubernetes-driver,{LABEL_PROFILE_UID}={source_uid}"); + let template = single_labelled_resource(&ctx.namespace, TEMPLATE_KIND, &selector).await?; + let warm_pool = single_labelled_resource(&ctx.namespace, WARM_POOL_KIND, &selector).await?; + + let template_name = object_name(&template)?; + let warm_pool_name = object_name(&warm_pool)?; + if template_name != template_annotation || warm_pool_name != template_annotation { + return Err(format!( + "generated resource names should match {ANNOTATION_TEMPLATE}={template_annotation}; template={template_name}, warm_pool={warm_pool_name}" + )); + } + + assert_generated_labels(&template, name)?; + assert_generated_labels(&warm_pool, name)?; + if warm_pool["spec"]["sandboxTemplateRef"]["name"].as_str() != Some(&template_name) { + return Err(format!( + "SandboxWarmPool/{warm_pool_name} does not reference SandboxTemplate/{template_name}: {warm_pool}" + )); + } + assert_template_invariants(&template, case)?; + wait_for_warm_pool_ready(&ctx.namespace, &warm_pool_name, Duration::from_secs(120)).await?; + + Ok(GeneratedProfile { + source_name: name.to_string(), + selector, + template_name, + warm_pool_name, + }) +} + +async fn single_labelled_resource( + namespace: &str, + kind: &str, + selector: &str, +) -> Result { + let list = + wait_for_resource_by_label(namespace, kind, selector, Duration::from_secs(120)).await?; + let list_items = items(&list); + if list_items.len() != 1 { + return Err(format!( + "expected exactly one {kind} with selector {selector}, got {}: {list}", + list_items.len() + )); + } + Ok(list_items[0].clone()) +} + +async fn wait_for_warm_pool_ready( + namespace: &str, + name: &str, + timeout: Duration, +) -> Result<(), String> { + let deadline = Instant::now() + timeout; + let mut last_observation: Option; + + loop { + match kubectl_json(&["get", WARM_POOL_KIND, name, "-n", namespace, "-o", "json"]).await { + Ok(value) if warm_pool_ready(&value) => return Ok(()), + Ok(value) => last_observation = Some(value), + Err(err) => last_observation = Some(serde_json::json!({ "error": err })), + } + + if Instant::now() >= deadline { + let last_observation = + last_observation.unwrap_or_else(|| serde_json::json!("")); + return Err(format!( + "SandboxWarmPool/{name} did not become ready within {}s. Last observation:\n{}", + timeout.as_secs(), + serde_json::to_string_pretty(&last_observation) + .unwrap_or_else(|_| last_observation.to_string()) + )); + } + tokio::time::sleep(Duration::from_secs(2)).await; + } +} + +fn warm_pool_ready(value: &Value) -> bool { + let desired = value["spec"]["replicas"].as_u64().unwrap_or(1); + if desired == 0 { + return true; + } + + if status_condition_is_true(value, "Ready") || status_condition_is_true(value, "Available") { + return true; + } + + for path in [ + "/status/readyReplicas", + "/status/availableReplicas", + "/status/currentReadyReplicas", + ] { + if value + .pointer(path) + .and_then(Value::as_u64) + .is_some_and(|ready| ready >= desired) + { + return true; + } + } + + false +} + +fn status_condition_is_true(value: &Value, condition_type: &str) -> bool { + value["status"]["conditions"] + .as_array() + .into_iter() + .flatten() + .any(|condition| { + condition["type"].as_str() == Some(condition_type) + && condition["status"] + .as_str() + .is_some_and(|status| status.eq_ignore_ascii_case("true")) + }) +} + +async fn create_and_assert_claimed( + ctx: &mut TestContext, + sandbox_name: &str, + extra_create_args: &[&str], + marker: &str, + profile: &GeneratedProfile, +) -> Result<(), String> { + ctx.sandboxes.push(sandbox_name.to_string()); + let mut args = vec!["sandbox", "create", "--name", sandbox_name, "--no-tty"]; + args.extend_from_slice(extra_create_args); + args.extend_from_slice(&["--", "echo", marker]); + let (output, code) = run_cli(&args).await; + if code != 0 || !output.contains(marker) { + return Err(format!( + "sandbox create for {sandbox_name} did not succeed through warm pool (exit {code}); expected marker {marker:?}. Output:\n{output}" + )); + } + + let claim = wait_for_claim(ctx, sandbox_name).await?; + if claim["metadata"]["labels"][LABEL_ALLOCATION].as_str() != Some("sandbox-claim") { + return Err(format!( + "SandboxClaim for {sandbox_name} is missing allocation label: {claim}" + )); + } + if claim["metadata"]["labels"][LABEL_SANDBOX_WORKSPACE].as_str() != Some("default") { + return Err(format!( + "SandboxClaim for {sandbox_name} is not in default workspace: {claim}" + )); + } + if claim["spec"]["warmPoolRef"]["name"].as_str() != Some(&profile.warm_pool_name) { + return Err(format!( + "SandboxClaim for {sandbox_name} used wrong warm pool; expected {}, claim: {claim}", + profile.warm_pool_name + )); + } + + let names = sandbox_names().await?; + if !names.iter().any(|name| name == sandbox_name) { + return Err(format!( + "sandbox list --names did not include claimed sandbox {sandbox_name}; names={names:?}" + )); + } + wait_for_sandbox_exec_contains( + sandbox_name, + &["echo", "claimed-ok"], + "claimed-ok", + Duration::from_secs(120), + ) + .await?; + + Ok(()) +} + +async fn wait_for_claim(ctx: &TestContext, sandbox_name: &str) -> Result { + let selector = format!("{LABEL_ALLOCATION}=sandbox-claim,{LABEL_SANDBOX_NAME}={sandbox_name}"); + single_labelled_resource(&ctx.namespace, CLAIM_KIND, &selector).await +} + +async fn create_and_assert_direct_fallback( + ctx: &mut TestContext, + sandbox_name: &str, +) -> Result<(), String> { + ctx.sandboxes.push(sandbox_name.to_string()); + let args = [ + "sandbox", + "create", + "--name", + sandbox_name, + "--no-tty", + "--env", + "FOO=baz", + "--", + "echo", + "fallback-ok", + ]; + let (output, code) = run_cli(&args).await; + if code != 0 || !output.contains("fallback-ok") { + return Err(format!( + "fallback sandbox create failed (exit {code}); output:\n{output}" + )); + } + + let claim_selector = + format!("{LABEL_ALLOCATION}=sandbox-claim,{LABEL_SANDBOX_NAME}={sandbox_name}"); + let claims = kubectl_json(&[ + "get", + CLAIM_KIND, + "-n", + &ctx.namespace, + "-l", + &claim_selector, + "-o", + "json", + ]) + .await?; + if !items(&claims).is_empty() { + return Err(format!( + "fallback sandbox {sandbox_name} unexpectedly allocated through SandboxClaim: {claims}" + )); + } + + let sandbox_selector = + format!("{LABEL_MANAGED_BY}=openshell,{LABEL_SANDBOX_NAME}={sandbox_name}"); + let sandbox = single_labelled_resource(&ctx.namespace, SANDBOX_KIND, &sandbox_selector).await?; + if object_name(&sandbox)? != format!("default--{sandbox_name}") { + return Err(format!( + "fallback sandbox used unexpected direct Sandbox name: {sandbox}" + )); + } + + Ok(()) +} + +async fn delete_profile_and_assert_gc( + ctx: &mut TestContext, + deleted: &GeneratedProfile, + remaining: &[&GeneratedProfile], +) -> Result<(), String> { + delete_resource(&ctx.namespace, "configmap", &deleted.source_name).await?; + wait_for_resource_absent_by_label( + &ctx.namespace, + WARM_POOL_KIND, + &deleted.selector, + Duration::from_secs(120), + ) + .await?; + wait_for_resource_absent_by_label( + &ctx.namespace, + TEMPLATE_KIND, + &deleted.selector, + Duration::from_secs(120), + ) + .await?; + + for profile in remaining { + let template = + single_labelled_resource(&ctx.namespace, TEMPLATE_KIND, &profile.selector).await?; + let warm_pool = + single_labelled_resource(&ctx.namespace, WARM_POOL_KIND, &profile.selector).await?; + if object_name(&template)? != profile.template_name + || object_name(&warm_pool)? != profile.warm_pool_name + { + return Err(format!( + "generated resources for profile {} changed while deleting {}", + profile.source_name, deleted.source_name + )); + } + } + + Ok(()) +} + +fn profile_configmap_yaml(namespace: &str, name: &str, case: ProfileCase) -> String { + let mut profile = String::from( + r#" version = 1 + workspace = "default" + replicas = 1 +"#, + ); + match case { + ProfileCase::Default => {} + ProfileCase::Env => profile.push_str( + r#" + [environment] + FOO = "bar" +"#, + ), + ProfileCase::Cpu => profile.push_str( + r#" + [resources] + cpu = "0.2" +"#, + ), + } + + format!( + r#"apiVersion: v1 +kind: ConfigMap +metadata: + name: {name} + namespace: {namespace} + labels: + openshell.ai/warm-pool-profile: "true" +data: + warm-pool.toml: | +{profile}"# + ) +} + +fn assert_generated_labels(obj: &Value, source_name: &str) -> Result<(), String> { + let labels = obj["metadata"]["labels"] + .as_object() + .ok_or_else(|| format!("generated object is missing metadata.labels: {obj}"))?; + if labels.get(LABEL_ENABLED).and_then(Value::as_str) != Some("true") { + return Err(format!( + "generated object is missing {LABEL_ENABLED}=true: {obj}" + )); + } + if labels.get(LABEL_MANAGED_BY).and_then(Value::as_str) != Some("openshell-kubernetes-driver") { + return Err(format!( + "generated object is missing driver manager label: {obj}" + )); + } + if labels + .get("openshell.ai/warm-pool-profile") + .and_then(Value::as_str) + .is_none_or(str::is_empty) + { + return Err(format!( + "generated object for source {source_name} is missing profile label: {obj}" + )); + } + Ok(()) +} + +fn assert_template_invariants(template: &Value, case: ProfileCase) -> Result<(), String> { + if template["spec"]["podTemplate"]["spec"]["dnsPolicy"].as_str() != Some("ClusterFirst") { + return Err(format!( + "generated SandboxTemplate must set pod dnsPolicy=ClusterFirst: {template}" + )); + } + + let agent = agent_container(template)?; + let image = agent["image"] + .as_str() + .ok_or_else(|| format!("agent container is missing image: {template}"))?; + if image.is_empty() { + return Err(format!( + "agent container image must not be empty: {template}" + )); + } + + let env = agent["env"].as_array().map_or(&[][..], Vec::as_slice); + for reserved in ["OPENSHELL_SANDBOX_ID", "OPENSHELL_SANDBOX"] { + if env + .iter() + .any(|entry| entry["name"].as_str() == Some(reserved)) + { + return Err(format!( + "warm-pool template must not include reserved env {reserved}: {template}" + )); + } + } + + match case { + ProfileCase::Default => {} + ProfileCase::Env => { + let foo = env + .iter() + .find(|entry| entry["name"].as_str() == Some("FOO")) + .and_then(|entry| entry["value"].as_str()); + if foo != Some("bar") { + return Err(format!( + "env warm-pool template did not include FOO=bar: {template}" + )); + } + } + ProfileCase::Cpu => { + let resources = &agent["resources"]; + if resources["limits"]["cpu"].as_str() != Some("0.2") + || resources["requests"]["cpu"].as_str() != Some("0.2") + { + return Err(format!( + "CPU warm-pool template did not render cpu requests/limits=0.2: {template}" + )); + } + } + } + + Ok(()) +} + +fn agent_container(template: &Value) -> Result<&Value, String> { + let containers = template["spec"]["podTemplate"]["spec"]["containers"] + .as_array() + .ok_or_else(|| format!("SandboxTemplate is missing pod containers: {template}"))?; + containers + .iter() + .find(|container| container["name"].as_str() == Some("agent")) + .ok_or_else(|| format!("SandboxTemplate is missing agent container: {template}")) +} + +fn object_name(obj: &Value) -> Result { + obj["metadata"]["name"] + .as_str() + .map(ToOwned::to_owned) + .ok_or_else(|| format!("object is missing metadata.name: {obj}")) +} + +fn unique_profile_suffix() -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + let value = (nanos ^ u128::from(std::process::id())) & 0xffff_ffff; + format!("{value:08x}") +} + +fn unique_sandbox_suffix() -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + let value = (nanos ^ u128::from(std::process::id())) & 0x00ff_ffff; + format!("{value:06x}") +} + +fn env_flag(name: &str) -> bool { + std::env::var(name).is_ok_and(|value| { + value == "1" || value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("yes") + }) +} diff --git a/e2e/with-kube-gateway.sh b/e2e/with-kube-gateway.sh index 0a114288e8..2e585abaf3 100755 --- a/e2e/with-kube-gateway.sh +++ b/e2e/with-kube-gateway.sh @@ -593,6 +593,10 @@ fi echo "Installing agent-sandbox CRDs and controller (${AGENT_SANDBOX_VERSION})..." _agent_sandbox_base="https://github.com/kubernetes-sigs/agent-sandbox/releases/download/${AGENT_SANDBOX_VERSION}" kctl apply -f "${_agent_sandbox_base}/manifest.yaml" +if [ "${OPENSHELL_E2E_KUBE_WARM_POOL:-0}" = "1" ]; then + echo "Installing agent-sandbox extension CRDs and controllers (${AGENT_SANDBOX_VERSION})..." + kctl apply -f "${_agent_sandbox_base}/extensions.yaml" +fi wait_for_agent_sandbox_crd kctl -n agent-sandbox-system rollout status deployment/agent-sandbox-controller --timeout=300s