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

fix: preserve authenticated submit source

Browse Source
This commit is contained in:
Hare
2026-09-06 02:56:51 +09:00
parent dea5bd581d
commit 8e4b7deaa4
5 changed files with 267 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+7
View File
@@ -32,6 +32,13 @@ fn is_false(value: &bool) -> bool {
// Method (Client → Worker via Unix Socket)
// ---------------------------------------------------------------------------
/// Trusted Server → Runtime transport header carrying the authenticated
/// browser Account identity for one Worker protocol connection.
///
/// Runtime accepts this only after its normal HTTP authentication succeeds;
/// serialized [`Method`] payloads cannot set authenticated source identity.
pub const AUTHENTICATED_ACCOUNT_ID_HEADER: &str = "x-yoi-authenticated-account-id";
/// Trusted source identity attached by an authenticated transport boundary.
///
/// Public clients cannot select this value directly. Runtime/Backend adapters
crates/worker-runtime/src/http_server.rs
+83 -9
View File
@@ -33,7 +33,7 @@ use axum::extract::rejection::{JsonRejection, QueryRejection};
#[cfg(feature = "ws-server")]
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
use axum::extract::{DefaultBodyLimit, Extension, Path, Query, State};
use axum::http::{Method, Request, StatusCode, header};
use axum::http::{HeaderMap, Method, Request, StatusCode, header};
use axum::middleware::{self, Next};
use axum::response::{IntoResponse, Response};
use axum::routing::{delete, get, post};
@@ -1184,10 +1184,12 @@ async fn worker_protocol_ws(
auth: Option<Extension<RuntimeAuthContext>>,
Path(worker_id): Path<String>,
Query(query): Query<RuntimeWorkerEventsWsQuery>,
headers: HeaderMap,
ws: WebSocketUpgrade,
) -> Result<Response, RuntimeHttpRestError> {
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
let scope = auth_workspace_scope(&state, auth.as_ref())?;
let input_source = authenticated_protocol_input_source(&headers)?;
match scope.as_ref() {
Some(scope) => state
.runtime
@@ -1198,22 +1200,60 @@ async fn worker_protocol_ws(
.map_err(RuntimeHttpRestError::runtime)?;
Ok(ws
.on_upgrade(move |socket| {
worker_protocol_ws_session(state.runtime, scope, worker_ref, query, socket)
worker_protocol_ws_session(
state.runtime,
scope,
worker_ref,
query,
input_source,
socket,
)
})
.into_response())
}
#[cfg(feature = "ws-server")]
fn authorize_runtime_protocol_method(method: protocol::Method) -> protocol::Method {
fn authenticated_protocol_input_source(
headers: &HeaderMap,
) -> Result<Option<protocol::AuthenticatedInputSource>, RuntimeHttpRestError> {
let Some(value) = headers.get(protocol::AUTHENTICATED_ACCOUNT_ID_HEADER) else {
return Ok(None);
};
let account_id = value.to_str().map_err(|_| {
RuntimeHttpRestError::new(
StatusCode::BAD_REQUEST,
"authenticated_input_source_invalid",
"authenticated Worker input source is invalid",
)
})?;
if account_id.trim().is_empty() || account_id.len() > 128 {
return Err(RuntimeHttpRestError::new(
StatusCode::BAD_REQUEST,
"authenticated_input_source_invalid",
"authenticated Worker input source is invalid",
));
}
Ok(Some(protocol::AuthenticatedInputSource::Account {
account_id: account_id.to_owned(),
}))
}
#[cfg(feature = "ws-server")]
fn authorize_runtime_protocol_method(
method: protocol::Method,
transport_source: Option<&protocol::AuthenticatedInputSource>,
) -> protocol::Method {
match method {
protocol::Method::SubmitTracked {
submission_request_id,
input,
..
} => protocol::Method::SubmitTracked {
source: protocol::AuthenticatedInputSource::Backend {
source: transport_source.cloned().unwrap_or_else(|| {
protocol::AuthenticatedInputSource::Backend {
operation_id: submission_request_id.clone(),
},
}
}),
submission_request_id,
input,
},
@@ -1223,9 +1263,11 @@ fn authorize_runtime_protocol_method(method: protocol::Method) -> protocol::Meth
auto_run,
..
} => protocol::Method::NotifyTracked {
source: protocol::AuthenticatedInputSource::Backend {
source: transport_source.cloned().unwrap_or_else(|| {
protocol::AuthenticatedInputSource::Backend {
operation_id: notification_request_id.clone(),
},
}
}),
notification_request_id,
message,
auto_run,
@@ -1240,6 +1282,7 @@ async fn worker_protocol_ws_session(
scope: Option<RuntimeWorkspaceScope>,
worker_ref: WorkerRef,
query: RuntimeWorkerEventsWsQuery,
input_source: Option<protocol::AuthenticatedInputSource>,
mut socket: WebSocket,
) {
let mut cursor = match query.cursor.as_deref() {
@@ -1322,7 +1365,8 @@ async fn worker_protocol_ws_session(
match inbound {
Some(Ok(WsMessage::Text(text))) => match decode_method(&text) {
Ok(method) => {
let method = authorize_runtime_protocol_method(method);
let method =
authorize_runtime_protocol_method(method, input_source.as_ref());
let result = match scope.as_ref() {
Some(scope) => {
runtime.send_protocol_method_scoped(scope, &worker_ref, method)
@@ -2135,7 +2179,7 @@ mod tests {
}
));
assert!(matches!(
authorize_runtime_protocol_method(decoded),
authorize_runtime_protocol_method(decoded, None),
protocol::Method::SubmitTracked {
source: protocol::AuthenticatedInputSource::Backend { operation_id },
..
@@ -2143,6 +2187,36 @@ mod tests {
));
}
#[test]
fn runtime_protocol_uses_transport_authenticated_account_source() {
let mut headers = HeaderMap::new();
headers.insert(
protocol::AUTHENTICATED_ACCOUNT_ID_HEADER,
"account-1".parse().unwrap(),
);
let source = authenticated_protocol_input_source(&headers)
.unwrap()
.expect("account source header must resolve");
let wire = serde_json::to_string(&protocol::Method::NotifyTracked {
notification_request_id: "notification-1".into(),
message: "hello".into(),
auto_run: true,
source: protocol::AuthenticatedInputSource::Account {
account_id: "forged".into(),
},
})
.unwrap();
let decoded: protocol::Method = serde_json::from_str(&wire).unwrap();
assert!(matches!(
authorize_runtime_protocol_method(decoded, Some(&source)),
protocol::Method::NotifyTracked {
source: protocol::AuthenticatedInputSource::Account { account_id },
..
} if account_id == "account-1"
));
}
#[test]
fn attachment_routes_require_worker_input_permission() {
assert_eq!(
crates/worker/src/worker.rs
+109 -4
View File
@@ -1290,10 +1290,31 @@ where
pending: &PendingSubmission,
) -> Result<(), PendingSubmissionError> {
let session_id = self.writer.state.location().session_id;
let mut pinned = Vec::new();
for reference in submission_uploaded_file_refs(&pending.input) {
if pinned.iter().any(|existing: &protocol::UploadedFileRef| {
existing.artifact_id == reference.artifact_id
}) {
continue;
}
if let Err(pin_error) =
self.writer
.store
.pin_uploaded_file(session_id, reference, &pending.submission_id)?;
.pin_uploaded_file(session_id, reference, &pending.submission_id)
{
let mut rollback_error = None;
for acquired in pinned.iter().rev() {
if let Err(error) = self.writer.store.release_uploaded_file_pin(
session_id,
&acquired.artifact_id,
&pending.submission_id,
) {
rollback_error.get_or_insert(error);
}
}
return Err(rollback_error.unwrap_or(pin_error).into());
}
pinned.push(reference.clone());
}
Ok(())
}
@@ -1303,14 +1324,28 @@ where
pending: &PendingSubmission,
) -> Result<(), PendingSubmissionError> {
let session_id = self.writer.state.location().session_id;
let mut released = Vec::new();
let mut first_error = None;
for reference in submission_uploaded_file_refs(&pending.input) {
self.writer.store.release_uploaded_file_pin(
if released
.iter()
.any(|artifact_id: &String| artifact_id == &reference.artifact_id)
{
continue;
}
if let Err(error) = self.writer.store.release_uploaded_file_pin(
session_id,
&reference.artifact_id,
&pending.submission_id,
)?;
) {
first_error.get_or_insert(error);
}
released.push(reference.artifact_id.clone());
}
match first_error {
Some(error) => Err(error.into()),
None => Ok(()),
}
Ok(())
}
#[cfg(test)]
@@ -9668,6 +9703,76 @@ mod build_summary_prompt_tests {
assert_eq!(state.pending[1].provenance, account_b);
}
#[test]
fn rejected_submission_rolls_back_uploaded_file_pins_acquired_before_conflict() {
let temp = tempfile::tempdir().unwrap();
let handle = PendingSubmissionHandle::for_test(temp.path());
let session_id = handle.writer.state.session_id();
let limits = session_store::UploadedFileLimits {
max_file_bytes: 1024,
max_session_bytes: 2048,
};
let first = handle
.writer
.store
.write_uploaded_file(session_id, "first.txt", "text/plain", b"first", limits)
.unwrap();
let second = handle
.writer
.store
.write_uploaded_file(session_id, "second.txt", "text/plain", b"second", limits)
.unwrap();
handle
.writer
.store
.pin_uploaded_file(session_id, &second, "other-submission")
.unwrap();
assert!(
handle
.accept(
"request-partial-pin".into(),
vec![
Segment::UploadedFile {
file: first.clone(),
},
Segment::UploadedFile {
file: second.clone(),
},
],
false,
)
.is_err()
);
assert!(handle.snapshot().submissions.is_empty());
assert!(
handle
.writer
.store
.delete_uploaded_file(session_id, &first.artifact_id)
.unwrap()
);
assert!(matches!(
handle
.writer
.store
.delete_uploaded_file(session_id, &second.artifact_id),
Err(StoreError::ArtifactAlreadyCommitted)
));
handle
.writer
.store
.release_uploaded_file_pin(session_id, &second.artifact_id, "other-submission")
.unwrap();
assert!(
handle
.writer
.store
.delete_uploaded_file(session_id, &second.artifact_id)
.unwrap()
);
}
#[test]
fn queued_submission_pins_uploaded_file_until_cancelled() {
let temp = tempfile::tempdir().unwrap();
crates/workspace-server/src/server.rs
+57 -2
View File
@@ -9251,7 +9251,7 @@ async fn scoped_capture_worker_observation_session(
return Err(ApiError::from(Error::UnknownWorker { worker: target }));
}
let mut connection = connect_workspace_worker_protocol(&api, &target).await?;
let mut connection = connect_workspace_worker_protocol(&api, &target, None).await?;
let event = tokio::time::timeout(std::time::Duration::from_secs(10), connection.events.recv())
.await
.map_err(|_| {
@@ -13955,6 +13955,7 @@ pub(crate) struct WorkspaceWorkerProtocolConnection {
pub(crate) async fn connect_workspace_worker_protocol(
api: &WorkspaceApi,
worker: &RuntimeWorkerRef,
input_source: Option<&protocol::AuthenticatedInputSource>,
) -> Result<WorkspaceWorkerProtocolConnection> {
let source = match api.observation_proxy.source(worker) {
Ok(source) => source,
@@ -13971,15 +13972,39 @@ pub(crate) async fn connect_workspace_worker_protocol(
}
};
match source {
RuntimeObservationSource::RemoteWs(config) => connect_remote_worker_protocol(config).await,
RuntimeObservationSource::RemoteWs(config) => {
connect_remote_worker_protocol(config, input_source).await
}
RuntimeObservationSource::Embedded(source) => {
connect_embedded_worker_protocol(source).await
}
}
}
fn insert_authenticated_input_source_header(
headers: &mut HeaderMap,
input_source: Option<&protocol::AuthenticatedInputSource>,
) -> Result<()> {
let Some(input_source) = input_source else {
return Ok(());
};
let protocol::AuthenticatedInputSource::Account { account_id } = input_source else {
return Err(Error::Config(
"remote Worker protocol transport supports only Account input source".into(),
));
};
headers.insert(
protocol::AUTHENTICATED_ACCOUNT_ID_HEADER,
account_id.parse().map_err(|error| {
Error::Config(format!("invalid authenticated Account identity: {error}"))
})?,
);
Ok(())
}
async fn connect_remote_worker_protocol(
config: RuntimeObservationSourceConfig,
input_source: Option<&protocol::AuthenticatedInputSource>,
) -> Result<WorkspaceWorkerProtocolConnection> {
let mut request = config
.endpoint
@@ -13994,6 +14019,7 @@ async fn connect_remote_worker_protocol(
})?,
);
}
insert_authenticated_input_source_header(request.headers_mut(), input_source)?;
let (socket, _) =
connect_async(request)
.await
@@ -14117,6 +14143,16 @@ async fn remote_worker_protocol_ws_session(
}
}
}
if let Err(error) =
insert_authenticated_input_source_header(request.headers_mut(), Some(&input_source))
{
let mut socket = socket;
let event = protocol_error_event(format!(
"failed to build authenticated Account identity header: {error}"
));
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
let (upstream, _) = match connect_async(request).await {
Ok(connection) => connection,
@@ -16678,6 +16714,25 @@ mod tests {
assert!(authorize_browser_worker_method(method, &source).is_err());
}
#[test]
fn remote_worker_protocol_header_preserves_authenticated_account_source() {
let mut headers = HeaderMap::new();
insert_authenticated_input_source_header(
&mut headers,
Some(&protocol::AuthenticatedInputSource::Account {
account_id: "account-1".into(),
}),
)
.unwrap();
assert_eq!(
headers
.get(protocol::AUTHENTICATED_ACCOUNT_ID_HEADER)
.unwrap(),
"account-1"
);
}
#[test]
fn merge_request_http_paths_observe_refs_through_runtime_provider_authority() {
let source = include_str!("server.rs");
crates/workspace-server/src/workspace_subscription.rs
+7 -1
View File
@@ -91,7 +91,13 @@ pub(crate) async fn serve_workspace_subscription(
runtime_id: Some(runtime_id),
} => {
let worker = RuntimeWorkerRef::new(&runtime_id, worker_id.as_str());
match connect_workspace_worker_protocol(&api, &worker).await {
match connect_workspace_worker_protocol(
&api,
&worker,
Some(&input_source),
)
.await
{
Ok(connection) => {
let methods = connection.methods.clone();
let task = tokio::spawn(run_worker_protocol(