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

server: fence Worker control operation identity

Browse Source
This commit is contained in:
Hare
2026-08-17 00:38:48 +09:00
parent 8be2cfd2a3
commit 97828f8bd5
2 changed files with 434 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/server.rs
+218 -30
View File
@@ -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<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>,
@@ -5986,7 +5994,11 @@ 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)
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:{}",
@@ -5996,10 +6008,9 @@ async fn spawn_known_worker(
.collect::<String>()
);
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 {
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::<String>()
);
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;
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,
)?;
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 {
crates/workspace-server/src/store.rs
+207 -9
View File
@@ -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<String>,
}
#[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<String>,
pub created_at: String,
pub completed_at: Option<String>,
}
#[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<bool>;
fn reserve_worker_control_delegation_operation(
&self,
record: &WorkerControlDelegationOperationRecord,
) -> Result<WorkerControlDelegationOperationRecord>;
fn complete_worker_control_delegation_operation(
&self,
workspace_id: &str,
source_controller: &RuntimeWorkerRef,
operation_id: &str,
delegated_grant_id: &str,
completed_at: &str,
) -> Result<WorkerControlDelegationOperationRecord>;
fn get_ticket_assignment_operation(
&self,
@@ -2549,6 +2578,95 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
})
}
fn reserve_worker_control_delegation_operation(
&self,
record: &WorkerControlDelegationOperationRecord,
) -> Result<WorkerControlDelegationOperationRecord> {
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<WorkerControlDelegationOperationRecord> {
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<WorkerControlDelegationOperationRecord> {
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<Option<WorkerControlDelegationOperationRecord>> {
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(),