Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: fence revoked runtime bindings

Browse Source
This commit is contained in:
Hare
2026-09-06 02:52:15 +09:00
parent 75b85b46d1
commit ec845cbc25
3 changed files with 199 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/hosts.rs
+60 -1
View File
@@ -1126,15 +1126,19 @@ pub enum RuntimeRegistryUnregisterResult {
},
}
type RuntimeBindingGate = Arc<dyn Fn(&str) -> bool + Send + Sync>;
#[derive(Clone)]
pub struct RuntimeRegistry {
runtimes: Arc<RwLock<Vec<Arc<dyn WorkspaceWorkerRuntime>>>>,
runtime_binding_gate: Arc<RwLock<Option<RuntimeBindingGate>>>,
}
impl RuntimeRegistry {
pub fn new(runtimes: Vec<Arc<dyn WorkspaceWorkerRuntime>>) -> Self {
Self {
runtimes: Arc::new(RwLock::new(runtimes)),
runtime_binding_gate: Arc::new(RwLock::new(None)),
}
}
@@ -1142,6 +1146,27 @@ impl RuntimeRegistry {
Self::new(vec![Arc::new(embedded_runtime)])
}
pub fn set_runtime_binding_gate<F>(&self, gate: F)
where
F: Fn(&str) -> bool + Send + Sync + 'static,
{
*self
.runtime_binding_gate
.write()
.expect("runtime binding gate lock poisoned") = Some(Arc::new(gate));
}
fn runtime_binding_is_active(&self, runtime_id: &str) -> bool {
if runtime_id == EMBEDDED_RUNTIME_ID {
return true;
}
self.runtime_binding_gate
.read()
.expect("runtime binding gate lock poisoned")
.as_ref()
.is_none_or(|gate| gate(runtime_id))
}
pub fn register<R>(&self, runtime: R)
where
R: WorkspaceWorkerRuntime + 'static,
@@ -1795,13 +1820,19 @@ impl RuntimeRegistry {
self.runtimes
.read()
.expect("runtime registry lock poisoned")
.clone()
.iter()
.filter(|runtime| self.runtime_binding_is_active(runtime.runtime_id()))
.cloned()
.collect()
}
fn runtime(
&self,
runtime_id: &str,
) -> Result<Arc<dyn WorkspaceWorkerRuntime>, RuntimeRegistryError> {
if !self.runtime_binding_is_active(runtime_id) {
return Err(RuntimeRegistryError::UnknownRuntime(runtime_id.to_string()));
}
self.runtimes
.read()
.expect("runtime registry lock poisoned")
@@ -5165,6 +5196,34 @@ mod tests {
assert_eq!(from_runtime_a.label, "worker from runtime a");
}
#[test]
fn registry_gate_rejects_cached_runtime_immediately_after_binding_revocation() {
let registry = RuntimeRegistry::new(vec![Arc::new(FixtureRuntime::with_worker(
"runtime-a",
"host-a",
"worker-a",
"worker from runtime a",
))]);
let active = Arc::new(Mutex::new(true));
let gate_state = active.clone();
registry.set_runtime_binding_gate(move |_| {
*gate_state.lock().expect("gate state lock poisoned")
});
assert_eq!(registry.list_runtimes(10).items.len(), 1);
assert!(
registry
.worker(&RuntimeWorkerRef::new("runtime-a", "worker-a"))
.is_ok()
);
*active.lock().expect("gate state lock poisoned") = false;
assert!(registry.list_runtimes(10).items.is_empty());
assert!(matches!(
registry.worker(&RuntimeWorkerRef::new("runtime-a", "worker-a")),
Err(RuntimeRegistryError::UnknownRuntime(runtime_id)) if runtime_id == "runtime-a"
));
}
#[test]
fn registry_broadcasts_workspace_prompt_projection_revisions() {
let runtime =
crates/workspace-server/src/server.rs
+11 -3
View File
@@ -1578,7 +1578,7 @@ impl WorkspaceApi {
updated_at: config.workspace_created_at.clone(),
revoked_at: None,
},
false,
true,
)
.await?;
let embedded_audience = format!("embedded:{}", config.workspace_id);
@@ -1604,14 +1604,22 @@ impl WorkspaceApi {
"failed to initialize embedded Worker backend: {err}"
))
})?;
Self::new_with_execution_backend_and_broker(
let runtime_binding_store = store.clone();
let runtime_binding_workspace_id = config.workspace_id.clone();
let api = Self::new_with_execution_backend_and_broker(
config,
store,
Arc::new(execution_backend),
resource_broker,
Some(worker_remove_dispatcher),
)
.await
.await?;
api.runtime.set_runtime_binding_gate(move |runtime_id| {
runtime_binding_store
.workspace_runtime_binding_is_active(&runtime_binding_workspace_id, runtime_id)
.unwrap_or(false)
});
Ok(api)
}
#[cfg(test)]
crates/workspace-server/src/store.rs
+128 -1
View File
@@ -544,6 +544,11 @@ pub trait ControlPlaneStore: Send + Sync {
&self,
record: &WorkspaceBootstrapRecord,
) -> Result<WorkspaceBootstrapResult>;
fn workspace_runtime_binding_is_active(
&self,
workspace_id: &str,
runtime_id: &str,
) -> Result<bool>;
async fn get_workspace_runtime_binding(
&self,
workspace_id: &str,
@@ -1835,6 +1840,17 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
})
}
fn workspace_runtime_binding_is_active(
&self,
workspace_id: &str,
runtime_id: &str,
) -> Result<bool> {
Ok(
SqliteWorkspaceStore::get_workspace_runtime_binding(self, workspace_id, runtime_id)?
.is_some_and(|binding| binding.revoked_at.is_none()),
)
}
async fn get_workspace_runtime_binding(
&self,
workspace_id: &str,
@@ -6967,12 +6983,20 @@ mod tests {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("server.db");
prepare_schema_v50(&path, Some("workspace-a"));
Connection::open(&path)
.unwrap()
.execute(
"UPDATE trusted_runtime_records SET revoked_at = '2' WHERE runtime_id = 'shared'",
[],
)
.unwrap();
let store = SqliteWorkspaceStore::open(&path).unwrap();
let binding = store
.get_workspace_runtime_binding("workspace-a", "shared")
.unwrap()
.unwrap();
assert_eq!(binding.revoked_at.as_deref(), Some("2"));
assert!(binding.public_key.is_some());
assert!(
binding
@@ -6988,6 +7012,24 @@ mod tests {
|row| row.get(0),
)?;
assert_eq!(jti_workspace, "workspace-a");
let workspace_foreign_keys: i64 = conn.query_row(
"SELECT COUNT(*) FROM pragma_foreign_key_list('workspace_runtime_bindings') WHERE \"table\" = 'workspaces' AND \"from\" = 'workspace_id'",
[],
|row| row.get(0),
)?;
assert_eq!(workspace_foreign_keys, 1);
let unique_indexes: i64 = conn.query_row(
"SELECT COUNT(*) FROM pragma_index_list('workspace_runtime_bindings') WHERE \"unique\" = 1",
[],
|row| row.get(0),
)?;
assert!(unique_indexes >= 2);
let lookup_index: i64 = conn.query_row(
"SELECT COUNT(*) FROM sqlite_master WHERE type = 'index' AND name = 'idx_workspace_runtime_bindings_workspace'",
[],
|row| row.get(0),
)?;
assert_eq!(lookup_index, 1);
let violations: i64 = conn.query_row(
"SELECT COUNT(*) FROM pragma_foreign_key_check",
[],
@@ -7033,7 +7075,9 @@ mod tests {
#[test]
fn runtime_binding_identity_and_trust_uniqueness_are_workspace_scoped() {
let store = SqliteWorkspaceStore::in_memory().unwrap();
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("server.db");
let store = SqliteWorkspaceStore::open(&path).unwrap();
store
.with_conn(|conn| {
conn.execute_batch(
@@ -7097,6 +7141,89 @@ mod tests {
.len(),
1
);
assert!(
store
.revoke_workspace_runtime_binding("workspace-a", "shared", "2")
.unwrap()
);
drop(store);
let reopened = SqliteWorkspaceStore::open(&path).unwrap();
assert!(
reopened
.list_workspace_runtime_bindings("workspace-a", false)
.unwrap()
.is_empty()
);
assert_eq!(
reopened
.list_workspace_runtime_bindings("workspace-b", false)
.unwrap()
.len(),
1
);
assert!(
reopened
.get_workspace_runtime_binding("workspace-a", "shared")
.unwrap()
.unwrap()
.revoked_at
.is_some()
);
}
#[test]
fn embedded_runtime_binding_can_explicitly_rotate_restart_identity() {
let store = SqliteWorkspaceStore::in_memory().unwrap();
store
.with_conn(|conn| {
conn.execute_batch(
r#"
INSERT INTO accounts(account_id, kind, handle, display_name, created_at, updated_at)
VALUES ('owner', 'user', 'owner', 'Owner', '1', '1');
INSERT INTO workspaces(workspace_id, owner_account_id, display_name, state, created_at, updated_at)
VALUES ('workspace-a', 'owner', 'Workspace A', 'active', '1', '1');
"#,
)?;
Ok(())
})
.unwrap();
let first =
worker_runtime::auth::RuntimeIdentityMaterial::generate("embedded-first").unwrap();
let second =
worker_runtime::auth::RuntimeIdentityMaterial::generate("embedded-second").unwrap();
let binding = |public_key: String| WorkspaceRuntimeBinding {
workspace_id: "workspace-a".to_string(),
runtime_id: crate::hosts::EMBEDDED_RUNTIME_ID.to_string(),
display_name: "Embedded Runtime".to_string(),
base_url: "in-process://embedded".to_string(),
public_key: Some(public_key),
public_key_fingerprint: None,
created_at: "1".to_string(),
updated_at: "1".to_string(),
revoked_at: None,
};
store
.upsert_workspace_runtime_binding(binding(first.public_key.clone()), false)
.unwrap();
assert!(matches!(
store.upsert_workspace_runtime_binding(binding(second.public_key.clone()), false),
Err(Error::RuntimeBindingConflict(_))
));
assert_eq!(
store
.upsert_workspace_runtime_binding(binding(second.public_key.clone()), true)
.unwrap(),
WorkspaceRuntimeBindingUpsert::Replaced
);
let persisted = store
.get_workspace_runtime_binding("workspace-a", crate::hosts::EMBEDDED_RUNTIME_ID)
.unwrap()
.unwrap();
assert_eq!(
persisted.public_key.as_deref(),
Some(second.public_key.as_str())
);
}
#[test]