fix: fence revoked runtime bindings
This commit is contained in:
@@ -1126,15 +1126,19 @@ pub enum RuntimeRegistryUnregisterResult {
|
|||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type RuntimeBindingGate = Arc<dyn Fn(&str) -> bool + Send + Sync>;
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
pub struct RuntimeRegistry {
|
pub struct RuntimeRegistry {
|
||||||
runtimes: Arc<RwLock<Vec<Arc<dyn WorkspaceWorkerRuntime>>>>,
|
runtimes: Arc<RwLock<Vec<Arc<dyn WorkspaceWorkerRuntime>>>>,
|
||||||
|
runtime_binding_gate: Arc<RwLock<Option<RuntimeBindingGate>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl RuntimeRegistry {
|
impl RuntimeRegistry {
|
||||||
pub fn new(runtimes: Vec<Arc<dyn WorkspaceWorkerRuntime>>) -> Self {
|
pub fn new(runtimes: Vec<Arc<dyn WorkspaceWorkerRuntime>>) -> Self {
|
||||||
Self {
|
Self {
|
||||||
runtimes: Arc::new(RwLock::new(runtimes)),
|
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)])
|
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)
|
pub fn register<R>(&self, runtime: R)
|
||||||
where
|
where
|
||||||
R: WorkspaceWorkerRuntime + 'static,
|
R: WorkspaceWorkerRuntime + 'static,
|
||||||
@@ -1795,13 +1820,19 @@ impl RuntimeRegistry {
|
|||||||
self.runtimes
|
self.runtimes
|
||||||
.read()
|
.read()
|
||||||
.expect("runtime registry lock poisoned")
|
.expect("runtime registry lock poisoned")
|
||||||
.clone()
|
.iter()
|
||||||
|
.filter(|runtime| self.runtime_binding_is_active(runtime.runtime_id()))
|
||||||
|
.cloned()
|
||||||
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn runtime(
|
fn runtime(
|
||||||
&self,
|
&self,
|
||||||
runtime_id: &str,
|
runtime_id: &str,
|
||||||
) -> Result<Arc<dyn WorkspaceWorkerRuntime>, RuntimeRegistryError> {
|
) -> Result<Arc<dyn WorkspaceWorkerRuntime>, RuntimeRegistryError> {
|
||||||
|
if !self.runtime_binding_is_active(runtime_id) {
|
||||||
|
return Err(RuntimeRegistryError::UnknownRuntime(runtime_id.to_string()));
|
||||||
|
}
|
||||||
self.runtimes
|
self.runtimes
|
||||||
.read()
|
.read()
|
||||||
.expect("runtime registry lock poisoned")
|
.expect("runtime registry lock poisoned")
|
||||||
@@ -5165,6 +5196,34 @@ mod tests {
|
|||||||
assert_eq!(from_runtime_a.label, "worker from runtime a");
|
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]
|
#[test]
|
||||||
fn registry_broadcasts_workspace_prompt_projection_revisions() {
|
fn registry_broadcasts_workspace_prompt_projection_revisions() {
|
||||||
let runtime =
|
let runtime =
|
||||||
|
|||||||
@@ -1578,7 +1578,7 @@ impl WorkspaceApi {
|
|||||||
updated_at: config.workspace_created_at.clone(),
|
updated_at: config.workspace_created_at.clone(),
|
||||||
revoked_at: None,
|
revoked_at: None,
|
||||||
},
|
},
|
||||||
false,
|
true,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
let embedded_audience = format!("embedded:{}", config.workspace_id);
|
let embedded_audience = format!("embedded:{}", config.workspace_id);
|
||||||
@@ -1604,14 +1604,22 @@ impl WorkspaceApi {
|
|||||||
"failed to initialize embedded Worker backend: {err}"
|
"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,
|
config,
|
||||||
store,
|
store,
|
||||||
Arc::new(execution_backend),
|
Arc::new(execution_backend),
|
||||||
resource_broker,
|
resource_broker,
|
||||||
Some(worker_remove_dispatcher),
|
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)]
|
#[cfg(test)]
|
||||||
|
|||||||
@@ -544,6 +544,11 @@ pub trait ControlPlaneStore: Send + Sync {
|
|||||||
&self,
|
&self,
|
||||||
record: &WorkspaceBootstrapRecord,
|
record: &WorkspaceBootstrapRecord,
|
||||||
) -> Result<WorkspaceBootstrapResult>;
|
) -> Result<WorkspaceBootstrapResult>;
|
||||||
|
fn workspace_runtime_binding_is_active(
|
||||||
|
&self,
|
||||||
|
workspace_id: &str,
|
||||||
|
runtime_id: &str,
|
||||||
|
) -> Result<bool>;
|
||||||
async fn get_workspace_runtime_binding(
|
async fn get_workspace_runtime_binding(
|
||||||
&self,
|
&self,
|
||||||
workspace_id: &str,
|
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(
|
async fn get_workspace_runtime_binding(
|
||||||
&self,
|
&self,
|
||||||
workspace_id: &str,
|
workspace_id: &str,
|
||||||
@@ -6967,12 +6983,20 @@ mod tests {
|
|||||||
let temp = tempfile::tempdir().unwrap();
|
let temp = tempfile::tempdir().unwrap();
|
||||||
let path = temp.path().join("server.db");
|
let path = temp.path().join("server.db");
|
||||||
prepare_schema_v50(&path, Some("workspace-a"));
|
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 store = SqliteWorkspaceStore::open(&path).unwrap();
|
||||||
let binding = store
|
let binding = store
|
||||||
.get_workspace_runtime_binding("workspace-a", "shared")
|
.get_workspace_runtime_binding("workspace-a", "shared")
|
||||||
.unwrap()
|
.unwrap()
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
assert_eq!(binding.revoked_at.as_deref(), Some("2"));
|
||||||
assert!(binding.public_key.is_some());
|
assert!(binding.public_key.is_some());
|
||||||
assert!(
|
assert!(
|
||||||
binding
|
binding
|
||||||
@@ -6988,6 +7012,24 @@ mod tests {
|
|||||||
|row| row.get(0),
|
|row| row.get(0),
|
||||||
)?;
|
)?;
|
||||||
assert_eq!(jti_workspace, "workspace-a");
|
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(
|
let violations: i64 = conn.query_row(
|
||||||
"SELECT COUNT(*) FROM pragma_foreign_key_check",
|
"SELECT COUNT(*) FROM pragma_foreign_key_check",
|
||||||
[],
|
[],
|
||||||
@@ -7033,7 +7075,9 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn runtime_binding_identity_and_trust_uniqueness_are_workspace_scoped() {
|
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
|
store
|
||||||
.with_conn(|conn| {
|
.with_conn(|conn| {
|
||||||
conn.execute_batch(
|
conn.execute_batch(
|
||||||
@@ -7097,6 +7141,89 @@ mod tests {
|
|||||||
.len(),
|
.len(),
|
||||||
1
|
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]
|
#[test]
|
||||||
|
|||||||
Reference in New Issue
Block a user