diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 47b57ac8..952e9e89 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -2460,7 +2460,10 @@ impl Worker { .map(ToOwned::to_owned) }) .unwrap_or_else(|| "Workspace rejected Flow source resolution".to_string()); - return Err(WorkerError::FlowInput(message)); + return Err(WorkerError::FlowInput(format!( + "{message} (HTTP {})", + response.status + ))); } let source: flow::ResolvedFlowSource = serde_json::from_str(&response.body) .map_err(|error| WorkerError::FlowInput(format!("decode Flow source: {error}")))?; diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 9f6e5e2c..b61284bf 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -20387,6 +20387,81 @@ mod tests { )); } + #[tokio::test] + async fn runtime_request_proof_accepts_active_worker_create_reservation() { + let workspace = tempfile::tempdir().unwrap(); + let mut api = test_api(workspace.path()).await; + let identity = + worker_runtime::auth::RuntimeIdentityMaterial::generate("runtime-test").unwrap(); + configure_runtime_request_auth(&mut api, &identity, "runtime-test"); + let store = SqliteWorkspaceStore::open(&api.config.database_path).unwrap(); + let memory_settings = store + .get_workspace_memory_settings(TEST_WORKSPACE_ID) + .unwrap(); + let reserved = store + .reserve_worker_create( + TEST_WORKSPACE_ID, + "runtime-test", + "spawn-flow-race", + &"f".repeat(64), + &memory_settings, + ) + .unwrap(); + let worker_id = reserved.worker_id.to_string(); + let path = format!("/api/w/{TEST_WORKSPACE_ID}/flows/resolve"); + let signer = worker_runtime::auth::RuntimeRequestSourceSigner::from_identity(&identity); + let issue = || { + signer + .issue( + "server-test", + TEST_WORKSPACE_ID, + Some(worker_id.as_str()), + worker_runtime::auth::WORKSPACE_REQUEST_PERMISSION, + "POST", + &path, + b"{}", + i64::try_from(worker_runtime::auth::unix_now_seconds()).unwrap_or(i64::MAX), + 30, + ) + .unwrap() + }; + + let proof = issue(); + let source = crate::worker_source::verify_runtime_request_source_proof_with_store( + api.store.as_ref(), + &api.config, + &proof, + TEST_WORKSPACE_ID, + worker_runtime::auth::WORKSPACE_REQUEST_PERMISSION, + "POST", + &path, + &worker_runtime::auth::request_body_digest(b"{}"), + ) + .await + .expect("active create reservation should establish provisional Worker membership"); + assert_eq!(source.worker_id, Some(worker_id.clone())); + + store + .complete_worker_create_reservation(TEST_WORKSPACE_ID, reserved.worker_id) + .unwrap(); + let proof = issue(); + let result = crate::worker_source::verify_runtime_request_source_proof_with_store( + api.store.as_ref(), + &api.config, + &proof, + TEST_WORKSPACE_ID, + worker_runtime::auth::WORKSPACE_REQUEST_PERMISSION, + "POST", + &path, + &worker_runtime::auth::request_body_digest(b"{}"), + ) + .await; + assert!(matches!( + result, + Err(crate::worker_source::WorkerMutationSourceProofError::WorkerCatalogMembership) + )); + } + #[tokio::test] async fn runtime_request_proof_verifies_path_and_query_for_ticket_search() { let workspace = tempfile::tempdir().unwrap(); diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 92fd6dfb..58bc0aee 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -1002,6 +1002,11 @@ pub trait ControlPlaneStore: Send + Sync { workspace_id: &str, worker: &RuntimeWorkerRef, ) -> Result>; + fn has_active_worker_create_reservation( + &self, + workspace_id: &str, + worker: &RuntimeWorkerRef, + ) -> Result; fn list_worker_registry( &self, workspace_id: &str, @@ -3268,6 +3273,28 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } + fn has_active_worker_create_reservation( + &self, + workspace_id: &str, + worker: &RuntimeWorkerRef, + ) -> Result { + self.with_conn(|conn| { + conn.query_row( + r#"SELECT EXISTS( + SELECT 1 + FROM worker_create_reservations + WHERE workspace_id = ?1 + AND runtime_id = ?2 + AND worker_id = ?3 + AND state = 'reserved' + )"#, + params![workspace_id, worker.runtime_id, worker.worker_id], + |row| row.get::<_, bool>(0), + ) + .map_err(Error::from) + }) + } + fn list_worker_registry( &self, workspace_id: &str, @@ -10476,6 +10503,12 @@ INSERT INTO workdir_registry ( ); assert_eq!(reserved.memory_settings.settings_revision, 1); assert_eq!(reserved.memory_settings.language, "English"); + let reserved_worker = RuntimeWorkerRef::new("arcadia", reserved.worker_id.to_string()); + assert!( + store + .has_active_worker_create_reservation("workspace-a", &reserved_worker) + .unwrap() + ); let unchanged_memory_settings = store .update_workspace_memory_settings("workspace-a", 1, " English ") .unwrap(); @@ -10549,6 +10582,11 @@ INSERT INTO workdir_registry ( store .complete_worker_create_reservation("workspace-a", reserved.worker_id) .unwrap(); + assert!( + !store + .has_active_worker_create_reservation("workspace-a", &reserved_worker) + .unwrap() + ); let state: String = store .with_conn(|conn| { conn.query_row( diff --git a/crates/workspace-server/src/worker_source.rs b/crates/workspace-server/src/worker_source.rs index 2affde46..1144a132 100644 --- a/crates/workspace-server/src/worker_source.rs +++ b/crates/workspace-server/src/worker_source.rs @@ -106,7 +106,10 @@ pub async fn verify_runtime_request_source_proof_with_store( let member = store .get_worker_registry(workspace_id, &worker) .map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?; - if member.is_none() { + let reserved = store + .has_active_worker_create_reservation(workspace_id, &worker) + .map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?; + if member.is_none() && !reserved { return Err(WorkerMutationSourceProofError::WorkerCatalogMembership); } }