diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index de5e8549..c28ebae9 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -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 diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index 5c849c14..196b0c09 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -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>, Path(worker_id): Path, Query(query): Query, + headers: HeaderMap, ws: WebSocketUpgrade, ) -> Result { 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, 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 { - operation_id: submission_request_id.clone(), - }, + 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 { - operation_id: notification_request_id.clone(), - }, + 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, worker_ref: WorkerRef, query: RuntimeWorkerEventsWsQuery, + input_source: Option, 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!( diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 2aecd80d..0b79c4f0 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -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) { - self.writer - .store - .pin_uploaded_file(session_id, reference, &pending.submission_id)?; + 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) + { + 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(); diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index d53b239d..50733cdb 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -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 { 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 { 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"); diff --git a/crates/workspace-server/src/workspace_subscription.rs b/crates/workspace-server/src/workspace_subscription.rs index d72017f1..018a9622 100644 --- a/crates/workspace-server/src/workspace_subscription.rs +++ b/crates/workspace-server/src/workspace_subscription.rs @@ -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(