diff --git a/crates/worker-runtime/src/auth.rs b/crates/worker-runtime/src/auth.rs index 3c90b176..15dc0fb2 100644 --- a/crates/worker-runtime/src/auth.rs +++ b/crates/worker-runtime/src/auth.rs @@ -10,6 +10,10 @@ const PUBLIC_KEY_PREFIX: &str = "yoi-ed25519-pub:v1:"; const PRIVATE_KEY_PREFIX: &str = "yoi-ed25519-pkcs8:v1:"; const TOKEN_PREFIX: &str = "yoi-cap-v1"; const SIGNING_INPUT_PREFIX: &str = "yoi-cap-v1."; +pub const WORKER_MUTATION_SOURCE_PROOF_HEADER: &str = "x-yoi-worker-mutation-proof"; +const WORKER_MUTATION_SOURCE_PROOF_PREFIX: &str = "yoi-worker-source-v1"; +const WORKER_MUTATION_SOURCE_SIGNING_INPUT_PREFIX: &str = "yoi-worker-source-v1."; +pub const WORKER_REMOVE_PERMISSION: &str = "workspace:worker-remove"; #[derive(Debug, thiserror::Error)] pub enum RuntimeAuthError { @@ -41,6 +45,16 @@ pub enum RuntimeAuthError { MissingWorkspaceScope, #[error("capability token is missing required permission `{0}`")] MissingPermission(String), + #[error("source proof workspace `{actual}` does not match `{expected}`")] + WrongWorkspace { expected: String, actual: String }, + #[error("source proof Worker `{actual}` does not match `{expected}`")] + WrongWorker { expected: String, actual: String }, + #[error("source proof actor kind is not allowed")] + WrongActorKind, + #[error("source proof operation is not allowed")] + WrongOperation, + #[error("source proof mutation target does not match the request")] + WrongMutationTarget, } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -210,6 +224,201 @@ pub fn verify_capability_token( }) } +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkerMutationSourceClaims { + pub iss: String, + pub aud: String, + pub workspace_id: String, + pub worker_id: String, + pub actor_kind: WorkerMutationActorKind, + pub operation: WorkerMutationOperation, + pub target_runtime_id: String, + pub target_worker_id: String, + pub permission: String, + pub iat: u64, + pub exp: u64, + pub jti: String, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WorkerMutationActorKind { + Worker, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WorkerMutationOperation { + WorkerRemove, +} + +#[derive(Clone)] +pub struct RuntimeWorkerMutationSourceSigner { + runtime_id: String, + private_key: String, +} + +impl fmt::Debug for RuntimeWorkerMutationSourceSigner { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("RuntimeWorkerMutationSourceSigner") + .field("runtime_id", &self.runtime_id) + .field("private_key", &"[redacted]") + .finish() + } +} + +impl RuntimeWorkerMutationSourceSigner { + pub fn from_identity(identity: &RuntimeIdentityMaterial) -> Self { + Self { + runtime_id: identity.identity_id.clone(), + private_key: identity.private_key.clone(), + } + } + + pub fn runtime_id(&self) -> &str { + &self.runtime_id + } + + pub fn issue_worker_remove( + &self, + audience: impl Into, + workspace_id: impl Into, + source_worker_id: impl Into, + target_runtime_id: impl Into, + target_worker_id: impl Into, + ttl_seconds: u64, + ) -> Result { + let issued_at = unix_now_seconds(); + let claims = WorkerMutationSourceClaims { + iss: self.runtime_id.clone(), + aud: audience.into(), + workspace_id: workspace_id.into(), + worker_id: source_worker_id.into(), + actor_kind: WorkerMutationActorKind::Worker, + operation: WorkerMutationOperation::WorkerRemove, + target_runtime_id: target_runtime_id.into(), + target_worker_id: target_worker_id.into(), + permission: WORKER_REMOVE_PERMISSION.to_string(), + iat: issued_at, + exp: issued_at.saturating_add(ttl_seconds), + jti: new_token_id()?, + }; + self.sign(&claims) + } + + pub fn sign(&self, claims: &WorkerMutationSourceClaims) -> Result { + if claims.iss != self.runtime_id { + return Err(RuntimeAuthError::UnknownIssuer(claims.iss.clone())); + } + let private = decode_private_key(&self.private_key)?; + let pair = Ed25519KeyPair::from_pkcs8(&private) + .map_err(|_| RuntimeAuthError::InvalidPrivateKey)?; + let payload = URL_SAFE_NO_PAD.encode(serde_json::to_vec(claims)?); + let signing_input = format!("{WORKER_MUTATION_SOURCE_SIGNING_INPUT_PREFIX}{payload}"); + let signature = pair.sign(signing_input.as_bytes()); + Ok(format!( + "{WORKER_MUTATION_SOURCE_PROOF_PREFIX}.{payload}.{}", + URL_SAFE_NO_PAD.encode(signature.as_ref()) + )) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct WorkerMutationSourceExpectation<'a> { + pub runtime_id: &'a str, + pub audience: &'a str, + pub workspace_id: &'a str, + pub worker_id: Option<&'a str>, + pub actor_kind: WorkerMutationActorKind, + pub operation: WorkerMutationOperation, + pub target_runtime_id: &'a str, + pub target_worker_id: &'a str, + pub permission: &'a str, +} + +pub fn decode_worker_mutation_source_claims( + token: &str, +) -> Result { + let (payload, _) = split_worker_mutation_source_proof(token)?; + Ok(serde_json::from_slice(&URL_SAFE_NO_PAD.decode(payload)?)?) +} + +pub fn verify_worker_mutation_source_proof( + trusted_runtime_public_key: &str, + token: &str, + expected: &WorkerMutationSourceExpectation<'_>, + now_seconds: u64, +) -> Result { + let (payload, signature) = split_worker_mutation_source_proof(token)?; + let claims_json = URL_SAFE_NO_PAD.decode(payload)?; + let claims: WorkerMutationSourceClaims = serde_json::from_slice(&claims_json)?; + let public_key = decode_public_key(trusted_runtime_public_key)?; + let signing_input = format!("{WORKER_MUTATION_SOURCE_SIGNING_INPUT_PREFIX}{payload}"); + UnparsedPublicKey::new(&ED25519, public_key) + .verify(signing_input.as_bytes(), &signature) + .map_err(|_| RuntimeAuthError::InvalidSignature)?; + + if claims.iss != expected.runtime_id { + return Err(RuntimeAuthError::UnknownIssuer(claims.iss)); + } + if claims.aud != expected.audience { + return Err(RuntimeAuthError::WrongAudience { + expected: expected.audience.to_string(), + actual: claims.aud, + }); + } + if claims.exp <= now_seconds || claims.iat > now_seconds.saturating_add(60) { + return Err(RuntimeAuthError::Expired); + } + if claims.workspace_id != expected.workspace_id { + return Err(RuntimeAuthError::WrongWorkspace { + expected: expected.workspace_id.to_string(), + actual: claims.workspace_id, + }); + } + if let Some(worker_id) = expected.worker_id { + if claims.worker_id != worker_id { + return Err(RuntimeAuthError::WrongWorker { + expected: worker_id.to_string(), + actual: claims.worker_id, + }); + } + } + if claims.actor_kind != expected.actor_kind { + return Err(RuntimeAuthError::WrongActorKind); + } + if claims.operation != expected.operation { + return Err(RuntimeAuthError::WrongOperation); + } + if claims.target_runtime_id != expected.target_runtime_id + || claims.target_worker_id != expected.target_worker_id + { + return Err(RuntimeAuthError::WrongMutationTarget); + } + if claims.permission != expected.permission { + return Err(RuntimeAuthError::MissingPermission( + expected.permission.to_string(), + )); + } + if claims.jti.trim().is_empty() { + return Err(RuntimeAuthError::InvalidTokenFormat); + } + Ok(claims) +} + +fn split_worker_mutation_source_proof(token: &str) -> Result<(&str, Vec), RuntimeAuthError> { + let mut parts = token.split('.'); + match (parts.next(), parts.next(), parts.next(), parts.next()) { + (Some(prefix), Some(payload), Some(signature), None) + if prefix == WORKER_MUTATION_SOURCE_PROOF_PREFIX => + { + Ok((payload, URL_SAFE_NO_PAD.decode(signature)?)) + } + _ => Err(RuntimeAuthError::InvalidTokenFormat), + } +} + fn split_token(token: &str) -> Result<(&str, Vec), RuntimeAuthError> { let mut parts = token.split('.'); match (parts.next(), parts.next(), parts.next(), parts.next()) { @@ -277,6 +486,112 @@ impl fmt::Display for RuntimeAuthContext { mod tests { use super::*; + #[test] + fn worker_mutation_source_proof_binds_all_source_authority_claims() { + let runtime = RuntimeIdentityMaterial::generate("runtime-main").unwrap(); + let signer = RuntimeWorkerMutationSourceSigner::from_identity(&runtime); + let claims = WorkerMutationSourceClaims { + iss: "runtime-main".to_string(), + aud: "server-main".to_string(), + workspace_id: "workspace-a".to_string(), + worker_id: "worker-7".to_string(), + actor_kind: WorkerMutationActorKind::Worker, + operation: WorkerMutationOperation::WorkerRemove, + target_runtime_id: "runtime-target".to_string(), + target_worker_id: "worker-target".to_string(), + permission: WORKER_REMOVE_PERMISSION.to_string(), + iat: 90, + exp: 100, + jti: "source-proof-1".to_string(), + }; + let token = signer.sign(&claims).unwrap(); + let expected = WorkerMutationSourceExpectation { + runtime_id: "runtime-main", + audience: "server-main", + workspace_id: "workspace-a", + worker_id: Some("worker-7"), + actor_kind: WorkerMutationActorKind::Worker, + operation: WorkerMutationOperation::WorkerRemove, + target_runtime_id: "runtime-target", + target_worker_id: "worker-target", + permission: WORKER_REMOVE_PERMISSION, + }; + + assert_eq!( + verify_worker_mutation_source_proof(&runtime.public_key, &token, &expected, 99) + .unwrap(), + claims + ); + + let wrong_worker = WorkerMutationSourceExpectation { + worker_id: Some("worker-8"), + ..expected.clone() + }; + assert!(matches!( + verify_worker_mutation_source_proof(&runtime.public_key, &token, &wrong_worker, 99), + Err(RuntimeAuthError::WrongWorker { .. }) + )); + let wrong_scope = WorkerMutationSourceExpectation { + workspace_id: "workspace-b", + ..expected.clone() + }; + assert!(matches!( + verify_worker_mutation_source_proof(&runtime.public_key, &token, &wrong_scope, 99), + Err(RuntimeAuthError::WrongWorkspace { .. }) + )); + let wrong_audience = WorkerMutationSourceExpectation { + audience: "server-other", + ..expected.clone() + }; + assert!(matches!( + verify_worker_mutation_source_proof(&runtime.public_key, &token, &wrong_audience, 99), + Err(RuntimeAuthError::WrongAudience { .. }) + )); + assert!(matches!( + verify_worker_mutation_source_proof(&runtime.public_key, &token, &expected, 101), + Err(RuntimeAuthError::Expired) + )); + } + + #[test] + fn worker_mutation_source_proof_rejects_spoofed_runtime_signature() { + let trusted = RuntimeIdentityMaterial::generate("runtime-main").unwrap(); + let spoofed = RuntimeIdentityMaterial::generate("runtime-main").unwrap(); + let claims = WorkerMutationSourceClaims { + iss: "runtime-main".to_string(), + aud: "server-main".to_string(), + workspace_id: "workspace-a".to_string(), + worker_id: "worker-7".to_string(), + actor_kind: WorkerMutationActorKind::Worker, + operation: WorkerMutationOperation::WorkerRemove, + target_runtime_id: "runtime-target".to_string(), + target_worker_id: "worker-target".to_string(), + permission: WORKER_REMOVE_PERMISSION.to_string(), + iat: 90, + exp: 100, + jti: "source-proof-2".to_string(), + }; + let token = RuntimeWorkerMutationSourceSigner::from_identity(&spoofed) + .sign(&claims) + .unwrap(); + let expected = WorkerMutationSourceExpectation { + runtime_id: "runtime-main", + audience: "server-main", + workspace_id: "workspace-a", + worker_id: Some("worker-7"), + actor_kind: WorkerMutationActorKind::Worker, + operation: WorkerMutationOperation::WorkerRemove, + target_runtime_id: "runtime-target", + target_worker_id: "worker-target", + permission: WORKER_REMOVE_PERMISSION, + }; + + assert!(matches!( + verify_worker_mutation_source_proof(&trusted.public_key, &token, &expected, 99), + Err(RuntimeAuthError::InvalidSignature) + )); + } + #[test] fn capability_token_verifies_signature_audience_expiry_and_permission() { let server = RuntimeIdentityMaterial::generate("server-main").unwrap(); diff --git a/crates/worker-runtime/src/catalog.rs b/crates/worker-runtime/src/catalog.rs index 3adbfb45..e32f450d 100644 --- a/crates/worker-runtime/src/catalog.rs +++ b/crates/worker-runtime/src/catalog.rs @@ -179,8 +179,6 @@ pub struct WorkingDirectoryStatus { pub struct WorkspaceApiRef { pub workspace_id: String, pub base_url: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub runtime_id: Option, } impl std::fmt::Debug for WorkspaceApiRef { @@ -189,7 +187,6 @@ impl std::fmt::Debug for WorkspaceApiRef { .debug_struct("WorkspaceApiRef") .field("workspace_id", &self.workspace_id) .field("base_url", &self.base_url) - .field("runtime_id", &self.runtime_id) .finish() } } @@ -292,3 +289,31 @@ pub struct WorkerLifecycleAck { pub worker_ref: WorkerRef, pub status: WorkerStatus, } + +#[cfg(test)] +mod tests { + use super::WorkspaceApiRef; + + #[test] + fn workspace_api_ref_public_schema_contains_no_source_credentials_or_claim_choices() { + let value = serde_json::to_value(WorkspaceApiRef { + workspace_id: "workspace-a".to_string(), + base_url: "https://server.invalid".to_string(), + }) + .unwrap(); + let object = value.as_object().unwrap(); + assert_eq!(object.len(), 2); + assert!(object.contains_key("workspace_id")); + assert!(object.contains_key("base_url")); + for forbidden in [ + "runtime_id", + "worker_id", + "permission", + "private_key", + "bearer_token", + "signing_handle", + ] { + assert!(!object.contains_key(forbidden), "unexpected {forbidden}"); + } + } +} diff --git a/crates/worker-runtime/src/execution.rs b/crates/worker-runtime/src/execution.rs index 5f43b43a..381e0c84 100644 --- a/crates/worker-runtime/src/execution.rs +++ b/crates/worker-runtime/src/execution.rs @@ -248,6 +248,7 @@ pub struct WorkerExecutionSpawnRequest { /// Monotonic execution generation reserved durably before launch. pub run_generation: u64, pub request: crate::catalog::CreateWorkerRequest, + pub workspace_scope: Option, pub context: WorkerExecutionContext, pub working_directory: Option, pub config_bundle: Option, @@ -260,6 +261,7 @@ pub struct WorkerExecutionRestoreRequest { /// Monotonic execution generation reserved durably before restore. pub run_generation: u64, pub request: crate::catalog::CreateWorkerRequest, + pub workspace_scope: Option, pub context: WorkerExecutionContext, pub previous_working_directory: Option, pub working_directory: Option, diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index f08f245b..dce58519 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -1753,7 +1753,6 @@ mod tests { request.workspace_api = Some(WorkspaceApiRef { workspace_id: workspace_id.to_string(), base_url: format!("https://workspace.example/{workspace_id}"), - runtime_id: None, }); request } @@ -2324,7 +2323,6 @@ mod tests { workspace_api: WorkspaceApiRef { workspace_id: "local".to_string(), base_url: "http://127.0.0.1:8787".to_string(), - runtime_id: None, }, }, ) diff --git a/crates/worker-runtime/src/lib.rs b/crates/worker-runtime/src/lib.rs index 001a102a..44817026 100644 --- a/crates/worker-runtime/src/lib.rs +++ b/crates/worker-runtime/src/lib.rs @@ -26,6 +26,7 @@ pub mod resource; pub mod retention; mod runtime; pub mod worker_backend; +pub mod worker_source; pub mod working_directory; #[cfg(feature = "fs-store")] diff --git a/crates/worker-runtime/src/main.rs b/crates/worker-runtime/src/main.rs index 448ffa42..0e4371e2 100644 --- a/crates/worker-runtime/src/main.rs +++ b/crates/worker-runtime/src/main.rs @@ -87,6 +87,9 @@ fn build_runtime(config: &ProcessConfig) -> Result { }; let mut factory = ProfileRuntimeWorkerFactory::new(fs_paths.worker_dir.join("worker-root")) .with_runtime_store_dir(runtime_store_dir); + if let Some(identity) = read_runtime_auth_file(&runtime_auth_path(config))?.identity { + factory = factory.with_remote_worker_mutation_identity(identity); + } if let Some(endpoint) = config.backend_resource_endpoint.clone() { factory = factory.with_resource_client(Arc::new( worker_runtime::resource::HttpBackendResourceClient::new( diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 2c9a4961..2b70adce 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -548,6 +548,7 @@ impl Runtime { worker_ref: worker_ref.clone(), run_generation: 1, request, + workspace_scope: scope.cloned(), context: self.execution_context(worker_ref.clone()), working_directory: None, config_bundle: None, @@ -826,13 +827,10 @@ impl Runtime { if let Some(existing) = worker.request.workspace_api.as_ref() && (existing.workspace_id != workspace_api.workspace_id || existing.base_url.trim_end_matches('/') - != workspace_api.base_url.trim_end_matches('/') - || existing.runtime_id.as_ref().is_some_and(|runtime_id| { - workspace_api.runtime_id.as_ref() != Some(runtime_id) - })) + != workspace_api.base_url.trim_end_matches('/')) { return Err(RuntimeError::InvalidRequest( - "Workspace API replacement cannot change Worker Workspace identity, Runtime identity, or base URL" + "Workspace API replacement cannot change Worker Workspace identity or base URL" .to_string(), )); } @@ -902,10 +900,17 @@ impl Runtime { })?; state.worker_mut(worker_ref)?.run_generation = run_generation; state.persist_worker(&worker_ref.worker_id)?; + let workspace_scope = worker_request.workspace_api.as_ref().and_then(|api| { + state + .workspace_owners + .get(&api.workspace_id) + .map(|server_id| RuntimeWorkspaceScope::new(&api.workspace_id, server_id)) + }); let request = WorkerExecutionRestoreRequest { worker_ref: worker_ref.clone(), run_generation, request: worker_request, + workspace_scope, context: self.execution_context(worker_ref.clone()), previous_working_directory, working_directory: None, @@ -1595,9 +1600,15 @@ impl Runtime { }; for candidate in candidates { - let backend = { + let (backend, workspace_scope) = { let state = self.lock()?; - state.execution_backend.clone() + let workspace_scope = candidate.request.workspace_api.as_ref().and_then(|api| { + state + .workspace_owners + .get(&api.workspace_id) + .map(|server_id| RuntimeWorkspaceScope::new(&api.workspace_id, server_id)) + }); + (state.execution_backend.clone(), workspace_scope) }; let Some(backend) = backend else { return Ok(()); @@ -1606,6 +1617,7 @@ impl Runtime { worker_ref: candidate.worker_ref.clone(), run_generation: candidate.run_generation, request: candidate.request, + workspace_scope, context: self.execution_context(candidate.worker_ref.clone()), previous_working_directory: candidate.previous_working_directory, working_directory: None, @@ -2790,7 +2802,6 @@ mod tests { request.workspace_api = Some(WorkspaceApiRef { workspace_id: workspace_id.to_string(), base_url: format!("https://workspace.example/{workspace_id}"), - runtime_id: None, }); request } @@ -3239,7 +3250,6 @@ mod tests { let replacement = WorkspaceApiRef { workspace_id: "workspace-a".to_string(), base_url: "https://workspace.example/workspace-a/".to_string(), - runtime_id: Some("runtime-a".to_string()), }; runtime diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 2825c162..e4fe6303 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -14,6 +14,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, mpsc}; use std::time::Duration; +use crate::auth::RuntimeIdentityMaterial; use crate::catalog::{ CreateWorkerRequest, ProfileSourceArchiveHttpRef, ProfileSourceArchiveSource, WorkingDirectoryRequest, WorkingDirectoryStatus, @@ -26,6 +27,9 @@ use crate::execution::{ use crate::identity::WorkerRef; use crate::interaction::{WorkerInput, WorkerInputKind}; use crate::resource::{BackendResourceClient, ProfileSourceArchiveCache}; +use crate::worker_source::{ + EmbeddedWorkerMutationDispatcher, RuntimeOwnedWorkspaceClient, RuntimeWorkerMutationForwarder, +}; use crate::working_directory::{ WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer, }; @@ -49,10 +53,9 @@ use worker::feature::builtin::{ #[cfg(feature = "ws-server")] use worker::ipc::protocol_session::{live_log_entry_event, subscribe_worker_protocol_session}; use worker::{ - PromptLoader, RuntimeWorkspaceHttpClient, SegmentLogSink, - WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerController, WorkerError, - WorkerFilesystemAuthority, WorkerHandle, WorkerSharedState, WorkerWorkspaceContext, - WorkspaceClient, WorkspaceId, + PromptLoader, SegmentLogSink, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, + WorkerController, WorkerError, WorkerFilesystemAuthority, WorkerHandle, WorkerSharedState, + WorkerWorkspaceContext, WorkspaceClient, WorkspaceId, }; const DEFAULT_BACKEND_ID: &str = "worker-crate"; @@ -215,6 +218,9 @@ pub struct ProfileRuntimeWorkerFactory { worker_aggregate_root: Option, resource_client: Option>, profile_archive_cache: Arc, + runtime_id: Option, + worker_mutation_identity: Option, + embedded_worker_mutation_dispatcher: Option>, } impl ProfileRuntimeWorkerFactory { @@ -226,9 +232,38 @@ impl ProfileRuntimeWorkerFactory { worker_aggregate_root: None, resource_client: None, profile_archive_cache: Arc::new(ProfileSourceArchiveCache::default()), + runtime_id: None, + worker_mutation_identity: None, + embedded_worker_mutation_dispatcher: None, } } + pub fn with_runtime_id(mut self, runtime_id: impl Into) -> Self { + self.runtime_id = Some(runtime_id.into()); + self + } + + pub fn with_remote_worker_mutation_identity( + mut self, + identity: RuntimeIdentityMaterial, + ) -> Self { + self.runtime_id = Some(identity.identity_id.clone()); + self.worker_mutation_identity = Some(identity); + self.embedded_worker_mutation_dispatcher = None; + self + } + + pub fn with_embedded_worker_mutation_dispatcher( + mut self, + runtime_id: impl Into, + dispatcher: Arc, + ) -> Self { + self.runtime_id = Some(runtime_id.into()); + self.worker_mutation_identity = None; + self.embedded_worker_mutation_dispatcher = Some(dispatcher); + self + } + pub fn with_runtime_store_dir(mut self, runtime_store_dir: impl Into) -> Self { self.worker_aggregate_root = Some(runtime_store_dir.into().join("workers")); self @@ -348,38 +383,59 @@ enum RuntimeWorkspaceBackendRef { } impl RuntimeWorkspaceBackendRef { - fn from_worker_request(request: &CreateWorkerRequest) -> Self { - if let Some(api) = request.workspace_api.as_ref() - && let Some(runtime_id) = api - .runtime_id - .as_ref() - .filter(|runtime_id| !runtime_id.trim().is_empty()) - { + fn from_worker_request(request: &CreateWorkerRequest, runtime_id: Option<&str>) -> Self { + if let (Some(api), Some(runtime_id)) = (request.workspace_api.as_ref(), runtime_id) { return Self::Http { workspace_id: api.workspace_id.clone(), base_url: api.base_url.clone(), - runtime_id: runtime_id.clone(), + runtime_id: runtime_id.to_string(), }; } Self::None } - fn worker_context(&self, worker_ref: &WorkerRef) -> WorkerWorkspaceContext { + fn worker_context( + &self, + worker_ref: &WorkerRef, + workspace_scope: Option<&crate::runtime::RuntimeWorkspaceScope>, + mutation_identity: Option<&RuntimeIdentityMaterial>, + embedded_dispatcher: Option<&Arc>, + ) -> WorkerWorkspaceContext { match self { Self::None => WorkerWorkspaceContext::no_workspace(), Self::Http { workspace_id, base_url, runtime_id, - } => WorkerWorkspaceContext::with_client( - WorkspaceId::new(workspace_id.clone()).ok(), - Arc::new(RuntimeWorkspaceHttpClient::new( + } => { + let mut client = RuntimeOwnedWorkspaceClient::new( workspace_id.clone(), base_url.clone(), runtime_id.clone(), worker_ref.worker_id.to_string(), - )), - ), + ); + if let (Some(scope), Some(identity)) = (workspace_scope, mutation_identity) { + client = client.with_worker_remove(RuntimeWorkerMutationForwarder::remote( + identity, + scope.clone(), + worker_ref.worker_id.to_string(), + base_url.clone(), + )); + } else if let (Some(scope), Some(dispatcher)) = + (workspace_scope, embedded_dispatcher) + { + client = client.with_worker_remove(RuntimeWorkerMutationForwarder::embedded( + runtime_id, + scope.clone(), + worker_ref.worker_id.to_string(), + (*dispatcher).clone(), + )); + } + WorkerWorkspaceContext::with_client( + WorkspaceId::new(workspace_id.clone()).ok(), + Arc::new(client), + ) + } } } } @@ -478,13 +534,11 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { ) }) .unwrap_or(WorkerFilesystemAuthority::None); - let workspace_backend_ref = - RuntimeWorkspaceBackendRef::from_worker_request(&request.request); - let observation_runtime_id = request - .request - .workspace_api - .as_ref() - .and_then(|api| api.runtime_id.clone()); + let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request( + &request.request, + self.runtime_id.as_deref(), + ); + let observation_runtime_id = self.runtime_id.clone(); let observation_workspace_id = request .request .workspace_api @@ -492,7 +546,12 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { .map(|api| api.workspace_id.clone()); let observation_grants = request.request.worker_observation_grants.clone(); let observation_enabled = request.request.worker_observation_enabled; - let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref); + let workspace_context = workspace_backend_ref.worker_context( + &request.worker_ref, + request.workspace_scope.as_ref(), + self.worker_mutation_identity.as_ref(), + self.embedded_worker_mutation_dispatcher.as_ref(), + ); let selector = profile.as_ref(); let archive = self .resolve_profile_source_archive(&request.request.profile_source) @@ -628,13 +687,11 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { ) }) .unwrap_or(WorkerFilesystemAuthority::None); - let workspace_backend_ref = - RuntimeWorkspaceBackendRef::from_worker_request(&request.request); - let observation_runtime_id = request - .request - .workspace_api - .as_ref() - .and_then(|api| api.runtime_id.clone()); + let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request( + &request.request, + self.runtime_id.as_deref(), + ); + let observation_runtime_id = self.runtime_id.clone(); let observation_workspace_id = request .request .workspace_api @@ -642,7 +699,12 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { .map(|api| api.workspace_id.clone()); let observation_grants = request.request.worker_observation_grants.clone(); let observation_enabled = request.request.worker_observation_enabled; - let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref); + let workspace_context = workspace_backend_ref.worker_context( + &request.worker_ref, + request.workspace_scope.as_ref(), + self.worker_mutation_identity.as_ref(), + self.embedded_worker_mutation_dispatcher.as_ref(), + ); let (manifest, loader) = Self::restore_fallback_manifest(&worker_name)?; let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?; @@ -1756,6 +1818,36 @@ mod tests { use manifest::{Scope, WorkerManifest}; use session_store::{LogEntry, WorkerMetadataStore}; + #[test] + fn restart_restore_reconstructs_runtime_owned_worker_mutation_client() { + let identity = RuntimeIdentityMaterial::generate("runtime-source").unwrap(); + let worker_ref = WorkerRef::new(crate::identity::WorkerId::new(17)); + let backend = RuntimeWorkspaceBackendRef::Http { + workspace_id: "workspace-a".to_string(), + base_url: "https://server.invalid".to_string(), + runtime_id: "runtime-source".to_string(), + }; + let scope = crate::runtime::RuntimeWorkspaceScope::new("workspace-a", "server-main"); + + let before_restart = + backend.worker_context(&worker_ref, Some(&scope), Some(&identity), None); + let after_restore = + backend.worker_context(&worker_ref, Some(&scope), Some(&identity), None); + + assert_eq!( + before_restart.client_handle().kind(), + "runtime-owned-workspace-client" + ); + assert_eq!( + after_restore.client_handle().kind(), + "runtime-owned-workspace-client" + ); + assert_eq!( + after_restore.client_handle().workspace_id(), + Some("workspace-a") + ); + } + #[test] fn notify_run_state_allows_running_worker_inbox_delivery() { assert_eq!( @@ -1878,9 +1970,16 @@ mod tests { .as_ref() .map(|binding| binding.root().to_path_buf()) .unwrap_or_else(|| self.cwd.clone()); - let workspace_backend_ref = - RuntimeWorkspaceBackendRef::from_worker_request(&request.request); - let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref); + let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request( + &request.request, + Some("runtime-test"), + ); + let workspace_context = workspace_backend_ref.worker_context( + &request.worker_ref, + request.workspace_scope.as_ref(), + None, + None, + ); let workspace_client = workspace_context.client_handle(); self.observed_workspace_clients.lock().unwrap().push(( workspace_client.kind().to_string(), @@ -1916,6 +2015,7 @@ mod tests { worker_ref: request.worker_ref, run_generation: request.run_generation, request: request.request, + workspace_scope: request.workspace_scope, context: request.context, working_directory: request.working_directory, config_bundle: request.config_bundle, @@ -2188,6 +2288,7 @@ mod tests { worker_ref: worker_ref.clone(), run_generation: 1, request: create_request("1"), + workspace_scope: None, context: test_execution_context(worker_ref), working_directory: None, config_bundle: None, @@ -2286,14 +2387,15 @@ mod tests { request.workspace_api = Some(crate::catalog::WorkspaceApiRef { workspace_id: "workspace-restore".to_string(), base_url: "http://workspace.invalid".to_string(), - runtime_id: Some("runtime-restore".to_string()), }); let controller = ProfileRuntimeWorkerFactory::new(root.path()) + .with_runtime_id("runtime-restore") .with_runtime_store_dir(&runtime_store_dir) .restore_controller(WorkerExecutionRestoreRequest { worker_ref: worker_ref.clone(), run_generation: 1, request, + workspace_scope: None, context: test_execution_context(worker_ref), previous_working_directory: None, working_directory: None, @@ -2420,7 +2522,6 @@ mod tests { request.workspace_api = Some(crate::catalog::WorkspaceApiRef { workspace_id: "ws-test".to_string(), base_url: "http://127.0.0.1:3999".to_string(), - runtime_id: Some("runtime-test".to_string()), }); let detail = runtime.create_worker(request).unwrap(); @@ -2453,7 +2554,7 @@ mod tests { assert_eq!( observed_workspace_clients.lock().unwrap().as_slice(), &[( - "runtime-http-proxy".to_string(), + "runtime-owned-workspace-client".to_string(), Some("ws-test".to_string()), true, )] diff --git a/crates/worker-runtime/src/worker_source.rs b/crates/worker-runtime/src/worker_source.rs new file mode 100644 index 00000000..b9fdd0c1 --- /dev/null +++ b/crates/worker-runtime/src/worker_source.rs @@ -0,0 +1,613 @@ +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use worker::{ + WorkspaceClient, WorkspaceClientError, WorkspaceRequest, WorkspaceRequestMethod, + WorkspaceResponse, +}; + +use crate::auth::{ + RuntimeAuthError, RuntimeIdentityMaterial, RuntimeWorkerMutationSourceSigner, + WORKER_REMOVE_PERMISSION, WorkerMutationActorKind, WorkerMutationOperation, + WorkerMutationSourceClaims, new_token_id, +}; +use crate::runtime::RuntimeWorkspaceScope; + +pub const DEFAULT_WORKER_MUTATION_SOURCE_TTL_SECONDS: u64 = 60; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RuntimeOwnedWorkerMutationProof { + Remote(String), + InProcess(InProcessWorkerMutationProof), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InProcessWorkerMutationProof { + claims: WorkerMutationSourceClaims, +} + +impl InProcessWorkerMutationProof { + pub fn claims(&self) -> &WorkerMutationSourceClaims { + &self.claims + } + + pub fn into_claims(self) -> WorkerMutationSourceClaims { + self.claims + } +} + +#[derive(Clone)] +pub struct RuntimeWorkerMutationSourceAuthority { + mode: RuntimeWorkerMutationSourceMode, +} + +#[derive(Clone)] +enum RuntimeWorkerMutationSourceMode { + Remote { + signer: RuntimeWorkerMutationSourceSigner, + }, + Embedded { + runtime_id: String, + audience: String, + }, +} + +impl RuntimeWorkerMutationSourceAuthority { + pub fn remote(identity: &RuntimeIdentityMaterial) -> Self { + Self { + mode: RuntimeWorkerMutationSourceMode::Remote { + signer: RuntimeWorkerMutationSourceSigner::from_identity(identity), + }, + } + } + + pub fn embedded(runtime_id: impl Into, workspace_id: impl AsRef) -> Self { + Self { + mode: RuntimeWorkerMutationSourceMode::Embedded { + runtime_id: runtime_id.into(), + audience: format!("embedded:{}", workspace_id.as_ref()), + }, + } + } + + pub fn issue_worker_remove( + &self, + scope: &RuntimeWorkspaceScope, + source_worker_id: &str, + target_runtime_id: &str, + target_worker_id: &str, + ) -> Result { + match &self.mode { + RuntimeWorkerMutationSourceMode::Remote { signer } => { + let token = signer.issue_worker_remove( + &scope.server_id, + &scope.workspace_id, + source_worker_id, + target_runtime_id, + target_worker_id, + DEFAULT_WORKER_MUTATION_SOURCE_TTL_SECONDS, + )?; + Ok(RuntimeOwnedWorkerMutationProof::Remote(token)) + } + RuntimeWorkerMutationSourceMode::Embedded { + runtime_id, + audience, + } => { + let issued_at = unix_now_seconds(); + Ok(RuntimeOwnedWorkerMutationProof::InProcess( + InProcessWorkerMutationProof { + claims: WorkerMutationSourceClaims { + iss: runtime_id.clone(), + aud: audience.clone(), + workspace_id: scope.workspace_id.clone(), + worker_id: source_worker_id.to_string(), + actor_kind: WorkerMutationActorKind::Worker, + operation: WorkerMutationOperation::WorkerRemove, + target_runtime_id: target_runtime_id.to_string(), + target_worker_id: target_worker_id.to_string(), + permission: WORKER_REMOVE_PERMISSION.to_string(), + iat: issued_at, + exp: issued_at + .saturating_add(DEFAULT_WORKER_MUTATION_SOURCE_TTL_SECONDS), + jti: new_token_id()?, + }, + }, + )) + } + } + } +} + +pub trait EmbeddedWorkerMutationDispatcher: Send + Sync { + fn execute_worker_remove( + &self, + proof: InProcessWorkerMutationProof, + target_runtime_id: &str, + target_worker_id: &str, + ) -> Result; +} + +#[derive(Clone)] +enum RuntimeWorkerMutationTransport { + Remote { + base_url: String, + client: reqwest::blocking::Client, + }, + Embedded { + dispatcher: Arc, + }, +} + +#[derive(Clone)] +pub struct RuntimeWorkerMutationForwarder { + authority: RuntimeWorkerMutationSourceAuthority, + scope: RuntimeWorkspaceScope, + source_worker_id: String, + transport: RuntimeWorkerMutationTransport, +} + +impl RuntimeWorkerMutationForwarder { + pub fn remote( + identity: &RuntimeIdentityMaterial, + scope: RuntimeWorkspaceScope, + source_worker_id: impl Into, + base_url: impl Into, + ) -> Self { + Self { + authority: RuntimeWorkerMutationSourceAuthority::remote(identity), + scope, + source_worker_id: source_worker_id.into(), + transport: RuntimeWorkerMutationTransport::Remote { + base_url: base_url.into().trim_end_matches('/').to_string(), + client: reqwest::blocking::Client::new(), + }, + } + } + + pub fn embedded( + runtime_id: impl Into, + scope: RuntimeWorkspaceScope, + source_worker_id: impl Into, + dispatcher: Arc, + ) -> Self { + let runtime_id = runtime_id.into(); + Self { + authority: RuntimeWorkerMutationSourceAuthority::embedded( + &runtime_id, + &scope.workspace_id, + ), + scope, + source_worker_id: source_worker_id.into(), + transport: RuntimeWorkerMutationTransport::Embedded { dispatcher }, + } + } + + pub fn execute_worker_remove( + &self, + target_runtime_id: &str, + target_worker_id: &str, + ) -> Result { + let proof = self.authority.issue_worker_remove( + &self.scope, + &self.source_worker_id, + target_runtime_id, + target_worker_id, + )?; + match (&self.transport, proof) { + ( + RuntimeWorkerMutationTransport::Remote { base_url, client }, + RuntimeOwnedWorkerMutationProof::Remote(token), + ) => { + let url = format!( + "{base_url}/api/w/{}/workers/remove", + self.scope.workspace_id + ); + let body = serde_json::json!({ + "target_runtime_id": target_runtime_id, + "target_worker_id": target_worker_id, + }); + let response = client + .post(url) + .header(crate::auth::WORKER_MUTATION_SOURCE_PROOF_HEADER, token) + .json(&body) + .send() + .map_err(|error| { + RuntimeWorkerMutationForwardError::Transport(error.to_string()) + })?; + let status = response.status().as_u16(); + let body = response.text().map_err(|error| { + RuntimeWorkerMutationForwardError::Transport(error.to_string()) + })?; + Ok(WorkspaceResponse { status, body }) + } + ( + RuntimeWorkerMutationTransport::Embedded { dispatcher }, + RuntimeOwnedWorkerMutationProof::InProcess(claims), + ) => dispatcher.execute_worker_remove(claims, target_runtime_id, target_worker_id), + _ => Err(RuntimeWorkerMutationForwardError::AuthorityTransportMismatch), + } + } +} + +pub struct RuntimeOwnedWorkspaceClient { + workspace_id: String, + base_url: String, + runtime_id: String, + worker_id: String, + worker_remove: Option, +} + +impl RuntimeOwnedWorkspaceClient { + pub fn new( + workspace_id: impl Into, + base_url: impl Into, + runtime_id: impl Into, + worker_id: impl Into, + ) -> Self { + Self { + workspace_id: workspace_id.into(), + base_url: base_url.into().trim_end_matches('/').to_string(), + runtime_id: runtime_id.into(), + worker_id: worker_id.into(), + worker_remove: None, + } + } + + pub fn with_worker_remove(mut self, worker_remove: RuntimeWorkerMutationForwarder) -> Self { + self.worker_remove = Some(worker_remove); + self + } +} + +impl std::fmt::Debug for RuntimeOwnedWorkspaceClient { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("RuntimeOwnedWorkspaceClient") + .field("workspace_id", &self.workspace_id) + .field("base_url", &self.base_url) + .field("source", &"Runtime-owned") + .field( + "worker_remove", + &self.worker_remove.as_ref().map(|_| "enabled"), + ) + .finish() + } +} + +impl WorkspaceClient for RuntimeOwnedWorkspaceClient { + fn workspace_id(&self) -> Option<&str> { + Some(&self.workspace_id) + } + + fn kind(&self) -> &str { + "runtime-owned-workspace-client" + } + + fn is_available(&self) -> bool { + true + } + + fn execute( + &self, + request: WorkspaceRequest, + ) -> Result { + let base_url = self.base_url.clone(); + let runtime_id = self.runtime_id.clone(); + let worker_id = self.worker_id.clone(); + if tokio::runtime::Handle::try_current().is_ok() { + std::thread::spawn(move || { + execute_runtime_owned_workspace_http(&base_url, &runtime_id, &worker_id, request) + }) + .join() + .map_err(|_| { + WorkspaceClientError::Request("workspace request thread panicked".to_string()) + })? + } else { + execute_runtime_owned_workspace_http( + &self.base_url, + &self.runtime_id, + &self.worker_id, + request, + ) + } + } + + fn execute_worker_remove( + &self, + target_runtime_id: &str, + target_worker_id: &str, + ) -> Result { + self.worker_remove + .as_ref() + .ok_or_else(|| { + WorkspaceClientError::Unavailable( + "Runtime-owned WorkerRemove forwarding is unavailable".to_string(), + ) + })? + .execute_worker_remove(target_runtime_id, target_worker_id) + .map_err(|error| WorkspaceClientError::Request(error.to_string())) + } +} + +fn execute_runtime_owned_workspace_http( + base_url: &str, + runtime_id: &str, + worker_id: &str, + request: WorkspaceRequest, +) -> Result { + if !request.path.starts_with('/') || request.path.starts_with("//") { + return Err(WorkspaceClientError::InvalidPath(request.path)); + } + let url = format!("{base_url}{}", request.path); + let method = match request.method { + WorkspaceRequestMethod::Get => reqwest::Method::GET, + WorkspaceRequestMethod::Post => reqwest::Method::POST, + WorkspaceRequestMethod::Put => reqwest::Method::PUT, + WorkspaceRequestMethod::Patch => reqwest::Method::PATCH, + WorkspaceRequestMethod::Delete => reqwest::Method::DELETE, + }; + let client = reqwest::blocking::Client::new(); + let mut request_builder = client + .request(method, url) + .header("x-yoi-runtime-id", runtime_id) + .header("x-yoi-worker-id", worker_id); + if let Some(body) = request.body { + request_builder = request_builder + .header(reqwest::header::CONTENT_TYPE, "application/json") + .body(body); + } + let response = request_builder + .send() + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + let status = response.status().as_u16(); + let body = response + .text() + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + Ok(WorkspaceResponse { status, body }) +} + +#[derive(Debug, thiserror::Error)] +pub enum RuntimeWorkerMutationForwardError { + #[error(transparent)] + Auth(#[from] RuntimeAuthError), + #[error("Worker mutation forwarding transport failed: {0}")] + Transport(String), + #[error("Worker mutation source authority does not match its forwarding transport")] + AuthorityTransportMismatch, + #[error("embedded Worker mutation dispatcher failed: {0}")] + Embedded(String), +} + +fn unix_now_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::auth::{ + WorkerMutationSourceExpectation, decode_worker_mutation_source_claims, + verify_worker_mutation_source_proof, + }; + + #[test] + fn ordinary_workspace_forwarding_stamps_legacy_source_only_inside_runtime() { + use std::io::{Read, Write}; + use std::net::TcpListener; + use std::sync::Mutex; + + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let received = Arc::new(Mutex::new(String::new())); + let received_for_server = received.clone(); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut bytes = [0_u8; 4096]; + let count = stream.read(&mut bytes).unwrap(); + *received_for_server.lock().unwrap() = + String::from_utf8_lossy(&bytes[..count]).into_owned(); + stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}") + .unwrap(); + }); + + let client = RuntimeOwnedWorkspaceClient::new( + "workspace-a", + format!("http://{address}"), + "runtime-a", + "worker-a", + ); + let response = client + .execute(WorkspaceRequest::get("/api/w/workspace-a/tickets/search")) + .unwrap(); + assert_eq!(response.status, 200); + server.join().unwrap(); + let request = received.lock().unwrap().to_ascii_lowercase(); + assert!(request.contains("x-yoi-runtime-id: runtime-a")); + assert!(request.contains("x-yoi-worker-id: worker-a")); + assert!(!request.contains("authorization:")); + } + + #[test] + fn remote_authority_stamps_and_signs_worker_remove_without_caller_claim_choices() { + let identity = RuntimeIdentityMaterial::generate("runtime-a").unwrap(); + let authority = RuntimeWorkerMutationSourceAuthority::remote(&identity); + let scope = RuntimeWorkspaceScope { + workspace_id: "workspace-a".to_string(), + server_id: "server-a".to_string(), + }; + + let RuntimeOwnedWorkerMutationProof::Remote(token) = authority + .issue_worker_remove(&scope, "worker-source", "runtime-b", "worker-target") + .unwrap() + else { + panic!("remote authority must produce a signed proof"); + }; + let claims = decode_worker_mutation_source_claims(&token).unwrap(); + let expected = WorkerMutationSourceExpectation { + runtime_id: "runtime-a", + audience: "server-a", + workspace_id: "workspace-a", + worker_id: Some("worker-source"), + actor_kind: WorkerMutationActorKind::Worker, + operation: WorkerMutationOperation::WorkerRemove, + target_runtime_id: "runtime-b", + target_worker_id: "worker-target", + permission: WORKER_REMOVE_PERMISSION, + }; + assert_eq!( + verify_worker_mutation_source_proof( + &identity.public_key, + &token, + &expected, + claims.iat + ) + .unwrap(), + claims + ); + } + + #[test] + fn remote_forwarder_stamps_signed_proof_inside_runtime_before_http_delivery() { + use std::io::{Read, Write}; + use std::net::TcpListener; + use std::sync::Mutex; + + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let received = Arc::new(Mutex::new(String::new())); + let received_for_server = received.clone(); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut bytes = [0_u8; 8192]; + let count = stream.read(&mut bytes).unwrap(); + *received_for_server.lock().unwrap() = + String::from_utf8_lossy(&bytes[..count]).into_owned(); + stream + .write_all(b"HTTP/1.1 204 No Content\r\nContent-Length: 0\r\n\r\n") + .unwrap(); + }); + + let identity = RuntimeIdentityMaterial::generate("runtime-a").unwrap(); + let scope = RuntimeWorkspaceScope::new("workspace-a", "server-a"); + let forwarder = RuntimeWorkerMutationForwarder::remote( + &identity, + scope, + "worker-source", + format!("http://{address}"), + ); + let response = forwarder + .execute_worker_remove("runtime-target", "worker-target") + .unwrap(); + assert_eq!(response.status, 204); + server.join().unwrap(); + + let request = received.lock().unwrap().clone(); + assert!(request.starts_with("POST /api/w/workspace-a/workers/remove HTTP/1.1")); + assert!(request.contains( + r#"{"target_runtime_id":"runtime-target","target_worker_id":"worker-target"}"# + )); + let token = request + .lines() + .find_map(|line| { + line.split_once(':').and_then(|(name, value)| { + name.eq_ignore_ascii_case(crate::auth::WORKER_MUTATION_SOURCE_PROOF_HEADER) + .then(|| value.trim()) + }) + }) + .expect("proof header"); + let claims = decode_worker_mutation_source_claims(token).unwrap(); + let expected = WorkerMutationSourceExpectation { + runtime_id: "runtime-a", + audience: "server-a", + workspace_id: "workspace-a", + worker_id: Some("worker-source"), + actor_kind: WorkerMutationActorKind::Worker, + operation: WorkerMutationOperation::WorkerRemove, + target_runtime_id: "runtime-target", + target_worker_id: "worker-target", + permission: WORKER_REMOVE_PERMISSION, + }; + verify_worker_mutation_source_proof(&identity.public_key, token, &expected, claims.iat) + .unwrap(); + } + + #[test] + fn embedded_forwarder_delivers_in_process_proof_with_the_request() { + use std::sync::Mutex; + + #[derive(Default)] + struct RecordingDispatcher { + seen: Mutex>, + } + impl EmbeddedWorkerMutationDispatcher for RecordingDispatcher { + fn execute_worker_remove( + &self, + proof: InProcessWorkerMutationProof, + target_runtime_id: &str, + target_worker_id: &str, + ) -> Result { + *self.seen.lock().unwrap() = Some(( + proof.into_claims(), + target_runtime_id.to_string(), + target_worker_id.to_string(), + )); + Ok(WorkspaceResponse { + status: 202, + body: "accepted".to_string(), + }) + } + } + + let dispatcher = Arc::new(RecordingDispatcher::default()); + let scope = RuntimeWorkspaceScope::new("workspace-a", "server-a"); + let forwarder = RuntimeWorkerMutationForwarder::embedded( + "runtime-embedded", + scope, + "worker-source", + dispatcher.clone(), + ); + let response = forwarder + .execute_worker_remove("runtime-target", "worker-target") + .unwrap(); + assert_eq!(response.status, 202); + let (claims, target_runtime_id, target_worker_id) = + dispatcher.seen.lock().unwrap().take().unwrap(); + assert_eq!(claims.iss, "runtime-embedded"); + assert_eq!(claims.worker_id, "worker-source"); + assert_eq!(claims.target_runtime_id, "runtime-target"); + assert_eq!(claims.target_worker_id, "worker-target"); + assert_eq!(target_runtime_id, "runtime-target"); + assert_eq!(target_worker_id, "worker-target"); + } + + #[test] + fn embedded_authority_uses_the_same_claim_contract_without_a_credential() { + let authority = + RuntimeWorkerMutationSourceAuthority::embedded("runtime-embedded", "workspace-a"); + let scope = RuntimeWorkspaceScope { + workspace_id: "workspace-a".to_string(), + server_id: "server-unused-for-embedded".to_string(), + }; + + let RuntimeOwnedWorkerMutationProof::InProcess(proof) = authority + .issue_worker_remove(&scope, "worker-source", "runtime-b", "worker-target") + .unwrap() + else { + panic!("embedded authority must produce an in-process proof"); + }; + let claims = proof.claims(); + assert_eq!(claims.iss, "runtime-embedded"); + assert_eq!(claims.aud, "embedded:workspace-a"); + assert_eq!(claims.workspace_id, "workspace-a"); + assert_eq!(claims.worker_id, "worker-source"); + assert_eq!(claims.operation, WorkerMutationOperation::WorkerRemove); + assert_eq!(claims.target_runtime_id, "runtime-b"); + assert_eq!(claims.target_worker_id, "worker-target"); + assert_eq!(claims.permission, WORKER_REMOVE_PERMISSION); + assert!(!claims.jti.is_empty()); + } +} diff --git a/crates/worker/src/feature/builtin/memory.rs b/crates/worker/src/feature/builtin/memory.rs index 774dc950..b4eec17e 100644 --- a/crates/worker/src/feature/builtin/memory.rs +++ b/crates/worker/src/feature/builtin/memory.rs @@ -346,11 +346,9 @@ mod tests { use llm_engine::tool::ToolDefinition; fn test_client() -> Arc { - Arc::new(crate::worker::RuntimeWorkspaceHttpClient::new( + Arc::new(crate::worker::TestWorkspaceHttpClient::new( "workspace", "http://backend", - "test-runtime", - "test-worker", )) } diff --git a/crates/worker/src/feature/builtin/objective.rs b/crates/worker/src/feature/builtin/objective.rs index ef2e828b..db5d86a9 100644 --- a/crates/worker/src/feature/builtin/objective.rs +++ b/crates/worker/src/feature/builtin/objective.rs @@ -628,12 +628,7 @@ mod tests { #[test] fn workspace_http_objective_tools_include_objective_crud_tools() { let names = tool_names(workspace_http_objective_tools(Arc::new( - crate::worker::RuntimeWorkspaceHttpClient::new( - "workspace", - "http://backend", - "test-runtime", - "test-worker", - ), + crate::worker::TestWorkspaceHttpClient::new("workspace", "http://backend"), ))); assert_eq!( diff --git a/crates/worker/src/feature/builtin/ticket.rs b/crates/worker/src/feature/builtin/ticket.rs index 32d3eb30..8a392b79 100644 --- a/crates/worker/src/feature/builtin/ticket.rs +++ b/crates/worker/src/feature/builtin/ticket.rs @@ -1376,12 +1376,7 @@ provider = "github" #[tokio::test(flavor = "multi_thread")] async fn workspace_http_backend_invoke_is_safe_inside_async_context() { let backend = WorkspaceHttpTicketBackend::new(Arc::new( - crate::worker::RuntimeWorkspaceHttpClient::new( - "workspace-a", - "not-a-url", - "test-runtime", - "test-worker", - ), + crate::worker::TestWorkspaceHttpClient::new("workspace-a", "not-a-url"), )); let error = backend @@ -1412,11 +1407,9 @@ provider = "github" .write_all(b"HTTP/1.1 204 No Content\r\nContent-Length: 0\r\n\r\n") .unwrap(); }); - let client = Arc::new(crate::worker::RuntimeWorkspaceHttpClient::new( + let client = Arc::new(crate::worker::TestWorkspaceHttpClient::new( "workspace-a", format!("http://{address}"), - "test-runtime", - "worker-a", )); let backend = WorkspaceHttpTicketBackend::new(client); @@ -1457,12 +1450,7 @@ provider = "github" }); let backend = WorkspaceHttpTicketBackend::new(Arc::new( - crate::worker::RuntimeWorkspaceHttpClient::new( - "workspace-a", - base_url, - "test-runtime", - "test-worker", - ), + crate::worker::TestWorkspaceHttpClient::new("workspace-a", base_url), )); let created = backend.create(NewTicket::new("HTTP ticket")).unwrap(); diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 3c0d1e12..7df10a5a 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -40,9 +40,9 @@ pub use runtime::dir::RuntimeDir; pub use segment_log_sink::SegmentLogSink; pub use shared_state::WorkerSharedState; pub use worker::{ - LocalWorkingDirectory, RuntimeWorkspaceHttpClient, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, - Worker, WorkerError, WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, - WorkspaceClient, WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspaceRequest, - WorkspaceRequestMethod, WorkspaceResponse, apply_worker_manifest, marker_workspace_client, + LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError, + WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceClient, + WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspaceRequest, WorkspaceRequestMethod, + WorkspaceResponse, apply_worker_manifest, marker_workspace_client, unavailable_workspace_client, }; diff --git a/crates/worker/src/skill.rs b/crates/worker/src/skill.rs index cef7a564..effd90c7 100644 --- a/crates/worker/src/skill.rs +++ b/crates/worker/src/skill.rs @@ -212,8 +212,8 @@ mod tests { break; } } - assert_eq!(runtime_header.as_deref(), Some("runtime-test")); - assert_eq!(worker_header.as_deref(), Some("test-worker")); + assert_eq!(runtime_header, None); + assert_eq!(worker_header, None); assert_eq!(authorization, None); let body = serde_json::json!({ "authority": "workspace-backend-skills-v0", @@ -236,12 +236,7 @@ mod tests { .unwrap(); }); - let client = crate::worker::RuntimeWorkspaceHttpClient::new( - "ws-1", - format!("http://{addr}"), - "runtime-test", - "test-worker", - ); + let client = crate::worker::TestWorkspaceHttpClient::new("ws-1", format!("http://{addr}")); let catalog = (&client as &dyn WorkspaceClient).list_skills().unwrap(); assert_eq!(catalog.entries[0].name, "triage-errors"); assert_eq!(catalog.entries[0].provenance.id, "workspace:triage-errors"); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index d9fc0f7f..565c785c 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -224,6 +224,18 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync { fn execute(&self, request: WorkspaceRequest) -> Result; + /// Executes the destructive WorkerRemove operation through Runtime-owned source proof. + /// Target identity is operation data; source identity and permission are never caller inputs. + fn execute_worker_remove( + &self, + _target_runtime_id: &str, + _target_worker_id: &str, + ) -> Result { + Err(WorkspaceClientError::Unavailable( + "Runtime-owned WorkerRemove forwarding is unavailable".to_string(), + )) + } + /// Trusted review-attempt context is injected by the Internal SubWorker spawn layer. /// It is never accepted from a model-visible tool argument. fn reviewer_attempt_context(&self) -> Option<&ReviewerAttemptContext> { @@ -313,53 +325,31 @@ impl WorkspaceClient for ReviewerChildWorkspaceClient { } } -/// HTTP forwarding client created by Runtime for one concrete Worker execution. -/// -/// The upstream endpoint and source headers are private implementation details; -/// model-visible tools can only submit [`WorkspaceRequest`] values through the -/// [`WorkspaceClient`] trait. -pub struct RuntimeWorkspaceHttpClient { +#[cfg(test)] +#[derive(Clone, Debug)] +pub(crate) struct TestWorkspaceHttpClient { workspace_id: String, base_url: String, - runtime_id: String, - worker_id: String, } -impl std::fmt::Debug for RuntimeWorkspaceHttpClient { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter - .debug_struct("RuntimeWorkspaceHttpClient") - .field("workspace_id", &self.workspace_id) - .field("base_url", &self.base_url) - .field("runtime_id", &self.runtime_id) - .field("worker_id", &self.worker_id) - .finish() - } -} - -impl RuntimeWorkspaceHttpClient { - pub fn new( - workspace_id: impl Into, - base_url: impl Into, - runtime_id: impl Into, - worker_id: impl Into, - ) -> Self { +#[cfg(test)] +impl TestWorkspaceHttpClient { + pub(crate) fn new(workspace_id: impl Into, base_url: impl Into) -> Self { Self { workspace_id: workspace_id.into(), base_url: base_url.into().trim_end_matches('/').to_string(), - runtime_id: runtime_id.into(), - worker_id: worker_id.into(), } } } -impl WorkspaceClient for RuntimeWorkspaceHttpClient { +#[cfg(test)] +impl WorkspaceClient for TestWorkspaceHttpClient { fn workspace_id(&self) -> Option<&str> { Some(&self.workspace_id) } fn kind(&self) -> &str { - "runtime-http-proxy" + "test-http" } fn is_available(&self) -> bool { @@ -371,32 +361,28 @@ impl WorkspaceClient for RuntimeWorkspaceHttpClient { request: WorkspaceRequest, ) -> Result { let base_url = self.base_url.clone(); - let runtime_id = self.runtime_id.clone(); - let worker_id = self.worker_id.clone(); if tokio::runtime::Handle::try_current().is_ok() { - std::thread::spawn(move || { - execute_runtime_workspace_http(&base_url, &runtime_id, &worker_id, request) - }) - .join() - .map_err(|_| { - WorkspaceClientError::Request("workspace request thread panicked".to_string()) - })? + std::thread::spawn(move || execute_test_workspace_http(&base_url, request)) + .join() + .map_err(|_| { + WorkspaceClientError::Request( + "test workspace request thread panicked".to_string(), + ) + })? } else { - execute_runtime_workspace_http(&base_url, &runtime_id, &worker_id, request) + execute_test_workspace_http(&base_url, request) } } } -fn execute_runtime_workspace_http( +#[cfg(test)] +fn execute_test_workspace_http( base_url: &str, - runtime_id: &str, - worker_id: &str, request: WorkspaceRequest, ) -> Result { if !request.path.starts_with('/') || request.path.starts_with("//") { return Err(WorkspaceClientError::InvalidPath(request.path)); } - let url = format!("{base_url}{}", request.path); let method = match request.method { WorkspaceRequestMethod::Get => reqwest::Method::GET, WorkspaceRequestMethod::Post => reqwest::Method::POST, @@ -405,16 +391,11 @@ fn execute_runtime_workspace_http( WorkspaceRequestMethod::Delete => reqwest::Method::DELETE, }; let client = reqwest::blocking::Client::new(); - let mut request_builder = client - .request(method, url) - .header("x-yoi-runtime-id", runtime_id) - .header("x-yoi-worker-id", worker_id); + let mut builder = client.request(method, format!("{base_url}{}", request.path)); if let Some(body) = request.body { - request_builder = request_builder - .header(reqwest::header::CONTENT_TYPE, "application/json") - .body(body); + builder = builder.body(body); } - let response = request_builder + let response = builder .send() .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; let status = response.status().as_u16(); @@ -6959,11 +6940,9 @@ mod build_summary_prompt_tests { }); WorkerWorkspaceContext::with_client( Some(WorkspaceId::new("test-memory").unwrap()), - Arc::new(RuntimeWorkspaceHttpClient::new( + Arc::new(TestWorkspaceHttpClient::new( "test-memory", format!("http://{addr}"), - "test-runtime", - "test-worker", )), ) } @@ -7096,11 +7075,9 @@ mod build_summary_prompt_tests { store, WorkerWorkspaceContext::with_client( Some(WorkspaceId::new("ws-skill").unwrap()), - Arc::new(RuntimeWorkspaceHttpClient::new( + Arc::new(TestWorkspaceHttpClient::new( "ws-skill", format!("http://{addr}"), - "test-runtime", - "test-worker", )), ), authority, @@ -7142,58 +7119,6 @@ mod build_summary_prompt_tests { })); } - #[test] - fn runtime_workspace_client_sends_runtime_worker_identity_without_bearer() { - use std::io::{BufRead, BufReader, Write}; - use std::net::TcpListener; - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let address = listener.local_addr().unwrap(); - let server = std::thread::spawn(move || { - let (mut stream, _) = listener.accept().unwrap(); - let mut reader = BufReader::new(stream.try_clone().unwrap()); - let mut first_line = String::new(); - reader.read_line(&mut first_line).unwrap(); - assert!(first_line.contains("/api/w/workspace-a/tickets/search")); - let mut runtime_id = String::new(); - let mut worker_id = String::new(); - let mut authorization = String::new(); - loop { - let mut line = String::new(); - reader.read_line(&mut line).unwrap(); - if let Some(value) = line.strip_prefix("x-yoi-runtime-id: ") { - runtime_id = value.trim().to_string(); - } - if let Some(value) = line.strip_prefix("x-yoi-worker-id: ") { - worker_id = value.trim().to_string(); - } - if let Some(value) = line.strip_prefix("authorization: ") { - authorization = value.trim().to_string(); - } - if line == "\r\n" || line.is_empty() { - break; - } - } - assert_eq!(runtime_id, "runtime-a"); - assert_eq!(worker_id, "worker-a"); - assert!(authorization.is_empty()); - stream - .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}") - .unwrap(); - }); - let client = RuntimeWorkspaceHttpClient::new( - "workspace-a", - format!("http://{address}"), - "runtime-a", - "worker-a", - ); - let response = client - .execute(WorkspaceRequest::get("/api/w/workspace-a/tickets/search")) - .unwrap(); - assert_eq!(response.status, 200); - server.join().unwrap(); - } - #[tokio::test] async fn cancelled_internal_extract_does_not_commit_pointer_or_completed_audit() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 46aecd87..6d9c40d1 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -55,7 +55,7 @@ use worker_runtime::interaction::{ use worker_runtime::management::{RuntimeOptions as EmbeddedRuntimeOptions, RuntimeStatus}; use worker_runtime::profile_archive::{ProfileSourceArchive, ProfileSourceArchiveInput}; -const EMBEDDED_RUNTIME_ID: &str = "embedded-worker-runtime"; +pub(crate) const EMBEDDED_RUNTIME_ID: &str = "embedded-worker-runtime"; const EMBEDDED_HOST_KIND: &str = "embedded-worker-runtime-host"; const REMOTE_HOST_KIND: &str = "remote-worker-runtime-host"; const MAX_DIAGNOSTICS: usize = 16; @@ -4039,7 +4039,6 @@ mod tests { WorkspaceApiRef { workspace_id: "workspace-test".to_string(), base_url: "http://127.0.0.1:8787".to_string(), - runtime_id: Some("runtime-test".to_string()), } } diff --git a/crates/workspace-server/src/lib.rs b/crates/workspace-server/src/lib.rs index 544f38fb..daa9f66e 100644 --- a/crates/workspace-server/src/lib.rs +++ b/crates/workspace-server/src/lib.rs @@ -24,6 +24,7 @@ pub mod runtime_subscription; pub mod server; pub mod skills; pub mod store; +pub mod worker_source; mod workspace_subscription; pub use authority::{ diff --git a/crates/workspace-server/src/retention.rs b/crates/workspace-server/src/retention.rs index 561d6e48..f6e41e01 100644 --- a/crates/workspace-server/src/retention.rs +++ b/crates/workspace-server/src/retention.rs @@ -1354,7 +1354,7 @@ mod tests { let path = temp.path().join("server.db"); { let s = SqliteWorkspaceStore::open(&path).unwrap(); - s.with_conn(|c|{c.execute("INSERT INTO workspaces(workspace_id,display_name,state,created_at,updated_at)VALUES('legacy','Legacy','active','old','old')",[])?;c.execute_batch("DROP TRIGGER seed_worker_retention_policy_after_workspace_insert;DROP TABLE worker_retention_audit_events;DROP TABLE worker_tombstones;DROP TABLE worker_session_archives;DROP TABLE worker_diagnostics_archives;DROP TABLE worker_orphan_diagnostics;DROP TABLE worker_removal_operations;DROP TABLE workspace_worker_retention_policies;DROP TABLE workspace_worker_retention_policy_revisions;DELETE FROM __yoi_schema_migrations WHERE version=28;")?;Ok(())}).unwrap(); + s.with_conn(|c|{c.execute("INSERT INTO workspaces(workspace_id,display_name,state,created_at,updated_at)VALUES('legacy','Legacy','active','old','old')",[])?;c.execute_batch("DROP TRIGGER seed_worker_retention_policy_after_workspace_insert;DROP TABLE worker_retention_audit_events;DROP TABLE worker_tombstones;DROP TABLE worker_session_archives;DROP TABLE worker_diagnostics_archives;DROP TABLE worker_orphan_diagnostics;DROP TABLE worker_removal_operations;DROP TABLE workspace_worker_retention_policies;DROP TABLE workspace_worker_retention_policy_revisions;DELETE FROM __yoi_schema_migrations WHERE version>=28;")?;Ok(())}).unwrap(); } let reopened = SqliteWorkspaceStore::open(&path).unwrap(); let p = reopened.worker_retention_policy("legacy").unwrap().unwrap(); diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 95ce0bf4..5f2f4d55 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -62,15 +62,15 @@ use crate::companion::{ }; use crate::config::{BackendRuntimesConfigFile, RemoteRuntimeConfigFile, resolve_remote_runtime}; use crate::hosts::{ - ConfigBundleCheckResult, ConfigBundleSyncResult, DiagnosticSeverity, EmbeddedWorkerRuntime, - HostSummary, RemoteRuntimeConfig, RemoteWorkerRuntime, RuntimeDiagnostic, RuntimeRegistry, - RuntimeRegistryError, RuntimeRegistryUnregisterResult, RuntimeSummary, TicketWorkerRole, - WorkerCapabilitySummary, WorkerCompletionsRequest, WorkerCompletionsResult, - WorkerImplementationSummary, WorkerInputKind, WorkerInputRequest, WorkerInputResult, - WorkerLifecycleRequest, WorkerLifecycleResult, WorkerOperationState, WorkerRestoreResult, - WorkerSpawnAcceptanceRequirement, WorkerSpawnIntent, WorkerSpawnRequest, WorkerSpawnResult, - WorkerSpawnWorkingDirectoryRequest, WorkerSummary, WorkerTicketAssignmentRequest, - WorkerWorkspaceSummary, + ConfigBundleCheckResult, ConfigBundleSyncResult, DiagnosticSeverity, EMBEDDED_RUNTIME_ID, + EmbeddedWorkerRuntime, HostSummary, RemoteRuntimeConfig, RemoteWorkerRuntime, + RuntimeDiagnostic, RuntimeRegistry, RuntimeRegistryError, RuntimeRegistryUnregisterResult, + RuntimeSummary, TicketWorkerRole, WorkerCapabilitySummary, WorkerCompletionsRequest, + WorkerCompletionsResult, WorkerImplementationSummary, WorkerInputKind, WorkerInputRequest, + WorkerInputResult, WorkerLifecycleRequest, WorkerLifecycleResult, WorkerOperationState, + WorkerRestoreResult, WorkerSpawnAcceptanceRequirement, WorkerSpawnIntent, WorkerSpawnRequest, + WorkerSpawnResult, WorkerSpawnWorkingDirectoryRequest, WorkerSummary, + WorkerTicketAssignmentRequest, WorkerWorkspaceSummary, }; use crate::identity::WorkspaceIdentity; use crate::memory_backend::execute_memory_backend_operation_with_authority; @@ -250,8 +250,8 @@ const ORCHESTRATOR_ATTENTION_PROMPT: &str = include_str!(concat!( #[derive(Clone)] pub struct WorkspaceApi { - config: ServerConfig, - store: Arc, + pub(crate) config: ServerConfig, + pub(crate) store: Arc, authority: SqliteWorkspaceAuthority, runtime: Arc, companion: Arc, @@ -269,6 +269,15 @@ impl WorkspaceApi { let resource_broker = BackendResourceBroker::default(); let execution_backend = WorkerRuntimeExecutionBackend::new( ProfileRuntimeWorkerFactory::new(config.workspace_root.clone()) + .with_embedded_worker_mutation_dispatcher( + EMBEDDED_RUNTIME_ID, + Arc::new( + crate::worker_source::EmbeddedServerWorkerMutationDispatcher::new( + config.clone(), + store.clone(), + ), + ), + ) .with_runtime_store_dir(config.embedded_runtime_store_root.clone()) .with_resource_client(Arc::new(resource_broker.clone())), ) @@ -394,7 +403,7 @@ impl WorkspaceApi { &self.runtime_subscription_broker } - fn workspace_api_ref(&self, runtime_id: &str) -> WorkspaceApiRef { + fn workspace_api_ref(&self, _runtime_id: &str) -> WorkspaceApiRef { WorkspaceApiRef { workspace_id: self.config.workspace_id.clone(), base_url: self @@ -404,7 +413,6 @@ impl WorkspaceApi { .unwrap_or_else(|| "http://127.0.0.1:8787".to_string()) .trim_end_matches('/') .to_string(), - runtime_id: Some(runtime_id.to_string()), } } @@ -1118,6 +1126,10 @@ pub fn build_router(api: WorkspaceApi) -> Router { "/api/w/{workspace_id}/companion/cancel", post(scoped_post_companion_cancel), ) + .route( + "/api/w/{workspace_id}/workers/remove", + post(scoped_worker_remove_source_boundary), + ) .route( "/api/runtimes/{runtime_id}/workers", get(list_runtime_workers).post(create_runtime_worker), @@ -4811,6 +4823,71 @@ async fn scoped_workspace_protocol_ws( .into_response()) } +#[derive(Debug, Deserialize)] +struct WorkerRemoveBoundaryRequest { + target_runtime_id: String, + target_worker_id: String, +} + +async fn scoped_worker_remove_source_boundary( + State(api): State, + AxumPath(path): AxumPath, + headers: HeaderMap, + Json(request): Json, +) -> Response { + if let Err(error) = validate_workspace_scope(&api, &path.workspace_id) { + return error.into_response(); + } + let proof = match crate::worker_source::presented_worker_remove_source(&headers, None) { + Ok(proof) => proof, + Err(error) => { + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({ "error": error.to_string() })), + ) + .into_response(); + } + }; + match crate::worker_source::verify_worker_remove_source( + &api, + proof, + &request.target_runtime_id, + &request.target_worker_id, + ) + .await + { + Ok(source) => ( + StatusCode::NOT_IMPLEMENTED, + Json(serde_json::json!({ + "error": "WorkerRemove lifecycle is deferred to its consumer Ticket", + "source": { + "runtime_id": source.runtime_id, + "worker_id": source.worker_id, + "actor_kind": source.actor_kind, + "permission": source.permission, + } + })), + ) + .into_response(), + Err(error) => { + let status = match error { + crate::worker_source::WorkerMutationSourceProofError::Replay => { + StatusCode::CONFLICT + } + crate::worker_source::WorkerMutationSourceProofError::Authority(_) => { + StatusCode::INTERNAL_SERVER_ERROR + } + _ => StatusCode::FORBIDDEN, + }; + ( + status, + Json(serde_json::json!({ "error": error.to_string() })), + ) + .into_response() + } + } +} + async fn scoped_list_workers( State(api): State, AxumPath(path): AxumPath, @@ -10561,12 +10638,19 @@ mod tests { use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; use tower::ServiceExt; + use worker_runtime::auth::{ + RuntimeIdentityMaterial, RuntimeWorkerMutationSourceSigner, WORKER_REMOVE_PERMISSION, + decode_worker_mutation_source_claims, + }; use worker_runtime::resource::BackendResourceClient; + use worker_runtime::worker_source::{ + RuntimeOwnedWorkerMutationProof, RuntimeWorkerMutationSourceAuthority, + }; use worker_runtime::working_directory::WorkingDirectoryMaterializer; use crate::hosts::{ - TicketWorkerRole, WorkerInputKind, WorkerOperationState, WorkerSpawnAcceptanceRequirement, - WorkerSpawnIntent, + RemoteRuntimeAuthConfig, RuntimeCapabilitySummary, TicketWorkerRole, WorkerInputKind, + WorkerOperationState, WorkerSpawnAcceptanceRequirement, WorkerSpawnIntent, }; use crate::store::{ MemoryDocumentRecord, MemoryStagingRecord, ObjectiveRecord, ObjectiveResourceRecord, @@ -10643,11 +10727,10 @@ mod tests { const TEST_REPOSITORY_ID: &str = "main"; const TEST_CREATED_AT: &str = "2026-06-23T06:43:28Z"; - fn test_worker_workspace_api(runtime_id: &str) -> WorkspaceApiRef { + fn test_worker_workspace_api(_runtime_id: &str) -> WorkspaceApiRef { WorkspaceApiRef { workspace_id: TEST_WORKSPACE_ID.to_string(), base_url: "http://127.0.0.1:8787".to_string(), - runtime_id: Some(runtime_id.to_string()), } } @@ -13141,6 +13224,359 @@ mod tests { .unwrap() } + #[tokio::test] + async fn destructive_worker_remove_rejects_browser_and_legacy_source_headers() { + let headers = HeaderMap::new(); + assert!(matches!( + crate::worker_source::presented_worker_remove_source(&headers, None), + Err(crate::worker_source::WorkerMutationSourceProofError::Missing) + )); + + let mut spoofed = HeaderMap::new(); + spoofed.insert("x-yoi-runtime-id", "runtime-spoofed".parse().unwrap()); + spoofed.insert("x-yoi-worker-id", "worker-spoofed".parse().unwrap()); + assert!(matches!( + crate::worker_source::presented_worker_remove_source(&spoofed, None), + Err(crate::worker_source::WorkerMutationSourceProofError::Missing) + )); + + let temp = tempfile::tempdir().unwrap(); + let app = build_router(test_api(temp.path()).await); + let body = r#"{"target_runtime_id":"runtime-target","target_worker_id":"target-worker"}"#; + let browser = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri(format!("/api/w/{TEST_WORKSPACE_ID}/workers/remove")) + .header(CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(browser.status(), StatusCode::UNAUTHORIZED); + let legacy = app + .oneshot( + Request::builder() + .method("POST") + .uri(format!("/api/w/{TEST_WORKSPACE_ID}/workers/remove")) + .header(CONTENT_TYPE, "application/json") + .header("x-yoi-runtime-id", "runtime-spoofed") + .header("x-yoi-worker-id", "worker-spoofed") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(legacy.status(), StatusCode::UNAUTHORIZED); + } + + #[tokio::test] + async fn embedded_worker_remove_proof_derives_source_and_rejects_replay_and_wrong_target() { + let temp = tempfile::tempdir().unwrap(); + let api = test_api(temp.path()).await; + seed_worker_source_member(&api, EMBEDDED_RUNTIME_ID, "7"); + let scope = worker_runtime::RuntimeWorkspaceScope::new( + api.config.workspace_id.clone(), + "server-unused-for-embedded", + ); + let authority = RuntimeWorkerMutationSourceAuthority::embedded( + EMBEDDED_RUNTIME_ID, + &api.config.workspace_id, + ); + let RuntimeOwnedWorkerMutationProof::InProcess(proof) = authority + .issue_worker_remove(&scope, "7", "runtime-target", "target-worker") + .unwrap() + else { + panic!("embedded Runtime must produce in-process claims"); + }; + + let wrong_target = crate::worker_source::verify_worker_remove_source( + &api, + crate::worker_source::PresentedWorkerMutationSourceProof::InProcess(proof.clone()), + "runtime-target", + "different-worker", + ) + .await; + assert!(matches!( + wrong_target, + Err(crate::worker_source::WorkerMutationSourceProofError::Invalid) + )); + + let verified = crate::worker_source::verify_worker_remove_source( + &api, + crate::worker_source::PresentedWorkerMutationSourceProof::InProcess(proof.clone()), + "runtime-target", + "target-worker", + ) + .await + .unwrap(); + assert_eq!(verified.runtime_id, EMBEDDED_RUNTIME_ID); + assert_eq!(verified.worker_id, "7"); + assert_eq!(verified.permission, WORKER_REMOVE_PERMISSION); + + let replay = crate::worker_source::verify_worker_remove_source( + &api, + crate::worker_source::PresentedWorkerMutationSourceProof::InProcess(proof), + "runtime-target", + "target-worker", + ) + .await; + assert!(matches!( + replay, + Err(crate::worker_source::WorkerMutationSourceProofError::Replay) + )); + + let RuntimeOwnedWorkerMutationProof::InProcess(fresh_proof) = authority + .issue_worker_remove(&scope, "7", "runtime-target", "target-worker") + .unwrap() + else { + panic!("embedded Runtime must produce in-process claims"); + }; + let dispatcher = crate::worker_source::EmbeddedServerWorkerMutationDispatcher::new( + api.config.clone(), + api.store.clone(), + ); + let response = + worker_runtime::worker_source::EmbeddedWorkerMutationDispatcher::execute_worker_remove( + &dispatcher, + fresh_proof, + "runtime-target", + "target-worker", + ) + .unwrap(); + assert_eq!(response.status, 501); + } + + #[tokio::test] + async fn remote_worker_remove_proof_requires_current_runtime_trust_scope_and_catalog_member() { + let temp = tempfile::tempdir().unwrap(); + let identity = RuntimeIdentityMaterial::generate("runtime-remote").unwrap(); + let mut config = test_server_config(temp.path()); + config.remote_runtime_sources.push(RemoteRuntimeConfig { + runtime_id: "runtime-remote".to_string(), + display_name: "Remote Runtime".to_string(), + base_url: "https://runtime.invalid".to_string(), + bearer_token: None, + auth: Some(RemoteRuntimeAuthConfig { + server_id: "server-main".to_string(), + server_private_key: identity.private_key.clone(), + }), + cached_capabilities: RuntimeCapabilitySummary { + can_list_hosts: true, + can_list_workers: true, + can_get_worker: true, + can_spawn_worker: true, + can_stop_worker: true, + has_workspace_fs: false, + has_shell: false, + has_git: false, + supports_worktrees: false, + supports_backend_internal_tools: false, + workspace_scope: TEST_WORKSPACE_ID.to_string(), + max_workers: 1, + os: "test".to_string(), + arch: "test".to_string(), + }, + cached_status: "connected".to_string(), + timeout: std::time::Duration::from_secs(1), + }); + let store = SqliteWorkspaceStore::open(config.database_path.clone()).unwrap(); + let trust = crate::store::TrustedRuntimeRecord { + runtime_id: "runtime-remote".to_string(), + display_name: "Remote Runtime".to_string(), + base_url: "https://runtime.invalid".to_string(), + public_key: identity.public_key.clone(), + created_at: "2026-08-11T00:00:00Z".to_string(), + updated_at: "2026-08-11T00:00:00Z".to_string(), + revoked_at: None, + }; + store.upsert_trusted_runtime(&trust).unwrap(); + let api = WorkspaceApi::new_with_execution_backend( + config, + Arc::new(store), + Arc::new(DeterministicExecutionBackend::default()), + ) + .await + .unwrap(); + seed_worker_source_member(&api, "runtime-remote", "7"); + + let signer = RuntimeWorkerMutationSourceSigner::from_identity(&identity); + let token = signer + .issue_worker_remove( + "server-main", + &api.config.workspace_id, + "7", + "runtime-target", + "target-worker", + 60, + ) + .unwrap(); + let verified = crate::worker_source::verify_worker_remove_source( + &api, + crate::worker_source::PresentedWorkerMutationSourceProof::Remote(&token), + "runtime-target", + "target-worker", + ) + .await + .unwrap(); + assert_eq!(verified.runtime_id, "runtime-remote"); + assert_eq!(verified.worker_id, "7"); + + let wrong_scope = signer + .issue_worker_remove( + "server-wrong", + &api.config.workspace_id, + "7", + "runtime-target", + "target-worker", + 60, + ) + .unwrap(); + assert!(matches!( + crate::worker_source::verify_worker_remove_source( + &api, + crate::worker_source::PresentedWorkerMutationSourceProof::Remote(&wrong_scope), + "runtime-target", + "target-worker", + ) + .await, + Err(crate::worker_source::WorkerMutationSourceProofError::WrongAudience) + )); + + let wrong_workspace = signer + .issue_worker_remove( + "server-main", + "workspace-other", + "7", + "runtime-target", + "target-worker", + 60, + ) + .unwrap(); + assert!(matches!( + crate::worker_source::verify_worker_remove_source( + &api, + crate::worker_source::PresentedWorkerMutationSourceProof::Remote(&wrong_workspace), + "runtime-target", + "target-worker", + ) + .await, + Err(crate::worker_source::WorkerMutationSourceProofError::WrongWorkspace) + )); + + let missing_worker = signer + .issue_worker_remove( + "server-main", + &api.config.workspace_id, + "999", + "runtime-target", + "target-worker", + 60, + ) + .unwrap(); + assert!(matches!( + crate::worker_source::verify_worker_remove_source( + &api, + crate::worker_source::PresentedWorkerMutationSourceProof::Remote(&missing_worker), + "runtime-target", + "target-worker", + ) + .await, + Err(crate::worker_source::WorkerMutationSourceProofError::WorkerCatalogMembership) + )); + + let mut expired_claims = decode_worker_mutation_source_claims(&token).unwrap(); + expired_claims.iat = 1; + expired_claims.exp = 2; + expired_claims.jti = "expired-proof".to_string(); + let expired = signer.sign(&expired_claims).unwrap(); + assert!(matches!( + crate::worker_source::verify_worker_remove_source( + &api, + crate::worker_source::PresentedWorkerMutationSourceProof::Remote(&expired), + "runtime-target", + "target-worker", + ) + .await, + Err(crate::worker_source::WorkerMutationSourceProofError::Expired) + )); + + let route_token = signer + .issue_worker_remove( + "server-main", + &api.config.workspace_id, + "7", + "runtime-target", + "target-worker", + 60, + ) + .unwrap(); + let route_response = build_router(api.clone()) + .oneshot( + Request::builder() + .method("POST") + .uri(format!("/api/w/{TEST_WORKSPACE_ID}/workers/remove")) + .header(CONTENT_TYPE, "application/json") + .header( + worker_runtime::auth::WORKER_MUTATION_SOURCE_PROOF_HEADER, + route_token, + ) + .body(Body::from( + r#"{"target_runtime_id":"runtime-target","target_worker_id":"target-worker"}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(route_response.status(), StatusCode::NOT_IMPLEMENTED); + + let mut revoked = trust; + revoked.revoked_at = Some("2026-08-11T00:01:00Z".to_string()); + let authority = SqliteWorkspaceStore::open(api.config.database_path.clone()).unwrap(); + authority.upsert_trusted_runtime(&revoked).unwrap(); + let revoked_token = signer + .issue_worker_remove( + "server-main", + &api.config.workspace_id, + "7", + "runtime-target", + "target-worker", + 60, + ) + .unwrap(); + assert!(matches!( + crate::worker_source::verify_worker_remove_source( + &api, + crate::worker_source::PresentedWorkerMutationSourceProof::Remote(&revoked_token), + "runtime-target", + "target-worker", + ) + .await, + Err(crate::worker_source::WorkerMutationSourceProofError::RevokedRuntimeTrust) + )); + } + + fn seed_worker_source_member(api: &WorkspaceApi, runtime_id: &str, worker_id: &str) { + let now = now_registry_timestamp(); + api.store + .upsert_worker_registry(&WorkerRegistryRecord { + workspace_id: api.config.workspace_id.clone(), + worker: RuntimeWorkerRef::new(runtime_id, worker_id), + display_name: worker_id.to_string(), + profile: None, + retention_state: "normal".to_string(), + transcript_ref: None, + session_ref: None, + summary_ref: None, + diagnostics_ref: None, + created_at: now.clone(), + updated_at: now, + }) + .unwrap(); + } + fn seed_cleanup_worker( api: &WorkspaceApi, runtime_worker_id: u64, diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 8b3786e8..f5fa9b83 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -161,6 +161,11 @@ const MIGRATIONS: &[Migration] = &[ name: "create Worker retention authority", apply: crate::retention::create_worker_retention_tables, }, + Migration { + version: 29, + name: "create Worker mutation source proof replay guard", + apply: create_worker_mutation_source_proof_replay_guard, + }, ]; struct Migration { @@ -460,6 +465,15 @@ pub trait ControlPlaneStore: Send + Sync { async fn schema_version(&self) -> Result; async fn upsert_workspace(&self, record: &WorkspaceRecord) -> Result<()>; async fn get_workspace(&self, workspace_id: &str) -> Result>; + async fn get_trusted_runtime(&self, runtime_id: &str) -> Result>; + async fn consume_worker_mutation_source_jti( + &self, + runtime_id: &str, + jti: &str, + expires_at: u64, + now_seconds: u64, + consumed_at: &str, + ) -> Result; fn list_workspaces(&self) -> Result>; fn upsert_repository(&self, record: &RepositoryRecord) -> Result<()>; fn get_repository( @@ -896,6 +910,44 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } + async fn get_trusted_runtime(&self, runtime_id: &str) -> Result> { + self.with_conn(|conn| { + conn.query_row( + r#"SELECT runtime_id, display_name, base_url, public_key, created_at, updated_at, revoked_at + FROM trusted_runtime_records WHERE runtime_id = ?1"#, + params![runtime_id], + read_trusted_runtime_record, + ) + .optional() + .map_err(Error::from) + }) + } + + async fn consume_worker_mutation_source_jti( + &self, + runtime_id: &str, + jti: &str, + expires_at: u64, + now_seconds: u64, + consumed_at: &str, + ) -> Result { + self.with_conn(|conn| { + let transaction = conn.unchecked_transaction()?; + transaction.execute( + "DELETE FROM worker_mutation_source_proof_jtis WHERE expires_at < ?1", + params![now_seconds], + )?; + let inserted = transaction.execute( + r#"INSERT OR IGNORE INTO worker_mutation_source_proof_jtis ( + runtime_id, jti, expires_at, consumed_at + ) VALUES (?1, ?2, ?3, ?4)"#, + params![runtime_id, jti, expires_at, consumed_at], + )?; + transaction.commit()?; + Ok(inserted == 1) + }) + } + fn list_workspaces(&self) -> Result> { self.with_conn(|conn| { let mut stmt = conn.prepare( @@ -4326,6 +4378,23 @@ fn current_schema_version(conn: &Connection) -> Result { .map_err(Error::from) } +fn create_worker_mutation_source_proof_replay_guard(conn: &Connection) -> Result<()> { + conn.execute_batch( + r#" + CREATE TABLE IF NOT EXISTS worker_mutation_source_proof_jtis ( + runtime_id TEXT NOT NULL, + jti TEXT NOT NULL, + expires_at INTEGER NOT NULL, + consumed_at TEXT NOT NULL, + PRIMARY KEY (runtime_id, jti) + ); + CREATE INDEX IF NOT EXISTS idx_worker_mutation_source_proof_jtis_expiry + ON worker_mutation_source_proof_jtis(expires_at); + "#, + )?; + Ok(()) +} + fn apply_migrations(conn: &Connection) -> Result<()> { let current = current_schema_version(conn)?; for migration in MIGRATIONS @@ -4915,7 +4984,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 28); + assert_eq!(current_schema_version(&conn).unwrap(), 29); assert!(table_exists(&conn, "worker_workdir_attachment_reservations").unwrap()); } @@ -4948,7 +5017,7 @@ CREATE TABLE flow_events (event_id TEXT PRIMARY KEY); apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 28); + assert_eq!(current_schema_version(&conn).unwrap(), 29); assert!(table_exists(&conn, "flow_sources").unwrap()); assert!(table_exists(&conn, "flow_source_revisions").unwrap()); assert!(!table_exists(&conn, "flow_instances").unwrap()); @@ -5015,7 +5084,7 @@ INSERT INTO worker_workdir_attachment_reservations ( apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 28); + assert_eq!(current_schema_version(&conn).unwrap(), 29); let repositories_sql: String = conn .query_row( "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'repositories'", @@ -5195,7 +5264,7 @@ INSERT INTO workdir_registry ( let db = dir.path().join("control-plane.sqlite"); let store = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 28); + assert_eq!(store.schema_version().await.unwrap(), 29); assert!( !store .with_conn(|conn| table_exists(conn, "worker_workspace_credentials")) @@ -5212,7 +5281,7 @@ INSERT INTO workdir_registry ( store.upsert_workspace(&record).await.unwrap(); let reopened = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(reopened.schema_version().await.unwrap(), 28); + assert_eq!(reopened.schema_version().await.unwrap(), 29); assert_eq!( reopened.get_workspace("local-dev").await.unwrap(), Some(record) @@ -5759,7 +5828,7 @@ INSERT INTO workdir_registry ( .unwrap(); let store = SqliteWorkspaceStore::from_connection(conn).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 28); + assert_eq!(store.schema_version().await.unwrap(), 29); store .with_conn(|conn| { @@ -5948,7 +6017,7 @@ CREATE TABLE ticket_assignment_operations ( #[tokio::test] async fn repository_records_round_trip() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 28); + assert_eq!(store.schema_version().await.unwrap(), 29); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -6014,7 +6083,7 @@ CREATE TABLE ticket_assignment_operations ( #[tokio::test] async fn memory_authority_records_round_trip_and_close_staging() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 28); + assert_eq!(store.schema_version().await.unwrap(), 29); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -6277,7 +6346,7 @@ CREATE TABLE ticket_assignment_operations ( #[tokio::test] async fn account_and_login_records_round_trip() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 28); + assert_eq!(store.schema_version().await.unwrap(), 29); let now = "2026-07-22T00:00:00Z".to_string(); let account = AccountRecord { account_id: "acct-user-alice".to_string(), @@ -6455,6 +6524,51 @@ CREATE TABLE ticket_assignment_operations ( ); } + #[tokio::test] + async fn worker_mutation_source_jti_replay_guard_survives_reopen() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("server.db"); + let store = SqliteWorkspaceStore::open(&path).unwrap(); + store + .upsert_trusted_runtime(&TrustedRuntimeRecord { + runtime_id: "runtime-a".to_string(), + display_name: "Runtime A".to_string(), + base_url: "https://runtime.invalid".to_string(), + public_key: "public-key".to_string(), + created_at: "2026-08-11T00:00:00Z".to_string(), + updated_at: "2026-08-11T00:00:00Z".to_string(), + revoked_at: None, + }) + .unwrap(); + assert!( + store + .consume_worker_mutation_source_jti( + "runtime-a", + "proof-1", + 2_000, + 1_000, + "2026-08-11T00:00:00Z", + ) + .await + .unwrap() + ); + drop(store); + + let reopened = SqliteWorkspaceStore::open(&path).unwrap(); + assert!( + !reopened + .consume_worker_mutation_source_jti( + "runtime-a", + "proof-1", + 2_000, + 1_001, + "2026-08-11T00:00:01Z", + ) + .await + .unwrap() + ); + } + fn table_names(conn: &Connection) -> BTreeSet { let mut stmt = conn .prepare( diff --git a/crates/workspace-server/src/worker_source.rs b/crates/workspace-server/src/worker_source.rs new file mode 100644 index 00000000..932250ad --- /dev/null +++ b/crates/workspace-server/src/worker_source.rs @@ -0,0 +1,288 @@ +use std::time::{SystemTime, UNIX_EPOCH}; + +use axum::http::HeaderMap; +use worker_runtime::auth::{ + WorkerMutationActorKind, WorkerMutationOperation, WorkerMutationSourceClaims, + WorkerMutationSourceExpectation, decode_worker_mutation_source_claims, + verify_worker_mutation_source_proof, +}; +use worker_runtime::worker_source::InProcessWorkerMutationProof; + +use crate::hosts::RemoteRuntimeConfig; +use crate::server::WorkspaceApi; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum PresentedWorkerMutationSourceProof<'a> { + Remote(&'a str), + InProcess(InProcessWorkerMutationProof), +} + +pub fn presented_worker_remove_source<'a>( + headers: &'a HeaderMap, + in_process: Option, +) -> Result, WorkerMutationSourceProofError> { + if let Some(claims) = in_process { + return Ok(PresentedWorkerMutationSourceProof::InProcess(claims)); + } + headers + .get(worker_runtime::auth::WORKER_MUTATION_SOURCE_PROOF_HEADER) + .and_then(|value| value.to_str().ok()) + .filter(|value| !value.trim().is_empty()) + .map(PresentedWorkerMutationSourceProof::Remote) + .ok_or(WorkerMutationSourceProofError::Missing) +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct VerifiedWorkerMutationSource { + pub runtime_id: String, + pub worker_id: String, + pub actor_kind: WorkerMutationActorKind, + pub permission: String, + pub jti: String, +} + +#[derive(Debug, thiserror::Error)] +pub enum WorkerMutationSourceProofError { + #[error("Worker mutation source proof is required")] + Missing, + #[error("Worker mutation source proof is invalid")] + Invalid, + #[error("Worker mutation source proof is not authorized for this Server")] + WrongAudience, + #[error("Worker mutation source proof is not authorized for this Workspace")] + WrongWorkspace, + #[error("Worker mutation source proof actor is not allowed")] + WrongActor, + #[error("Worker mutation source proof lacks `{0}` permission")] + MissingPermission(String), + #[error("Worker mutation source proof is expired")] + Expired, + #[error("Runtime trust is missing or revoked")] + RevokedRuntimeTrust, + #[error("Worker mutation source proof was already consumed")] + Replay, + #[error("source Worker is not a current member of this Workspace Runtime catalog")] + WorkerCatalogMembership, + #[error("source proof authority failed: {0}")] + Authority(String), +} + +pub async fn verify_worker_remove_source( + api: &WorkspaceApi, + proof: PresentedWorkerMutationSourceProof<'_>, + target_runtime_id: &str, + target_worker_id: &str, +) -> Result { + verify_worker_remove_source_with( + &api.config, + &api.store, + proof, + target_runtime_id, + target_worker_id, + ) + .await +} + +async fn verify_worker_remove_source_with( + config: &crate::server::ServerConfig, + store: &std::sync::Arc, + proof: PresentedWorkerMutationSourceProof<'_>, + target_runtime_id: &str, + target_worker_id: &str, +) -> Result { + let required_permission = worker_runtime::auth::WORKER_REMOVE_PERMISSION; + let now = unix_now_seconds(); + let claims = match proof { + PresentedWorkerMutationSourceProof::Remote(token) => { + let unverified = decode_worker_mutation_source_claims(token) + .map_err(|_| WorkerMutationSourceProofError::Invalid)?; + let audience = remote_audience(config, &unverified.iss)?; + let trusted = store + .get_trusted_runtime(&unverified.iss) + .await + .map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))? + .filter(|record| record.revoked_at.is_none()) + .ok_or(WorkerMutationSourceProofError::RevokedRuntimeTrust)?; + let expected = WorkerMutationSourceExpectation { + runtime_id: &unverified.iss, + audience, + workspace_id: &config.workspace_id, + worker_id: None, + actor_kind: WorkerMutationActorKind::Worker, + operation: WorkerMutationOperation::WorkerRemove, + target_runtime_id, + target_worker_id, + permission: required_permission, + }; + verify_worker_mutation_source_proof(&trusted.public_key, token, &expected, now) + .map_err(map_auth_error)? + } + PresentedWorkerMutationSourceProof::InProcess(proof) => { + let claims = proof.into_claims(); + if config + .remote_runtime_sources + .iter() + .any(|runtime| runtime.runtime_id == claims.iss) + { + return Err(WorkerMutationSourceProofError::Invalid); + } + validate_in_process_claims( + &claims, + &format!("embedded:{}", config.workspace_id), + &config.workspace_id, + target_runtime_id, + target_worker_id, + required_permission, + now, + )?; + claims + } + }; + + let worker = worker_runtime::identity::RuntimeWorkerRef { + runtime_id: claims.iss.clone(), + worker_id: claims.worker_id.clone(), + }; + let member = store + .get_worker_registry(&config.workspace_id, &worker) + .map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?; + if member.is_none() { + return Err(WorkerMutationSourceProofError::WorkerCatalogMembership); + } + + let consumed_at = chrono::Utc::now().to_rfc3339(); + let consumed = store + .consume_worker_mutation_source_jti(&claims.iss, &claims.jti, claims.exp, now, &consumed_at) + .await + .map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?; + if !consumed { + return Err(WorkerMutationSourceProofError::Replay); + } + + Ok(VerifiedWorkerMutationSource { + runtime_id: claims.iss, + worker_id: claims.worker_id, + actor_kind: claims.actor_kind, + permission: claims.permission, + jti: claims.jti, + }) +} + +#[derive(Clone)] +pub(crate) struct EmbeddedServerWorkerMutationDispatcher { + config: crate::server::ServerConfig, + store: std::sync::Arc, +} + +impl EmbeddedServerWorkerMutationDispatcher { + pub(crate) fn new( + config: crate::server::ServerConfig, + store: std::sync::Arc, + ) -> Self { + Self { config, store } + } +} + +impl worker_runtime::worker_source::EmbeddedWorkerMutationDispatcher + for EmbeddedServerWorkerMutationDispatcher +{ + fn execute_worker_remove( + &self, + proof: InProcessWorkerMutationProof, + target_runtime_id: &str, + target_worker_id: &str, + ) -> Result< + worker::WorkspaceResponse, + worker_runtime::worker_source::RuntimeWorkerMutationForwardError, + > { + futures::executor::block_on(verify_worker_remove_source_with( + &self.config, + &self.store, + PresentedWorkerMutationSourceProof::InProcess(proof), + target_runtime_id, + target_worker_id, + )) + .map_err(|error| { + worker_runtime::worker_source::RuntimeWorkerMutationForwardError::Embedded( + error.to_string(), + ) + })?; + Ok(worker::WorkspaceResponse { + status: 501, + body: "WorkerRemove lifecycle is not implemented by this operation boundary" + .to_string(), + }) + } +} + +fn remote_audience<'a>( + config: &'a crate::server::ServerConfig, + runtime_id: &str, +) -> Result<&'a str, WorkerMutationSourceProofError> { + config + .remote_runtime_sources + .iter() + .find(|runtime| runtime.runtime_id == runtime_id) + .and_then(|runtime: &RemoteRuntimeConfig| runtime.auth.as_ref()) + .map(|auth| auth.server_id.as_str()) + .ok_or(WorkerMutationSourceProofError::RevokedRuntimeTrust) +} + +fn validate_in_process_claims( + claims: &WorkerMutationSourceClaims, + audience: &str, + workspace_id: &str, + target_runtime_id: &str, + target_worker_id: &str, + permission: &str, + now: u64, +) -> Result<(), WorkerMutationSourceProofError> { + if claims.aud != audience { + return Err(WorkerMutationSourceProofError::WrongAudience); + } + if claims.workspace_id != workspace_id { + return Err(WorkerMutationSourceProofError::WrongWorkspace); + } + if claims.actor_kind != WorkerMutationActorKind::Worker { + return Err(WorkerMutationSourceProofError::WrongActor); + } + if claims.operation != WorkerMutationOperation::WorkerRemove + || claims.target_runtime_id != target_runtime_id + || claims.target_worker_id != target_worker_id + { + return Err(WorkerMutationSourceProofError::Invalid); + } + if claims.permission != permission { + return Err(WorkerMutationSourceProofError::MissingPermission( + permission.to_string(), + )); + } + if claims.exp <= now || claims.iat > now.saturating_add(60) || claims.jti.trim().is_empty() { + return Err(WorkerMutationSourceProofError::Expired); + } + Ok(()) +} + +fn map_auth_error(error: worker_runtime::auth::RuntimeAuthError) -> WorkerMutationSourceProofError { + use worker_runtime::auth::RuntimeAuthError; + match error { + RuntimeAuthError::WrongAudience { .. } => WorkerMutationSourceProofError::WrongAudience, + RuntimeAuthError::WrongWorkspace { .. } => WorkerMutationSourceProofError::WrongWorkspace, + RuntimeAuthError::WrongActorKind => WorkerMutationSourceProofError::WrongActor, + RuntimeAuthError::WrongOperation | RuntimeAuthError::WrongMutationTarget => { + WorkerMutationSourceProofError::Invalid + } + RuntimeAuthError::MissingPermission(permission) => { + WorkerMutationSourceProofError::MissingPermission(permission) + } + RuntimeAuthError::Expired => WorkerMutationSourceProofError::Expired, + _ => WorkerMutationSourceProofError::Invalid, + } +} + +fn unix_now_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +}