runtime: prove Worker mutation source authority

This commit is contained in:
2026-08-12 04:03:09 +09:00
parent 86be3a6865
commit 8cc0aaf8d2
21 changed files with 2037 additions and 230 deletions
+315
View File
@@ -10,6 +10,10 @@ const PUBLIC_KEY_PREFIX: &str = "yoi-ed25519-pub:v1:";
const PRIVATE_KEY_PREFIX: &str = "yoi-ed25519-pkcs8:v1:"; const PRIVATE_KEY_PREFIX: &str = "yoi-ed25519-pkcs8:v1:";
const TOKEN_PREFIX: &str = "yoi-cap-v1"; const TOKEN_PREFIX: &str = "yoi-cap-v1";
const SIGNING_INPUT_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)] #[derive(Debug, thiserror::Error)]
pub enum RuntimeAuthError { pub enum RuntimeAuthError {
@@ -41,6 +45,16 @@ pub enum RuntimeAuthError {
MissingWorkspaceScope, MissingWorkspaceScope,
#[error("capability token is missing required permission `{0}`")] #[error("capability token is missing required permission `{0}`")]
MissingPermission(String), 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)] #[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<String>,
workspace_id: impl Into<String>,
source_worker_id: impl Into<String>,
target_runtime_id: impl Into<String>,
target_worker_id: impl Into<String>,
ttl_seconds: u64,
) -> Result<String, RuntimeAuthError> {
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<String, RuntimeAuthError> {
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<WorkerMutationSourceClaims, RuntimeAuthError> {
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<WorkerMutationSourceClaims, RuntimeAuthError> {
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<u8>), 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<u8>), RuntimeAuthError> { fn split_token(token: &str) -> Result<(&str, Vec<u8>), RuntimeAuthError> {
let mut parts = token.split('.'); let mut parts = token.split('.');
match (parts.next(), parts.next(), parts.next(), parts.next()) { match (parts.next(), parts.next(), parts.next(), parts.next()) {
@@ -277,6 +486,112 @@ impl fmt::Display for RuntimeAuthContext {
mod tests { mod tests {
use super::*; 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] #[test]
fn capability_token_verifies_signature_audience_expiry_and_permission() { fn capability_token_verifies_signature_audience_expiry_and_permission() {
let server = RuntimeIdentityMaterial::generate("server-main").unwrap(); let server = RuntimeIdentityMaterial::generate("server-main").unwrap();
+28 -3
View File
@@ -179,8 +179,6 @@ pub struct WorkingDirectoryStatus {
pub struct WorkspaceApiRef { pub struct WorkspaceApiRef {
pub workspace_id: String, pub workspace_id: String,
pub base_url: String, pub base_url: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub runtime_id: Option<String>,
} }
impl std::fmt::Debug for WorkspaceApiRef { impl std::fmt::Debug for WorkspaceApiRef {
@@ -189,7 +187,6 @@ impl std::fmt::Debug for WorkspaceApiRef {
.debug_struct("WorkspaceApiRef") .debug_struct("WorkspaceApiRef")
.field("workspace_id", &self.workspace_id) .field("workspace_id", &self.workspace_id)
.field("base_url", &self.base_url) .field("base_url", &self.base_url)
.field("runtime_id", &self.runtime_id)
.finish() .finish()
} }
} }
@@ -292,3 +289,31 @@ pub struct WorkerLifecycleAck {
pub worker_ref: WorkerRef, pub worker_ref: WorkerRef,
pub status: WorkerStatus, 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}");
}
}
}
+2
View File
@@ -248,6 +248,7 @@ pub struct WorkerExecutionSpawnRequest {
/// Monotonic execution generation reserved durably before launch. /// Monotonic execution generation reserved durably before launch.
pub run_generation: u64, pub run_generation: u64,
pub request: crate::catalog::CreateWorkerRequest, pub request: crate::catalog::CreateWorkerRequest,
pub workspace_scope: Option<crate::runtime::RuntimeWorkspaceScope>,
pub context: WorkerExecutionContext, pub context: WorkerExecutionContext,
pub working_directory: Option<WorkingDirectoryBinding>, pub working_directory: Option<WorkingDirectoryBinding>,
pub config_bundle: Option<ConfigBundle>, pub config_bundle: Option<ConfigBundle>,
@@ -260,6 +261,7 @@ pub struct WorkerExecutionRestoreRequest {
/// Monotonic execution generation reserved durably before restore. /// Monotonic execution generation reserved durably before restore.
pub run_generation: u64, pub run_generation: u64,
pub request: crate::catalog::CreateWorkerRequest, pub request: crate::catalog::CreateWorkerRequest,
pub workspace_scope: Option<crate::runtime::RuntimeWorkspaceScope>,
pub context: WorkerExecutionContext, pub context: WorkerExecutionContext,
pub previous_working_directory: Option<WorkingDirectoryStatus>, pub previous_working_directory: Option<WorkingDirectoryStatus>,
pub working_directory: Option<WorkingDirectoryBinding>, pub working_directory: Option<WorkingDirectoryBinding>,
-2
View File
@@ -1753,7 +1753,6 @@ mod tests {
request.workspace_api = Some(WorkspaceApiRef { request.workspace_api = Some(WorkspaceApiRef {
workspace_id: workspace_id.to_string(), workspace_id: workspace_id.to_string(),
base_url: format!("https://workspace.example/{workspace_id}"), base_url: format!("https://workspace.example/{workspace_id}"),
runtime_id: None,
}); });
request request
} }
@@ -2324,7 +2323,6 @@ mod tests {
workspace_api: WorkspaceApiRef { workspace_api: WorkspaceApiRef {
workspace_id: "local".to_string(), workspace_id: "local".to_string(),
base_url: "http://127.0.0.1:8787".to_string(), base_url: "http://127.0.0.1:8787".to_string(),
runtime_id: None,
}, },
}, },
) )
+1
View File
@@ -26,6 +26,7 @@ pub mod resource;
pub mod retention; pub mod retention;
mod runtime; mod runtime;
pub mod worker_backend; pub mod worker_backend;
pub mod worker_source;
pub mod working_directory; pub mod working_directory;
#[cfg(feature = "fs-store")] #[cfg(feature = "fs-store")]
+3
View File
@@ -87,6 +87,9 @@ fn build_runtime(config: &ProcessConfig) -> Result<Runtime, ProcessError> {
}; };
let mut factory = ProfileRuntimeWorkerFactory::new(fs_paths.worker_dir.join("worker-root")) let mut factory = ProfileRuntimeWorkerFactory::new(fs_paths.worker_dir.join("worker-root"))
.with_runtime_store_dir(runtime_store_dir); .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() { if let Some(endpoint) = config.backend_resource_endpoint.clone() {
factory = factory.with_resource_client(Arc::new( factory = factory.with_resource_client(Arc::new(
worker_runtime::resource::HttpBackendResourceClient::new( worker_runtime::resource::HttpBackendResourceClient::new(
+19 -9
View File
@@ -548,6 +548,7 @@ impl Runtime {
worker_ref: worker_ref.clone(), worker_ref: worker_ref.clone(),
run_generation: 1, run_generation: 1,
request, request,
workspace_scope: scope.cloned(),
context: self.execution_context(worker_ref.clone()), context: self.execution_context(worker_ref.clone()),
working_directory: None, working_directory: None,
config_bundle: None, config_bundle: None,
@@ -826,13 +827,10 @@ impl Runtime {
if let Some(existing) = worker.request.workspace_api.as_ref() if let Some(existing) = worker.request.workspace_api.as_ref()
&& (existing.workspace_id != workspace_api.workspace_id && (existing.workspace_id != workspace_api.workspace_id
|| existing.base_url.trim_end_matches('/') || existing.base_url.trim_end_matches('/')
!= workspace_api.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)
}))
{ {
return Err(RuntimeError::InvalidRequest( 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(), .to_string(),
)); ));
} }
@@ -902,10 +900,17 @@ impl Runtime {
})?; })?;
state.worker_mut(worker_ref)?.run_generation = run_generation; state.worker_mut(worker_ref)?.run_generation = run_generation;
state.persist_worker(&worker_ref.worker_id)?; 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 { let request = WorkerExecutionRestoreRequest {
worker_ref: worker_ref.clone(), worker_ref: worker_ref.clone(),
run_generation, run_generation,
request: worker_request, request: worker_request,
workspace_scope,
context: self.execution_context(worker_ref.clone()), context: self.execution_context(worker_ref.clone()),
previous_working_directory, previous_working_directory,
working_directory: None, working_directory: None,
@@ -1595,9 +1600,15 @@ impl Runtime {
}; };
for candidate in candidates { for candidate in candidates {
let backend = { let (backend, workspace_scope) = {
let state = self.lock()?; 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 { let Some(backend) = backend else {
return Ok(()); return Ok(());
@@ -1606,6 +1617,7 @@ impl Runtime {
worker_ref: candidate.worker_ref.clone(), worker_ref: candidate.worker_ref.clone(),
run_generation: candidate.run_generation, run_generation: candidate.run_generation,
request: candidate.request, request: candidate.request,
workspace_scope,
context: self.execution_context(candidate.worker_ref.clone()), context: self.execution_context(candidate.worker_ref.clone()),
previous_working_directory: candidate.previous_working_directory, previous_working_directory: candidate.previous_working_directory,
working_directory: None, working_directory: None,
@@ -2790,7 +2802,6 @@ mod tests {
request.workspace_api = Some(WorkspaceApiRef { request.workspace_api = Some(WorkspaceApiRef {
workspace_id: workspace_id.to_string(), workspace_id: workspace_id.to_string(),
base_url: format!("https://workspace.example/{workspace_id}"), base_url: format!("https://workspace.example/{workspace_id}"),
runtime_id: None,
}); });
request request
} }
@@ -3239,7 +3250,6 @@ mod tests {
let replacement = WorkspaceApiRef { let replacement = WorkspaceApiRef {
workspace_id: "workspace-a".to_string(), workspace_id: "workspace-a".to_string(),
base_url: "https://workspace.example/workspace-a/".to_string(), base_url: "https://workspace.example/workspace-a/".to_string(),
runtime_id: Some("runtime-a".to_string()),
}; };
runtime runtime
+141 -40
View File
@@ -14,6 +14,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, mpsc}; use std::sync::{Arc, Mutex, mpsc};
use std::time::Duration; use std::time::Duration;
use crate::auth::RuntimeIdentityMaterial;
use crate::catalog::{ use crate::catalog::{
CreateWorkerRequest, ProfileSourceArchiveHttpRef, ProfileSourceArchiveSource, CreateWorkerRequest, ProfileSourceArchiveHttpRef, ProfileSourceArchiveSource,
WorkingDirectoryRequest, WorkingDirectoryStatus, WorkingDirectoryRequest, WorkingDirectoryStatus,
@@ -26,6 +27,9 @@ use crate::execution::{
use crate::identity::WorkerRef; use crate::identity::WorkerRef;
use crate::interaction::{WorkerInput, WorkerInputKind}; use crate::interaction::{WorkerInput, WorkerInputKind};
use crate::resource::{BackendResourceClient, ProfileSourceArchiveCache}; use crate::resource::{BackendResourceClient, ProfileSourceArchiveCache};
use crate::worker_source::{
EmbeddedWorkerMutationDispatcher, RuntimeOwnedWorkspaceClient, RuntimeWorkerMutationForwarder,
};
use crate::working_directory::{ use crate::working_directory::{
WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer, WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer,
}; };
@@ -49,10 +53,9 @@ use worker::feature::builtin::{
#[cfg(feature = "ws-server")] #[cfg(feature = "ws-server")]
use worker::ipc::protocol_session::{live_log_entry_event, subscribe_worker_protocol_session}; use worker::ipc::protocol_session::{live_log_entry_event, subscribe_worker_protocol_session};
use worker::{ use worker::{
PromptLoader, RuntimeWorkspaceHttpClient, SegmentLogSink, PromptLoader, SegmentLogSink, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker,
WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerController, WorkerError, WorkerController, WorkerError, WorkerFilesystemAuthority, WorkerHandle, WorkerSharedState,
WorkerFilesystemAuthority, WorkerHandle, WorkerSharedState, WorkerWorkspaceContext, WorkerWorkspaceContext, WorkspaceClient, WorkspaceId,
WorkspaceClient, WorkspaceId,
}; };
const DEFAULT_BACKEND_ID: &str = "worker-crate"; const DEFAULT_BACKEND_ID: &str = "worker-crate";
@@ -215,6 +218,9 @@ pub struct ProfileRuntimeWorkerFactory {
worker_aggregate_root: Option<PathBuf>, worker_aggregate_root: Option<PathBuf>,
resource_client: Option<Arc<dyn BackendResourceClient>>, resource_client: Option<Arc<dyn BackendResourceClient>>,
profile_archive_cache: Arc<ProfileSourceArchiveCache>, profile_archive_cache: Arc<ProfileSourceArchiveCache>,
runtime_id: Option<String>,
worker_mutation_identity: Option<RuntimeIdentityMaterial>,
embedded_worker_mutation_dispatcher: Option<Arc<dyn EmbeddedWorkerMutationDispatcher>>,
} }
impl ProfileRuntimeWorkerFactory { impl ProfileRuntimeWorkerFactory {
@@ -226,9 +232,38 @@ impl ProfileRuntimeWorkerFactory {
worker_aggregate_root: None, worker_aggregate_root: None,
resource_client: None, resource_client: None,
profile_archive_cache: Arc::new(ProfileSourceArchiveCache::default()), 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<String>) -> 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<String>,
dispatcher: Arc<dyn EmbeddedWorkerMutationDispatcher>,
) -> 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<PathBuf>) -> Self { pub fn with_runtime_store_dir(mut self, runtime_store_dir: impl Into<PathBuf>) -> Self {
self.worker_aggregate_root = Some(runtime_store_dir.into().join("workers")); self.worker_aggregate_root = Some(runtime_store_dir.into().join("workers"));
self self
@@ -348,38 +383,59 @@ enum RuntimeWorkspaceBackendRef {
} }
impl RuntimeWorkspaceBackendRef { impl RuntimeWorkspaceBackendRef {
fn from_worker_request(request: &CreateWorkerRequest) -> Self { fn from_worker_request(request: &CreateWorkerRequest, runtime_id: Option<&str>) -> Self {
if let Some(api) = request.workspace_api.as_ref() if let (Some(api), Some(runtime_id)) = (request.workspace_api.as_ref(), runtime_id) {
&& let Some(runtime_id) = api
.runtime_id
.as_ref()
.filter(|runtime_id| !runtime_id.trim().is_empty())
{
return Self::Http { return Self::Http {
workspace_id: api.workspace_id.clone(), workspace_id: api.workspace_id.clone(),
base_url: api.base_url.clone(), base_url: api.base_url.clone(),
runtime_id: runtime_id.clone(), runtime_id: runtime_id.to_string(),
}; };
} }
Self::None 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<dyn EmbeddedWorkerMutationDispatcher>>,
) -> WorkerWorkspaceContext {
match self { match self {
Self::None => WorkerWorkspaceContext::no_workspace(), Self::None => WorkerWorkspaceContext::no_workspace(),
Self::Http { Self::Http {
workspace_id, workspace_id,
base_url, base_url,
runtime_id, runtime_id,
} => WorkerWorkspaceContext::with_client( } => {
WorkspaceId::new(workspace_id.clone()).ok(), let mut client = RuntimeOwnedWorkspaceClient::new(
Arc::new(RuntimeWorkspaceHttpClient::new(
workspace_id.clone(), workspace_id.clone(),
base_url.clone(), base_url.clone(),
runtime_id.clone(), runtime_id.clone(),
worker_ref.worker_id.to_string(), 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); .unwrap_or(WorkerFilesystemAuthority::None);
let workspace_backend_ref = let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request(
RuntimeWorkspaceBackendRef::from_worker_request(&request.request); &request.request,
let observation_runtime_id = request self.runtime_id.as_deref(),
.request );
.workspace_api let observation_runtime_id = self.runtime_id.clone();
.as_ref()
.and_then(|api| api.runtime_id.clone());
let observation_workspace_id = request let observation_workspace_id = request
.request .request
.workspace_api .workspace_api
@@ -492,7 +546,12 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
.map(|api| api.workspace_id.clone()); .map(|api| api.workspace_id.clone());
let observation_grants = request.request.worker_observation_grants.clone(); let observation_grants = request.request.worker_observation_grants.clone();
let observation_enabled = request.request.worker_observation_enabled; 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 selector = profile.as_ref();
let archive = self let archive = self
.resolve_profile_source_archive(&request.request.profile_source) .resolve_profile_source_archive(&request.request.profile_source)
@@ -628,13 +687,11 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
) )
}) })
.unwrap_or(WorkerFilesystemAuthority::None); .unwrap_or(WorkerFilesystemAuthority::None);
let workspace_backend_ref = let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request(
RuntimeWorkspaceBackendRef::from_worker_request(&request.request); &request.request,
let observation_runtime_id = request self.runtime_id.as_deref(),
.request );
.workspace_api let observation_runtime_id = self.runtime_id.clone();
.as_ref()
.and_then(|api| api.runtime_id.clone());
let observation_workspace_id = request let observation_workspace_id = request
.request .request
.workspace_api .workspace_api
@@ -642,7 +699,12 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
.map(|api| api.workspace_id.clone()); .map(|api| api.workspace_id.clone());
let observation_grants = request.request.worker_observation_grants.clone(); let observation_grants = request.request.worker_observation_grants.clone();
let observation_enabled = request.request.worker_observation_enabled; 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 (manifest, loader) = Self::restore_fallback_manifest(&worker_name)?;
let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?; let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?;
@@ -1756,6 +1818,36 @@ mod tests {
use manifest::{Scope, WorkerManifest}; use manifest::{Scope, WorkerManifest};
use session_store::{LogEntry, WorkerMetadataStore}; 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] #[test]
fn notify_run_state_allows_running_worker_inbox_delivery() { fn notify_run_state_allows_running_worker_inbox_delivery() {
assert_eq!( assert_eq!(
@@ -1878,9 +1970,16 @@ mod tests {
.as_ref() .as_ref()
.map(|binding| binding.root().to_path_buf()) .map(|binding| binding.root().to_path_buf())
.unwrap_or_else(|| self.cwd.clone()); .unwrap_or_else(|| self.cwd.clone());
let workspace_backend_ref = let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request(
RuntimeWorkspaceBackendRef::from_worker_request(&request.request); &request.request,
let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref); 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(); let workspace_client = workspace_context.client_handle();
self.observed_workspace_clients.lock().unwrap().push(( self.observed_workspace_clients.lock().unwrap().push((
workspace_client.kind().to_string(), workspace_client.kind().to_string(),
@@ -1916,6 +2015,7 @@ mod tests {
worker_ref: request.worker_ref, worker_ref: request.worker_ref,
run_generation: request.run_generation, run_generation: request.run_generation,
request: request.request, request: request.request,
workspace_scope: request.workspace_scope,
context: request.context, context: request.context,
working_directory: request.working_directory, working_directory: request.working_directory,
config_bundle: request.config_bundle, config_bundle: request.config_bundle,
@@ -2188,6 +2288,7 @@ mod tests {
worker_ref: worker_ref.clone(), worker_ref: worker_ref.clone(),
run_generation: 1, run_generation: 1,
request: create_request("1"), request: create_request("1"),
workspace_scope: None,
context: test_execution_context(worker_ref), context: test_execution_context(worker_ref),
working_directory: None, working_directory: None,
config_bundle: None, config_bundle: None,
@@ -2286,14 +2387,15 @@ mod tests {
request.workspace_api = Some(crate::catalog::WorkspaceApiRef { request.workspace_api = Some(crate::catalog::WorkspaceApiRef {
workspace_id: "workspace-restore".to_string(), workspace_id: "workspace-restore".to_string(),
base_url: "http://workspace.invalid".to_string(), base_url: "http://workspace.invalid".to_string(),
runtime_id: Some("runtime-restore".to_string()),
}); });
let controller = ProfileRuntimeWorkerFactory::new(root.path()) let controller = ProfileRuntimeWorkerFactory::new(root.path())
.with_runtime_id("runtime-restore")
.with_runtime_store_dir(&runtime_store_dir) .with_runtime_store_dir(&runtime_store_dir)
.restore_controller(WorkerExecutionRestoreRequest { .restore_controller(WorkerExecutionRestoreRequest {
worker_ref: worker_ref.clone(), worker_ref: worker_ref.clone(),
run_generation: 1, run_generation: 1,
request, request,
workspace_scope: None,
context: test_execution_context(worker_ref), context: test_execution_context(worker_ref),
previous_working_directory: None, previous_working_directory: None,
working_directory: None, working_directory: None,
@@ -2420,7 +2522,6 @@ mod tests {
request.workspace_api = Some(crate::catalog::WorkspaceApiRef { request.workspace_api = Some(crate::catalog::WorkspaceApiRef {
workspace_id: "ws-test".to_string(), workspace_id: "ws-test".to_string(),
base_url: "http://127.0.0.1:3999".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(); let detail = runtime.create_worker(request).unwrap();
@@ -2453,7 +2554,7 @@ mod tests {
assert_eq!( assert_eq!(
observed_workspace_clients.lock().unwrap().as_slice(), observed_workspace_clients.lock().unwrap().as_slice(),
&[( &[(
"runtime-http-proxy".to_string(), "runtime-owned-workspace-client".to_string(),
Some("ws-test".to_string()), Some("ws-test".to_string()),
true, true,
)] )]
+613
View File
@@ -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<String>, workspace_id: impl AsRef<str>) -> 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<RuntimeOwnedWorkerMutationProof, RuntimeAuthError> {
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<WorkspaceResponse, RuntimeWorkerMutationForwardError>;
}
#[derive(Clone)]
enum RuntimeWorkerMutationTransport {
Remote {
base_url: String,
client: reqwest::blocking::Client,
},
Embedded {
dispatcher: Arc<dyn EmbeddedWorkerMutationDispatcher>,
},
}
#[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<String>,
base_url: impl Into<String>,
) -> 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<String>,
scope: RuntimeWorkspaceScope,
source_worker_id: impl Into<String>,
dispatcher: Arc<dyn EmbeddedWorkerMutationDispatcher>,
) -> 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<WorkspaceResponse, RuntimeWorkerMutationForwardError> {
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<RuntimeWorkerMutationForwarder>,
}
impl RuntimeOwnedWorkspaceClient {
pub fn new(
workspace_id: impl Into<String>,
base_url: impl Into<String>,
runtime_id: impl Into<String>,
worker_id: impl Into<String>,
) -> 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<WorkspaceResponse, WorkspaceClientError> {
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<WorkspaceResponse, WorkspaceClientError> {
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<WorkspaceResponse, WorkspaceClientError> {
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<Option<(WorkerMutationSourceClaims, String, String)>>,
}
impl EmbeddedWorkerMutationDispatcher for RecordingDispatcher {
fn execute_worker_remove(
&self,
proof: InProcessWorkerMutationProof,
target_runtime_id: &str,
target_worker_id: &str,
) -> Result<WorkspaceResponse, RuntimeWorkerMutationForwardError> {
*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());
}
}
+1 -3
View File
@@ -346,11 +346,9 @@ mod tests {
use llm_engine::tool::ToolDefinition; use llm_engine::tool::ToolDefinition;
fn test_client() -> Arc<dyn WorkspaceClient> { fn test_client() -> Arc<dyn WorkspaceClient> {
Arc::new(crate::worker::RuntimeWorkspaceHttpClient::new( Arc::new(crate::worker::TestWorkspaceHttpClient::new(
"workspace", "workspace",
"http://backend", "http://backend",
"test-runtime",
"test-worker",
)) ))
} }
@@ -628,12 +628,7 @@ mod tests {
#[test] #[test]
fn workspace_http_objective_tools_include_objective_crud_tools() { fn workspace_http_objective_tools_include_objective_crud_tools() {
let names = tool_names(workspace_http_objective_tools(Arc::new( let names = tool_names(workspace_http_objective_tools(Arc::new(
crate::worker::RuntimeWorkspaceHttpClient::new( crate::worker::TestWorkspaceHttpClient::new("workspace", "http://backend"),
"workspace",
"http://backend",
"test-runtime",
"test-worker",
),
))); )));
assert_eq!( assert_eq!(
+3 -15
View File
@@ -1376,12 +1376,7 @@ provider = "github"
#[tokio::test(flavor = "multi_thread")] #[tokio::test(flavor = "multi_thread")]
async fn workspace_http_backend_invoke_is_safe_inside_async_context() { async fn workspace_http_backend_invoke_is_safe_inside_async_context() {
let backend = WorkspaceHttpTicketBackend::new(Arc::new( let backend = WorkspaceHttpTicketBackend::new(Arc::new(
crate::worker::RuntimeWorkspaceHttpClient::new( crate::worker::TestWorkspaceHttpClient::new("workspace-a", "not-a-url"),
"workspace-a",
"not-a-url",
"test-runtime",
"test-worker",
),
)); ));
let error = backend 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") .write_all(b"HTTP/1.1 204 No Content\r\nContent-Length: 0\r\n\r\n")
.unwrap(); .unwrap();
}); });
let client = Arc::new(crate::worker::RuntimeWorkspaceHttpClient::new( let client = Arc::new(crate::worker::TestWorkspaceHttpClient::new(
"workspace-a", "workspace-a",
format!("http://{address}"), format!("http://{address}"),
"test-runtime",
"worker-a",
)); ));
let backend = WorkspaceHttpTicketBackend::new(client); let backend = WorkspaceHttpTicketBackend::new(client);
@@ -1457,12 +1450,7 @@ provider = "github"
}); });
let backend = WorkspaceHttpTicketBackend::new(Arc::new( let backend = WorkspaceHttpTicketBackend::new(Arc::new(
crate::worker::RuntimeWorkspaceHttpClient::new( crate::worker::TestWorkspaceHttpClient::new("workspace-a", base_url),
"workspace-a",
base_url,
"test-runtime",
"test-worker",
),
)); ));
let created = backend.create(NewTicket::new("HTTP ticket")).unwrap(); let created = backend.create(NewTicket::new("HTTP ticket")).unwrap();
+4 -4
View File
@@ -40,9 +40,9 @@ pub use runtime::dir::RuntimeDir;
pub use segment_log_sink::SegmentLogSink; pub use segment_log_sink::SegmentLogSink;
pub use shared_state::WorkerSharedState; pub use shared_state::WorkerSharedState;
pub use worker::{ pub use worker::{
LocalWorkingDirectory, RuntimeWorkspaceHttpClient, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError,
Worker, WorkerError, WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceClient,
WorkspaceClient, WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspaceRequest, WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspaceRequest, WorkspaceRequestMethod,
WorkspaceRequestMethod, WorkspaceResponse, apply_worker_manifest, marker_workspace_client, WorkspaceResponse, apply_worker_manifest, marker_workspace_client,
unavailable_workspace_client, unavailable_workspace_client,
}; };
+3 -8
View File
@@ -212,8 +212,8 @@ mod tests {
break; break;
} }
} }
assert_eq!(runtime_header.as_deref(), Some("runtime-test")); assert_eq!(runtime_header, None);
assert_eq!(worker_header.as_deref(), Some("test-worker")); assert_eq!(worker_header, None);
assert_eq!(authorization, None); assert_eq!(authorization, None);
let body = serde_json::json!({ let body = serde_json::json!({
"authority": "workspace-backend-skills-v0", "authority": "workspace-backend-skills-v0",
@@ -236,12 +236,7 @@ mod tests {
.unwrap(); .unwrap();
}); });
let client = crate::worker::RuntimeWorkspaceHttpClient::new( let client = crate::worker::TestWorkspaceHttpClient::new("ws-1", format!("http://{addr}"));
"ws-1",
format!("http://{addr}"),
"runtime-test",
"test-worker",
);
let catalog = (&client as &dyn WorkspaceClient).list_skills().unwrap(); let catalog = (&client as &dyn WorkspaceClient).list_skills().unwrap();
assert_eq!(catalog.entries[0].name, "triage-errors"); assert_eq!(catalog.entries[0].name, "triage-errors");
assert_eq!(catalog.entries[0].provenance.id, "workspace:triage-errors"); assert_eq!(catalog.entries[0].provenance.id, "workspace:triage-errors");
+33 -108
View File
@@ -224,6 +224,18 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync {
fn execute(&self, request: WorkspaceRequest) fn execute(&self, request: WorkspaceRequest)
-> Result<WorkspaceResponse, WorkspaceClientError>; -> Result<WorkspaceResponse, WorkspaceClientError>;
/// 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<WorkspaceResponse, WorkspaceClientError> {
Err(WorkspaceClientError::Unavailable(
"Runtime-owned WorkerRemove forwarding is unavailable".to_string(),
))
}
/// Trusted review-attempt context is injected by the Internal SubWorker spawn layer. /// Trusted review-attempt context is injected by the Internal SubWorker spawn layer.
/// It is never accepted from a model-visible tool argument. /// It is never accepted from a model-visible tool argument.
fn reviewer_attempt_context(&self) -> Option<&ReviewerAttemptContext> { 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. #[cfg(test)]
/// #[derive(Clone, Debug)]
/// The upstream endpoint and source headers are private implementation details; pub(crate) struct TestWorkspaceHttpClient {
/// model-visible tools can only submit [`WorkspaceRequest`] values through the
/// [`WorkspaceClient`] trait.
pub struct RuntimeWorkspaceHttpClient {
workspace_id: String, workspace_id: String,
base_url: String, base_url: String,
runtime_id: String,
worker_id: String,
} }
impl std::fmt::Debug for RuntimeWorkspaceHttpClient { #[cfg(test)]
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { impl TestWorkspaceHttpClient {
formatter pub(crate) fn new(workspace_id: impl Into<String>, base_url: impl Into<String>) -> Self {
.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<String>,
base_url: impl Into<String>,
runtime_id: impl Into<String>,
worker_id: impl Into<String>,
) -> Self {
Self { Self {
workspace_id: workspace_id.into(), workspace_id: workspace_id.into(),
base_url: base_url.into().trim_end_matches('/').to_string(), 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> { fn workspace_id(&self) -> Option<&str> {
Some(&self.workspace_id) Some(&self.workspace_id)
} }
fn kind(&self) -> &str { fn kind(&self) -> &str {
"runtime-http-proxy" "test-http"
} }
fn is_available(&self) -> bool { fn is_available(&self) -> bool {
@@ -371,32 +361,28 @@ impl WorkspaceClient for RuntimeWorkspaceHttpClient {
request: WorkspaceRequest, request: WorkspaceRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> { ) -> Result<WorkspaceResponse, WorkspaceClientError> {
let base_url = self.base_url.clone(); 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() { if tokio::runtime::Handle::try_current().is_ok() {
std::thread::spawn(move || { std::thread::spawn(move || execute_test_workspace_http(&base_url, request))
execute_runtime_workspace_http(&base_url, &runtime_id, &worker_id, request)
})
.join() .join()
.map_err(|_| { .map_err(|_| {
WorkspaceClientError::Request("workspace request thread panicked".to_string()) WorkspaceClientError::Request(
"test workspace request thread panicked".to_string(),
)
})? })?
} else { } 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, base_url: &str,
runtime_id: &str,
worker_id: &str,
request: WorkspaceRequest, request: WorkspaceRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> { ) -> Result<WorkspaceResponse, WorkspaceClientError> {
if !request.path.starts_with('/') || request.path.starts_with("//") { if !request.path.starts_with('/') || request.path.starts_with("//") {
return Err(WorkspaceClientError::InvalidPath(request.path)); return Err(WorkspaceClientError::InvalidPath(request.path));
} }
let url = format!("{base_url}{}", request.path);
let method = match request.method { let method = match request.method {
WorkspaceRequestMethod::Get => reqwest::Method::GET, WorkspaceRequestMethod::Get => reqwest::Method::GET,
WorkspaceRequestMethod::Post => reqwest::Method::POST, WorkspaceRequestMethod::Post => reqwest::Method::POST,
@@ -405,16 +391,11 @@ fn execute_runtime_workspace_http(
WorkspaceRequestMethod::Delete => reqwest::Method::DELETE, WorkspaceRequestMethod::Delete => reqwest::Method::DELETE,
}; };
let client = reqwest::blocking::Client::new(); let client = reqwest::blocking::Client::new();
let mut request_builder = client let mut builder = client.request(method, format!("{base_url}{}", request.path));
.request(method, url)
.header("x-yoi-runtime-id", runtime_id)
.header("x-yoi-worker-id", worker_id);
if let Some(body) = request.body { if let Some(body) = request.body {
request_builder = request_builder builder = builder.body(body);
.header(reqwest::header::CONTENT_TYPE, "application/json")
.body(body);
} }
let response = request_builder let response = builder
.send() .send()
.map_err(|error| WorkspaceClientError::Request(error.to_string()))?; .map_err(|error| WorkspaceClientError::Request(error.to_string()))?;
let status = response.status().as_u16(); let status = response.status().as_u16();
@@ -6959,11 +6940,9 @@ mod build_summary_prompt_tests {
}); });
WorkerWorkspaceContext::with_client( WorkerWorkspaceContext::with_client(
Some(WorkspaceId::new("test-memory").unwrap()), Some(WorkspaceId::new("test-memory").unwrap()),
Arc::new(RuntimeWorkspaceHttpClient::new( Arc::new(TestWorkspaceHttpClient::new(
"test-memory", "test-memory",
format!("http://{addr}"), format!("http://{addr}"),
"test-runtime",
"test-worker",
)), )),
) )
} }
@@ -7096,11 +7075,9 @@ mod build_summary_prompt_tests {
store, store,
WorkerWorkspaceContext::with_client( WorkerWorkspaceContext::with_client(
Some(WorkspaceId::new("ws-skill").unwrap()), Some(WorkspaceId::new("ws-skill").unwrap()),
Arc::new(RuntimeWorkspaceHttpClient::new( Arc::new(TestWorkspaceHttpClient::new(
"ws-skill", "ws-skill",
format!("http://{addr}"), format!("http://{addr}"),
"test-runtime",
"test-worker",
)), )),
), ),
authority, 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] #[tokio::test]
async fn cancelled_internal_extract_does_not_commit_pointer_or_completed_audit() { async fn cancelled_internal_extract_does_not_commit_pointer_or_completed_audit() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
+1 -2
View File
@@ -55,7 +55,7 @@ use worker_runtime::interaction::{
use worker_runtime::management::{RuntimeOptions as EmbeddedRuntimeOptions, RuntimeStatus}; use worker_runtime::management::{RuntimeOptions as EmbeddedRuntimeOptions, RuntimeStatus};
use worker_runtime::profile_archive::{ProfileSourceArchive, ProfileSourceArchiveInput}; 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 EMBEDDED_HOST_KIND: &str = "embedded-worker-runtime-host";
const REMOTE_HOST_KIND: &str = "remote-worker-runtime-host"; const REMOTE_HOST_KIND: &str = "remote-worker-runtime-host";
const MAX_DIAGNOSTICS: usize = 16; const MAX_DIAGNOSTICS: usize = 16;
@@ -4039,7 +4039,6 @@ mod tests {
WorkspaceApiRef { WorkspaceApiRef {
workspace_id: "workspace-test".to_string(), workspace_id: "workspace-test".to_string(),
base_url: "http://127.0.0.1:8787".to_string(), base_url: "http://127.0.0.1:8787".to_string(),
runtime_id: Some("runtime-test".to_string()),
} }
} }
+1
View File
@@ -24,6 +24,7 @@ pub mod runtime_subscription;
pub mod server; pub mod server;
pub mod skills; pub mod skills;
pub mod store; pub mod store;
pub mod worker_source;
mod workspace_subscription; mod workspace_subscription;
pub use authority::{ pub use authority::{
+1 -1
View File
@@ -1354,7 +1354,7 @@ mod tests {
let path = temp.path().join("server.db"); let path = temp.path().join("server.db");
{ {
let s = SqliteWorkspaceStore::open(&path).unwrap(); 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 reopened = SqliteWorkspaceStore::open(&path).unwrap();
let p = reopened.worker_retention_policy("legacy").unwrap().unwrap(); let p = reopened.worker_retention_policy("legacy").unwrap().unwrap();
+453 -17
View File
@@ -62,15 +62,15 @@ use crate::companion::{
}; };
use crate::config::{BackendRuntimesConfigFile, RemoteRuntimeConfigFile, resolve_remote_runtime}; use crate::config::{BackendRuntimesConfigFile, RemoteRuntimeConfigFile, resolve_remote_runtime};
use crate::hosts::{ use crate::hosts::{
ConfigBundleCheckResult, ConfigBundleSyncResult, DiagnosticSeverity, EmbeddedWorkerRuntime, ConfigBundleCheckResult, ConfigBundleSyncResult, DiagnosticSeverity, EMBEDDED_RUNTIME_ID,
HostSummary, RemoteRuntimeConfig, RemoteWorkerRuntime, RuntimeDiagnostic, RuntimeRegistry, EmbeddedWorkerRuntime, HostSummary, RemoteRuntimeConfig, RemoteWorkerRuntime,
RuntimeRegistryError, RuntimeRegistryUnregisterResult, RuntimeSummary, TicketWorkerRole, RuntimeDiagnostic, RuntimeRegistry, RuntimeRegistryError, RuntimeRegistryUnregisterResult,
WorkerCapabilitySummary, WorkerCompletionsRequest, WorkerCompletionsResult, RuntimeSummary, TicketWorkerRole, WorkerCapabilitySummary, WorkerCompletionsRequest,
WorkerImplementationSummary, WorkerInputKind, WorkerInputRequest, WorkerInputResult, WorkerCompletionsResult, WorkerImplementationSummary, WorkerInputKind, WorkerInputRequest,
WorkerLifecycleRequest, WorkerLifecycleResult, WorkerOperationState, WorkerRestoreResult, WorkerInputResult, WorkerLifecycleRequest, WorkerLifecycleResult, WorkerOperationState,
WorkerSpawnAcceptanceRequirement, WorkerSpawnIntent, WorkerSpawnRequest, WorkerSpawnResult, WorkerRestoreResult, WorkerSpawnAcceptanceRequirement, WorkerSpawnIntent, WorkerSpawnRequest,
WorkerSpawnWorkingDirectoryRequest, WorkerSummary, WorkerTicketAssignmentRequest, WorkerSpawnResult, WorkerSpawnWorkingDirectoryRequest, WorkerSummary,
WorkerWorkspaceSummary, WorkerTicketAssignmentRequest, WorkerWorkspaceSummary,
}; };
use crate::identity::WorkspaceIdentity; use crate::identity::WorkspaceIdentity;
use crate::memory_backend::execute_memory_backend_operation_with_authority; use crate::memory_backend::execute_memory_backend_operation_with_authority;
@@ -250,8 +250,8 @@ const ORCHESTRATOR_ATTENTION_PROMPT: &str = include_str!(concat!(
#[derive(Clone)] #[derive(Clone)]
pub struct WorkspaceApi { pub struct WorkspaceApi {
config: ServerConfig, pub(crate) config: ServerConfig,
store: Arc<dyn ControlPlaneStore>, pub(crate) store: Arc<dyn ControlPlaneStore>,
authority: SqliteWorkspaceAuthority, authority: SqliteWorkspaceAuthority,
runtime: Arc<RuntimeRegistry>, runtime: Arc<RuntimeRegistry>,
companion: Arc<CompanionConsole>, companion: Arc<CompanionConsole>,
@@ -269,6 +269,15 @@ impl WorkspaceApi {
let resource_broker = BackendResourceBroker::default(); let resource_broker = BackendResourceBroker::default();
let execution_backend = WorkerRuntimeExecutionBackend::new( let execution_backend = WorkerRuntimeExecutionBackend::new(
ProfileRuntimeWorkerFactory::new(config.workspace_root.clone()) 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_runtime_store_dir(config.embedded_runtime_store_root.clone())
.with_resource_client(Arc::new(resource_broker.clone())), .with_resource_client(Arc::new(resource_broker.clone())),
) )
@@ -394,7 +403,7 @@ impl WorkspaceApi {
&self.runtime_subscription_broker &self.runtime_subscription_broker
} }
fn workspace_api_ref(&self, runtime_id: &str) -> WorkspaceApiRef { fn workspace_api_ref(&self, _runtime_id: &str) -> WorkspaceApiRef {
WorkspaceApiRef { WorkspaceApiRef {
workspace_id: self.config.workspace_id.clone(), workspace_id: self.config.workspace_id.clone(),
base_url: self base_url: self
@@ -404,7 +413,6 @@ impl WorkspaceApi {
.unwrap_or_else(|| "http://127.0.0.1:8787".to_string()) .unwrap_or_else(|| "http://127.0.0.1:8787".to_string())
.trim_end_matches('/') .trim_end_matches('/')
.to_string(), .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", "/api/w/{workspace_id}/companion/cancel",
post(scoped_post_companion_cancel), post(scoped_post_companion_cancel),
) )
.route(
"/api/w/{workspace_id}/workers/remove",
post(scoped_worker_remove_source_boundary),
)
.route( .route(
"/api/runtimes/{runtime_id}/workers", "/api/runtimes/{runtime_id}/workers",
get(list_runtime_workers).post(create_runtime_worker), get(list_runtime_workers).post(create_runtime_worker),
@@ -4811,6 +4823,71 @@ async fn scoped_workspace_protocol_ws(
.into_response()) .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<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>,
headers: HeaderMap,
Json(request): Json<WorkerRemoveBoundaryRequest>,
) -> 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( async fn scoped_list_workers(
State(api): State<WorkspaceApi>, State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>, AxumPath(path): AxumPath<ScopedWorkspacePath>,
@@ -10561,12 +10638,19 @@ mod tests {
use tokio_tungstenite::connect_async; use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::Message;
use tower::ServiceExt; 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::resource::BackendResourceClient;
use worker_runtime::worker_source::{
RuntimeOwnedWorkerMutationProof, RuntimeWorkerMutationSourceAuthority,
};
use worker_runtime::working_directory::WorkingDirectoryMaterializer; use worker_runtime::working_directory::WorkingDirectoryMaterializer;
use crate::hosts::{ use crate::hosts::{
TicketWorkerRole, WorkerInputKind, WorkerOperationState, WorkerSpawnAcceptanceRequirement, RemoteRuntimeAuthConfig, RuntimeCapabilitySummary, TicketWorkerRole, WorkerInputKind,
WorkerSpawnIntent, WorkerOperationState, WorkerSpawnAcceptanceRequirement, WorkerSpawnIntent,
}; };
use crate::store::{ use crate::store::{
MemoryDocumentRecord, MemoryStagingRecord, ObjectiveRecord, ObjectiveResourceRecord, MemoryDocumentRecord, MemoryStagingRecord, ObjectiveRecord, ObjectiveResourceRecord,
@@ -10643,11 +10727,10 @@ mod tests {
const TEST_REPOSITORY_ID: &str = "main"; const TEST_REPOSITORY_ID: &str = "main";
const TEST_CREATED_AT: &str = "2026-06-23T06:43:28Z"; 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 { WorkspaceApiRef {
workspace_id: TEST_WORKSPACE_ID.to_string(), workspace_id: TEST_WORKSPACE_ID.to_string(),
base_url: "http://127.0.0.1:8787".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() .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( fn seed_cleanup_worker(
api: &WorkspaceApi, api: &WorkspaceApi,
runtime_worker_id: u64, runtime_worker_id: u64,
+123 -9
View File
@@ -161,6 +161,11 @@ const MIGRATIONS: &[Migration] = &[
name: "create Worker retention authority", name: "create Worker retention authority",
apply: crate::retention::create_worker_retention_tables, 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 { struct Migration {
@@ -460,6 +465,15 @@ pub trait ControlPlaneStore: Send + Sync {
async fn schema_version(&self) -> Result<i64>; async fn schema_version(&self) -> Result<i64>;
async fn upsert_workspace(&self, record: &WorkspaceRecord) -> Result<()>; async fn upsert_workspace(&self, record: &WorkspaceRecord) -> Result<()>;
async fn get_workspace(&self, workspace_id: &str) -> Result<Option<WorkspaceRecord>>; async fn get_workspace(&self, workspace_id: &str) -> Result<Option<WorkspaceRecord>>;
async fn get_trusted_runtime(&self, runtime_id: &str) -> Result<Option<TrustedRuntimeRecord>>;
async fn consume_worker_mutation_source_jti(
&self,
runtime_id: &str,
jti: &str,
expires_at: u64,
now_seconds: u64,
consumed_at: &str,
) -> Result<bool>;
fn list_workspaces(&self) -> Result<Vec<WorkspaceRecord>>; fn list_workspaces(&self) -> Result<Vec<WorkspaceRecord>>;
fn upsert_repository(&self, record: &RepositoryRecord) -> Result<()>; fn upsert_repository(&self, record: &RepositoryRecord) -> Result<()>;
fn get_repository( fn get_repository(
@@ -896,6 +910,44 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
}) })
} }
async fn get_trusted_runtime(&self, runtime_id: &str) -> Result<Option<TrustedRuntimeRecord>> {
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<bool> {
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<Vec<WorkspaceRecord>> { fn list_workspaces(&self) -> Result<Vec<WorkspaceRecord>> {
self.with_conn(|conn| { self.with_conn(|conn| {
let mut stmt = conn.prepare( let mut stmt = conn.prepare(
@@ -4326,6 +4378,23 @@ fn current_schema_version(conn: &Connection) -> Result<i64> {
.map_err(Error::from) .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<()> { fn apply_migrations(conn: &Connection) -> Result<()> {
let current = current_schema_version(conn)?; let current = current_schema_version(conn)?;
for migration in MIGRATIONS for migration in MIGRATIONS
@@ -4915,7 +4984,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT);
apply_migrations(&conn).unwrap(); 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()); 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(); 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_sources").unwrap());
assert!(table_exists(&conn, "flow_source_revisions").unwrap()); assert!(table_exists(&conn, "flow_source_revisions").unwrap());
assert!(!table_exists(&conn, "flow_instances").unwrap()); assert!(!table_exists(&conn, "flow_instances").unwrap());
@@ -5015,7 +5084,7 @@ INSERT INTO worker_workdir_attachment_reservations (
apply_migrations(&conn).unwrap(); 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 let repositories_sql: String = conn
.query_row( .query_row(
"SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'repositories'", "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 db = dir.path().join("control-plane.sqlite");
let store = SqliteWorkspaceStore::open(&db).unwrap(); let store = SqliteWorkspaceStore::open(&db).unwrap();
assert_eq!(store.schema_version().await.unwrap(), 28); assert_eq!(store.schema_version().await.unwrap(), 29);
assert!( assert!(
!store !store
.with_conn(|conn| table_exists(conn, "worker_workspace_credentials")) .with_conn(|conn| table_exists(conn, "worker_workspace_credentials"))
@@ -5212,7 +5281,7 @@ INSERT INTO workdir_registry (
store.upsert_workspace(&record).await.unwrap(); store.upsert_workspace(&record).await.unwrap();
let reopened = SqliteWorkspaceStore::open(&db).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!( assert_eq!(
reopened.get_workspace("local-dev").await.unwrap(), reopened.get_workspace("local-dev").await.unwrap(),
Some(record) Some(record)
@@ -5759,7 +5828,7 @@ INSERT INTO workdir_registry (
.unwrap(); .unwrap();
let store = SqliteWorkspaceStore::from_connection(conn).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 store
.with_conn(|conn| { .with_conn(|conn| {
@@ -5948,7 +6017,7 @@ CREATE TABLE ticket_assignment_operations (
#[tokio::test] #[tokio::test]
async fn repository_records_round_trip() { async fn repository_records_round_trip() {
let store = SqliteWorkspaceStore::in_memory().unwrap(); 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 { let workspace = WorkspaceRecord {
workspace_id: "local-dev".to_string(), workspace_id: "local-dev".to_string(),
owner_account_id: None, owner_account_id: None,
@@ -6014,7 +6083,7 @@ CREATE TABLE ticket_assignment_operations (
#[tokio::test] #[tokio::test]
async fn memory_authority_records_round_trip_and_close_staging() { async fn memory_authority_records_round_trip_and_close_staging() {
let store = SqliteWorkspaceStore::in_memory().unwrap(); 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 { let workspace = WorkspaceRecord {
workspace_id: "local-dev".to_string(), workspace_id: "local-dev".to_string(),
owner_account_id: None, owner_account_id: None,
@@ -6277,7 +6346,7 @@ CREATE TABLE ticket_assignment_operations (
#[tokio::test] #[tokio::test]
async fn account_and_login_records_round_trip() { async fn account_and_login_records_round_trip() {
let store = SqliteWorkspaceStore::in_memory().unwrap(); 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 now = "2026-07-22T00:00:00Z".to_string();
let account = AccountRecord { let account = AccountRecord {
account_id: "acct-user-alice".to_string(), 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<String> { fn table_names(conn: &Connection) -> BTreeSet<String> {
let mut stmt = conn let mut stmt = conn
.prepare( .prepare(
@@ -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<InProcessWorkerMutationProof>,
) -> Result<PresentedWorkerMutationSourceProof<'a>, 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<VerifiedWorkerMutationSource, WorkerMutationSourceProofError> {
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<dyn crate::store::ControlPlaneStore>,
proof: PresentedWorkerMutationSourceProof<'_>,
target_runtime_id: &str,
target_worker_id: &str,
) -> Result<VerifiedWorkerMutationSource, WorkerMutationSourceProofError> {
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<dyn crate::store::ControlPlaneStore>,
}
impl EmbeddedServerWorkerMutationDispatcher {
pub(crate) fn new(
config: crate::server::ServerConfig,
store: std::sync::Arc<dyn crate::store::ControlPlaneStore>,
) -> 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()
}