diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index e7708b27..f2bbd16e 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -1126,15 +1126,19 @@ pub enum RuntimeRegistryUnregisterResult { }, } +type RuntimeBindingGate = Arc bool + Send + Sync>; + #[derive(Clone)] pub struct RuntimeRegistry { runtimes: Arc>>>, + runtime_binding_gate: Arc>>, } impl RuntimeRegistry { pub fn new(runtimes: Vec>) -> 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(&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(&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, 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 = diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index f674cc9f..2eb438ab 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -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)] diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 37bc6f9b..6d04b945 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -544,6 +544,11 @@ pub trait ControlPlaneStore: Send + Sync { &self, record: &WorkspaceBootstrapRecord, ) -> Result; + fn workspace_runtime_binding_is_active( + &self, + workspace_id: &str, + runtime_id: &str, + ) -> Result; 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 { + 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]