runtime: prove Worker mutation source authority
This commit is contained in:
@@ -10,6 +10,10 @@ const PUBLIC_KEY_PREFIX: &str = "yoi-ed25519-pub:v1:";
|
||||
const PRIVATE_KEY_PREFIX: &str = "yoi-ed25519-pkcs8:v1:";
|
||||
const TOKEN_PREFIX: &str = "yoi-cap-v1";
|
||||
const SIGNING_INPUT_PREFIX: &str = "yoi-cap-v1.";
|
||||
pub const WORKER_MUTATION_SOURCE_PROOF_HEADER: &str = "x-yoi-worker-mutation-proof";
|
||||
const WORKER_MUTATION_SOURCE_PROOF_PREFIX: &str = "yoi-worker-source-v1";
|
||||
const WORKER_MUTATION_SOURCE_SIGNING_INPUT_PREFIX: &str = "yoi-worker-source-v1.";
|
||||
pub const WORKER_REMOVE_PERMISSION: &str = "workspace:worker-remove";
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum RuntimeAuthError {
|
||||
@@ -41,6 +45,16 @@ pub enum RuntimeAuthError {
|
||||
MissingWorkspaceScope,
|
||||
#[error("capability token is missing required permission `{0}`")]
|
||||
MissingPermission(String),
|
||||
#[error("source proof workspace `{actual}` does not match `{expected}`")]
|
||||
WrongWorkspace { expected: String, actual: String },
|
||||
#[error("source proof Worker `{actual}` does not match `{expected}`")]
|
||||
WrongWorker { expected: String, actual: String },
|
||||
#[error("source proof actor kind is not allowed")]
|
||||
WrongActorKind,
|
||||
#[error("source proof operation is not allowed")]
|
||||
WrongOperation,
|
||||
#[error("source proof mutation target does not match the request")]
|
||||
WrongMutationTarget,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
@@ -210,6 +224,201 @@ pub fn verify_capability_token(
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct WorkerMutationSourceClaims {
|
||||
pub iss: String,
|
||||
pub aud: String,
|
||||
pub workspace_id: String,
|
||||
pub worker_id: String,
|
||||
pub actor_kind: WorkerMutationActorKind,
|
||||
pub operation: WorkerMutationOperation,
|
||||
pub target_runtime_id: String,
|
||||
pub target_worker_id: String,
|
||||
pub permission: String,
|
||||
pub iat: u64,
|
||||
pub exp: u64,
|
||||
pub jti: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum WorkerMutationActorKind {
|
||||
Worker,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum WorkerMutationOperation {
|
||||
WorkerRemove,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RuntimeWorkerMutationSourceSigner {
|
||||
runtime_id: String,
|
||||
private_key: String,
|
||||
}
|
||||
|
||||
impl fmt::Debug for RuntimeWorkerMutationSourceSigner {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("RuntimeWorkerMutationSourceSigner")
|
||||
.field("runtime_id", &self.runtime_id)
|
||||
.field("private_key", &"[redacted]")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl RuntimeWorkerMutationSourceSigner {
|
||||
pub fn from_identity(identity: &RuntimeIdentityMaterial) -> Self {
|
||||
Self {
|
||||
runtime_id: identity.identity_id.clone(),
|
||||
private_key: identity.private_key.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn runtime_id(&self) -> &str {
|
||||
&self.runtime_id
|
||||
}
|
||||
|
||||
pub fn issue_worker_remove(
|
||||
&self,
|
||||
audience: impl Into<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> {
|
||||
let mut parts = token.split('.');
|
||||
match (parts.next(), parts.next(), parts.next(), parts.next()) {
|
||||
@@ -277,6 +486,112 @@ impl fmt::Display for RuntimeAuthContext {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn worker_mutation_source_proof_binds_all_source_authority_claims() {
|
||||
let runtime = RuntimeIdentityMaterial::generate("runtime-main").unwrap();
|
||||
let signer = RuntimeWorkerMutationSourceSigner::from_identity(&runtime);
|
||||
let claims = WorkerMutationSourceClaims {
|
||||
iss: "runtime-main".to_string(),
|
||||
aud: "server-main".to_string(),
|
||||
workspace_id: "workspace-a".to_string(),
|
||||
worker_id: "worker-7".to_string(),
|
||||
actor_kind: WorkerMutationActorKind::Worker,
|
||||
operation: WorkerMutationOperation::WorkerRemove,
|
||||
target_runtime_id: "runtime-target".to_string(),
|
||||
target_worker_id: "worker-target".to_string(),
|
||||
permission: WORKER_REMOVE_PERMISSION.to_string(),
|
||||
iat: 90,
|
||||
exp: 100,
|
||||
jti: "source-proof-1".to_string(),
|
||||
};
|
||||
let token = signer.sign(&claims).unwrap();
|
||||
let expected = WorkerMutationSourceExpectation {
|
||||
runtime_id: "runtime-main",
|
||||
audience: "server-main",
|
||||
workspace_id: "workspace-a",
|
||||
worker_id: Some("worker-7"),
|
||||
actor_kind: WorkerMutationActorKind::Worker,
|
||||
operation: WorkerMutationOperation::WorkerRemove,
|
||||
target_runtime_id: "runtime-target",
|
||||
target_worker_id: "worker-target",
|
||||
permission: WORKER_REMOVE_PERMISSION,
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
verify_worker_mutation_source_proof(&runtime.public_key, &token, &expected, 99)
|
||||
.unwrap(),
|
||||
claims
|
||||
);
|
||||
|
||||
let wrong_worker = WorkerMutationSourceExpectation {
|
||||
worker_id: Some("worker-8"),
|
||||
..expected.clone()
|
||||
};
|
||||
assert!(matches!(
|
||||
verify_worker_mutation_source_proof(&runtime.public_key, &token, &wrong_worker, 99),
|
||||
Err(RuntimeAuthError::WrongWorker { .. })
|
||||
));
|
||||
let wrong_scope = WorkerMutationSourceExpectation {
|
||||
workspace_id: "workspace-b",
|
||||
..expected.clone()
|
||||
};
|
||||
assert!(matches!(
|
||||
verify_worker_mutation_source_proof(&runtime.public_key, &token, &wrong_scope, 99),
|
||||
Err(RuntimeAuthError::WrongWorkspace { .. })
|
||||
));
|
||||
let wrong_audience = WorkerMutationSourceExpectation {
|
||||
audience: "server-other",
|
||||
..expected.clone()
|
||||
};
|
||||
assert!(matches!(
|
||||
verify_worker_mutation_source_proof(&runtime.public_key, &token, &wrong_audience, 99),
|
||||
Err(RuntimeAuthError::WrongAudience { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
verify_worker_mutation_source_proof(&runtime.public_key, &token, &expected, 101),
|
||||
Err(RuntimeAuthError::Expired)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_mutation_source_proof_rejects_spoofed_runtime_signature() {
|
||||
let trusted = RuntimeIdentityMaterial::generate("runtime-main").unwrap();
|
||||
let spoofed = RuntimeIdentityMaterial::generate("runtime-main").unwrap();
|
||||
let claims = WorkerMutationSourceClaims {
|
||||
iss: "runtime-main".to_string(),
|
||||
aud: "server-main".to_string(),
|
||||
workspace_id: "workspace-a".to_string(),
|
||||
worker_id: "worker-7".to_string(),
|
||||
actor_kind: WorkerMutationActorKind::Worker,
|
||||
operation: WorkerMutationOperation::WorkerRemove,
|
||||
target_runtime_id: "runtime-target".to_string(),
|
||||
target_worker_id: "worker-target".to_string(),
|
||||
permission: WORKER_REMOVE_PERMISSION.to_string(),
|
||||
iat: 90,
|
||||
exp: 100,
|
||||
jti: "source-proof-2".to_string(),
|
||||
};
|
||||
let token = RuntimeWorkerMutationSourceSigner::from_identity(&spoofed)
|
||||
.sign(&claims)
|
||||
.unwrap();
|
||||
let expected = WorkerMutationSourceExpectation {
|
||||
runtime_id: "runtime-main",
|
||||
audience: "server-main",
|
||||
workspace_id: "workspace-a",
|
||||
worker_id: Some("worker-7"),
|
||||
actor_kind: WorkerMutationActorKind::Worker,
|
||||
operation: WorkerMutationOperation::WorkerRemove,
|
||||
target_runtime_id: "runtime-target",
|
||||
target_worker_id: "worker-target",
|
||||
permission: WORKER_REMOVE_PERMISSION,
|
||||
};
|
||||
|
||||
assert!(matches!(
|
||||
verify_worker_mutation_source_proof(&trusted.public_key, &token, &expected, 99),
|
||||
Err(RuntimeAuthError::InvalidSignature)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_token_verifies_signature_audience_expiry_and_permission() {
|
||||
let server = RuntimeIdentityMaterial::generate("server-main").unwrap();
|
||||
|
||||
@@ -179,8 +179,6 @@ pub struct WorkingDirectoryStatus {
|
||||
pub struct WorkspaceApiRef {
|
||||
pub workspace_id: String,
|
||||
pub base_url: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub runtime_id: Option<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for WorkspaceApiRef {
|
||||
@@ -189,7 +187,6 @@ impl std::fmt::Debug for WorkspaceApiRef {
|
||||
.debug_struct("WorkspaceApiRef")
|
||||
.field("workspace_id", &self.workspace_id)
|
||||
.field("base_url", &self.base_url)
|
||||
.field("runtime_id", &self.runtime_id)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -292,3 +289,31 @@ pub struct WorkerLifecycleAck {
|
||||
pub worker_ref: WorkerRef,
|
||||
pub status: WorkerStatus,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::WorkspaceApiRef;
|
||||
|
||||
#[test]
|
||||
fn workspace_api_ref_public_schema_contains_no_source_credentials_or_claim_choices() {
|
||||
let value = serde_json::to_value(WorkspaceApiRef {
|
||||
workspace_id: "workspace-a".to_string(),
|
||||
base_url: "https://server.invalid".to_string(),
|
||||
})
|
||||
.unwrap();
|
||||
let object = value.as_object().unwrap();
|
||||
assert_eq!(object.len(), 2);
|
||||
assert!(object.contains_key("workspace_id"));
|
||||
assert!(object.contains_key("base_url"));
|
||||
for forbidden in [
|
||||
"runtime_id",
|
||||
"worker_id",
|
||||
"permission",
|
||||
"private_key",
|
||||
"bearer_token",
|
||||
"signing_handle",
|
||||
] {
|
||||
assert!(!object.contains_key(forbidden), "unexpected {forbidden}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -248,6 +248,7 @@ pub struct WorkerExecutionSpawnRequest {
|
||||
/// Monotonic execution generation reserved durably before launch.
|
||||
pub run_generation: u64,
|
||||
pub request: crate::catalog::CreateWorkerRequest,
|
||||
pub workspace_scope: Option<crate::runtime::RuntimeWorkspaceScope>,
|
||||
pub context: WorkerExecutionContext,
|
||||
pub working_directory: Option<WorkingDirectoryBinding>,
|
||||
pub config_bundle: Option<ConfigBundle>,
|
||||
@@ -260,6 +261,7 @@ pub struct WorkerExecutionRestoreRequest {
|
||||
/// Monotonic execution generation reserved durably before restore.
|
||||
pub run_generation: u64,
|
||||
pub request: crate::catalog::CreateWorkerRequest,
|
||||
pub workspace_scope: Option<crate::runtime::RuntimeWorkspaceScope>,
|
||||
pub context: WorkerExecutionContext,
|
||||
pub previous_working_directory: Option<WorkingDirectoryStatus>,
|
||||
pub working_directory: Option<WorkingDirectoryBinding>,
|
||||
|
||||
@@ -1753,7 +1753,6 @@ mod tests {
|
||||
request.workspace_api = Some(WorkspaceApiRef {
|
||||
workspace_id: workspace_id.to_string(),
|
||||
base_url: format!("https://workspace.example/{workspace_id}"),
|
||||
runtime_id: None,
|
||||
});
|
||||
request
|
||||
}
|
||||
@@ -2324,7 +2323,6 @@ mod tests {
|
||||
workspace_api: WorkspaceApiRef {
|
||||
workspace_id: "local".to_string(),
|
||||
base_url: "http://127.0.0.1:8787".to_string(),
|
||||
runtime_id: None,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
@@ -26,6 +26,7 @@ pub mod resource;
|
||||
pub mod retention;
|
||||
mod runtime;
|
||||
pub mod worker_backend;
|
||||
pub mod worker_source;
|
||||
pub mod working_directory;
|
||||
|
||||
#[cfg(feature = "fs-store")]
|
||||
|
||||
@@ -87,6 +87,9 @@ fn build_runtime(config: &ProcessConfig) -> Result<Runtime, ProcessError> {
|
||||
};
|
||||
let mut factory = ProfileRuntimeWorkerFactory::new(fs_paths.worker_dir.join("worker-root"))
|
||||
.with_runtime_store_dir(runtime_store_dir);
|
||||
if let Some(identity) = read_runtime_auth_file(&runtime_auth_path(config))?.identity {
|
||||
factory = factory.with_remote_worker_mutation_identity(identity);
|
||||
}
|
||||
if let Some(endpoint) = config.backend_resource_endpoint.clone() {
|
||||
factory = factory.with_resource_client(Arc::new(
|
||||
worker_runtime::resource::HttpBackendResourceClient::new(
|
||||
|
||||
@@ -548,6 +548,7 @@ impl Runtime {
|
||||
worker_ref: worker_ref.clone(),
|
||||
run_generation: 1,
|
||||
request,
|
||||
workspace_scope: scope.cloned(),
|
||||
context: self.execution_context(worker_ref.clone()),
|
||||
working_directory: None,
|
||||
config_bundle: None,
|
||||
@@ -826,13 +827,10 @@ impl Runtime {
|
||||
if let Some(existing) = worker.request.workspace_api.as_ref()
|
||||
&& (existing.workspace_id != workspace_api.workspace_id
|
||||
|| existing.base_url.trim_end_matches('/')
|
||||
!= workspace_api.base_url.trim_end_matches('/')
|
||||
|| existing.runtime_id.as_ref().is_some_and(|runtime_id| {
|
||||
workspace_api.runtime_id.as_ref() != Some(runtime_id)
|
||||
}))
|
||||
!= workspace_api.base_url.trim_end_matches('/'))
|
||||
{
|
||||
return Err(RuntimeError::InvalidRequest(
|
||||
"Workspace API replacement cannot change Worker Workspace identity, Runtime identity, or base URL"
|
||||
"Workspace API replacement cannot change Worker Workspace identity or base URL"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
@@ -902,10 +900,17 @@ impl Runtime {
|
||||
})?;
|
||||
state.worker_mut(worker_ref)?.run_generation = run_generation;
|
||||
state.persist_worker(&worker_ref.worker_id)?;
|
||||
let workspace_scope = worker_request.workspace_api.as_ref().and_then(|api| {
|
||||
state
|
||||
.workspace_owners
|
||||
.get(&api.workspace_id)
|
||||
.map(|server_id| RuntimeWorkspaceScope::new(&api.workspace_id, server_id))
|
||||
});
|
||||
let request = WorkerExecutionRestoreRequest {
|
||||
worker_ref: worker_ref.clone(),
|
||||
run_generation,
|
||||
request: worker_request,
|
||||
workspace_scope,
|
||||
context: self.execution_context(worker_ref.clone()),
|
||||
previous_working_directory,
|
||||
working_directory: None,
|
||||
@@ -1595,9 +1600,15 @@ impl Runtime {
|
||||
};
|
||||
|
||||
for candidate in candidates {
|
||||
let backend = {
|
||||
let (backend, workspace_scope) = {
|
||||
let state = self.lock()?;
|
||||
state.execution_backend.clone()
|
||||
let workspace_scope = candidate.request.workspace_api.as_ref().and_then(|api| {
|
||||
state
|
||||
.workspace_owners
|
||||
.get(&api.workspace_id)
|
||||
.map(|server_id| RuntimeWorkspaceScope::new(&api.workspace_id, server_id))
|
||||
});
|
||||
(state.execution_backend.clone(), workspace_scope)
|
||||
};
|
||||
let Some(backend) = backend else {
|
||||
return Ok(());
|
||||
@@ -1606,6 +1617,7 @@ impl Runtime {
|
||||
worker_ref: candidate.worker_ref.clone(),
|
||||
run_generation: candidate.run_generation,
|
||||
request: candidate.request,
|
||||
workspace_scope,
|
||||
context: self.execution_context(candidate.worker_ref.clone()),
|
||||
previous_working_directory: candidate.previous_working_directory,
|
||||
working_directory: None,
|
||||
@@ -2790,7 +2802,6 @@ mod tests {
|
||||
request.workspace_api = Some(WorkspaceApiRef {
|
||||
workspace_id: workspace_id.to_string(),
|
||||
base_url: format!("https://workspace.example/{workspace_id}"),
|
||||
runtime_id: None,
|
||||
});
|
||||
request
|
||||
}
|
||||
@@ -3239,7 +3250,6 @@ mod tests {
|
||||
let replacement = WorkspaceApiRef {
|
||||
workspace_id: "workspace-a".to_string(),
|
||||
base_url: "https://workspace.example/workspace-a/".to_string(),
|
||||
runtime_id: Some("runtime-a".to_string()),
|
||||
};
|
||||
|
||||
runtime
|
||||
|
||||
@@ -14,6 +14,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex, mpsc};
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::auth::RuntimeIdentityMaterial;
|
||||
use crate::catalog::{
|
||||
CreateWorkerRequest, ProfileSourceArchiveHttpRef, ProfileSourceArchiveSource,
|
||||
WorkingDirectoryRequest, WorkingDirectoryStatus,
|
||||
@@ -26,6 +27,9 @@ use crate::execution::{
|
||||
use crate::identity::WorkerRef;
|
||||
use crate::interaction::{WorkerInput, WorkerInputKind};
|
||||
use crate::resource::{BackendResourceClient, ProfileSourceArchiveCache};
|
||||
use crate::worker_source::{
|
||||
EmbeddedWorkerMutationDispatcher, RuntimeOwnedWorkspaceClient, RuntimeWorkerMutationForwarder,
|
||||
};
|
||||
use crate::working_directory::{
|
||||
WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer,
|
||||
};
|
||||
@@ -49,10 +53,9 @@ use worker::feature::builtin::{
|
||||
#[cfg(feature = "ws-server")]
|
||||
use worker::ipc::protocol_session::{live_log_entry_event, subscribe_worker_protocol_session};
|
||||
use worker::{
|
||||
PromptLoader, RuntimeWorkspaceHttpClient, SegmentLogSink,
|
||||
WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerController, WorkerError,
|
||||
WorkerFilesystemAuthority, WorkerHandle, WorkerSharedState, WorkerWorkspaceContext,
|
||||
WorkspaceClient, WorkspaceId,
|
||||
PromptLoader, SegmentLogSink, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker,
|
||||
WorkerController, WorkerError, WorkerFilesystemAuthority, WorkerHandle, WorkerSharedState,
|
||||
WorkerWorkspaceContext, WorkspaceClient, WorkspaceId,
|
||||
};
|
||||
|
||||
const DEFAULT_BACKEND_ID: &str = "worker-crate";
|
||||
@@ -215,6 +218,9 @@ pub struct ProfileRuntimeWorkerFactory {
|
||||
worker_aggregate_root: Option<PathBuf>,
|
||||
resource_client: Option<Arc<dyn BackendResourceClient>>,
|
||||
profile_archive_cache: Arc<ProfileSourceArchiveCache>,
|
||||
runtime_id: Option<String>,
|
||||
worker_mutation_identity: Option<RuntimeIdentityMaterial>,
|
||||
embedded_worker_mutation_dispatcher: Option<Arc<dyn EmbeddedWorkerMutationDispatcher>>,
|
||||
}
|
||||
|
||||
impl ProfileRuntimeWorkerFactory {
|
||||
@@ -226,9 +232,38 @@ impl ProfileRuntimeWorkerFactory {
|
||||
worker_aggregate_root: None,
|
||||
resource_client: None,
|
||||
profile_archive_cache: Arc::new(ProfileSourceArchiveCache::default()),
|
||||
runtime_id: None,
|
||||
worker_mutation_identity: None,
|
||||
embedded_worker_mutation_dispatcher: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_runtime_id(mut self, runtime_id: impl Into<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 {
|
||||
self.worker_aggregate_root = Some(runtime_store_dir.into().join("workers"));
|
||||
self
|
||||
@@ -348,38 +383,59 @@ enum RuntimeWorkspaceBackendRef {
|
||||
}
|
||||
|
||||
impl RuntimeWorkspaceBackendRef {
|
||||
fn from_worker_request(request: &CreateWorkerRequest) -> Self {
|
||||
if let Some(api) = request.workspace_api.as_ref()
|
||||
&& let Some(runtime_id) = api
|
||||
.runtime_id
|
||||
.as_ref()
|
||||
.filter(|runtime_id| !runtime_id.trim().is_empty())
|
||||
{
|
||||
fn from_worker_request(request: &CreateWorkerRequest, runtime_id: Option<&str>) -> Self {
|
||||
if let (Some(api), Some(runtime_id)) = (request.workspace_api.as_ref(), runtime_id) {
|
||||
return Self::Http {
|
||||
workspace_id: api.workspace_id.clone(),
|
||||
base_url: api.base_url.clone(),
|
||||
runtime_id: runtime_id.clone(),
|
||||
runtime_id: runtime_id.to_string(),
|
||||
};
|
||||
}
|
||||
Self::None
|
||||
}
|
||||
|
||||
fn worker_context(&self, worker_ref: &WorkerRef) -> WorkerWorkspaceContext {
|
||||
fn worker_context(
|
||||
&self,
|
||||
worker_ref: &WorkerRef,
|
||||
workspace_scope: Option<&crate::runtime::RuntimeWorkspaceScope>,
|
||||
mutation_identity: Option<&RuntimeIdentityMaterial>,
|
||||
embedded_dispatcher: Option<&Arc<dyn EmbeddedWorkerMutationDispatcher>>,
|
||||
) -> WorkerWorkspaceContext {
|
||||
match self {
|
||||
Self::None => WorkerWorkspaceContext::no_workspace(),
|
||||
Self::Http {
|
||||
workspace_id,
|
||||
base_url,
|
||||
runtime_id,
|
||||
} => WorkerWorkspaceContext::with_client(
|
||||
WorkspaceId::new(workspace_id.clone()).ok(),
|
||||
Arc::new(RuntimeWorkspaceHttpClient::new(
|
||||
} => {
|
||||
let mut client = RuntimeOwnedWorkspaceClient::new(
|
||||
workspace_id.clone(),
|
||||
base_url.clone(),
|
||||
runtime_id.clone(),
|
||||
worker_ref.worker_id.to_string(),
|
||||
)),
|
||||
),
|
||||
);
|
||||
if let (Some(scope), Some(identity)) = (workspace_scope, mutation_identity) {
|
||||
client = client.with_worker_remove(RuntimeWorkerMutationForwarder::remote(
|
||||
identity,
|
||||
scope.clone(),
|
||||
worker_ref.worker_id.to_string(),
|
||||
base_url.clone(),
|
||||
));
|
||||
} else if let (Some(scope), Some(dispatcher)) =
|
||||
(workspace_scope, embedded_dispatcher)
|
||||
{
|
||||
client = client.with_worker_remove(RuntimeWorkerMutationForwarder::embedded(
|
||||
runtime_id,
|
||||
scope.clone(),
|
||||
worker_ref.worker_id.to_string(),
|
||||
(*dispatcher).clone(),
|
||||
));
|
||||
}
|
||||
WorkerWorkspaceContext::with_client(
|
||||
WorkspaceId::new(workspace_id.clone()).ok(),
|
||||
Arc::new(client),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -478,13 +534,11 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
|
||||
)
|
||||
})
|
||||
.unwrap_or(WorkerFilesystemAuthority::None);
|
||||
let workspace_backend_ref =
|
||||
RuntimeWorkspaceBackendRef::from_worker_request(&request.request);
|
||||
let observation_runtime_id = request
|
||||
.request
|
||||
.workspace_api
|
||||
.as_ref()
|
||||
.and_then(|api| api.runtime_id.clone());
|
||||
let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request(
|
||||
&request.request,
|
||||
self.runtime_id.as_deref(),
|
||||
);
|
||||
let observation_runtime_id = self.runtime_id.clone();
|
||||
let observation_workspace_id = request
|
||||
.request
|
||||
.workspace_api
|
||||
@@ -492,7 +546,12 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
|
||||
.map(|api| api.workspace_id.clone());
|
||||
let observation_grants = request.request.worker_observation_grants.clone();
|
||||
let observation_enabled = request.request.worker_observation_enabled;
|
||||
let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref);
|
||||
let workspace_context = workspace_backend_ref.worker_context(
|
||||
&request.worker_ref,
|
||||
request.workspace_scope.as_ref(),
|
||||
self.worker_mutation_identity.as_ref(),
|
||||
self.embedded_worker_mutation_dispatcher.as_ref(),
|
||||
);
|
||||
let selector = profile.as_ref();
|
||||
let archive = self
|
||||
.resolve_profile_source_archive(&request.request.profile_source)
|
||||
@@ -628,13 +687,11 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
|
||||
)
|
||||
})
|
||||
.unwrap_or(WorkerFilesystemAuthority::None);
|
||||
let workspace_backend_ref =
|
||||
RuntimeWorkspaceBackendRef::from_worker_request(&request.request);
|
||||
let observation_runtime_id = request
|
||||
.request
|
||||
.workspace_api
|
||||
.as_ref()
|
||||
.and_then(|api| api.runtime_id.clone());
|
||||
let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request(
|
||||
&request.request,
|
||||
self.runtime_id.as_deref(),
|
||||
);
|
||||
let observation_runtime_id = self.runtime_id.clone();
|
||||
let observation_workspace_id = request
|
||||
.request
|
||||
.workspace_api
|
||||
@@ -642,7 +699,12 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
|
||||
.map(|api| api.workspace_id.clone());
|
||||
let observation_grants = request.request.worker_observation_grants.clone();
|
||||
let observation_enabled = request.request.worker_observation_enabled;
|
||||
let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref);
|
||||
let workspace_context = workspace_backend_ref.worker_context(
|
||||
&request.worker_ref,
|
||||
request.workspace_scope.as_ref(),
|
||||
self.worker_mutation_identity.as_ref(),
|
||||
self.embedded_worker_mutation_dispatcher.as_ref(),
|
||||
);
|
||||
let (manifest, loader) = Self::restore_fallback_manifest(&worker_name)?;
|
||||
|
||||
let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?;
|
||||
@@ -1756,6 +1818,36 @@ mod tests {
|
||||
use manifest::{Scope, WorkerManifest};
|
||||
use session_store::{LogEntry, WorkerMetadataStore};
|
||||
|
||||
#[test]
|
||||
fn restart_restore_reconstructs_runtime_owned_worker_mutation_client() {
|
||||
let identity = RuntimeIdentityMaterial::generate("runtime-source").unwrap();
|
||||
let worker_ref = WorkerRef::new(crate::identity::WorkerId::new(17));
|
||||
let backend = RuntimeWorkspaceBackendRef::Http {
|
||||
workspace_id: "workspace-a".to_string(),
|
||||
base_url: "https://server.invalid".to_string(),
|
||||
runtime_id: "runtime-source".to_string(),
|
||||
};
|
||||
let scope = crate::runtime::RuntimeWorkspaceScope::new("workspace-a", "server-main");
|
||||
|
||||
let before_restart =
|
||||
backend.worker_context(&worker_ref, Some(&scope), Some(&identity), None);
|
||||
let after_restore =
|
||||
backend.worker_context(&worker_ref, Some(&scope), Some(&identity), None);
|
||||
|
||||
assert_eq!(
|
||||
before_restart.client_handle().kind(),
|
||||
"runtime-owned-workspace-client"
|
||||
);
|
||||
assert_eq!(
|
||||
after_restore.client_handle().kind(),
|
||||
"runtime-owned-workspace-client"
|
||||
);
|
||||
assert_eq!(
|
||||
after_restore.client_handle().workspace_id(),
|
||||
Some("workspace-a")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notify_run_state_allows_running_worker_inbox_delivery() {
|
||||
assert_eq!(
|
||||
@@ -1878,9 +1970,16 @@ mod tests {
|
||||
.as_ref()
|
||||
.map(|binding| binding.root().to_path_buf())
|
||||
.unwrap_or_else(|| self.cwd.clone());
|
||||
let workspace_backend_ref =
|
||||
RuntimeWorkspaceBackendRef::from_worker_request(&request.request);
|
||||
let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref);
|
||||
let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request(
|
||||
&request.request,
|
||||
Some("runtime-test"),
|
||||
);
|
||||
let workspace_context = workspace_backend_ref.worker_context(
|
||||
&request.worker_ref,
|
||||
request.workspace_scope.as_ref(),
|
||||
None,
|
||||
None,
|
||||
);
|
||||
let workspace_client = workspace_context.client_handle();
|
||||
self.observed_workspace_clients.lock().unwrap().push((
|
||||
workspace_client.kind().to_string(),
|
||||
@@ -1916,6 +2015,7 @@ mod tests {
|
||||
worker_ref: request.worker_ref,
|
||||
run_generation: request.run_generation,
|
||||
request: request.request,
|
||||
workspace_scope: request.workspace_scope,
|
||||
context: request.context,
|
||||
working_directory: request.working_directory,
|
||||
config_bundle: request.config_bundle,
|
||||
@@ -2188,6 +2288,7 @@ mod tests {
|
||||
worker_ref: worker_ref.clone(),
|
||||
run_generation: 1,
|
||||
request: create_request("1"),
|
||||
workspace_scope: None,
|
||||
context: test_execution_context(worker_ref),
|
||||
working_directory: None,
|
||||
config_bundle: None,
|
||||
@@ -2286,14 +2387,15 @@ mod tests {
|
||||
request.workspace_api = Some(crate::catalog::WorkspaceApiRef {
|
||||
workspace_id: "workspace-restore".to_string(),
|
||||
base_url: "http://workspace.invalid".to_string(),
|
||||
runtime_id: Some("runtime-restore".to_string()),
|
||||
});
|
||||
let controller = ProfileRuntimeWorkerFactory::new(root.path())
|
||||
.with_runtime_id("runtime-restore")
|
||||
.with_runtime_store_dir(&runtime_store_dir)
|
||||
.restore_controller(WorkerExecutionRestoreRequest {
|
||||
worker_ref: worker_ref.clone(),
|
||||
run_generation: 1,
|
||||
request,
|
||||
workspace_scope: None,
|
||||
context: test_execution_context(worker_ref),
|
||||
previous_working_directory: None,
|
||||
working_directory: None,
|
||||
@@ -2420,7 +2522,6 @@ mod tests {
|
||||
request.workspace_api = Some(crate::catalog::WorkspaceApiRef {
|
||||
workspace_id: "ws-test".to_string(),
|
||||
base_url: "http://127.0.0.1:3999".to_string(),
|
||||
runtime_id: Some("runtime-test".to_string()),
|
||||
});
|
||||
let detail = runtime.create_worker(request).unwrap();
|
||||
|
||||
@@ -2453,7 +2554,7 @@ mod tests {
|
||||
assert_eq!(
|
||||
observed_workspace_clients.lock().unwrap().as_slice(),
|
||||
&[(
|
||||
"runtime-http-proxy".to_string(),
|
||||
"runtime-owned-workspace-client".to_string(),
|
||||
Some("ws-test".to_string()),
|
||||
true,
|
||||
)]
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user