fix: accept reserved workers during initial flow resolution
This commit is contained in:
@@ -2460,7 +2460,10 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
|
|||||||
.map(ToOwned::to_owned)
|
.map(ToOwned::to_owned)
|
||||||
})
|
})
|
||||||
.unwrap_or_else(|| "Workspace rejected Flow source resolution".to_string());
|
.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)
|
let source: flow::ResolvedFlowSource = serde_json::from_str(&response.body)
|
||||||
.map_err(|error| WorkerError::FlowInput(format!("decode Flow source: {error}")))?;
|
.map_err(|error| WorkerError::FlowInput(format!("decode Flow source: {error}")))?;
|
||||||
|
|||||||
@@ -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]
|
#[tokio::test]
|
||||||
async fn runtime_request_proof_verifies_path_and_query_for_ticket_search() {
|
async fn runtime_request_proof_verifies_path_and_query_for_ticket_search() {
|
||||||
let workspace = tempfile::tempdir().unwrap();
|
let workspace = tempfile::tempdir().unwrap();
|
||||||
|
|||||||
@@ -1002,6 +1002,11 @@ pub trait ControlPlaneStore: Send + Sync {
|
|||||||
workspace_id: &str,
|
workspace_id: &str,
|
||||||
worker: &RuntimeWorkerRef,
|
worker: &RuntimeWorkerRef,
|
||||||
) -> Result<Option<WorkerRegistryRecord>>;
|
) -> Result<Option<WorkerRegistryRecord>>;
|
||||||
|
fn has_active_worker_create_reservation(
|
||||||
|
&self,
|
||||||
|
workspace_id: &str,
|
||||||
|
worker: &RuntimeWorkerRef,
|
||||||
|
) -> Result<bool>;
|
||||||
fn list_worker_registry(
|
fn list_worker_registry(
|
||||||
&self,
|
&self,
|
||||||
workspace_id: &str,
|
workspace_id: &str,
|
||||||
@@ -3268,6 +3273,28 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn has_active_worker_create_reservation(
|
||||||
|
&self,
|
||||||
|
workspace_id: &str,
|
||||||
|
worker: &RuntimeWorkerRef,
|
||||||
|
) -> Result<bool> {
|
||||||
|
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(
|
fn list_worker_registry(
|
||||||
&self,
|
&self,
|
||||||
workspace_id: &str,
|
workspace_id: &str,
|
||||||
@@ -10476,6 +10503,12 @@ INSERT INTO workdir_registry (
|
|||||||
);
|
);
|
||||||
assert_eq!(reserved.memory_settings.settings_revision, 1);
|
assert_eq!(reserved.memory_settings.settings_revision, 1);
|
||||||
assert_eq!(reserved.memory_settings.language, "English");
|
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
|
let unchanged_memory_settings = store
|
||||||
.update_workspace_memory_settings("workspace-a", 1, " English ")
|
.update_workspace_memory_settings("workspace-a", 1, " English ")
|
||||||
.unwrap();
|
.unwrap();
|
||||||
@@ -10549,6 +10582,11 @@ INSERT INTO workdir_registry (
|
|||||||
store
|
store
|
||||||
.complete_worker_create_reservation("workspace-a", reserved.worker_id)
|
.complete_worker_create_reservation("workspace-a", reserved.worker_id)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
assert!(
|
||||||
|
!store
|
||||||
|
.has_active_worker_create_reservation("workspace-a", &reserved_worker)
|
||||||
|
.unwrap()
|
||||||
|
);
|
||||||
let state: String = store
|
let state: String = store
|
||||||
.with_conn(|conn| {
|
.with_conn(|conn| {
|
||||||
conn.query_row(
|
conn.query_row(
|
||||||
|
|||||||
@@ -106,7 +106,10 @@ pub async fn verify_runtime_request_source_proof_with_store(
|
|||||||
let member = store
|
let member = store
|
||||||
.get_worker_registry(workspace_id, &worker)
|
.get_worker_registry(workspace_id, &worker)
|
||||||
.map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?;
|
.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);
|
return Err(WorkerMutationSourceProofError::WorkerCatalogMembership);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user