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

fix: await owned cold protocol signing and propagate authorization failure

Browse Source
This commit is contained in:
Hare
2026-10-12 08:58:34 +09:00
parent d3fcfb2977
commit c0223a0dea
3 changed files with 128 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/hosts.rs
+32 -21
View File
@@ -1297,8 +1297,8 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
fn observation_source(
&self,
_worker_id: &str,
) -> Option<crate::observation::RuntimeObservationSource> {
None
) -> Result<Option<crate::observation::RuntimeObservationSource>, RuntimeDiagnostic> {
Ok(None)
}
fn send_input(&self, worker_id: &str, _request: WorkerInputRequest) -> WorkerInputResult {
@@ -2236,6 +2236,11 @@ impl RuntimeRegistry {
let runtime = self.runtime(runtime_id)?;
runtime
.observation_source(worker_id)
.map_err(|diagnostic| RuntimeRegistryError::RuntimeOperationFailed {
runtime_id: runtime_id.to_owned(),
code: diagnostic.code,
message: diagnostic.message,
})?
.ok_or_else(|| RuntimeRegistryError::UnknownWorker {
worker: worker.clone(),
})
@@ -3203,17 +3208,21 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
fn observation_source(
&self,
worker_id: &str,
) -> Option<crate::observation::RuntimeObservationSource> {
let worker_ref = self.worker_ref(worker_id)?;
) -> Result<Option<crate::observation::RuntimeObservationSource>, RuntimeDiagnostic> {
let Some(worker_ref) = self.worker_ref(worker_id) else {
return Ok(None);
};
if self.runtime.worker_detail(&worker_ref).is_err() {
return None;
return Ok(None);
}
Some(crate::observation::RuntimeObservationSource::embedded(
crate::observation::EmbeddedRuntimeObservationSource {
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id),
runtime: self.runtime.clone(),
worker_ref,
},
Ok(Some(
crate::observation::RuntimeObservationSource::embedded(
crate::observation::EmbeddedRuntimeObservationSource {
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id),
runtime: self.runtime.clone(),
worker_ref,
},
),
))
}
@@ -5324,20 +5333,22 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
fn observation_source(
&self,
worker_id: &str,
) -> Option<crate::observation::RuntimeObservationSource> {
) -> Result<Option<crate::observation::RuntimeObservationSource>, RuntimeDiagnostic> {
let path = format!("/v1/workers/{worker_id}/protocol/ws");
let bearer_token = match &self.workspace_authorization {
Some(authorization) => authorization
.issue("GET", &path, "workers:protocol", Some(worker_id), &[])
.ok(),
Some(authorization) => {
Some(authorization.issue("GET", &path, "workers:protocol", Some(worker_id), &[])?)
}
None => self.bearer_token.clone(),
};
Some(crate::observation::RuntimeObservationSource::remote_ws(
crate::observation::RuntimeObservationSourceConfig {
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id),
endpoint: self.ws_endpoint(worker_id),
bearer_token,
},
Ok(Some(
crate::observation::RuntimeObservationSource::remote_ws(
crate::observation::RuntimeObservationSourceConfig {
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id),
endpoint: self.ws_endpoint(worker_id),
bearer_token,
},
),
))
}
crates/workspace-server/src/server.rs
+32 -14
View File
@@ -31468,20 +31468,7 @@ pub(crate) async fn connect_workspace_worker_protocol(
worker: &RuntimeWorkerRef,
input_source: Option<&protocol::AuthenticatedInputSource>,
) -> Result<WorkspaceWorkerProtocolConnection> {
let source = match api.observation_proxy.source(worker) {
Ok(source) => source,
Err(ObservationProxyError::WorkerNotFound(_)) => api
.runtime
.observation_source(worker)
.map_err(|error| error.into_error())?,
Err(error) => {
return Err(Error::RuntimeOperationFailed {
runtime_id: worker.runtime_id.clone(),
code: error.code().to_string(),
message: error.message().to_string(),
});
}
};
let source = workspace_worker_protocol_source(api, worker).await?;
match source {
RuntimeObservationSource::RemoteWs(config) => {
connect_remote_worker_protocol(config, input_source).await
@@ -31494,6 +31481,37 @@ pub(crate) async fn connect_workspace_worker_protocol(
}
}
// Cold source discovery includes live authority/signing reads. Keep the whole
// synchronous callback on its admitted DB owner, not the Tokio WS task.
async fn workspace_worker_protocol_source(
api: &WorkspaceApi,
worker: &RuntimeWorkerRef,
) -> Result<RuntimeObservationSource> {
let source = match api.observation_proxy.source(worker) {
Ok(source) => source,
Err(ObservationProxyError::WorkerNotFound(_)) => {
let runtime = api.runtime.clone();
let worker = worker.clone();
api.store
.db_executor()
.run(move || {
runtime
.observation_source(&worker)
.map_err(|error| error.into_error())
})
.await?
}
Err(error) => {
return Err(Error::RuntimeOperationFailed {
runtime_id: worker.runtime_id.clone(),
code: error.code().to_string(),
message: error.message().to_string(),
});
}
};
Ok(source)
}
fn insert_authenticated_input_source_header(
headers: &mut HeaderMap,
input_source: Option<&protocol::AuthenticatedInputSource>,
crates/workspace-server/src/server_db_execution_tests.rs
+64 -1
View File
@@ -848,7 +848,70 @@ fn install_remote_fixture(
binding
}
#[tokio::test(flavor = "current_thread")]
async fn cold_protocol_source_signing_waits_off_tokio_and_never_falls_back_on_overload_or_revocation()
{
let fixture = Fixture::new().await;
let binding = install_remote_fixture(&fixture, "protocol-runtime", "http://127.0.0.1:9", true);
let worker = RuntimeWorkerRef::new(&binding.runtime_id, "source-worker");
let (mut busy, waiting) = BusyObservation::install_store(fixture.store.clone());
let mut lock = WriteLock::acquire(fixture.api.config.database_path.clone());
let writer = fixture.store.clone();
fixture
.executor
.dispatch(move || {
writer
.mark_api_token_used("missing-token", "2026-01-01T00:00:00Z")
.unwrap()
})
.unwrap();
observed(waiting).await.unwrap(); // SQL holds the actual shared connection mutex.
let mut source = Box::pin(workspace_worker_protocol_source(&fixture.api, &worker));
assert!(futures::poll!(&mut source).is_pending());
runtime_progress_while_locked(&lock).await;
assert!(futures::poll!(&mut source).is_pending());
assert!(
matches!(
workspace_worker_protocol_source(&fixture.api, &worker).await,
Err(Error::DbCapacityExceeded)
),
"overload must not return an unsigned or legacy-bearer source"
);
lock.release();
busy.resume();
let RuntimeObservationSource::RemoteWs(config) = observed(source).await.unwrap() else {
panic!("remote source expected")
};
let first = config
.bearer_token
.expect("fresh Workspace capability is required");
assert_ne!(first, "must-not-fall-back");
let RuntimeObservationSource::RemoteWs(second) =
workspace_worker_protocol_source(&fixture.api, &worker)
.await
.unwrap()
else {
panic!("remote source expected")
};
assert_ne!(
first,
second.bearer_token.unwrap(),
"cold lookup must sign a fresh request, not cache authentication success"
);
fixture
.store
.revoke_workspace_runtime_binding(WORKSPACE, &binding.runtime_id, "2026-01-02T00:00:00Z")
.unwrap();
let error = workspace_worker_protocol_source(&fixture.api, &worker)
.await
.err()
.expect("revoked authority cannot produce any source");
assert!(
matches!(error, Error::RuntimeOperationFailed { ref code, .. } if code == "workspace_runtime_authorization_stale"),
"{error:?}"
);
fixture.executor.drain().await.unwrap();
}
// Loopback fake at the Session/Worker Runtime HTTP boundary. It owns one reply,
// released by an explicit event; Drop unblocks accept/read and joins on failure.