From 97828f8bd58bfc2227e36db53e0370d5d2ae5b5a Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 17 Aug 2026 00:38:48 +0900 Subject: [PATCH] server: fence Worker control operation identity --- crates/workspace-server/src/server.rs | 266 ++++++++++++++++++++++---- crates/workspace-server/src/store.rs | 216 ++++++++++++++++++++- 2 files changed, 434 insertions(+), 48 deletions(-) diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 5e856c87..570508da 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -103,8 +103,9 @@ use crate::skills; use crate::store::{ AccountRecord, ApiTokenRecord, AuthChallengeRecord, BrowserSessionRecord, ControlPlaneStore, DeviceLoginFlowRecord, FlowSourceRecord, PasskeyCredentialRecord, RepositoryRecord, - TicketWorkerAssignmentRecord, UserRecord, WorkdirRegistryRecord, WorkerControlGrantRecord, - WorkerRegistryRecord, WorkerWorkdirLinkRecord, WorkspaceRecord, + TicketWorkerAssignmentRecord, UserRecord, WorkdirRegistryRecord, + WorkerControlDelegationOperationRecord, WorkerControlGrantRecord, WorkerRegistryRecord, + WorkerWorkdirLinkRecord, WorkspaceRecord, }; use crate::{Error, Result}; use worker_runtime::catalog::{ @@ -5966,6 +5967,13 @@ async fn list_known_workers( })) } +fn scoped_worker_control_operation_id(controller: &RuntimeWorkerRef, operation_id: &str) -> String { + format!( + "worker-control:{}:{}:{operation_id}", + controller.runtime_id, controller.worker_id + ) +} + async fn spawn_known_worker( State(api): State, AxumPath(path): AxumPath, @@ -5986,8 +5994,12 @@ async fn spawn_known_worker( .filter(|value| !value.is_empty()) .ok_or_else(|| Error::InvalidInput("control_operation_id is required".to_string()))? .to_string(); - let fingerprint_input = serde_json::to_vec(&request) - .map_err(|error| Error::InvalidInput(format!("invalid Worker spawn input: {error}")))?; + let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id); + let fingerprint_input = serde_json::to_vec(&serde_json::json!({ + "controller": &controller, + "request": &request, + })) + .map_err(|error| Error::InvalidInput(format!("invalid Worker spawn input: {error}")))?; let input_fingerprint = format!( "sha256:{}", Sha256::digest(&fingerprint_input) @@ -5996,10 +6008,9 @@ async fn spawn_known_worker( .collect::() ); request.resolved_control_operation = Some(WorkerControlOperation { - operation_id: operation_id.clone(), + operation_id: scoped_worker_control_operation_id(&controller, &operation_id), input_fingerprint, }); - let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id); let response = create_workspace_worker(State(api.clone()), headers, Json(request)).await?; if let Err(error) = api .store @@ -6123,51 +6134,84 @@ async fn delegate_worker_control_grant( "target_controller must differ from the current Worker".to_string(), ))); } - api.store - .get_worker_registry(&path.workspace_id, &request.target_controller)? - .ok_or_else(|| Error::UnknownWorker { - worker: request.target_controller.clone(), - })?; let operation_id = request.operation_id.trim(); if operation_id.is_empty() || operation_id.len() > 200 { return Err(ApiError::from(Error::InvalidInput( "operation_id must contain 1..=200 bytes".to_string(), ))); } - let expected_origin = format!("worker_control_{permission}:{}", grant.grant_id); - if let Some(existing) = api.store.get_worker_control_grant_by_operation( - &path.workspace_id, - &request.target_controller, - operation_id, - )? { - if existing.subject == grant.subject && existing.origin == expected_origin { - if transfer && grant.revoked_at.is_none() { - let lock = worker_control_lock(&api, &grant.grant_id); - let _guard = lock.lock().await; - let now = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); - if api - .store - .get_worker_control_grant(&path.workspace_id, &grant.grant_id)? - .is_some_and(|current| current.revoked_at.is_none()) - { - api.store.revoke_worker_control_grant( - &path.workspace_id, - &grant.grant_id, - &now, - )?; - } + let operation_input = serde_json::json!({ + "source_controller": &controller, + "source_grant_id": &grant.grant_id, + "action": permission, + "target_controller": &request.target_controller, + "subject": &grant.subject, + "permissions": &grant.permissions, + }); + let operation_bytes = serde_json::to_vec(&operation_input).map_err(|error| { + Error::InvalidInput(format!("invalid Worker delegation input: {error}")) + })?; + let input_fingerprint = format!( + "sha256:{}", + Sha256::digest(&operation_bytes) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + ); + let now = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let operation = api.store.reserve_worker_control_delegation_operation( + &WorkerControlDelegationOperationRecord { + workspace_id: path.workspace_id.clone(), + source_controller: controller.clone(), + source_grant_id: grant.grant_id.clone(), + operation_id: operation_id.to_string(), + input_fingerprint, + delegated_grant_id: None, + created_at: now.clone(), + completed_at: None, + }, + )?; + if let Some(delegated_grant_id) = operation.delegated_grant_id.as_deref() { + let delegated = api + .store + .get_worker_control_grant(&path.workspace_id, delegated_grant_id)? + .ok_or_else(|| { + Error::Store("completed Worker delegation references a missing grant".to_string()) + })?; + if transfer && grant.revoked_at.is_none() { + let lock = worker_control_lock(&api, &grant.grant_id); + let _guard = lock.lock().await; + if api + .store + .get_worker_control_grant(&path.workspace_id, &grant.grant_id)? + .is_some_and(|current| current.revoked_at.is_none()) + { + api.store + .revoke_worker_control_grant(&path.workspace_id, &grant.grant_id, &now)?; } - return Ok(Json(existing)); } - return Err(ApiError::from(Error::InvalidInput( - "operation_id was already used with different grant input".to_string(), - ))); + return Ok(Json(delegated)); } if grant.revoked_at.is_some() { return Err(ApiError::from(Error::UnknownWorker { worker: grant.subject, })); } + api.store + .get_active_worker_control_grant( + &path.workspace_id, + &controller, + &request.target_controller, + )? + .ok_or_else(|| Error::UnknownWorker { + worker: request.target_controller.clone(), + })?; + api.store + .get_worker_registry(&path.workspace_id, &request.target_controller)? + .ok_or_else(|| Error::UnknownWorker { + worker: request.target_controller.clone(), + })?; + let lock = worker_control_lock(&api, &grant.grant_id); let _guard = lock.lock().await; let current = api @@ -6184,7 +6228,20 @@ async fn delegate_worker_control_grant( .ok_or_else(|| Error::UnknownWorker { worker: grant.subject.clone(), })?; - let now = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + api.store + .get_active_worker_control_grant( + &path.workspace_id, + &controller, + &request.target_controller, + )? + .ok_or_else(|| Error::UnknownWorker { + worker: request.target_controller.clone(), + })?; + let delegated_operation_id = format!( + "worker-control-delegate:{}:{}:{}:{}:{}", + controller.runtime_id, controller.worker_id, current.grant_id, permission, operation_id + ); + let expected_origin = format!("worker_control_{permission}:{}", current.grant_id); let delegated = api .store .create_worker_control_grant(&WorkerControlGrantRecord { @@ -6195,7 +6252,7 @@ async fn delegate_worker_control_grant( relation: if transfer { "transferred" } else { "shared" }.to_string(), origin: expected_origin, permissions: current.permissions.clone(), - operation_id: operation_id.to_string(), + operation_id: delegated_operation_id, created_at: now.clone(), revoked_at: None, })?; @@ -6208,6 +6265,13 @@ async fn delegate_worker_control_grant( worker: current.subject, })); } + api.store.complete_worker_control_delegation_operation( + &path.workspace_id, + &controller, + operation_id, + &delegated.grant_id, + &now, + )?; Ok(Json(delegated)) } @@ -13278,6 +13342,14 @@ mod tests { .await .unwrap(); let controller = controller_worker.worker_ref; + assert_ne!( + scoped_worker_control_operation_id(&controller, "same-operation"), + scoped_worker_control_operation_id( + &RuntimeWorkerRef::new(&controller.runtime_id, "different-controller"), + "same-operation", + ), + "Runtime idempotency keys are scoped to the authenticated controller" + ); let mut headers = HeaderMap::new(); headers.insert( "x-yoi-runtime-id", @@ -13570,6 +13642,122 @@ mod tests { .unwrap(); assert_eq!(shared.controller, target_controller); assert_eq!(shared.relation, "shared"); + + let alternate_target = RuntimeWorkerRef::new(EMBEDDED_WORKER_RUNTIME_ID, "1000"); + let now = now_registry_timestamp(); + api.store + .upsert_worker_registry(&WorkerRegistryRecord { + workspace_id: workspace_id.clone(), + worker: alternate_target.clone(), + display_name: "Alternate known target".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.clone(), + }) + .unwrap(); + let registry_only_target = RuntimeWorkerRef::new(EMBEDDED_WORKER_RUNTIME_ID, "1001"); + api.store + .upsert_worker_registry(&WorkerRegistryRecord { + workspace_id: workspace_id.clone(), + worker: registry_only_target.clone(), + display_name: "Registry-only target".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.clone(), + }) + .unwrap(); + let unknown_target = share_worker_control_grant( + State(api.clone()), + AxumPath(ScopedWorkerControlGrantPath { + workspace_id: workspace_id.clone(), + grant_id: "orchestrator-share-source".to_string(), + }), + observation_headers.clone(), + Json(DelegateWorkerControlGrantRequest { + target_controller: registry_only_target, + operation_id: "share-registry-only".to_string(), + }), + ) + .await + .unwrap_err(); + assert_eq!( + unknown_target.into_response().status(), + StatusCode::NOT_FOUND + ); + for (grant_id, operation_id, subject) in [ + ( + "orchestrator-knows-alternate", + "seed-known-alternate", + alternate_target.clone(), + ), + ( + "orchestrator-second-share-source", + "seed-second-share", + generic.worker_ref.clone(), + ), + ] { + api.store + .create_worker_control_grant(&WorkerControlGrantRecord { + workspace_id: workspace_id.clone(), + grant_id: grant_id.to_string(), + controller: dedicated.worker.clone(), + subject, + relation: "spawned".to_string(), + origin: "test-delegation-conflict".to_string(), + permissions: vec!["share".to_string()], + operation_id: operation_id.to_string(), + created_at: now.clone(), + revoked_at: None, + }) + .unwrap(); + } + let changed_target = share_worker_control_grant( + State(api.clone()), + AxumPath(ScopedWorkerControlGrantPath { + workspace_id: workspace_id.clone(), + grant_id: "orchestrator-share-source".to_string(), + }), + observation_headers.clone(), + Json(DelegateWorkerControlGrantRequest { + target_controller: alternate_target, + operation_id: "share-operation".to_string(), + }), + ) + .await + .unwrap_err(); + assert_eq!( + changed_target.into_response().status(), + StatusCode::BAD_REQUEST + ); + let changed_source_grant = share_worker_control_grant( + State(api.clone()), + AxumPath(ScopedWorkerControlGrantPath { + workspace_id: workspace_id.clone(), + grant_id: "orchestrator-second-share-source".to_string(), + }), + observation_headers.clone(), + Json(DelegateWorkerControlGrantRequest { + target_controller: generic.worker_ref.clone(), + operation_id: "share-operation".to_string(), + }), + ) + .await + .unwrap_err(); + assert_eq!( + changed_source_grant.into_response().status(), + StatusCode::BAD_REQUEST + ); + let Json(transferred) = transfer_worker_control_grant( State(api.clone()), AxumPath(ScopedWorkerControlGrantPath { diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index f9dd91cd..eb89c8f8 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -186,6 +186,11 @@ const MIGRATIONS: &[Migration] = &[ name: "create durable Runtime Worker control grants", apply: create_worker_control_grant_authority, }, + Migration { + version: 34, + name: "create Worker control delegation operation authority", + apply: create_worker_control_delegation_operation_authority, + }, ]; struct Migration { @@ -347,6 +352,18 @@ pub struct WorkerControlGrantRecord { pub revoked_at: Option, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct WorkerControlDelegationOperationRecord { + pub workspace_id: String, + pub source_controller: RuntimeWorkerRef, + pub source_grant_id: String, + pub operation_id: String, + pub input_fingerprint: String, + pub delegated_grant_id: Option, + pub created_at: String, + pub completed_at: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct TicketWorkerAssignmentRecord { pub workspace_id: String, @@ -787,6 +804,18 @@ pub trait ControlPlaneStore: Send + Sync { grant_id: &str, revoked_at: &str, ) -> Result; + fn reserve_worker_control_delegation_operation( + &self, + record: &WorkerControlDelegationOperationRecord, + ) -> Result; + fn complete_worker_control_delegation_operation( + &self, + workspace_id: &str, + source_controller: &RuntimeWorkerRef, + operation_id: &str, + delegated_grant_id: &str, + completed_at: &str, + ) -> Result; fn get_ticket_assignment_operation( &self, @@ -2549,6 +2578,95 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } + fn reserve_worker_control_delegation_operation( + &self, + record: &WorkerControlDelegationOperationRecord, + ) -> Result { + self.with_conn(|conn| { + conn.execute( + r#"INSERT INTO worker_control_delegation_operations ( + workspace_id, source_controller_runtime_id, source_controller_worker_id, + source_grant_id, operation_id, input_fingerprint, + delegated_grant_id, created_at, completed_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9) + ON CONFLICT ( + workspace_id, source_controller_runtime_id, + source_controller_worker_id, operation_id + ) DO NOTHING"#, + params![ + record.workspace_id, + record.source_controller.runtime_id, + record.source_controller.worker_id, + record.source_grant_id, + record.operation_id, + record.input_fingerprint, + record.delegated_grant_id, + record.created_at, + record.completed_at, + ], + )?; + let persisted = read_worker_control_delegation_operation_by_key( + conn, + &record.workspace_id, + &record.source_controller, + &record.operation_id, + )? + .ok_or_else(|| { + Error::Store("worker control delegation operation was not persisted".to_string()) + })?; + if persisted.source_grant_id != record.source_grant_id + || persisted.input_fingerprint != record.input_fingerprint + { + return Err(Error::InvalidInput(format!( + "worker control delegation operation `{}` was already used with different input", + record.operation_id + ))); + } + Ok(persisted) + }) + } + + fn complete_worker_control_delegation_operation( + &self, + workspace_id: &str, + source_controller: &RuntimeWorkerRef, + operation_id: &str, + delegated_grant_id: &str, + completed_at: &str, + ) -> Result { + self.with_conn(|conn| { + conn.execute( + r#"UPDATE worker_control_delegation_operations + SET delegated_grant_id = ?5, completed_at = ?6 + WHERE workspace_id = ?1 + AND source_controller_runtime_id = ?2 + AND source_controller_worker_id = ?3 + AND operation_id = ?4 + AND (delegated_grant_id IS NULL OR delegated_grant_id = ?5)"#, + params![ + workspace_id, + source_controller.runtime_id, + source_controller.worker_id, + operation_id, + delegated_grant_id, + completed_at, + ], + )?; + read_worker_control_delegation_operation_by_key( + conn, + workspace_id, + source_controller, + operation_id, + )? + .filter(|record| record.delegated_grant_id.as_deref() == Some(delegated_grant_id)) + .ok_or_else(|| { + Error::InvalidInput(format!( + "worker control delegation operation `{operation_id}` completed with a different grant" + )) + }) + }) + } + fn get_ticket_assignment_operation( &self, workspace_id: &str, @@ -3906,6 +4024,52 @@ fn read_worker_control_grant_by_operation( .map_err(Error::from) } +fn read_worker_control_delegation_operation_record( + row: &rusqlite::Row<'_>, +) -> rusqlite::Result { + Ok(WorkerControlDelegationOperationRecord { + workspace_id: row.get(0)?, + source_controller: RuntimeWorkerRef::new( + row.get::<_, String>(1)?, + row.get::<_, u64>(2)?.to_string(), + ), + source_grant_id: row.get(3)?, + operation_id: row.get(4)?, + input_fingerprint: row.get(5)?, + delegated_grant_id: row.get(6)?, + created_at: row.get(7)?, + completed_at: row.get(8)?, + }) +} + +fn read_worker_control_delegation_operation_by_key( + conn: &Connection, + workspace_id: &str, + source_controller: &RuntimeWorkerRef, + operation_id: &str, +) -> Result> { + conn.query_row( + r#"SELECT workspace_id, + source_controller_runtime_id, source_controller_worker_id, + source_grant_id, operation_id, input_fingerprint, + delegated_grant_id, created_at, completed_at + FROM worker_control_delegation_operations + WHERE workspace_id = ?1 + AND source_controller_runtime_id = ?2 + AND source_controller_worker_id = ?3 + AND operation_id = ?4"#, + params![ + workspace_id, + source_controller.runtime_id, + source_controller.worker_id, + operation_id, + ], + read_worker_control_delegation_operation_record, + ) + .optional() + .map_err(Error::from) +} + fn current_ticket_worker_assignment_select_sql() -> String { "SELECT a.workspace_id, a.ticket_id, a.assignment_id, a.runtime_id, a.worker_id, \ a.assigned_by, a.assigned_at \ @@ -4970,6 +5134,40 @@ fn create_worker_control_grant_authority(conn: &Connection) -> Result<()> { Ok(()) } +fn create_worker_control_delegation_operation_authority(conn: &Connection) -> Result<()> { + conn.execute_batch( + r#" + CREATE TABLE worker_control_delegation_operations ( + workspace_id TEXT NOT NULL, + source_controller_runtime_id TEXT NOT NULL, + source_controller_worker_id INTEGER NOT NULL, + source_grant_id TEXT NOT NULL, + operation_id TEXT NOT NULL, + input_fingerprint TEXT NOT NULL, + delegated_grant_id TEXT, + created_at TEXT NOT NULL, + completed_at TEXT, + PRIMARY KEY ( + workspace_id, + source_controller_runtime_id, + source_controller_worker_id, + operation_id + ), + FOREIGN KEY (workspace_id, source_controller_runtime_id, source_controller_worker_id) + REFERENCES worker_registry (workspace_id, runtime_id, runtime_worker_id) + ON DELETE CASCADE, + FOREIGN KEY (workspace_id, source_grant_id) + REFERENCES worker_control_grants (workspace_id, grant_id) + ON DELETE CASCADE, + FOREIGN KEY (workspace_id, delegated_grant_id) + REFERENCES worker_control_grants (workspace_id, grant_id) + ON DELETE SET NULL + ); + "#, + )?; + Ok(()) +} + fn create_worker_mutation_source_proof_replay_guard(conn: &Connection) -> Result<()> { conn.execute_batch( r#" @@ -5643,7 +5841,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 33); + assert_eq!(current_schema_version(&conn).unwrap(), 34); assert!(table_exists(&conn, "worker_workdir_attachment_reservations").unwrap()); } @@ -5676,7 +5874,7 @@ CREATE TABLE flow_events (event_id TEXT PRIMARY KEY); apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 33); + assert_eq!(current_schema_version(&conn).unwrap(), 34); assert!(table_exists(&conn, "flow_sources").unwrap()); assert!(table_exists(&conn, "flow_source_revisions").unwrap()); assert!(!table_exists(&conn, "flow_instances").unwrap()); @@ -5743,7 +5941,7 @@ INSERT INTO worker_workdir_attachment_reservations ( apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 33); + assert_eq!(current_schema_version(&conn).unwrap(), 34); let repositories_sql: String = conn .query_row( "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'repositories'", @@ -5923,7 +6121,7 @@ INSERT INTO workdir_registry ( let db = dir.path().join("control-plane.sqlite"); let store = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 33); + assert_eq!(store.schema_version().await.unwrap(), 34); assert!( !store .with_conn(|conn| table_exists(conn, "worker_workspace_credentials")) @@ -5940,7 +6138,7 @@ INSERT INTO workdir_registry ( store.upsert_workspace(&record).await.unwrap(); let reopened = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(reopened.schema_version().await.unwrap(), 33); + assert_eq!(reopened.schema_version().await.unwrap(), 34); assert_eq!( reopened.get_workspace("local-dev").await.unwrap(), Some(record) @@ -6487,7 +6685,7 @@ INSERT INTO workdir_registry ( .unwrap(); let store = SqliteWorkspaceStore::from_connection(conn).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 33); + assert_eq!(store.schema_version().await.unwrap(), 34); store .with_conn(|conn| { @@ -6676,7 +6874,7 @@ CREATE TABLE ticket_assignment_operations ( #[tokio::test] async fn repository_records_round_trip() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 33); + assert_eq!(store.schema_version().await.unwrap(), 34); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -6742,7 +6940,7 @@ CREATE TABLE ticket_assignment_operations ( #[tokio::test] async fn memory_authority_records_round_trip_and_close_staging() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 33); + assert_eq!(store.schema_version().await.unwrap(), 34); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -7133,7 +7331,7 @@ CREATE TABLE ticket_assignment_operations ( #[tokio::test] async fn account_and_login_records_round_trip() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 33); + assert_eq!(store.schema_version().await.unwrap(), 34); let now = "2026-07-22T00:00:00Z".to_string(); let account = AccountRecord { account_id: "acct-user-alice".to_string(),