From ef04de78aeedf1ef7e5d1bcea44e8110886d31a1 Mon Sep 17 00:00:00 2001 From: Nancy Kwong Date: Mon, 19 Oct 2026 23:35:56 +0000 Subject: [PATCH] data-plane: open an API-relayed session stream on the route the API hands over --- Cargo.lock | 1 + crates/apps/data/plane/svc/src/main.rs | 8 + crates/libs/api-route/src/ed25519.rs | 27 + .../src/ed25519/data_plane_handover.rs | 226 +++++++++ .../src/ed25519/data_plane_handover/tests.rs | 244 +++++++++ crates/libs/api/src/command_stream.rs | 42 +- .../libs/api/src/command_stream/handover.rs | 271 ++++++++++ .../api/src/command_stream/handover/tests.rs | 353 +++++++++++++ crates/libs/api/src/routes/mod.rs | 17 +- .../common/grpc/src/internal_transport.rs | 18 + crates/libs/data/plane/server/Cargo.toml | 1 + crates/libs/data/plane/server/src/drain.rs | 23 +- .../libs/data/plane/server/src/drain/tests.rs | 2 + crates/libs/data/plane/server/src/lib.rs | 50 ++ crates/libs/data/plane/server/src/native.rs | 8 +- .../data/plane/server/src/native/handover.rs | 180 +++++++ .../plane/server/src/native/handover/tests.rs | 477 ++++++++++++++++++ .../libs/data/plane/server/src/native/sse.rs | 5 +- .../data/plane/server/src/native/websocket.rs | 23 +- crates/libs/data/plane/server/src/tests.rs | 1 + .../src/data_plane_sessions/route_rows.rs | 14 +- crates/libs/metrics/src/data_plane.rs | 10 + docs/configuration.md | 4 +- 23 files changed, 1976 insertions(+), 29 deletions(-) create mode 100644 crates/libs/api-route/src/ed25519/data_plane_handover.rs create mode 100644 crates/libs/api-route/src/ed25519/data_plane_handover/tests.rs create mode 100644 crates/libs/api/src/command_stream/handover.rs create mode 100644 crates/libs/api/src/command_stream/handover/tests.rs create mode 100644 crates/libs/data/plane/server/src/native/handover.rs create mode 100644 crates/libs/data/plane/server/src/native/handover/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 8603afb754..1a58d78e66 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5733,6 +5733,7 @@ dependencies = [ "serde_json", "sha2 0.11.0", "socket2", + "subtle", "tempfile", "thiserror 2.0.20", "time", diff --git a/crates/apps/data/plane/svc/src/main.rs b/crates/apps/data/plane/svc/src/main.rs index ee527b363a..23cf3e0fad 100644 --- a/crates/apps/data/plane/svc/src/main.rs +++ b/crates/apps/data/plane/svc/src/main.rs @@ -51,6 +51,11 @@ struct Cli { #[arg(long, default_value = DEFAULT_LISTEN)] listen: String, + /// Listen address for in-cluster callers. Keep it off public ingress: + /// requests on it may skip route resolution with a route handover. + #[arg(long, env = "SPACE_DATA_PLANE_INTERNAL_LISTEN")] + internal_listen: Option, + /// First-frame attach deadline for upgraded sockets without a header token. #[arg(long, default_value_t = DEFAULT_ATTACH_DEADLINE_MS)] attach_deadline_ms: u64, @@ -108,6 +113,9 @@ impl From for DataPlaneConfig { fn from(cli: Cli) -> Self { Self { listen: cli.listen, + internal_listen: cli + .internal_listen + .filter(|listen| !listen.trim().is_empty()), attach_deadline: Duration::from_millis(cli.attach_deadline_ms), unauthenticated_upgrade_cap_per_ip: cli.unauthenticated_upgrade_cap_per_ip, max_attached_streams: cli.max_attached_streams, diff --git a/crates/libs/api-route/src/ed25519.rs b/crates/libs/api-route/src/ed25519.rs index 56bf700a3c..5c4114e9d6 100644 --- a/crates/libs/api-route/src/ed25519.rs +++ b/crates/libs/api-route/src/ed25519.rs @@ -348,6 +348,30 @@ pub struct SpaceApiRouteEd25519Authority { } impl SpaceApiRouteEd25519Authority { + pub fn new( + issuer: SpaceApiRouteEd25519Issuer, + cluster_id: ClusterId, + signer: SpaceApiRouteEd25519Signer, + ) -> Self { + Self { + issuer, + cluster_id, + signer, + } + } + + pub fn issuer(&self) -> &SpaceApiRouteEd25519Issuer { + &self.issuer + } + + pub fn cluster_id(&self) -> &ClusterId { + &self.cluster_id + } + + pub fn verification_key(&self) -> SpaceApiRouteEd25519VerificationKey { + self.signer.public_key() + } + pub fn load_from_env() -> Result, SpaceApiRouteEd25519Error> { Self::load_from_environment_impl(&pplx_space_common_core::env::ProcessEnvironment) } @@ -923,5 +947,8 @@ fn unix_now_secs() -> u64 { u64::try_from(jiff::Timestamp::now().as_second()).unwrap_or_default() } +mod data_plane_handover; +pub use data_plane_handover::*; + #[cfg(test)] mod tests; diff --git a/crates/libs/api-route/src/ed25519/data_plane_handover.rs b/crates/libs/api-route/src/ed25519/data_plane_handover.rs new file mode 100644 index 0000000000..706b10f931 --- /dev/null +++ b/crates/libs/api-route/src/ed25519/data_plane_handover.rs @@ -0,0 +1,226 @@ +//! A route that the SPACE API resolved for a session it just created, signed so +//! the data plane can open that session's stream without resolving it again. +//! +//! The API mints a handover only for the in-cluster request it makes itself +//! with the session's fresh attach token. The data plane accepts one only on +//! its internal listener, and only when every bound field matches the request: +//! the session, the SHA-256 of the presented attach token, the session kind, +//! and the sandbox and execution generation that the embedded node route token +//! names. Anything else falls back to the ordinary route resolution. +//! +//! Handovers are signed with the route signing key under their own token +//! prefix and audience. Each prefix is part of the signed bytes, so a node +//! never accepts a handover as a route token and the data plane never accepts +//! a route token as a handover. + +use super::*; +use pplx_space_common_core::types::{ExecutionGeneration, NodeId, OrgId, SandboxId}; + +/// Request header that carries a handover from the API to the data plane. +pub const DATA_PLANE_ROUTE_HANDOVER_HEADER: &str = "x-space-route-handover"; +pub const DATA_PLANE_ROUTE_HANDOVER_AUDIENCE: &str = "space-data-plane-route-handover"; +pub const DATA_PLANE_ROUTE_HANDOVER_SCHEMA_VERSION: u8 = 1; +/// The API sends a handover on a request it is making right away. +pub const DATA_PLANE_ROUTE_HANDOVER_TTL_SECS: u64 = 10; +const DATA_PLANE_ROUTE_HANDOVER_MAX_LIFETIME_SECS: u64 = 30; +const DATA_PLANE_ROUTE_HANDOVER_CLOCK_SKEW_SECS: u64 = 5; +const HANDOVER_TOKEN_PREFIX: &str = "dp-handover-ed25519"; +/// Bounds the header; a larger route (a long command line or environment) +/// is resolved by the data plane instead. +pub const DATA_PLANE_ROUTE_HANDOVER_MAX_BYTES: usize = 16 * 1024; + +/// The route fields the data plane would otherwise get from route resolution. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +pub struct DataPlaneRouteHandoverRoute { + pub session_id: String, + pub token_sha256: String, + pub org_id: OrgId, + pub sandbox_id: SandboxId, + pub kind: String, + pub route_generation: u64, + pub current_node_id: NodeId, + pub backend_endpoint: String, + pub execution_generation: ExecutionGeneration, + pub revocation_epoch: u64, + pub org_revocation_epoch: u64, + pub adapter_state_json: String, + pub space_api_route_token: String, +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +pub struct DataPlaneRouteHandoverClaims { + #[serde(rename = "version")] + pub schema_version: u8, + pub issuer: SpaceApiRouteEd25519Issuer, + pub audience: String, + pub cluster_id: ClusterId, + #[serde(flatten)] + pub route: DataPlaneRouteHandoverRoute, + pub issued_at_unix_secs: u64, + pub expires_at_unix_secs: u64, +} + +#[derive(Debug, thiserror::Error)] +pub enum DataPlaneRouteHandoverError { + #[error("route handover is malformed")] + Malformed, + #[error("route handover exceeds {DATA_PLANE_ROUTE_HANDOVER_MAX_BYTES} bytes")] + TooLarge, + #[error("route handover key id is unknown")] + UnknownKeyId, + #[error("route handover signature is invalid")] + InvalidSignature, + #[error("route handover claims are invalid: {0}")] + InvalidClaims(#[from] serde_json::Error), + #[error("route handover schema version is unsupported")] + UnsupportedVersion, + #[error("route handover issuer is not trusted")] + WrongIssuer, + #[error("route handover audience is wrong")] + WrongAudience, + #[error("route handover was issued for another cluster")] + WrongCluster, + #[error("route handover is expired")] + Expired, + #[error("route handover lifetime is invalid")] + InvalidLifetime, + #[error("route handover could not be signed: {0}")] + Sign(#[from] SpaceApiRouteEd25519Error), +} + +impl DataPlaneRouteHandoverError { + pub fn label(&self) -> &'static str { + match self { + Self::Malformed => "malformed", + Self::TooLarge => "too_large", + Self::UnknownKeyId => "unknown_key_id", + Self::InvalidSignature => "invalid_signature", + Self::InvalidClaims(_) => "invalid_claims", + Self::UnsupportedVersion => "unsupported_version", + Self::WrongIssuer => "wrong_issuer", + Self::WrongAudience => "wrong_audience", + Self::WrongCluster => "wrong_cluster", + Self::Expired => "expired", + Self::InvalidLifetime => "invalid_lifetime", + Self::Sign(_) => "sign_failed", + } + } +} + +impl SpaceApiRouteEd25519Authority { + pub fn mint_data_plane_route_handover( + &self, + route: DataPlaneRouteHandoverRoute, + ) -> Result { + let issued_at_unix_secs = unix_now_secs(); + let claims = DataPlaneRouteHandoverClaims { + schema_version: DATA_PLANE_ROUTE_HANDOVER_SCHEMA_VERSION, + issuer: self.issuer.clone(), + audience: DATA_PLANE_ROUTE_HANDOVER_AUDIENCE.to_string(), + cluster_id: self.cluster_id.clone(), + route, + issued_at_unix_secs, + expires_at_unix_secs: issued_at_unix_secs + DATA_PLANE_ROUTE_HANDOVER_TTL_SECS, + }; + self.signer.mint_data_plane_route_handover(&claims) + } + + /// Signs `claims` as given, lifetime included; for exercising verifiers. + pub fn sign_data_plane_route_handover_claims( + &self, + claims: &DataPlaneRouteHandoverClaims, + ) -> Result { + self.signer.mint_data_plane_route_handover(claims) + } +} + +impl SpaceApiRouteEd25519Signer { + pub fn mint_data_plane_route_handover( + &self, + claims: &DataPlaneRouteHandoverClaims, + ) -> Result { + let payload = URL_SAFE_NO_PAD.encode(serde_json::to_vec(claims)?); + let signing_input = format!("{HANDOVER_TOKEN_PREFIX}.{}.{payload}", self.key_id); + let signature = self.key_pair.sign(signing_input.as_bytes()); + let handover = format!( + "{signing_input}.{}", + URL_SAFE_NO_PAD.encode(signature.as_ref()) + ); + if handover.len() > DATA_PLANE_ROUTE_HANDOVER_MAX_BYTES { + return Err(DataPlaneRouteHandoverError::TooLarge); + } + Ok(handover) + } +} + +impl SpaceApiRouteEd25519Verifier { + /// Checks the signature, issuer, audience, cluster, and lifetime. The + /// caller still has to match the route against the request it arrived on. + pub fn verify_data_plane_route_handover( + &self, + handover: &str, + ) -> Result { + if handover.len() > DATA_PLANE_ROUTE_HANDOVER_MAX_BYTES { + return Err(DataPlaneRouteHandoverError::TooLarge); + } + let mut parts = handover.split('.'); + if parts.next() != Some(HANDOVER_TOKEN_PREFIX) { + return Err(DataPlaneRouteHandoverError::Malformed); + } + let key_id = parts + .next() + .and_then(|value| value.parse::().ok()) + .ok_or(DataPlaneRouteHandoverError::Malformed)?; + let encoded_payload = parts.next().ok_or(DataPlaneRouteHandoverError::Malformed)?; + let encoded_signature = parts.next().ok_or(DataPlaneRouteHandoverError::Malformed)?; + if parts.next().is_some() { + return Err(DataPlaneRouteHandoverError::Malformed); + } + let signature = URL_SAFE_NO_PAD + .decode(encoded_signature) + .map_err(|_| DataPlaneRouteHandoverError::Malformed)?; + let public_key = self + .keys + .get(&key_id) + .ok_or(DataPlaneRouteHandoverError::UnknownKeyId)?; + let signing_input = format!("{HANDOVER_TOKEN_PREFIX}.{key_id}.{encoded_payload}"); + ED25519 + .verify_sig(public_key, signing_input.as_bytes(), &signature) + .map_err(|_| DataPlaneRouteHandoverError::InvalidSignature)?; + let payload = URL_SAFE_NO_PAD + .decode(encoded_payload) + .map_err(|_| DataPlaneRouteHandoverError::Malformed)?; + let claims = serde_json::from_slice::(&payload)?; + if claims.schema_version != DATA_PLANE_ROUTE_HANDOVER_SCHEMA_VERSION { + return Err(DataPlaneRouteHandoverError::UnsupportedVersion); + } + if claims.issuer != self.expected_issuer { + return Err(DataPlaneRouteHandoverError::WrongIssuer); + } + if claims.audience != DATA_PLANE_ROUTE_HANDOVER_AUDIENCE { + return Err(DataPlaneRouteHandoverError::WrongAudience); + } + if claims.cluster_id != self.expected_cluster_id { + return Err(DataPlaneRouteHandoverError::WrongCluster); + } + let lifetime = claims + .expires_at_unix_secs + .checked_sub(claims.issued_at_unix_secs) + .ok_or(DataPlaneRouteHandoverError::InvalidLifetime)?; + if lifetime == 0 || lifetime > DATA_PLANE_ROUTE_HANDOVER_MAX_LIFETIME_SECS { + return Err(DataPlaneRouteHandoverError::InvalidLifetime); + } + let now = unix_now_secs(); + let skew = DATA_PLANE_ROUTE_HANDOVER_CLOCK_SKEW_SECS; + if claims.issued_at_unix_secs > now.saturating_add(skew) { + return Err(DataPlaneRouteHandoverError::InvalidLifetime); + } + if claims.expires_at_unix_secs.saturating_add(skew) <= now { + return Err(DataPlaneRouteHandoverError::Expired); + } + Ok(claims.route) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/libs/api-route/src/ed25519/data_plane_handover/tests.rs b/crates/libs/api-route/src/ed25519/data_plane_handover/tests.rs new file mode 100644 index 0000000000..8a61af62e3 --- /dev/null +++ b/crates/libs/api-route/src/ed25519/data_plane_handover/tests.rs @@ -0,0 +1,244 @@ +use aws_lc_rs::rand::SystemRandom; + +use super::*; + +fn authority(issuer: &str, cluster: &str) -> SpaceApiRouteEd25519Authority { + let pkcs8 = Ed25519KeyPair::generate_pkcs8(&SystemRandom::new()).unwrap(); + SpaceApiRouteEd25519Authority::new( + SpaceApiRouteEd25519Issuer::parse(issuer).unwrap(), + ClusterId::parse(cluster).unwrap(), + SpaceApiRouteEd25519Signer::from_pkcs8( + SpaceApiRouteEd25519KeyId::parse("test-key").unwrap(), + pkcs8.as_ref(), + ) + .unwrap(), + ) +} + +fn verifier_for(authority: &SpaceApiRouteEd25519Authority) -> SpaceApiRouteEd25519Verifier { + SpaceApiRouteEd25519Verifier::new( + authority.issuer().clone(), + authority.cluster_id().clone(), + [authority.verification_key()], + ) + .unwrap() +} + +fn route() -> DataPlaneRouteHandoverRoute { + DataPlaneRouteHandoverRoute { + session_id: "proc_1".to_string(), + token_sha256: "a".repeat(64), + org_id: OrgId::new(), + sandbox_id: SandboxId::new(), + kind: "process".to_string(), + route_generation: 2, + current_node_id: NodeId::new(), + backend_endpoint: "https://10.0.0.1:9800".to_string(), + execution_generation: ExecutionGeneration::new(3).unwrap(), + revocation_epoch: 0, + org_revocation_epoch: 4, + adapter_state_json: r#"{"command":["true"]}"#.to_string(), + space_api_route_token: "ed25519.test-key.payload.signature".to_string(), + } +} + +fn sign_claims( + authority: &SpaceApiRouteEd25519Authority, + edit: impl FnOnce(&mut DataPlaneRouteHandoverClaims), +) -> String { + let now = unix_now_secs(); + let mut claims = DataPlaneRouteHandoverClaims { + schema_version: DATA_PLANE_ROUTE_HANDOVER_SCHEMA_VERSION, + issuer: authority.issuer().clone(), + audience: DATA_PLANE_ROUTE_HANDOVER_AUDIENCE.to_string(), + cluster_id: authority.cluster_id().clone(), + route: route(), + issued_at_unix_secs: now, + expires_at_unix_secs: now + DATA_PLANE_ROUTE_HANDOVER_TTL_SECS, + }; + edit(&mut claims); + authority + .signer + .mint_data_plane_route_handover(&claims) + .unwrap() +} + +#[test] +fn handover_round_trips_through_its_verifier() { + let authority = authority("space-api", "testing"); + let route = route(); + let handover = authority + .mint_data_plane_route_handover(route.clone()) + .unwrap(); + assert_eq!( + verifier_for(&authority) + .verify_data_plane_route_handover(&handover) + .unwrap(), + route + ); +} + +#[test] +fn handover_and_route_tokens_do_not_verify_as_each_other() { + let authority = authority("space-api", "testing"); + let verifier = verifier_for(&authority); + let handover = authority.mint_data_plane_route_handover(route()).unwrap(); + let as_route_token = SpaceApiRouteEd25519Token::parse(handover.clone()); + assert!( + as_route_token.is_err() || verifier.verify(&as_route_token.unwrap()).is_err(), + "a node must not accept a handover as a route token" + ); + + let route_token = authority + .mint_space_api_bridge( + pplx_space_common_core::types::SandboxRouteIdentity::try_new( + SandboxId::new(), + pplx_space_common_core::types::RouteGeneration::new(3).unwrap(), + NodeId::new(), + pplx_space_common_core::types::NodeManagedBy::Space, + ExecutionGeneration::new(3).unwrap(), + pplx_space_common_core::types::SandboxActivationAttemptId::new(), + None, + ) + .unwrap(), + ) + .unwrap(); + assert!(matches!( + verifier.verify_data_plane_route_handover(route_token.secret()), + Err(DataPlaneRouteHandoverError::Malformed) + )); +} + +#[test] +fn handover_rejects_forgery_and_tampering() { + let authority = authority("space-api", "testing"); + let verifier = verifier_for(&authority); + + let other = authority_with_same_identity(&authority); + let forged = other.mint_data_plane_route_handover(route()).unwrap(); + assert!(matches!( + verifier.verify_data_plane_route_handover(&forged), + Err(DataPlaneRouteHandoverError::InvalidSignature) + )); + + let handover = authority.mint_data_plane_route_handover(route()).unwrap(); + let mut parts = handover.split('.').map(str::to_string).collect::>(); + let mut claims: serde_json::Value = + serde_json::from_slice(&URL_SAFE_NO_PAD.decode(&parts[2]).unwrap()).unwrap(); + claims["backend_endpoint"] = "https://203.0.113.9:9800".into(); + parts[2] = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).unwrap()); + assert!(matches!( + verifier.verify_data_plane_route_handover(&parts.join(".")), + Err(DataPlaneRouteHandoverError::InvalidSignature) + )); + + let mut unknown_key = handover.split('.').map(str::to_string).collect::>(); + unknown_key[1] = "other-key".to_string(); + assert!(matches!( + verifier.verify_data_plane_route_handover(&unknown_key.join(".")), + Err(DataPlaneRouteHandoverError::UnknownKeyId) + )); +} + +fn authority_with_same_identity( + authority: &SpaceApiRouteEd25519Authority, +) -> SpaceApiRouteEd25519Authority { + let pkcs8 = Ed25519KeyPair::generate_pkcs8(&SystemRandom::new()).unwrap(); + SpaceApiRouteEd25519Authority::new( + authority.issuer().clone(), + authority.cluster_id().clone(), + SpaceApiRouteEd25519Signer::from_pkcs8( + SpaceApiRouteEd25519KeyId::parse("test-key").unwrap(), + pkcs8.as_ref(), + ) + .unwrap(), + ) +} + +#[test] +fn handover_rejects_expiry_long_lifetimes_and_foreign_issuers() { + let authority = authority("space-api", "testing"); + let verifier = verifier_for(&authority); + let now = unix_now_secs(); + + let expired = sign_claims(&authority, |claims| { + claims.issued_at_unix_secs = now - 60; + claims.expires_at_unix_secs = now - 50; + }); + assert!(matches!( + verifier.verify_data_plane_route_handover(&expired), + Err(DataPlaneRouteHandoverError::Expired) + )); + + let long_lived = sign_claims(&authority, |claims| { + claims.expires_at_unix_secs = now + 3600; + }); + assert!(matches!( + verifier.verify_data_plane_route_handover(&long_lived), + Err(DataPlaneRouteHandoverError::InvalidLifetime) + )); + + let future = sign_claims(&authority, |claims| { + claims.issued_at_unix_secs = now + 60; + claims.expires_at_unix_secs = now + 65; + }); + assert!(matches!( + verifier.verify_data_plane_route_handover(&future), + Err(DataPlaneRouteHandoverError::InvalidLifetime) + )); + + for (edit, expected) in [ + ( + Box::new(|claims: &mut DataPlaneRouteHandoverClaims| { + claims.issuer = SpaceApiRouteEd25519Issuer::parse("someone-else").unwrap(); + }) as Box, + "wrong_issuer", + ), + ( + Box::new(|claims: &mut DataPlaneRouteHandoverClaims| { + claims.audience = SPACE_API_ROUTE_ED25519_AUDIENCE.to_string(); + }), + "wrong_audience", + ), + ( + Box::new(|claims: &mut DataPlaneRouteHandoverClaims| { + claims.cluster_id = ClusterId::parse("elsewhere").unwrap(); + }), + "wrong_cluster", + ), + ( + Box::new(|claims: &mut DataPlaneRouteHandoverClaims| { + claims.schema_version = 9; + }), + "unsupported_version", + ), + ] { + let handover = sign_claims(&authority, edit); + assert_eq!( + verifier + .verify_data_plane_route_handover(&handover) + .unwrap_err() + .label(), + expected + ); + } +} + +#[test] +fn oversized_handovers_are_neither_minted_nor_parsed() { + let authority = authority("space-api", "testing"); + let mut large = route(); + large.adapter_state_json = format!( + r#"{{"command":["{}"]}}"#, + "x".repeat(DATA_PLANE_ROUTE_HANDOVER_MAX_BYTES) + ); + assert!(matches!( + authority.mint_data_plane_route_handover(large), + Err(DataPlaneRouteHandoverError::TooLarge) + )); + assert!(matches!( + verifier_for(&authority) + .verify_data_plane_route_handover(&"x".repeat(DATA_PLANE_ROUTE_HANDOVER_MAX_BYTES + 1)), + Err(DataPlaneRouteHandoverError::TooLarge) + )); +} diff --git a/crates/libs/api/src/command_stream.rs b/crates/libs/api/src/command_stream.rs index 9e83325c66..90200fb3f1 100644 --- a/crates/libs/api/src/command_stream.rs +++ b/crates/libs/api/src/command_stream.rs @@ -20,9 +20,14 @@ use hyper_util::{ client::legacy::{Client, connect::HttpConnector}, rt::TokioExecutor, }; +use pplx_space_common_svc::space_api_route_ed25519::DATA_PLANE_ROUTE_HANDOVER_HEADER; use pplx_space_public_api::{PublicCommand, data_plane_session_events_path}; use tokio_util::sync::CancellationToken; +mod handover; +pub use handover::DATA_PLANE_ROUTE_HANDOVER_ENV; +pub(crate) use handover::RouteHandoverIssuer; + pub const DATA_PLANE_INTERNAL_URL_ENV: &str = "SPACE_DATA_PLANE_INTERNAL_URL"; /// Event name and frame `type` of the first frame, which carries the command @@ -32,6 +37,9 @@ pub const DATA_PLANE_INTERNAL_URL_ENV: &str = "SPACE_DATA_PLANE_INTERNAL_URL"; pub const COMMAND_EVENT: &str = "command"; const CONNECT_TIMEOUT: Duration = Duration::from_secs(2); +// Past this, the data plane resolving the route is faster than waiting on +// the API's own read. +const HANDOVER_TIMEOUT: Duration = Duration::from_millis(500); // The data plane answers once the process has started. Past this, the client // is better served by the handle and its own attach than by a stalled request. const OPEN_TIMEOUT: Duration = Duration::from_secs(30); @@ -66,6 +74,7 @@ pub struct CommandStreamProxy { client: Client>, origin: String, shutdown: CancellationToken, + route_handover: Option, } impl CommandStreamProxy { @@ -96,9 +105,33 @@ impl CommandStreamProxy { client: Client::builder(TokioExecutor::new()).build(connector), origin: format!("http://{authority}"), shutdown, + route_handover: None, }) } + /// Hands the data plane the session's route with each stream it opens, so + /// the data plane does not resolve it again. + pub(crate) fn with_route_handover(mut self, issuer: RouteHandoverIssuer) -> Self { + self.route_handover = Some(issuer); + self + } + + async fn route_handover(&self, command: &PublicCommand, token: &str) -> Option { + let issuer = self.route_handover.as_ref()?; + let session_id = command.session_id.to_string(); + match tokio::time::timeout(HANDOVER_TIMEOUT, issuer.mint(&session_id, token)).await { + Ok(Ok(handover)) => Some(handover), + Ok(Err(error)) => { + tracing::debug!(%session_id, %error, "no route handover for command stream"); + None + } + Err(_) => { + tracing::debug!(%session_id, "route handover timed out"); + None + } + } + } + /// Returns a `201 text/event-stream` response for `command`, or `None` when /// the stream could not be opened and the caller should answer with JSON. pub(crate) async fn stream(&self, command: &PublicCommand) -> Option { @@ -117,14 +150,17 @@ impl CommandStreamProxy { async fn open(&self, command: &PublicCommand) -> Result { let token = &command.attach.as_ref().ok_or(OpenError::NoAttach)?.token; - let request = Request::get(format!( + let mut request = Request::get(format!( "{}{}", self.origin, data_plane_session_events_path(command.session_id) )) .header(header::AUTHORIZATION, format!("Bearer {token}")) - .header(header::ACCEPT, "text/event-stream") - .body(Empty::new())?; + .header(header::ACCEPT, "text/event-stream"); + if let Some(handover) = self.route_handover(command, token).await { + request = request.header(DATA_PLANE_ROUTE_HANDOVER_HEADER, handover); + } + let request = request.body(Empty::new())?; let upstream = tokio::time::timeout(OPEN_TIMEOUT, self.client.request(request)) .await .map_err(|_| OpenError::Timeout)??; diff --git a/crates/libs/api/src/command_stream/handover.rs b/crates/libs/api/src/command_stream/handover.rs new file mode 100644 index 0000000000..7e1a24d7b1 --- /dev/null +++ b/crates/libs/api/src/command_stream/handover.rs @@ -0,0 +1,271 @@ +//! The route handover the API attaches to its in-cluster request for a +//! session it just created, so the data plane can open that session's stream +//! without resolving the route again. +//! +//! The handover carries what route resolution would return for this attach, +//! and is minted only after the same checks route resolution makes: the +//! presented token matches the stored digest, is unexpired, and grants the +//! kind's connect scope; the session is active and has never bound a process; +//! and the sandbox is started and placed on a node with a node-API endpoint. +//! Anything else sends no handover, and the data plane resolves the route +//! itself: resolution refuses a restoring sandbox's route, and holds the route +//! of one not running while a wake is about to activate it. + +use pplx_space_common_core::types::{ + ExecutionGeneration, NodeId, NodeManagedBy, OrgId, PodUid, RouteGeneration, + SandboxActivationAttemptId, SandboxId, SandboxLifecycle, SandboxRouteIdentity, +}; +use pplx_space_common_svc::space_api_route_ed25519::{ + DataPlaneRouteHandoverRoute, SpaceApiRouteEd25519Authority, +}; +use pplx_space_common_svc::sql; +use pplx_space_public_api::SessionKind; +use sha2::{Digest, Sha256}; +use sqlx::PgPool; +use std::sync::Arc; +use subtle::ConstantTimeEq as _; + +/// Turns route handovers on for the in-cluster command stream. Enable it only +/// where the data plane's internal URL reaches its internal listener; +/// elsewhere the data plane ignores handovers and the route read is wasted. +pub const DATA_PLANE_ROUTE_HANDOVER_ENV: &str = "SPACE_DATA_PLANE_ROUTE_HANDOVER"; +const BACKEND_TRANSPORT_GRPC: &str = "grpc"; + +#[derive(Debug, thiserror::Error)] +pub(crate) enum HandoverError { + #[error("read the session route: {0}")] + Read(#[from] sqlx::Error), + #[error("session not found")] + NotFound, + #[error("session cannot take a handover: {0}")] + Ineligible(&'static str), + #[error("route identity is incomplete: {0}")] + Identity(String), + #[error("mint the route token: {0}")] + RouteToken(#[from] pplx_space_common_svc::space_api_route_ed25519::SpaceApiRouteEd25519Error), + #[error("mint the handover: {0}")] + Handover(#[from] pplx_space_common_svc::space_api_route_ed25519::DataPlaneRouteHandoverError), +} + +#[derive(Clone)] +pub(crate) struct RouteHandoverIssuer { + pool: PgPool, + authority: Arc, +} + +#[derive(sqlx::FromRow)] +struct HandoverRouteRow { + org_id: OrgId, + sandbox_id: SandboxId, + kind: String, + adapter_state_json: String, + route_generation: i64, + current_node_id: Option, + backend_endpoint: Option, + backend_transport: Option, + revocation_epoch: i64, + revocation_status: String, + org_revocation_epoch: i64, + scopes: Vec, + token_hash: Option, + token_unexpired: bool, + backend_process_id: Option, + sandbox_lifecycle: SandboxLifecycle, + sandbox_active_node_id: Option, + node_address: Option, + node_managed_by: Option, + direct_node_api_endpoint: Option, + sandbox_pod_uid: Option, + execution_generation: Option, + activation_attempt_id: Option, +} + +impl RouteHandoverIssuer { + pub(crate) fn new(pool: PgPool, authority: Arc) -> Self { + Self { pool, authority } + } + + /// Mints the handover for attaching `session_id` with `token`. + pub(crate) async fn mint( + &self, + session_id: &str, + token: &str, + ) -> Result { + // The same joins route resolution reads, so both see one placement. + let row = sqlx::query_as::<_, HandoverRouteRow>(sql!( + "SELECT d.org_id, + d.sandbox_id, + d.kind, + d.adapter_state::TEXT AS adapter_state_json, + d.route_generation, + d.current_node_id, + d.backend_endpoint, + d.backend_transport, + d.revocation_epoch, + d.revocation_status, + COALESCE(o.epoch, 0) AS org_revocation_epoch, + d.scopes, + d.token_hash, + COALESCE(d.token_expires_at > NOW(), FALSE) AS token_unexpired, + d.backend_process_id, + s.lifecycle AS sandbox_lifecycle, + COALESCE(direct_route.current_node_id, s.active_node_id) + AS sandbox_active_node_id, + HOST(n.address) AS node_address, + COALESCE(direct_route.substrate_kind, n.managed_by) AS node_managed_by, + nd.node_api_endpoint AS direct_node_api_endpoint, + COALESCE(NULLIF(BTRIM(e.pod_uid), ''), direct_route.pod_uid) + AS sandbox_pod_uid, + COALESCE(e.generation, direct_route.execution_generation) + AS execution_generation, + COALESCE(e.route_activation_id, direct_route.activation_attempt_id) + AS activation_attempt_id + FROM data_plane_sessions d + JOIN sandboxes s ON s.id = d.sandbox_id + LEFT JOIN data_plane_org_revocation_epochs o ON o.org_id = d.org_id + LEFT JOIN sandbox_direct_route_identities direct_route + ON direct_route.sandbox_id = s.id + AND s.active_node_id IS NOT NULL + LEFT JOIN nodes n + ON n.id = COALESCE(direct_route.current_node_id, s.active_node_id) + LEFT JOIN node_direct nd ON nd.node_id = n.id AND nd.deleted_at IS NULL + LEFT JOIN sandbox_k8s_executions e ON e.sandbox_id = s.id AND e.active + WHERE d.session_id = $1 + AND (d.expires_at IS NULL OR d.expires_at > NOW())" + )) + .bind(session_id) + .fetch_optional(&self.pool) + .await? + .ok_or(HandoverError::NotFound)?; + let route = handover_route(session_id, token, row)?; + let route_token = self.authority.mint_space_api_bridge(route.identity)?; + Ok(self + .authority + .mint_data_plane_route_handover(DataPlaneRouteHandoverRoute { + space_api_route_token: route_token.secret().to_string(), + ..route.handover + })?) + } +} + +struct HandoverRoute { + handover: DataPlaneRouteHandoverRoute, + identity: SandboxRouteIdentity, +} + +fn handover_route( + session_id: &str, + token: &str, + row: HandoverRouteRow, +) -> Result { + let kind = row + .kind + .parse::() + .map_err(|_| HandoverError::Ineligible("unknown_kind"))?; + if !matches!( + kind, + SessionKind::Terminal | SessionKind::Process | SessionKind::FileWatch + ) { + return Err(HandoverError::Ineligible("unsupported_kind")); + } + let token_sha256 = hex::encode(Sha256::digest(token.as_bytes())); + let stored = row.token_hash.as_deref().map(str::trim).unwrap_or_default(); + if stored.is_empty() || stored.as_bytes().ct_eq(token_sha256.as_bytes()).unwrap_u8() != 1 { + return Err(HandoverError::Ineligible("token_mismatch")); + } + if !row.token_unexpired { + return Err(HandoverError::Ineligible("token_expired")); + } + if !row + .scopes + .iter() + .any(|scope| scope == kind.connect_scope().as_str()) + { + return Err(HandoverError::Ineligible("scope_missing")); + } + if row.revocation_status != "active" { + return Err(HandoverError::Ineligible("revoked")); + } + if row.backend_process_id.is_some() { + return Err(HandoverError::Ineligible("process_bound")); + } + if row.sandbox_lifecycle != SandboxLifecycle::Started { + return Err(HandoverError::Ineligible("not_started")); + } + let current_node_id = row + .sandbox_active_node_id + .ok_or(HandoverError::Ineligible("unplaced"))?; + let backend_endpoint = pplx_space_common_grpc::internal_transport::node_api_endpoint( + row.direct_node_api_endpoint.as_deref(), + row.node_address.as_deref(), + ) + .ok_or(HandoverError::Ineligible("no_endpoint"))?; + let (Some(substrate_kind), Some(execution_generation), Some(activation_attempt_id)) = ( + row.node_managed_by, + row.execution_generation, + row.activation_attempt_id, + ) else { + return Err(HandoverError::Identity( + "route identity is incomplete".to_string(), + )); + }; + let execution_generation = ExecutionGeneration::new(execution_generation) + .map_err(|error| HandoverError::Identity(error.to_string()))?; + let identity = SandboxRouteIdentity::try_new( + row.sandbox_id, + u64::try_from(execution_generation.get()) + .ok() + .and_then(|value| RouteGeneration::new(value).ok()) + .ok_or_else(|| HandoverError::Identity("execution generation".to_string()))?, + current_node_id, + substrate_kind, + execution_generation, + activation_attempt_id, + row.sandbox_pod_uid, + ) + .map_err(|error| HandoverError::Identity(error.to_string()))?; + + // Route resolution stores this route on the session and bumps its + // generation when it differs from the stored one. The data plane compares + // generations when it later refreshes the route, so hand over the one + // that resolution will store. + let stored_generation = + u64::try_from(row.route_generation).map_err(|_| HandoverError::Ineligible("generation"))?; + let backend_transport = row + .backend_transport + .as_deref() + .filter(|value| !value.trim().is_empty()) + .unwrap_or(BACKEND_TRANSPORT_GRPC); + let unchanged = row.current_node_id == Some(current_node_id) + && row.backend_endpoint.as_deref() == Some(backend_endpoint.as_str()) + && row.backend_transport.as_deref() == Some(backend_transport); + let route_generation = if unchanged { + stored_generation + } else { + stored_generation + 1 + }; + + Ok(HandoverRoute { + handover: DataPlaneRouteHandoverRoute { + session_id: session_id.to_string(), + token_sha256, + org_id: row.org_id, + sandbox_id: row.sandbox_id, + kind: kind.as_str().to_string(), + route_generation, + current_node_id, + backend_endpoint, + execution_generation, + revocation_epoch: u64::try_from(row.revocation_epoch) + .map_err(|_| HandoverError::Ineligible("revocation_epoch"))?, + org_revocation_epoch: u64::try_from(row.org_revocation_epoch) + .map_err(|_| HandoverError::Ineligible("org_revocation_epoch"))?, + adapter_state_json: row.adapter_state_json, + space_api_route_token: String::new(), + }, + identity, + }) +} + +#[cfg(test)] +mod tests; diff --git a/crates/libs/api/src/command_stream/handover/tests.rs b/crates/libs/api/src/command_stream/handover/tests.rs new file mode 100644 index 0000000000..432a7670f1 --- /dev/null +++ b/crates/libs/api/src/command_stream/handover/tests.rs @@ -0,0 +1,353 @@ +use super::*; +use crate::command_stream::CommandStreamProxy; +use crate::routes::AppState; +use axum::{ + Router, + body::Body, + extract::State, + http::{HeaderMap, Method, Request, StatusCode, header}, + response::Response, + routing::get, +}; +use pplx_space_common_svc::space_api_route_ed25519::{ + DATA_PLANE_ROUTE_HANDOVER_HEADER, SPACE_API_ROUTE_ED25519_ISSUER_ENV, + SPACE_API_ROUTE_ED25519_SIGNING_KEY_ID_ENV, SPACE_API_ROUTE_ED25519_SIGNING_KEY_PATH_ENV, + SpaceApiRouteEd25519Token, SpaceApiRouteEd25519Verifier, + decode_unverified_space_api_route_ed25519_claims, +}; +use pplx_space_db_test_fixtures::{SEED_ORG_ID, SEED_USER_ID, TestDbResult, pg_test_db}; +use pplx_space_internal_api_server::ClusterApiService; +use pplx_space_internal_api_server::proto::{ + ResolveDataPlaneSessionRouteRequest, cluster_api_server::ClusterApi, +}; +use pplx_space_public_api::{DATA_PLANE_SESSION_EVENTS_PATH_TEMPLATE, PublicSandboxId}; +use std::sync::Mutex; +use tokio_util::sync::CancellationToken; +use tower::ServiceExt; +use uuid::Uuid; + +const KEY: &str = "route-handover-test-key"; + +fn route_signing_test_process() -> bool { + let key_path = pplx_space_common_svc::space_api_route_ed25519::test_signing_key_path() + .to_str() + .unwrap() + .to_string(); + pplx_space_test_env::enter_test_process(&[ + ("SPACE_CLUSTER_ID", Some(std::ffi::OsStr::new("testing"))), + ( + SPACE_API_ROUTE_ED25519_SIGNING_KEY_PATH_ENV, + Some(std::ffi::OsStr::new(&key_path)), + ), + ( + SPACE_API_ROUTE_ED25519_SIGNING_KEY_ID_ENV, + Some(std::ffi::OsStr::new("handover-test-key")), + ), + ( + SPACE_API_ROUTE_ED25519_ISSUER_ENV, + Some(std::ffi::OsStr::new("space-api")), + ), + ]) +} + +/// A data plane that records the handover each events request carries. +async fn spawn_data_plane() -> (String, Arc>>>) { + let seen = Arc::new(Mutex::new(Vec::new())); + let app = Router::new() + .route( + DATA_PLANE_SESSION_EVENTS_PATH_TEMPLATE, + get( + |State(seen): State>>>>, headers: HeaderMap| async move { + seen.lock().unwrap().push( + headers + .get(DATA_PLANE_ROUTE_HANDOVER_HEADER) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned), + ); + Response::builder() + .header(header::CONTENT_TYPE, "text/event-stream") + .body(Body::from("event: attached\ndata: {}\n\n")) + .unwrap() + }, + ), + ) + .with_state(seen.clone()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let origin = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + (origin, seen) +} + +struct Placed { + sandbox: PublicSandboxId, + node_id: NodeId, +} + +/// A started sandbox on a ready node with a running execution. +async fn place_sandbox(pool: &PgPool, address: &str, generation: i64) -> TestDbResult { + let node_id = NodeId::new(); + sqlx::query(sql!( + "INSERT INTO nodes (id, address, role, managed_by, state) + VALUES ($1, $2::INET, 'worker', 'k8s', 'ready')" + )) + .bind(node_id) + .bind(address) + .execute(pool) + .await?; + let sandbox = PublicSandboxId::from_uuid(Uuid::now_v7()); + sqlx::query(sql!( + "INSERT INTO sandboxes (id, user_id, org_id, managed_by, lifecycle, placement, + active_node_id, target_node_id, network_policy_override_json, + network_policy_generation) + VALUES ($1, $2, $3, 'k8s', 'started', 'placed', $4, $4, $5, 1)" + )) + .bind(sandbox.as_uuid()) + .bind(SEED_USER_ID) + .bind(SEED_ORG_ID) + .bind(node_id) + .bind(sqlx::types::Json( + pplx_space_net_common::NetworkPolicyV1::default(), + )) + .execute(pool) + .await?; + sqlx::query(sql!( + "INSERT INTO sandbox_k8s_executions ( + id, sandbox_id, generation, pod_name, pod_namespace, pod_uid, + pod_ip, route_activation_id, state, active + ) + VALUES ($1, $2, $3, $4, 'space', $5, '10.80.0.9'::INET, $6, 'running', true)" + )) + .bind(Uuid::now_v7()) + .bind(sandbox.as_uuid()) + .bind(generation) + .bind(format!( + "sandbox-pod-{}-{generation}", + sandbox.as_uuid().simple() + )) + .bind(format!("pod-uid-{generation}")) + .bind(Uuid::now_v7()) + .execute(pool) + .await?; + Ok(Placed { sandbox, node_id }) +} + +async fn post_streamed_command(app: &Router, sandbox: PublicSandboxId) -> TestDbResult { + let response = app + .clone() + .oneshot( + Request::builder() + .method(Method::POST) + .uri(format!("/v1/sandboxes/{sandbox}/commands")) + .header(header::AUTHORIZATION, format!("Bearer {KEY}")) + .header(header::CONTENT_TYPE, "application/json") + .header(header::ACCEPT, "text/event-stream") + .header("idempotency-key", Uuid::now_v7().to_string()) + .body(Body::from(serde_json::to_vec( + &serde_json::json!({"command": ["echo", "hello"]}), + )?))?, + ) + .await?; + assert_eq!(response.status(), StatusCode::CREATED); + let body = axum::body::to_bytes(response.into_body(), 1024 * 1024).await?; + let body = String::from_utf8(body.to_vec())?; + let (first, _) = body.split_once("\n\n").expect("command event"); + let data = first + .strip_prefix("event: command\ndata: ") + .expect("command event first"); + Ok(serde_json::from_str::(data)?["command"].clone()) +} + +use serde_json::Value; + +#[tokio::test] +async fn the_handover_carries_the_route_resolution_would_return() -> TestDbResult<()> { + if !route_signing_test_process() { + return Ok(()); + } + let db = pg_test_db().await?; + let pool = db.pool.clone(); + sqlx::query(sql!( + "INSERT INTO api_keys (id, org_id, name, key_hash, allowed_surfaces) + VALUES ($1, $2, 'route-handover-test', $3, ARRAY['native'])" + )) + .bind(Uuid::now_v7()) + .bind(SEED_ORG_ID) + .bind(crate::auth::hash_token(KEY)) + .execute(&pool) + .await?; + let authority = Arc::new(SpaceApiRouteEd25519Authority::load_required_from_env()?); + let verifier = SpaceApiRouteEd25519Verifier::new( + authority.issuer().clone(), + authority.cluster_id().clone(), + [authority.verification_key()], + )?; + let (origin, seen) = spawn_data_plane().await; + let mut state = AppState::from_pool(pool.clone()).await?; + state.command_stream = Some( + CommandStreamProxy::new(&origin, CancellationToken::new())?.with_route_handover( + RouteHandoverIssuer::new(pool.clone(), Arc::clone(&authority)), + ), + ); + let app = Router::new() + .nest("/v1", crate::routes::public::v1_router(state.clone())) + .with_state(state); + let cluster_api = ClusterApiService::new(pool.clone()); + + for (address, generation) in [("10.0.0.21", 4), ("fd00::21", 9)] { + let placed = place_sandbox(&pool, address, generation).await?; + let command = post_streamed_command(&app, placed.sandbox).await?; + let session_id = command["session_id"].as_str().unwrap().to_string(); + let token = command["attach"]["token"].as_str().unwrap().to_string(); + let handover = seen + .lock() + .unwrap() + .last() + .cloned() + .flatten() + .expect("the stream request carried a handover"); + let handed = verifier.verify_data_plane_route_handover(&handover)?; + + let resolved = cluster_api + .resolve_data_plane_session_route(tonic::Request::new( + ResolveDataPlaneSessionRouteRequest { + session_id: session_id.clone(), + protocol_session_id: None, + token_sha256: Some(hex::encode(Sha256::digest(token.as_bytes()))), + required_scope: Some(SessionKind::Process.connect_scope().as_str().to_string()), + requested_transport: None, + }, + )) + .await? + .into_inner(); + let revocation = resolved.revocation.unwrap(); + let backend = resolved.backend.unwrap(); + assert_eq!(handed.session_id, resolved.session_id, "{address}"); + assert_eq!( + handed.sandbox_id.to_string(), + resolved.sandbox_id, + "{address}" + ); + assert_eq!(handed.org_id.to_string(), resolved.org_id, "{address}"); + assert_eq!(handed.kind, resolved.kind, "{address}"); + assert_eq!(handed.current_node_id, placed.node_id, "{address}"); + assert_eq!( + handed.current_node_id.to_string(), + resolved.current_node_id, + "{address}" + ); + assert_eq!(handed.backend_endpoint, backend.endpoint, "{address}"); + assert_eq!( + handed.route_generation, resolved.route_generation, + "{address}: the generation resolution stores" + ); + assert_eq!(handed.revocation_epoch, revocation.epoch, "{address}"); + assert_eq!( + handed.org_revocation_epoch, revocation.org_epoch, + "{address}" + ); + assert_eq!( + serde_json::from_str::(&handed.adapter_state_json)?, + serde_json::from_str::(&resolved.adapter_state_json)?, + "{address}" + ); + assert_eq!(handed.execution_generation.get(), generation, "{address}"); + let route_identity = |token: &str| { + decode_unverified_space_api_route_ed25519_claims( + &SpaceApiRouteEd25519Token::parse(token.to_string()).unwrap(), + ) + .unwrap() + .route_identity + }; + assert_eq!( + route_identity(&handed.space_api_route_token), + route_identity(&resolved.space_api_route_token.unwrap().token), + "{address}: the same node route identity" + ); + } + Ok(()) +} + +fn eligible_row() -> HandoverRouteRow { + HandoverRouteRow { + org_id: OrgId::new(), + sandbox_id: SandboxId::new(), + kind: "process".to_string(), + adapter_state_json: "{}".to_string(), + route_generation: 1, + current_node_id: None, + backend_endpoint: None, + backend_transport: None, + revocation_epoch: 0, + revocation_status: "active".to_string(), + org_revocation_epoch: 0, + scopes: vec!["process:connect".to_string()], + token_hash: Some(hex::encode(Sha256::digest(b"attach-token"))), + token_unexpired: true, + backend_process_id: None, + sandbox_lifecycle: SandboxLifecycle::Started, + sandbox_active_node_id: Some(NodeId::new()), + node_address: Some("10.0.0.1".to_string()), + node_managed_by: Some(NodeManagedBy::Kubernetes), + direct_node_api_endpoint: None, + sandbox_pod_uid: Some(PodUid::parse("pod-uid").unwrap()), + execution_generation: Some(3), + activation_attempt_id: Some(SandboxActivationAttemptId::new()), + } +} + +#[test] +fn only_attaches_that_resolution_would_authorize_get_a_handover() { + let handed = handover_route("proc_1", "attach-token", eligible_row()).unwrap(); + assert_eq!( + handed.handover.route_generation, 2, + "the first route bumps it" + ); + let mut unchanged = eligible_row(); + unchanged.current_node_id = unchanged.sandbox_active_node_id; + unchanged.backend_endpoint = Some(handed.handover.backend_endpoint.clone()); + unchanged.backend_transport = Some("grpc".to_string()); + assert_eq!( + handover_route("proc_1", "attach-token", unchanged) + .unwrap() + .handover + .route_generation, + 1, + "a stored route is not bumped again" + ); + + type Edit = fn(&mut HandoverRouteRow); + let cases: [(&str, Edit); 11] = [ + ("token_mismatch", |row| { + row.token_hash = Some(hex::encode(Sha256::digest(b"another-token"))); + }), + ("token_mismatch", |row| row.token_hash = None), + ("token_expired", |row| row.token_unexpired = false), + ("scope_missing", |row| { + row.scopes = vec!["terminal:connect".to_string()] + }), + ("revoked", |row| { + row.revocation_status = "revoked".to_string() + }), + ("process_bound", |row| { + row.backend_process_id = Some("proc-1".to_string()); + }), + ("not_started", |row| { + row.sandbox_lifecycle = SandboxLifecycle::Restoring; + }), + ("not_started", |row| { + row.sandbox_lifecycle = SandboxLifecycle::Paused; + }), + ("unplaced", |row| row.sandbox_active_node_id = None), + ("no_endpoint", |row| row.node_address = None), + ("unsupported_kind", |row| row.kind = "port".to_string()), + ]; + for (expected, edit) in cases { + let mut row = eligible_row(); + edit(&mut row); + match handover_route("proc_1", "attach-token", row) { + Err(HandoverError::Ineligible(reason)) => assert_eq!(reason, expected), + Err(other) => panic!("{expected}: {other}"), + Ok(_) => panic!("{expected}: minted a handover"), + } + } +} diff --git a/crates/libs/api/src/routes/mod.rs b/crates/libs/api/src/routes/mod.rs index 8229f0f8ad..2468c12b80 100644 --- a/crates/libs/api/src/routes/mod.rs +++ b/crates/libs/api/src/routes/mod.rs @@ -103,6 +103,11 @@ pub enum AppStateError { NativeHiddenOperations(#[from] public::hidden_operations::HiddenOperationsConfigError), } +fn route_handover_enabled() -> bool { + std::env::var(crate::command_stream::DATA_PLANE_ROUTE_HANDOVER_ENV) + .is_ok_and(|value| matches!(value.trim(), "1" | "true" | "yes" | "on")) +} + impl AppState { #[cfg(test)] pub async fn from_pool(pool: PgPool) -> Result { @@ -155,6 +160,7 @@ impl AppState { feature_flags.clone(), ); task_tracker.spawn(cluster_admission.liveness_heartbeat(shutdown_token.child_token())); + let route_handover_authority = route_ed25519_authority.clone(); let sandbox_controller = crate::controllers::sandbox::SandboxController::with_cluster_admission( substrate.clone(), @@ -181,8 +187,15 @@ impl AppState { let template_build_state = crate::routes::e2b::TemplateBuildState::new(pool.clone()); let data_plane_base_url = crate::public_sessions::data_plane_base_url_from_env(); - let command_stream = - crate::command_stream::CommandStreamProxy::from_env(shutdown_token.child_token())?; + let command_stream = crate::command_stream::CommandStreamProxy::from_env( + shutdown_token.child_token(), + )? + .map(|proxy| match route_handover_authority { + Some(authority) if route_handover_enabled() => proxy.with_route_handover( + crate::command_stream::RouteHandoverIssuer::new(pool.clone(), authority), + ), + _ => proxy, + }); // Startup check: without a route-signing authority, nothing runs // inline. let inline_operations = crate::inline_operations::InlineOperations::new( diff --git a/crates/libs/common/grpc/src/internal_transport.rs b/crates/libs/common/grpc/src/internal_transport.rs index 0eb9136c70..26220d23a0 100644 --- a/crates/libs/common/grpc/src/internal_transport.rs +++ b/crates/libs/common/grpc/src/internal_transport.rs @@ -94,6 +94,24 @@ fn advertised_endpoint_prefix_from_environment(env: &dyn Environment) -> &'stati } } +/// The node API port a node's registered address implies. +pub const DEFAULT_NODE_API_PORT: u16 = 9800; + +/// A node's node-API endpoint: its explicitly registered endpoint, or else its +/// registered address on [`DEFAULT_NODE_API_PORT`] behind +/// [`advertised_endpoint_prefix`]. Database INET host values need IPv6 +/// brackets, which `SocketAddr` formatting adds. +pub fn node_api_endpoint(direct: Option<&str>, address: Option<&str>) -> Option { + direct + .filter(|value| !value.trim().is_empty()) + .map(ToOwned::to_owned) + .or_else(|| { + let ip = address?.trim().parse::().ok()?; + let address = std::net::SocketAddr::new(ip, DEFAULT_NODE_API_PORT); + Some(format!("{}{address}", advertised_endpoint_prefix())) + }) +} + /// Add the configured TCP scheme only to bare addresses. Explicit schemes, /// including unsupported ones, remain intact for transport validation. pub fn qualify_endpoint(address: &str) -> String { diff --git a/crates/libs/data/plane/server/Cargo.toml b/crates/libs/data/plane/server/Cargo.toml index 52f7a6435b..67eb32bf14 100644 --- a/crates/libs/data/plane/server/Cargo.toml +++ b/crates/libs/data/plane/server/Cargo.toml @@ -41,6 +41,7 @@ sha2 = { workspace = true } # `all` gates TcpKeepalive::with_retries, which the listener sets, and the # socket option getters the listener test reads. socket2 = { workspace = true, features = ["all"] } +subtle = { workspace = true } tempfile = { workspace = true } thiserror = { workspace = true } time = { workspace = true, features = ["formatting"] } diff --git a/crates/libs/data/plane/server/src/drain.rs b/crates/libs/data/plane/server/src/drain.rs index 46d07efa72..6084da7876 100644 --- a/crates/libs/data/plane/server/src/drain.rs +++ b/crates/libs/data/plane/server/src/drain.rs @@ -51,9 +51,13 @@ enum DrainEnding { /// `shutdown_grace` it waits for the streams until the process is killed or a /// second stop signal cancels `force_token`, which ends the drain the same way /// the deadline does. +/// +/// `internal`, when set, is a second listener for in-cluster callers. It +/// stops accepting and drains together with the public one. pub(super) async fn serve_and_drain( listener: DataPlaneListener, app: Router, + internal: Option<(DataPlaneListener, Router)>, lifecycle: DrainLifecycle, ) -> std::io::Result<()> { let DrainLifecycle { @@ -76,15 +80,32 @@ pub(super) async fn serve_and_drain( ); } }; + let internal_stop_accepting = { + let shutdown_token = shutdown_token.clone(); + async move { shutdown_token.cancelled().await } + }; let server = axum::serve( listener, app.into_make_service_with_connect_info::(), ) .with_graceful_shutdown(stop_accepting); + let internal_server = async move { + match internal { + Some((listener, app)) => { + axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(internal_stop_accepting) + .await + } + None => Ok(()), + } + }; let drained = { let tasks = tasks.clone(); async move { - server.await?; + tokio::try_join!(async move { server.await }, internal_server)?; // The server has stopped, so nothing new can be tracked; closing // lets `wait` finish once the upgraded connections end. tasks.close(); diff --git a/crates/libs/data/plane/server/src/drain/tests.rs b/crates/libs/data/plane/server/src/drain/tests.rs index e2a5c7e654..f7e85a5966 100644 --- a/crates/libs/data/plane/server/src/drain/tests.rs +++ b/crates/libs/data/plane/server/src/drain/tests.rs @@ -30,6 +30,7 @@ async fn start_drain( let server = tokio::spawn(serve_and_drain( data_plane_listener(listener), app(stream_shutdown_token.clone()), + None, DrainLifecycle { shutdown_token: shutdown_token.clone(), stream_shutdown_token: stream_shutdown_token.clone(), @@ -389,6 +390,7 @@ async fn unattached_websocket_is_told_to_reconnect_when_the_drain_starts() { let server = tokio::spawn(serve_and_drain( data_plane_listener(listener), router(state), + None, DrainLifecycle { shutdown_token: shutdown_token.clone(), stream_shutdown_token, diff --git a/crates/libs/data/plane/server/src/lib.rs b/crates/libs/data/plane/server/src/lib.rs index e4f39fab57..f0d0b52956 100644 --- a/crates/libs/data/plane/server/src/lib.rs +++ b/crates/libs/data/plane/server/src/lib.rs @@ -10,6 +10,7 @@ use std::sync::{ }; use std::time::Duration; +use axum::Extension; use axum::body::Body; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::extract::{ConnectInfo, Path, State}; @@ -31,6 +32,7 @@ use pplx_space_common_core::types::cluster::process_cluster_id; use pplx_space_common_core::types::{NodeId, OrgId, ProcessId, SandboxId, SimpleDuration}; use pplx_space_common_svc::bearer_token::{SpaceBearerToken, insert_space_bearer_token}; use pplx_space_common_svc::bootstrap::{GrpcClientTrace, wrap_grpc_client_channel}; +use pplx_space_common_svc::space_api_route_ed25519::SpaceApiRouteEd25519Verifier; use pplx_space_common_svc::{HttpServerMakeSpan, HttpSpanMetadata}; use pplx_space_feature_flags::{ Attributes, FeatureFlagClient, SPACE_CLUSTER_API_CLIENT_SIDE_BALANCING_FLAG, @@ -214,6 +216,9 @@ fn record_ecw_backend_rpc_latency( #[derive(Debug, Clone)] pub struct DataPlaneConfig { pub listen: String, + /// A second listener for in-cluster callers, which public ingress must + /// not route to. Only requests on it may carry a route handover. + pub internal_listen: Option, pub attach_deadline: Duration, pub unauthenticated_upgrade_cap_per_ip: usize, pub max_attached_streams: usize, @@ -249,6 +254,7 @@ impl Default for DataPlaneConfig { route_cache_ttl: Duration::from_millis(DEFAULT_ROUTE_CACHE_TTL_MS), upload_admission: None, shutdown_grace: None, + internal_listen: None, } } } @@ -285,6 +291,9 @@ struct DataPlaneState { /// would be cancelled by the very signal the drain outlives. stream_shutdown_token: CancellationToken, upgrade_tasks: TaskTracker, + /// Verifies routes the API hands over on the internal listener; `None` + /// resolves every route. + route_handover: Option>, } #[derive(Clone)] @@ -464,6 +473,10 @@ pub async fn run( drain_started, stream_shutdown_token: stream_shutdown_token.clone(), upgrade_tasks: upgrade_tasks.clone(), + route_handover: config + .internal_listen + .as_ref() + .and_then(|_| route_handover_verifier()), }; tracing::info!( @@ -472,6 +485,26 @@ pub async fn run( "attached streams keep a confirmed revocation verdict through failed checks \ for up to the revocation tolerance" ); + let internal = match &config.internal_listen { + Some(listen) => { + let listener = tokio::net::TcpListener::bind(listen) + .await + .map_err(|source| RunDataPlaneError::Bind { + listen: listen.clone(), + source, + })?; + tracing::info!( + listen = %listen, + route_handover = state.route_handover.is_some(), + "space-data-plane-svc internal listener" + ); + Some(( + data_plane_listener(listener), + router(state.clone()).layer(Extension(native::InternalListener)), + )) + } + None => None, + }; let app = router(state); let listener = tokio::net::TcpListener::bind(&config.listen) .await @@ -484,6 +517,7 @@ pub async fn run( serve_and_drain( data_plane_listener(listener), app, + internal, DrainLifecycle { shutdown_token, stream_shutdown_token, @@ -540,6 +574,22 @@ fn data_plane_listener(listener: tokio::net::TcpListener) -> DataPlaneListener { ) } +/// Route handovers verify against the keys nodes use for route tokens. Without +/// them, every route is resolved. +fn route_handover_verifier() -> Option> { + match pplx_space_common_svc::space_api_route_ed25519::load_space_api_route_ed25519_verification_bundle() { + Ok(Some(bundle)) => Some(Arc::new(bundle.into_verifier())), + Ok(None) => { + tracing::warn!("no route verification keys configured; route handovers are disabled"); + None + } + Err(error) => { + tracing::warn!(%error, "route verification keys are invalid; route handovers are disabled"); + None + } + } +} + async fn connect_cluster_api_clients_for_startup( shutdown_token: &CancellationToken, ) -> Result { diff --git a/crates/libs/data/plane/server/src/native.rs b/crates/libs/data/plane/server/src/native.rs index 479f2aed72..70515c3178 100644 --- a/crates/libs/data/plane/server/src/native.rs +++ b/crates/libs/data/plane/server/src/native.rs @@ -6,6 +6,7 @@ use pplx_space_metrics::data_plane::{ }; mod auth; +mod handover; mod protocol; mod proxy; mod sse; @@ -16,6 +17,7 @@ const NATIVE_ROUTE_CACHE_SWEEP_INTERVAL: Duration = Duration::from_secs(60); const NATIVE_ROUTE_CACHE_SWEEP_MAX_ENTRIES: usize = 16 * 1024; pub(super) use auth::*; +pub(super) use handover::*; pub(super) use protocol::*; pub(super) use proxy::*; pub(super) use sse::*; @@ -52,14 +54,16 @@ const SESSION_NOT_READY_POLL: Duration = Duration::from_millis(500); /// [`resolve_session_route`] for a client attach, which waits up to the /// attach deadline while the session's sandbox has no route: a wake started /// with the session, such as the one `createCommand` starts for a sandbox at -/// rest, places the sandbox again in about that time. +/// rest, places the sandbox again in about that time. A route `handover` is +/// offered once; the retries resolve without it. pub(super) async fn resolve_attach_route( state: &DataPlaneState, session_id: DataPlaneSessionId, token: &str, + mut handover: Option, ) -> Result { retry_not_ready_until(tokio::time::Instant::now() + state.attach_deadline, || { - resolve_session_route(state, session_id, token) + resolve_session_route_with_handover(state, session_id, token, handover.take()) }) .await } diff --git a/crates/libs/data/plane/server/src/native/handover.rs b/crates/libs/data/plane/server/src/native/handover.rs new file mode 100644 index 0000000000..4c7bf19c6c --- /dev/null +++ b/crates/libs/data/plane/server/src/native/handover.rs @@ -0,0 +1,180 @@ +//! Opening a session's stream on the route the API handed over, instead of +//! resolving it again, when the API itself relays the stream of a session it +//! just created. +//! +//! A handover replaces route resolution only when all of these hold: +//! - the request arrived on the internal listener, which public ingress never +//! reaches; +//! - the handover verifies against the route verification keys, is unexpired, +//! and names this cluster; +//! - it names this session and its kind, and was minted for the SHA-256 of the +//! attach token on this request; +//! - its node route token, verified with the same keys, names the handover's +//! sandbox, node, and execution generation. +//! +//! Any other request resolves the route as before. Revocation is unchanged: +//! the stream remembers the epochs the handover carries and checks them +//! periodically. A route that went stale between the handover and the open +//! (a wake, restore, or migration) is refused by the node, and the stream +//! refreshes it as it would any rejected route. + +use super::*; +use pplx_space_common_svc::space_api_route_ed25519::{ + DATA_PLANE_ROUTE_HANDOVER_HEADER, DataPlaneRouteHandoverRoute, SpaceApiRouteEd25519Token, +}; +use pplx_space_node_api_client::proto::SpaceApiRouteToken as ProtoSpaceApiRouteToken; +use subtle::ConstantTimeEq as _; + +/// Marks requests that arrived on the internal listener. +#[derive(Clone, Copy, Debug)] +pub(crate) struct InternalListener; + +/// The handover a request carries, and whether the listener it arrived on may +/// honor one. +pub(crate) struct OfferedRouteHandover { + handover: String, + internal_listener: bool, +} + +impl OfferedRouteHandover { + pub(crate) fn from_request( + headers: &HeaderMap, + internal_listener: Option, + ) -> Option { + let handover = headers.get(DATA_PLANE_ROUTE_HANDOVER_HEADER)?; + Some(Self { + handover: handover.to_str().unwrap_or_default().to_string(), + internal_listener: internal_listener.is_some(), + }) + } +} + +/// Resolves the session's route, or takes it from a valid handover. +pub(super) async fn resolve_session_route_with_handover( + state: &DataPlaneState, + session_id: DataPlaneSessionId, + token: &str, + handover: Option, +) -> Result { + if let Some(handover) = handover { + let kind = session_id.kind().as_str(); + match accept_route_handover(state, session_id, token, &handover) { + Ok(attached) => { + pplx_space_metrics::data_plane::route_handover(kind, "accepted"); + state.revocation_cache.remember_active( + session_id, + attached.revocation_epoch, + attached.org_id, + attached.org_revocation_epoch, + start_revocation_read(state), + ); + return Ok(attached); + } + Err(reason) => { + pplx_space_metrics::data_plane::route_handover(kind, reason); + tracing::debug!( + session_id = %session_id, + reason, + "route handover refused; resolving the session route" + ); + } + } + } + resolve_session_route(state, session_id, token).await +} + +pub(super) fn accept_route_handover( + state: &DataPlaneState, + session_id: DataPlaneSessionId, + token: &str, + offered: &OfferedRouteHandover, +) -> Result { + if !offered.internal_listener { + return Err("public_listener"); + } + let verifier = state.route_handover.as_deref().ok_or("not_configured")?; + if !matches!( + session_id.kind(), + SessionKind::Terminal | SessionKind::Process | SessionKind::FileWatch + ) { + return Err("unsupported_kind"); + } + let route = verifier + .verify_data_plane_route_handover(&offered.handover) + .map_err(|error| error.label())?; + let DataPlaneRouteHandoverRoute { + session_id: handed_session_id, + token_sha256: handed_token_sha256, + org_id, + sandbox_id, + kind, + route_generation, + current_node_id, + backend_endpoint, + execution_generation, + revocation_epoch, + org_revocation_epoch, + adapter_state_json, + space_api_route_token, + } = route; + if handed_session_id != session_id.to_string() { + return Err("wrong_session"); + } + if handed_token_sha256 + .as_bytes() + .ct_eq(token_sha256(token).as_bytes()) + .unwrap_u8() + != 1 + { + return Err("wrong_token"); + } + if kind.parse::().ok() != Some(session_id.kind()) { + return Err("wrong_kind"); + } + + let route_token = SpaceApiRouteEd25519Token::parse(space_api_route_token.clone()) + .map_err(|_| "invalid_route_token")?; + let route_claims = verifier + .verify(&route_token) + .map_err(|_| "invalid_route_token")?; + let identity = &route_claims.route_identity; + if identity.sandbox_id() != sandbox_id + || identity.current_node_id() != current_node_id + || identity.execution_generation() != execution_generation + { + return Err("route_token_mismatch"); + } + let space_api_route_token = parse_proto_space_api_route_token( + SpaceApiRouteTokenRouteSource::DataPlaneSession, + ProtoSpaceApiRouteToken { + token: space_api_route_token, + }, + ) + .ok_or("invalid_route_token")?; + + Ok(AttachedSession { + cache_epoch: state.route_cache.next_cache_epoch(), + cache_key: None, + session_id, + sandbox_id, + org_id, + kind: session_id.kind(), + route_generation, + revocation_epoch, + org_revocation_epoch, + current_node_id, + backend_endpoint: backend_endpoint + .parse() + .map_err(|_| "invalid_backend_endpoint")?, + target_port: None, + target_protocol: None, + adapter_state_json, + // A handover is minted only for a session that has never bound a + // process; any later attach resolves the bound one. + backend_process_id: None, + space_api_route_token: Some(space_api_route_token), + }) +} + +#[cfg(test)] +mod tests; diff --git a/crates/libs/data/plane/server/src/native/handover/tests.rs b/crates/libs/data/plane/server/src/native/handover/tests.rs new file mode 100644 index 0000000000..ba7a5ef74a --- /dev/null +++ b/crates/libs/data/plane/server/src/native/handover/tests.rs @@ -0,0 +1,477 @@ +use super::*; +use aws_lc_rs::rand::SystemRandom; +use aws_lc_rs::signature::Ed25519KeyPair; +use pplx_space_common_core::types::{ + ClusterId, ExecutionGeneration, NodeManagedBy, RouteGeneration, SandboxActivationAttemptId, + SandboxRouteIdentity, +}; +use pplx_space_common_svc::space_api_route::{ + RouteIdentityRemediation, SpaceApiRouteIdentityDimension, set_route_identity_rejection_metadata, +}; +use pplx_space_common_svc::space_api_route_ed25519::{ + SPACE_API_ROUTE_ED25519_TOKEN_METADATA_KEY, SpaceApiRouteEd25519Authority, + SpaceApiRouteEd25519Issuer, SpaceApiRouteEd25519KeyId, SpaceApiRouteEd25519Signer, +}; +use pplx_space_node_api_client::proto::DataPlaneBackendEndpoint; +use std::sync::Mutex; +use std::sync::atomic::{AtomicUsize, Ordering}; + +fn authority() -> SpaceApiRouteEd25519Authority { + let pkcs8 = Ed25519KeyPair::generate_pkcs8(&SystemRandom::new()).unwrap(); + SpaceApiRouteEd25519Authority::new( + SpaceApiRouteEd25519Issuer::parse("space-api").unwrap(), + ClusterId::parse("testing").unwrap(), + SpaceApiRouteEd25519Signer::from_pkcs8( + SpaceApiRouteEd25519KeyId::parse("test-key").unwrap(), + pkcs8.as_ref(), + ) + .unwrap(), + ) +} + +fn verifier(authority: &SpaceApiRouteEd25519Authority) -> Arc { + Arc::new( + SpaceApiRouteEd25519Verifier::new( + authority.issuer().clone(), + authority.cluster_id().clone(), + [authority.verification_key()], + ) + .unwrap(), + ) +} + +fn route_token( + authority: &SpaceApiRouteEd25519Authority, + sandbox_id: SandboxId, + node_id: NodeId, + execution_generation: i64, +) -> String { + authority + .mint_space_api_bridge( + SandboxRouteIdentity::try_new( + sandbox_id, + RouteGeneration::new(execution_generation.cast_unsigned()).unwrap(), + node_id, + NodeManagedBy::Space, + ExecutionGeneration::new(execution_generation).unwrap(), + SandboxActivationAttemptId::new(), + None, + ) + .unwrap(), + ) + .unwrap() + .secret() + .to_string() +} + +fn attach_token() -> String { + format!("{DATA_PLANE_TOKEN_PREFIX}{}", "d".repeat(64)) +} + +/// A fake cluster API and worker for one process session. The worker records +/// the route token on every start, and refuses the handed-over one when +/// `refuse` holds it, as a worker refuses a route that a restore or migration +/// replaced. +struct Fake { + endpoint: String, + sandbox_id: SandboxId, + node_id: NodeId, + resolved_route_token: String, + resolves: Arc, + started_with: Arc>>, + refuse: Arc>>, +} + +async fn spawn_fake( + authority: &SpaceApiRouteEd25519Authority, + session_id: DataPlaneSessionId, +) -> Fake { + use prost::Message as _; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let sandbox_id = SandboxId::new(); + let node_id = NodeId::new(); + let resolved_route_token = route_token(authority, sandbox_id, node_id, 6); + let mut route = crate::tests::route_response_with_token( + session_id, + Some(ProtoSpaceApiRouteToken { + token: resolved_route_token.clone(), + }), + ); + route.sandbox_id = sandbox_id.to_string(); + route.current_node_id = node_id.to_string(); + route.backend = Some(DataPlaneBackendEndpoint { + node_id: node_id.to_string(), + endpoint: endpoint.clone(), + transport: None, + }); + route.adapter_state_json = r#"{"command":["true"]}"#.to_string(); + let route = route.encode_to_vec(); + let process_id = pplx_space_public_api::process_id_for_session(session_id).to_string(); + let mut started = vec![0]; + let encoded = ProcessEvent { + process_id, + pid: Some(42), + event: Some(spaced_process_event::Event::Started( + pplx_space_sbx_mgr_client::operations::ProcessStarted { start_time: None }, + )), + } + .encode_to_vec(); + started.extend_from_slice(&(encoded.len() as u32).to_be_bytes()); + started.extend(encoded); + let started = bytes::Bytes::from(started); + + let fake = Fake { + endpoint, + sandbox_id, + node_id, + resolved_route_token, + resolves: Arc::new(AtomicUsize::new(0)), + started_with: Arc::new(Mutex::new(Vec::new())), + refuse: Arc::new(Mutex::new(None)), + }; + let resolves = Arc::clone(&fake.resolves); + let started_with = Arc::clone(&fake.started_with); + let refuse = Arc::clone(&fake.refuse); + let service = axum::Router::new().fallback(move |request: Request| { + let route = route.clone(); + let started = started.clone(); + let resolves = Arc::clone(&resolves); + let started_with = Arc::clone(&started_with); + let refuse = Arc::clone(&refuse); + async move { + let path = request.uri().path().to_string(); + if path.ends_with("/Check") { + return crate::tests::grpc_test_response(vec![8, 1]); + } + if path.ends_with("/ResolveDataPlaneSessionRoute") { + resolves.fetch_add(1, Ordering::SeqCst); + return crate::tests::grpc_test_response(route); + } + if path.ends_with("/BindDataPlaneSessionProcess") { + return crate::tests::grpc_test_response(Vec::new()); + } + assert!( + path.ends_with("/SpacedStartProcess"), + "unexpected call {path}" + ); + let token = request + .headers() + .get(SPACE_API_ROUTE_ED25519_TOKEN_METADATA_KEY) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + started_with.lock().unwrap().push(token.clone()); + if refuse.lock().unwrap().as_deref() == Some(token.as_str()) { + let mut status = tonic::Status::permission_denied("stale execution generation"); + set_route_identity_rejection_metadata( + &mut status, + RouteIdentityRemediation::RefreshRoute, + Some(SpaceApiRouteIdentityDimension::ExecutionGeneration), + ); + return status.into_http::(); + } + let frames = futures_util::stream::iter([Ok::<_, Infallible>( + hyper::body::Frame::data(started), + )]) + .chain(futures_util::stream::pending()); + Response::builder() + .header("content-type", "application/grpc") + .body(Body::new(http_body_util::StreamBody::new(frames))) + .unwrap() + } + }); + tokio::spawn(async move { + axum::serve(listener, service).await.unwrap(); + }); + fake +} + +/// The handover the API mints for this session's fresh route. +fn handover_route( + authority: &SpaceApiRouteEd25519Authority, + fake: &Fake, + session_id: DataPlaneSessionId, +) -> DataPlaneRouteHandoverRoute { + DataPlaneRouteHandoverRoute { + session_id: session_id.to_string(), + token_sha256: token_sha256(&attach_token()), + org_id: OrgId::new(), + sandbox_id: fake.sandbox_id, + kind: session_id.kind().as_str().to_string(), + route_generation: 2, + current_node_id: fake.node_id, + backend_endpoint: fake.endpoint.clone(), + execution_generation: ExecutionGeneration::new(5).unwrap(), + revocation_epoch: 0, + org_revocation_epoch: 0, + adapter_state_json: r#"{"command":["true"]}"#.to_string(), + space_api_route_token: route_token(authority, fake.sandbox_id, fake.node_id, 5), + } +} + +fn test_process() -> bool { + pplx_space_test_env::enter_test_process(&[ + ( + "SPACE_INTERNAL_RPC_INSECURE_DEV", + Some(std::ffi::OsStr::new("true")), + ), + ( + "SPACE_INTERNAL_TOKEN", + Some(std::ffi::OsStr::new("test-token")), + ), + ]) +} + +enum Listener { + Internal, + Public, +} + +/// Opens the session's event stream with `handover` and waits for the start. +async fn open_with_handover( + fake: &Fake, + verifier: Arc, + session_id: DataPlaneSessionId, + handover: Option, + listener: Listener, +) { + let mut state = crate::tests::test_data_plane_state().await; + state.cluster_api = ClusterApiClients::service_only( + pplx_space_node_api_client::cluster::connect(fake.endpoint.clone(), "test") + .await + .unwrap(), + ); + state.revocation_check_interval = Duration::from_secs(3600); + state.route_handover = Some(verifier); + let app = match listener { + Listener::Internal => router(state).layer(Extension(InternalListener)), + Listener::Public => router(state), + }; + let mut request = Request::get(format!("/v1/sessions/{session_id}/events")) + .header(header::AUTHORIZATION, format!("Bearer {}", attach_token())); + if let Some(handover) = handover { + request = request.header(DATA_PLANE_ROUTE_HANDOVER_HEADER, handover); + } + let response = tower::ServiceExt::oneshot(app, request.body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let mut body = response.into_body(); + let mut seen = String::new(); + tokio::time::timeout(Duration::from_secs(5), async { + while !seen.contains("process_started") { + let frame = body.frame().await.expect("stream open").unwrap(); + if let Ok(data) = frame.into_data() { + seen.push_str(&String::from_utf8_lossy(&data)); + } + } + }) + .await + .unwrap_or_else(|_| panic!("process_started within 5s; saw {seen}")); +} + +#[tokio::test] +async fn a_valid_handover_opens_the_stream_without_resolving() { + if !test_process() { + return; + } + let authority = authority(); + let session_id = DataPlaneSessionId::new_for_kind(SessionKind::Process); + let fake = spawn_fake(&authority, session_id).await; + let route = handover_route(&authority, &fake, session_id); + let handed_token = route.space_api_route_token.clone(); + let handover = authority.mint_data_plane_route_handover(route).unwrap(); + + open_with_handover( + &fake, + verifier(&authority), + session_id, + Some(handover), + Listener::Internal, + ) + .await; + + assert_eq!( + fake.resolves.load(Ordering::SeqCst), + 0, + "no route resolution" + ); + assert_eq!( + *fake.started_with.lock().unwrap(), + vec![handed_token], + "the process started on the handed-over route" + ); +} + +#[tokio::test] +async fn refused_handovers_fall_back_to_resolving_the_route() { + if !test_process() { + return; + } + let authority = authority(); + let impostor = authority_with_other_key(); + for (case, listener, handover) in [ + ("public listener", Listener::Public, Handover::Valid), + ("forged", Listener::Internal, Handover::SignedBy(&impostor)), + ("expired", Listener::Internal, Handover::Expired), + ("other session", Listener::Internal, Handover::OtherSession), + ( + "other attach token", + Listener::Internal, + Handover::OtherToken, + ), + ( + "mismatched route token", + Listener::Internal, + Handover::OtherExecution, + ), + ("garbage", Listener::Internal, Handover::Garbage), + ] { + let session_id = DataPlaneSessionId::new_for_kind(SessionKind::Process); + let fake = spawn_fake(&authority, session_id).await; + let handover = handover.mint(&authority, &fake, session_id); + + open_with_handover( + &fake, + verifier(&authority), + session_id, + Some(handover), + listener, + ) + .await; + + assert_eq!( + fake.resolves.load(Ordering::SeqCst), + 1, + "{case}: the route was resolved" + ); + assert_eq!( + *fake.started_with.lock().unwrap(), + vec![fake.resolved_route_token.clone()], + "{case}: the process started on the resolved route" + ); + } +} + +#[tokio::test] +async fn a_route_that_went_stale_after_its_handover_is_refreshed() { + if !test_process() { + return; + } + let authority = authority(); + let session_id = DataPlaneSessionId::new_for_kind(SessionKind::Process); + let fake = spawn_fake(&authority, session_id).await; + let route = handover_route(&authority, &fake, session_id); + let handed_token = route.space_api_route_token.clone(); + // A restore replaced the execution the handover names. + *fake.refuse.lock().unwrap() = Some(handed_token.clone()); + let handover = authority.mint_data_plane_route_handover(route).unwrap(); + + open_with_handover( + &fake, + verifier(&authority), + session_id, + Some(handover), + Listener::Internal, + ) + .await; + + assert_eq!(fake.resolves.load(Ordering::SeqCst), 1, "resolved once"); + assert_eq!( + *fake.started_with.lock().unwrap(), + vec![handed_token, fake.resolved_route_token.clone()], + "the refused start was retried on the resolved route" + ); +} + +#[tokio::test] +async fn handovers_are_ignored_where_none_are_verified() { + if !test_process() { + return; + } + let authority = authority(); + let session_id = DataPlaneSessionId::new_for_kind(SessionKind::Process); + let fake = spawn_fake(&authority, session_id).await; + let handover = authority + .mint_data_plane_route_handover(handover_route(&authority, &fake, session_id)) + .unwrap(); + let state = crate::tests::test_data_plane_state().await; + let offered = OfferedRouteHandover { + handover, + internal_listener: true, + }; + assert_eq!( + accept_route_handover(&state, session_id, &attach_token(), &offered).unwrap_err(), + "not_configured" + ); +} + +fn authority_with_other_key() -> SpaceApiRouteEd25519Authority { + authority() +} + +enum Handover<'a> { + Valid, + SignedBy(&'a SpaceApiRouteEd25519Authority), + Expired, + OtherSession, + OtherToken, + OtherExecution, + Garbage, +} + +impl Handover<'_> { + fn mint( + &self, + authority: &SpaceApiRouteEd25519Authority, + fake: &Fake, + session_id: DataPlaneSessionId, + ) -> String { + let mut route = handover_route(authority, fake, session_id); + match self { + Self::Valid => {} + Self::SignedBy(other) => return other.mint_data_plane_route_handover(route).unwrap(), + Self::Expired => return expired_handover(authority, route), + Self::OtherSession => { + route.session_id = + DataPlaneSessionId::new_for_kind(SessionKind::Process).to_string(); + } + Self::OtherToken => { + route.token_sha256 = + token_sha256(&format!("{DATA_PLANE_TOKEN_PREFIX}{}", "e".repeat(64))); + } + Self::OtherExecution => { + route.space_api_route_token = + route_token(authority, fake.sandbox_id, fake.node_id, 4); + } + Self::Garbage => return "dp-handover-ed25519.test-key.e30.AAAA".to_string(), + } + authority.mint_data_plane_route_handover(route).unwrap() + } +} + +fn expired_handover( + authority: &SpaceApiRouteEd25519Authority, + route: DataPlaneRouteHandoverRoute, +) -> String { + use pplx_space_common_svc::space_api_route_ed25519::{ + DATA_PLANE_ROUTE_HANDOVER_AUDIENCE, DATA_PLANE_ROUTE_HANDOVER_SCHEMA_VERSION, + DataPlaneRouteHandoverClaims, + }; + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + authority + .sign_data_plane_route_handover_claims(&DataPlaneRouteHandoverClaims { + schema_version: DATA_PLANE_ROUTE_HANDOVER_SCHEMA_VERSION, + issuer: authority.issuer().clone(), + audience: DATA_PLANE_ROUTE_HANDOVER_AUDIENCE.to_string(), + cluster_id: authority.cluster_id().clone(), + route, + issued_at_unix_secs: now - 60, + expires_at_unix_secs: now - 50, + }) + .unwrap() +} diff --git a/crates/libs/data/plane/server/src/native/sse.rs b/crates/libs/data/plane/server/src/native/sse.rs index a6565fbd0e..4feadb6229 100644 --- a/crates/libs/data/plane/server/src/native/sse.rs +++ b/crates/libs/data/plane/server/src/native/sse.rs @@ -5,6 +5,7 @@ use super::*; pub(crate) async fn session_events( State(state): State, Path(session_id): Path, + internal_listener: Option>, headers: HeaderMap, uri: Uri, ) -> Response { @@ -68,7 +69,9 @@ pub(crate) async fn session_events( } }; - let attached = match resolve_attach_route(&state, session_id, &token).await { + let handover = + OfferedRouteHandover::from_request(&headers, internal_listener.map(|Extension(l)| l)); + let attached = match resolve_attach_route(&state, session_id, &token, handover).await { Ok(attached) => attached, Err(error) => { pplx_space_metrics::data_plane::attach( diff --git a/crates/libs/data/plane/server/src/native/websocket.rs b/crates/libs/data/plane/server/src/native/websocket.rs index 7854206365..e4fa376552 100644 --- a/crates/libs/data/plane/server/src/native/websocket.rs +++ b/crates/libs/data/plane/server/src/native/websocket.rs @@ -7,6 +7,7 @@ pub(crate) async fn attach_session( State(state): State, ConnectInfo(peer): ConnectInfo, Path(session_id): Path, + internal_listener: Option>, headers: HeaderMap, upgrade: WebSocketUpgrade, uri: axum::http::Uri, @@ -16,6 +17,11 @@ pub(crate) async fn attach_session( .parse::() .map_err(|_| AttachError::BadRequest("invalid session id"))?; let header_token = token_from_headers(&headers)?; + // A handover vouches for the token it was minted for, so it only rides + // along with a header token, never a first-frame one. + let handover = header_token.as_ref().and_then(|_| { + OfferedRouteHandover::from_request(&headers, internal_listener.map(|Extension(l)| l)) + }); if state.drain_started.is_cancelled() { return Ok(native_draining_response(session_id.kind().as_str())); } @@ -25,7 +31,7 @@ pub(crate) async fn attach_session( let task = state.upgrade_tasks.token(); Ok(upgrade.on_upgrade(move |socket| async move { let _task = task; - handle_socket(state, peer.ip(), session_id, header_token, socket).await; + handle_socket(state, peer.ip(), session_id, header_token, handover, socket).await; })) } @@ -34,6 +40,7 @@ pub(crate) async fn handle_socket( source_ip: IpAddr, session_id: DataPlaneSessionId, header_token: Option, + handover: Option, mut socket: WebSocket, ) { let attach_deadline = tokio::time::Instant::now() + state.attach_deadline; @@ -62,6 +69,7 @@ pub(crate) async fn handle_socket( &state, session_id, header_token, + handover, &mut socket, attach_deadline, ) => { @@ -129,14 +137,17 @@ pub(super) async fn authenticate_native_socket( state: &DataPlaneState, session_id: DataPlaneSessionId, header_token: Option, + mut handover: Option, socket: &mut WebSocket, deadline: tokio::time::Instant, ) -> Option { - let attached = - attach_token_and_route(header_token, socket.recv(), deadline, |token| async move { - resolve_session_route(state, session_id, &token).await - }) - .await; + // The handover is offered once; a resolution waiting for the sandbox's + // placement retries without it. + let attached = attach_token_and_route(header_token, socket.recv(), deadline, |token| { + let handover = handover.take(); + async move { resolve_session_route_with_handover(state, session_id, &token, handover).await } + }) + .await; let frame = match attached { Ok(attached) => return Some(attached), Err(AttachRefused::Token(error)) => { diff --git a/crates/libs/data/plane/server/src/tests.rs b/crates/libs/data/plane/server/src/tests.rs index 91af349cc6..7db8cf3845 100644 --- a/crates/libs/data/plane/server/src/tests.rs +++ b/crates/libs/data/plane/server/src/tests.rs @@ -951,6 +951,7 @@ pub(crate) async fn test_data_plane_state() -> DataPlaneState { drain_started: CancellationToken::new(), stream_shutdown_token: CancellationToken::new(), upgrade_tasks: TaskTracker::new(), + route_handover: None, } } diff --git a/crates/libs/internal/api/server/src/data_plane_sessions/route_rows.rs b/crates/libs/internal/api/server/src/data_plane_sessions/route_rows.rs index 27944b8b41..04968f0c4d 100644 --- a/crates/libs/internal/api/server/src/data_plane_sessions/route_rows.rs +++ b/crates/libs/internal/api/server/src/data_plane_sessions/route_rows.rs @@ -2,8 +2,6 @@ use super::*; -pub(super) const DEFAULT_NODE_API_PORT: u16 = 9800; - // Each read here places the sandbox on // `COALESCE(direct_route.current_node_id, s.active_node_id)`, and joins the // Direct route identity only while `s.active_node_id` is set. A worker's own @@ -174,17 +172,7 @@ pub(super) fn node_api_backend_endpoint( direct: Option<&str>, address: Option<&str>, ) -> Option { - direct - .filter(|value| !value.trim().is_empty()) - .map(ToOwned::to_owned) - .or_else(|| { - let ip = address?.trim().parse::().ok()?; - let address = std::net::SocketAddr::new(ip, DEFAULT_NODE_API_PORT); - Some(format!( - "{}{address}", - pplx_space_common_grpc::internal_transport::advertised_endpoint_prefix() - )) - }) + pplx_space_common_grpc::internal_transport::node_api_endpoint(direct, address) } pub(super) fn resolve_backend_endpoint(row: &SessionRouteRow) -> Option { diff --git a/crates/libs/metrics/src/data_plane.rs b/crates/libs/metrics/src/data_plane.rs index d02fab5b79..0f3a7c105b 100644 --- a/crates/libs/metrics/src/data_plane.rs +++ b/crates/libs/metrics/src/data_plane.rs @@ -362,6 +362,16 @@ pub fn attach(kind: &str, result: &str) { ); } +/// Whether an attach used the route the API handed over (`accepted`) or fell +/// back to resolving it (every other `result`). +pub fn route_handover(kind: &str, result: &str) { + counter( + "data_plane.route_handover", + 1, + &[tag("kind", kind), tag("result", result)], + ); +} + pub fn stream_closed(kind: &str, result: &str) { counter( "data_plane.stream.closed", diff --git a/docs/configuration.md b/docs/configuration.md index 244cd1ba25..1b41042d36 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -173,6 +173,7 @@ because disabling it by mistake would leave every create queued forever. | `SPACE_PUBLIC_CURSOR_SECRET` | **required** | | | `SPACE_DATA_PLANE_PUBLIC_URL` | empty | Empty yields relative URLs. | | `SPACE_DATA_PLANE_INTERNAL_URL` | unset | `http://` origin of the data plane as the API reaches it. Enables streaming command creation (`Accept: text/event-stream`); unset answers those requests with the JSON handle. Streams relayed this way end when the API shuts down. | +| `SPACE_DATA_PLANE_ROUTE_HANDOVER` | off | API: with the command stream on, also send the data plane the session's route, signed with the route signing key, so the data plane opens the stream without resolving the route. Turn it on only where `SPACE_DATA_PLANE_INTERNAL_URL` reaches the data plane's internal listener: anywhere else the data plane ignores the route and resolves it anyway. Accepts `true`/`1`/`yes`/`on`. | | `SPACE_FEDERATED_GRANT_AUDIENCE`, `SPACE_FEDERATED_IDENTITY_ISSUER`, `SPACE_FEDERATED_REJECTED_SUBJECTS` | `credential-provider`, `space`, `[]` | | | `SPACE_E2B_HOME_WORKDIR_TEMPLATES`, `SPACE_ECW_HOSTS` | `code-runtime`, none | E2B routes return 404 without hosts. | | `SPACE_CLUSTER_ID` | unset | Required for route signing. | @@ -185,6 +186,7 @@ because disabling it by mistake would leave every create queued forever. | `SPACE_DATA_PLANE_ORG_TRANSFER_BYTES_PER_SECOND`, `SPACE_DATA_PLANE_ORG_TRANSFER_BURST_BYTES` | 100000000, 1000000000 | Per-org token bucket on each replica for `POST /files` request bodies and `GET /files` response bodies together. An HTTP/1.x upload over the rate slows through TCP backpressure rather than failing, but one that arrives while its org owes more than the admission wait could sleep off gets 429 before it takes any permit, and an admitted one that meets more than 15 s of debt, which the org's other transfers can run up meanwhile, ends with 429 and frees its permits rather than hold them through the sleep, unless that frame is the last of a body that declared its length; bytes are counted as the data plane reads them, before the throttle holds a frame, so what crosses the wire can run ahead by about a frame per upload plus socket buffering. HTTP/2 uploads and downloads are never slowed, since their streams share an HTTP/2 connection's flow control with other requests (a download's with other sandboxes' calls on a pooled backend connection). Their bytes are charged to the bucket as they move, and a new one gets 429 while the org owes the bucket, which is most of the time while the org's HTTP/1.x uploads are being slowed. So the rate binds a download only as it starts: downloads started while the org owes nothing all run as fast as their clients read them, however many there are. An org's HTTP/2 uploads can run ahead of the rate by at most its in-flight byte share, or one upload's 256 MiB body limit if the share is smaller. Starting values for testing, not measured ones. The per-org buckets are a bounded cache, swept in the background at most once a minute, started by a transfer: unreferenced buckets that are full or have sat idle for 10 minutes are dropped, and when a new org would take a replica past 65,536 tracked orgs the fullest unreferenced ones go, at most once a second. The org's next transfer then starts from a full burst, whatever the dropped bucket still had or owed. The rate accepts `off`; every other value, the burst included, must be non-zero. Applies only with upload admission on. | | `SPACE_DATA_PLANE_UPLOAD_MIN_BODY_BYTES_PER_SECOND`, `SPACE_DATA_PLANE_UPLOAD_MIN_BODY_RATE_WINDOW_MS` | 65536, 15000 | An admitted upload must deliver rate × window bytes in each consecutive window of time spent waiting on the client, or it ends with 408 and frees its permits. Time in the byte-rate bucket or spent spooling and writing to the guest doesn't count. At the defaults that's 960 KiB per 15 s, so an upload smaller than 960 KiB must finish before the pod has waited 15 s on its client in total. It applies to HTTP/1.x uploads only: an HTTP/2 upload shares its connection's flow control with other requests, so it isn't held to the floor; the time limit below bounds how long it holds its permits instead. The rate accepts `off`; every other value, the window included, must be non-zero. Applies only with upload admission on. | | `SPACE_DATA_PLANE_UPLOAD_HTTP2_MAX_DURATION_MS` | 600000 | How long after admission an HTTP/2 upload may take to finish before it ends with 408 and frees its permits; a body whose declared length has all arrived gets 5 s more for its end, which an HTTP/2 client may send in a frame of its own. A time rather than a rate, so another stream's slowness on the connection can't trip it early: at the default a full-size 256 MiB upload has it at about 450 KB/s. Accepts `off`, which lets a stalled HTTP/2 upload hold its permits until its client gives up; otherwise it must be non-zero. Applies only with upload admission on. | +| `SPACE_DATA_PLANE_INTERNAL_LISTEN` | unset | Data plane: a second listen address for in-cluster callers. Keep it off public ingress. Only requests on it may carry a route handover from the API, and only if `SPACE_API_ROUTE_ED25519_VERIFICATION_BUNDLE_PATH` (the route verification keys nodes use) and `SPACE_CLUSTER_ID` are set too. A handover that fails any check falls back to resolving the route. | | `SPACE_DATA_PLANE_SHUTDOWN_GRACE_SECS` | unset | Data plane: how long live streams keep running after SIGTERM before they're cancelled. Unset, they run until they end or the termination grace period kills the process; a second stop signal ends the drain either way, without waiting for a `Process.Start` that's still opening. Keep it plus 5 seconds within the pod's termination grace period minus its preStop delay. Values under 8 seconds, the stream-open budget, are rejected, so an ECW `Process.Start` already opening at SIGTERM can report itself. See [architecture](architecture.md) for what each stream kind sees. | ### Sizing dispatcher execution concurrency @@ -452,7 +454,7 @@ calls. | `SPACE_NODE_API_AUTH_PEER_NODE_SERVICE_ACCOUNTS` | empty | Accounts another node's node-api presents. They may call only `NodeExportTemplate` and `GetNodeInfo`, after admission and the transfer grant. | | `SPACE_NODE_API_NAMESPACE`, `SPACE_NODE_API_TLS_SERVER_NAME`, `SPACE_CLUSTER_API_TLS_SERVER_NAME` | namespace chain, from endpoint | | | `SPACE_API_ROUTE_ED25519_ISSUER`, `SPACE_API_ROUTE_ED25519_SIGNING_KEY_ID`, `SPACE_API_ROUTE_ED25519_SIGNING_KEY_PATH` | unset = no signing | Set all three. | -| `SPACE_API_ROUTE_ED25519_VERIFICATION_BUNDLE_PATH` | required on sandbox managers | | +| `SPACE_API_ROUTE_ED25519_VERIFICATION_BUNDLE_PATH` | required on sandbox managers | Also verifies route handovers on the data plane's internal listener. | | `SPACE_API_ROUTE_GENERATION_FENCE_ENABLED` | on | | | `SPACE_ECW_CREDENTIAL_KEY_PATH` | **required** on api and internal-api | | | `SPACE_ECW_CREDENTIAL_PREVIOUS_KEY_PATH` | unset | Set only while rotating. |