diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 11b1b3dfb0..c5840e4d4c 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -348,8 +348,20 @@ impl InProcessAppServerClient { // this loop can keep draining runtime events // while the request is blocked on client input. tokio::spawn(async move { - let result = request_sender.request(*request).await; - let _ = response_tx.send(result); + // Device ceremonies belong to the waiting UI. Preserve + // its cancellation through this buffering task. + let cancellable = matches!(*request, + ClientRequest::UserVerificationStatus { .. } + | ClientRequest::UserVerificationEnroll { .. } + | ClientRequest::UserVerificationDelete { .. } + | ClientRequest::UserVerificationVerify { .. }); + let mut response_tx = response_tx; + tokio::select! { + _ = response_tx.closed(), if cancellable => {} + result = request_sender.request(*request) => { + let _ = response_tx.send(result); + } + } }); } Some(ClientCommand::Notify { diff --git a/codex-rs/app-server/src/connection_rpc_gate.rs b/codex-rs/app-server/src/connection_rpc_gate.rs index fb2aedd352..dc1dc5307b 100644 --- a/codex-rs/app-server/src/connection_rpc_gate.rs +++ b/codex-rs/app-server/src/connection_rpc_gate.rs @@ -38,6 +38,10 @@ impl ConnectionRpcGate { drop(token); } + pub(crate) fn is_closed(&self) -> bool { + self.tasks.is_closed() + } + pub(crate) async fn close(&self) { let mut accepting = self.accepting.lock().await; *accepting = false; diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index bc012ce7c9..65a7d4e1db 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -188,6 +188,7 @@ enum InProcessClientMessage { Request { request: Box, response_tx: oneshot::Sender, + cancellation: tokio_util::sync::CancellationToken, }, Notification { notification: ClientNotification, @@ -206,7 +207,7 @@ enum InProcessClientMessage { } enum ProcessorCommand { - Request(Box), + Request(Box, tokio_util::sync::CancellationToken), Notification(ClientNotification), } @@ -218,9 +219,12 @@ pub struct InProcessClientSender { impl InProcessClientSender { pub async fn request(&self, request: ClientRequest) -> IoResult { let (response_tx, response_rx) = oneshot::channel(); + let cancellation = tokio_util::sync::CancellationToken::new(); + let _cancel_on_drop = cancellation.clone().drop_guard(); self.try_send_client_message(InProcessClientMessage::Request { request: Box::new(request), response_tx, + cancellation, })?; response_rx.await.map_err(|err| { IoError::new( @@ -482,6 +486,9 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult IoResult { match command { - Some(ProcessorCommand::Request(request)) => { + Some(ProcessorCommand::Request(request, cancellation)) => { let was_initialized = session.initialized(); processor .process_client_request( @@ -505,6 +514,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult IoResult { match message { - Some(InProcessClientMessage::Request { request, response_tx }) => { + Some(InProcessClientMessage::Request { request, response_tx, cancellation }) => { let request = *request; let request_id = request.id().clone(); match pending_request_responses.entry(request_id.clone()) { @@ -591,7 +601,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult {} Err(mpsc::error::TrySendError::Full(_)) => { if let Some(response_tx) = diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index c4ec80c9bf..7bb96f6ad0 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -938,6 +938,9 @@ pub async fn run_main_with_transport_options( state_db: state_db.clone(), config_warnings, session_source, + user_verification: Arc::new(crate::user_verification::Service::new(Arc::clone( + &auth_manager, + ))), auth_manager, installation_id, code_mode_session_provider, diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index da84311484..3557370eee 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -135,6 +135,7 @@ fn reject_removed_permission_profile(request: &JSONRPCRequest) -> Result<(), JSO } pub(crate) struct MessageProcessor { + user_verification: Arc, outgoing: Arc, models_refresh_worker: ModelsRefreshWorker, turn_cost_worker: Option, @@ -168,6 +169,7 @@ pub(crate) struct MessageProcessor { #[derive(Debug)] pub(crate) struct ConnectionSessionState { pub(crate) rpc_gate: Arc, + pub(crate) origin: crate::transport::ConnectionOrigin, pub(crate) mcp_event_streams: McpEventStreams, initialized: OnceLock, } @@ -182,15 +184,10 @@ pub(crate) struct InitializedConnectionSessionState { pub(crate) client_mcp_extensions: ClientMcpExtensions, } -impl Default for ConnectionSessionState { - fn default() -> Self { - Self::new() - } -} - impl ConnectionSessionState { - pub(crate) fn new() -> Self { + pub(crate) fn new(origin: crate::transport::ConnectionOrigin) -> Self { Self { + origin, rpc_gate: Arc::new(ConnectionRpcGate::new()), mcp_event_streams: McpEventStreams::default(), initialized: OnceLock::new(), @@ -256,6 +253,7 @@ pub(crate) struct MessageProcessorArgs { pub(crate) config_warnings: Vec, pub(crate) session_source: SessionSource, pub(crate) auth_manager: Arc, + pub(crate) user_verification: Arc, pub(crate) installation_id: String, pub(crate) code_mode_session_provider: Option>, pub(crate) rpc_transport: AppServerRpcTransport, @@ -281,6 +279,7 @@ impl MessageProcessor { config_warnings, session_source, auth_manager, + user_verification, installation_id, code_mode_session_provider, rpc_transport, @@ -448,6 +447,7 @@ impl MessageProcessor { Arc::clone(&config), config_warnings.clone(), rpc_transport, + Arc::clone(&user_verification), ); let marketplace_processor = MarketplaceRequestProcessor::new( Arc::clone(&config), @@ -569,6 +569,7 @@ impl MessageProcessor { ); Self { + user_verification, outgoing, models_refresh_worker, turn_cost_worker, @@ -671,6 +672,7 @@ impl MessageProcessor { request: ClientRequest, session: Arc, outbound_initialized: &AtomicBool, + cancellation: tokio_util::sync::CancellationToken, ) { let request_id = ConnectionRequestId { connection_id, @@ -678,8 +680,9 @@ impl MessageProcessor { }; let request_span = crate::app_server_tracing::typed_request_span(&request, connection_id, &session); - let request_context = + let mut request_context = RequestContext::new(request_id.clone(), request_span, /*parent_trace*/ None); + request_context.cancellation = cancellation; tracing::trace!( ?connection_id, request_id = ?request_id.request_id, @@ -992,11 +995,43 @@ impl MessageProcessor { ClientRequest::Initialize { .. } => { panic!("Initialize should be handled before initialized request dispatch"); } - ClientRequest::UserVerificationStatus { .. } + request @ (ClientRequest::UserVerificationStatus { .. } | ClientRequest::UserVerificationEnroll { .. } | ClientRequest::UserVerificationDelete { .. } - | ClientRequest::UserVerificationVerify { .. } => { - Err(crate::user_verification::unavailable()) + | ClientRequest::UserVerificationVerify { .. }) => { + return Box::pin(async { + // Keep verification temporaries out of unrelated RPCs' stack frames. + let operation = match request { + ClientRequest::UserVerificationStatus { .. } => { + crate::user_verification::Operation::Status + } + ClientRequest::UserVerificationEnroll { .. } => { + crate::user_verification::Operation::Enroll + } + ClientRequest::UserVerificationDelete { .. } => { + crate::user_verification::Operation::Delete + } + ClientRequest::UserVerificationVerify { params, .. } => { + crate::user_verification::Operation::Verify(params) + } + _ => unreachable!("matched a user-verification request"), + }; + let response = self + .user_verification + .handle( + operation, + Arc::clone(&session.rpc_gate), + request_context.cancellation.clone(), + session.origin, + ) + .await?; + let (payload, check) = response.into_parts(); + self.outgoing + .send_response_as_checked(request_id.clone(), payload, check) + .await; + Ok(()) + }) + .await; } ClientRequest::ServerDiagnostics { .. } => Ok(Some(read_server_diagnostics().into())), ClientRequest::ConfigRead { params, .. } => self diff --git a/codex-rs/app-server/src/message_processor_tracing_tests.rs b/codex-rs/app-server/src/message_processor_tracing_tests.rs index 82c1ad0127..260092a4e7 100644 --- a/codex-rs/app-server/src/message_processor_tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor_tracing_tests.rs @@ -129,7 +129,9 @@ impl TracingHarness { _codex_home: codex_home, processor, outgoing_rx, - session: Arc::new(ConnectionSessionState::new()), + session: Arc::new(ConnectionSessionState::new( + crate::transport::ConnectionOrigin::Stdio, + )), tracing, }; @@ -265,6 +267,9 @@ async fn build_test_processor( state_db: None, config_warnings: Vec::new(), session_source: SessionSource::VSCode, + user_verification: Arc::new(crate::user_verification::Service::new(Arc::clone( + &auth_manager, + ))), auth_manager, installation_id: "11111111-1111-4111-8111-111111111111".to_string(), code_mode_session_provider: None, diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index e6e0ae87a4..5c20dad5d7 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -61,6 +61,7 @@ pub(crate) struct ConnectionRequestId { #[derive(Clone)] pub(crate) struct RequestContext { request_id: ConnectionRequestId, + pub(crate) cancellation: tokio_util::sync::CancellationToken, span: Span, parent_trace: Option, _diagnostics_guard: Arc, @@ -74,6 +75,7 @@ impl RequestContext { ) -> Self { Self { request_id, + cancellation: tokio_util::sync::CancellationToken::new(), span, parent_trace, _diagnostics_guard: Arc::new(IN_FLIGHT_REQUESTS.track()), @@ -638,6 +640,42 @@ impl OutgoingMessageSender { .await; } + /// Revalidates a sensitive result after reserving queue capacity, with no + /// suspension between the identity check and handing off the response. + pub(crate) async fn send_response_as_checked( + &self, + request_id: ConnectionRequestId, + response: ClientResponsePayload, + check: impl FnOnce() -> std::result::Result<(), JSONRPCErrorError>, + ) { + let _context = self.take_request_context(&request_id).await; + let Ok(permit) = self.sender.reserve().await else { + return; + }; + let message = match check() { + Ok(()) => { + self.analytics_events_client.track_response( + request_id.connection_id.0, + request_id.request_id.clone(), + &response, + ); + OutgoingMessage::Response(OutgoingResponse { + id: request_id.request_id, + result: Box::new(response), + }) + } + Err(error) => OutgoingMessage::Error(OutgoingError { + id: request_id.request_id, + error, + }), + }; + permit.send(OutgoingEnvelope::ToConnection { + connection_id: request_id.connection_id, + message, + write_complete_tx: None, + }); + } + async fn send_response_as_inner( &self, request_id: ConnectionRequestId, diff --git a/codex-rs/app-server/src/request_processors/initialize_processor.rs b/codex-rs/app-server/src/request_processors/initialize_processor.rs index 442b71f335..d6c77f12eb 100644 --- a/codex-rs/app-server/src/request_processors/initialize_processor.rs +++ b/codex-rs/app-server/src/request_processors/initialize_processor.rs @@ -8,6 +8,8 @@ use codex_login::default_client::USER_AGENT_SUFFIX; use codex_login::default_client::get_codex_user_agent; use codex_login::default_client::set_default_client_residency_requirement; use codex_login::default_client::set_default_originator; +use codex_protocol::mcp::ClientMcpExtensions; +use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID; use super::*; use crate::message_processor::ConnectionSessionState; @@ -22,6 +24,7 @@ pub(crate) struct InitializeRequestProcessor { config: Arc, config_warnings: Arc>, rpc_transport: AppServerRpcTransport, + user_verification: Arc, } impl InitializeRequestProcessor { @@ -31,6 +34,7 @@ impl InitializeRequestProcessor { config: Arc, config_warnings: Vec, rpc_transport: AppServerRpcTransport, + user_verification: Arc, ) -> Self { Self { outgoing, @@ -38,6 +42,7 @@ impl InitializeRequestProcessor { config, config_warnings: Arc::new(config_warnings), rpc_transport, + user_verification, } } @@ -71,7 +76,7 @@ impl InitializeRequestProcessor { let experimental_api_enabled = capabilities.experimental_api; let request_attestation = capabilities.request_attestation; let extensions = capabilities.extensions.as_ref(); - let client_mcp_extensions = codex_mcp::client_mcp_extensions( + let mut client_mcp_extensions = codex_mcp::client_mcp_extensions( extensions, capabilities.mcp_server_openai_form_elicitation, ); @@ -90,6 +95,31 @@ impl InitializeRequestProcessor { "Invalid clientInfo.name: '{name}'. Must be a valid HTTP header value." ))); } + // The bundled TUI shares this build and implements the typed verification UI. + // Independently deployed UIs need their own rollout before receiving this mode. + let user_verification_enabled = experimental_api_enabled + && matches!( + session.origin, + crate::transport::ConnectionOrigin::InProcess + ) + && name == "codex-tui" + && tokio::task::spawn_blocking(self.user_verification.device_supported) + .await + .unwrap_or(false); + if user_verification_enabled { + let mut extensions = client_mcp_extensions + .iter() + .map(|(id, value)| (id.to_string(), value.clone())) + .collect::>(); + let settings = extensions + .entry(OPENAI_ELICITATION_EXTENSION_ID.to_string()) + .or_insert_with(|| serde_json::json!({})); + if !settings.is_object() { + *settings = serde_json::json!({}); + } + settings["userVerification"] = serde_json::json!({}); + client_mcp_extensions = ClientMcpExtensions::new(extensions); + } let originator = name.clone(); let user_agent_suffix = format!("{name}; {version}"); let mutates_global_identity = !NON_ORIGINATING_CLIENT_NAMES.contains(&name.as_str()); @@ -107,6 +137,11 @@ impl InitializeRequestProcessor { { return Err(invalid_request("Already initialized")); } + if user_verification_enabled { + self.outgoing + .enable_user_verification_connection(connection_id) + .await; + } if mutates_global_identity { // Only real client initialization may mutate process-global client metadata. diff --git a/codex-rs/app-server/src/transport.rs b/codex-rs/app-server/src/transport.rs index 3203065331..e94df28833 100644 --- a/codex-rs/app-server/src/transport.rs +++ b/codex-rs/app-server/src/transport.rs @@ -57,7 +57,7 @@ impl ConnectionState { outbound_initialized, outbound_experimental_api_enabled, outbound_opted_out_notification_methods, - session: Arc::new(ConnectionSessionState::new()), + session: Arc::new(ConnectionSessionState::new(origin)), } } } diff --git a/codex-rs/app-server/src/user_verification.rs b/codex-rs/app-server/src/user_verification.rs index 6cc5f9b36d..3f8112b899 100644 --- a/codex-rs/app-server/src/user_verification.rs +++ b/codex-rs/app-server/src/user_verification.rs @@ -1,7 +1,259 @@ -//! Local verification value conversion; provider dispatch follows in the next change. +//! Dispatches local verification onto one bounded native worker per service. +//! Request guards remain active through response delivery; registration is a later integration. + +use crate::connection_rpc_gate::ConnectionRpcGate; +use crate::transport::ConnectionOrigin; +use codex_app_server_protocol as rpc; +use codex_login::AuthManager; +use codex_user_verification as native; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Semaphore; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; -#[allow(dead_code)] #[path = "user_verification_adapter.rs"] mod adapter; +use adapter::error; +use adapter::native_error; pub(crate) use adapter::unavailable; +use adapter::unavailable_reason; +use adapter::validate; + +pub(crate) enum Operation { + Status, + Enroll, + Delete, + Verify(rpc::UserVerificationVerifyParams), +} + +/// Shared local provider dependencies and a bounded native worker slot. +/// The slot stays with the worker until it exits, including after RPC cancellation. +pub(crate) struct Service { + auth_manager: Arc, + provider: Arc< + dyn Fn(native::UserVerificationKeyNamespace) -> Arc + + Send + + Sync, + >, + platform_supported: bool, + pub(crate) device_supported: fn() -> bool, + worker: Arc, +} + +impl Service { + pub(crate) fn new(auth_manager: Arc) -> Self { + Self { + auth_manager, + provider: Arc::new(native::platform_provider), + platform_supported: native::platform_supported(), + // Host activation follows the bundled TUI integration. + device_supported: || false, + worker: Arc::new(Semaphore::new(/*permits*/ 1)), + } + } + + pub(crate) async fn handle( + &self, + operation: Operation, + connection: Arc, + cancellation: CancellationToken, + origin: ConnectionOrigin, + ) -> Result { + // A network peer must sign on its own device. Stdio belongs to the local + // parent process; in-process calls belong to the embedded application. + if matches!( + origin, + ConnectionOrigin::WebSocket | ConnectionOrigin::RemoteControl + ) && !matches!(operation, Operation::Status) + { + return Err(unavailable()); + } + let operation = match operation { + Operation::Verify(params) => NativeOperation::Verify(validate(params)?), + Operation::Status => NativeOperation::Status, + Operation::Enroll => NativeOperation::Enroll, + Operation::Delete => NativeOperation::Delete, + }; + let auth_manager = Arc::clone(&self.auth_manager); + let auth_changes = auth_manager.auth_change_receiver(); + let auth_revision = *auth_changes.borrow(); + let Some(identity) = identity(&auth_manager).filter(|_| self.platform_supported) else { + return match operation { + NativeOperation::Status => Ok(GuardedResponse { + payload: rpc::UserVerificationStatusResponse { + credential_id: None, + unavailable_reason: Some( + rpc::UserVerificationUnavailableReason::ProviderUnavailable, + ), + unavailable_message: Some( + "User verification is not available in this build or account.".into(), + ), + } + .into(), + guard: CancelOnDrop(native::UserVerificationRequestGuard::default()), + }), + NativeOperation::Enroll | NativeOperation::Delete | NativeOperation::Verify(_) => { + Err(unavailable()) + } + }; + }; + let provider = (self.provider)(native::UserVerificationKeyNamespace::new(&identity)); + let guard = native::UserVerificationRequestGuard::with_activity_check(move || { + !connection.is_closed() + && !cancellation.is_cancelled() + && *auth_changes.borrow() == auth_revision + && self::identity(&auth_manager).as_ref() == Some(&identity) + }); + run(operation, provider, guard, Arc::clone(&self.worker)).await + } +} + +pub(crate) struct GuardedResponse { + pub(crate) payload: rpc::ClientResponsePayload, + guard: CancelOnDrop, +} + +impl GuardedResponse { + pub(crate) fn into_parts( + self, + ) -> ( + rpc::ClientResponsePayload, + impl FnOnce() -> Result<(), rpc::JSONRPCErrorError>, + ) { + (self.payload, move || { + self.guard.0.check().map_err(native_error) + }) + } +} + +fn identity(auth_manager: &AuthManager) -> Option { + auth_manager.auth_cached()?.get_chatgpt_account_user_id() +} + +enum NativeOperation { + Status, + Enroll, + Delete, + Verify(native::UserVerificationRequest), +} + +/// Cancels blocking work when its RPC future is dropped or times out. +struct CancelOnDrop(native::UserVerificationRequestGuard); +impl Drop for CancelOnDrop { + fn drop(&mut self) { + self.0.cancel(); + } +} + +async fn run( + operation: NativeOperation, + provider: Arc, + guard: native::UserVerificationRequestGuard, + worker: Arc, +) -> Result { + let cancel = CancelOnDrop(guard.clone()); + let permit = worker.try_acquire_owned().map_err(|_| { + error( + rpc::UserVerificationErrorDetails::Failed { + reason: rpc::UserVerificationFailureReason::ProviderError, + }, + "Another user verification operation is still running.", + ) + })?; + let (result_tx, work) = oneshot::channel(); + // Native APIs may block even after cancellation. They must not occupy Tokio's + // blocking pool (which runtime shutdown joins), or create unbounded workers. + std::thread::Builder::new() + .name("user-verification".into()) + .spawn(move || { + let result = (|| { + guard.check().map_err(native_error)?; + let result: rpc::ClientResponsePayload = match operation { + NativeOperation::Status => { + let status = provider.status(&guard).map_err(native_error)?; + rpc::UserVerificationStatusResponse { + credential_id: status.credential.map(|key| key.credential_id), + unavailable_reason: status.unavailable_reason.map(unavailable_reason), + unavailable_message: status.unavailable_message, + } + .into() + } + NativeOperation::Enroll => { + let key = provider.ensure_key(&guard).map_err(native_error)?; + // TODO: start enrollment, sign proof of possession, then finish registration. + // This implementation establishes the local key only. + rpc::UserVerificationEnrollResponse { + credential_id: key.credential.credential_id, + } + .into() + } + NativeOperation::Delete => { + // TODO: retain credential references and revoke server registrations when + // backend enrollment is connected. Current enroll only creates local keys. + provider.delete(&guard).map_err(native_error)?; + rpc::UserVerificationDeleteResponse {}.into() + } + NativeOperation::Verify(request) => { + let proof = provider.verify(&request, &guard).map_err(native_error)?; + rpc::UserVerificationVerifyResponse { + proof: rpc::UserVerificationProof { + credential_id: proof.credential_id, + signature: proof.signature, + }, + } + .into() + } + }; + guard.check().map_err(native_error)?; + Ok(result) + })(); + drop(permit); + let _ = result_tx.send(result); + }) + .map_err(|_| { + error( + rpc::UserVerificationErrorDetails::Failed { + reason: rpc::UserVerificationFailureReason::ProviderError, + }, + "User verification could not start.", + ) + })?; + let result = match tokio::time::timeout(Duration::from_secs(/*secs*/ 120), work).await { + Ok(Ok(result)) => result, + Ok(Err(_)) => Err(error( + rpc::UserVerificationErrorDetails::Failed { + reason: rpc::UserVerificationFailureReason::ProviderError, + }, + "User verification could not complete.", + )), + Err(_) => Err(error( + rpc::UserVerificationErrorDetails::Failed { + reason: rpc::UserVerificationFailureReason::Timeout, + }, + "User verification timed out.", + )), + }; + // The UI may disconnect or change accounts after the blocking worker completes + // but before this task receives its result. + if result.is_ok() { + cancel.0.check().map_err(native_error)?; + } + result.map(|payload| GuardedResponse { + payload, + guard: cancel, + }) +} + +#[cfg(test)] +#[path = "user_verification_tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "user_verification_rpc_tests.rs"] +mod rpc_tests; + +#[cfg(test)] +#[path = "user_verification_test_support.rs"] +mod test_support; diff --git a/codex-rs/app-server/src/user_verification_rpc_tests.rs b/codex-rs/app-server/src/user_verification_rpc_tests.rs new file mode 100644 index 0000000000..591d77ae50 --- /dev/null +++ b/codex-rs/app-server/src/user_verification_rpc_tests.rs @@ -0,0 +1,188 @@ +//! Exercises cancellation and local-device security through the client JSON-RPC dispatcher. + +use super::test_support::Harness; +use super::test_support::write_auth; +use super::*; +use crate::outgoing_message::ConnectionId; +use crate::outgoing_message::OutgoingMessage; +use anyhow::Result; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::sync::atomic::Ordering; +use tokio::time::timeout; + +#[tokio::test] +async fn user_verification_rpc_auth_revision_discards_native_proof() -> Result<()> { + let mut h = Harness::new(ConnectionOrigin::Stdio, || true).await?; + h.initialize("test-local-ui", /*opt_in*/ true).await; + h.start_verify().await?; + write_auth(h.home.path(), "second")?; + h.auth.reload().await; + h.provider.released.store(/*val*/ true, Ordering::Release); + let OutgoingMessage::Error(error) = h.response().await else { + panic!("auth change must reject proof") + }; + assert_eq!( + error.error.data, + Some(json!({"type": "cancelled", "reason": "interrupted"})) + ); + h.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn user_verification_rpc_disconnect_cancels_worker_and_releases_slot() -> Result<()> { + let mut h = Harness::new(ConnectionOrigin::Stdio, || true).await?; + h.initialize("test-local-ui", /*opt_in*/ true).await; + h.start_verify().await?; + timeout( + Duration::from_secs(/*secs*/ 5), + h.processor.connection_closed(ConnectionId(1), &h.session), + ) + .await?; + assert!(matches!(h.response().await, OutgoingMessage::Error(_))); + assert_eq!(h.service.worker.available_permits(), 1); + h.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn user_verification_rpc_network_peers_cannot_use_server_keys() -> Result<()> { + for origin in [ConnectionOrigin::WebSocket, ConnectionOrigin::RemoteControl] { + let mut h = Harness::new(origin, || true).await?; + h.initialize("codex-tui", /*opt_in*/ true).await; + for method in [ + "userVerification/enroll", + "userVerification/delete", + "userVerification/verify", + ] { + let params = if method.ends_with("verify") { + json!({"challenge":"AQ", "title":"Approve", "description":""}) + } else { + json!({}) + }; + h.send(/*id*/ 1, method, params).await; + let OutgoingMessage::Error(error) = h.response().await else { + panic!("remote native operation must fail") + }; + assert_eq!( + error.error.data, + Some(json!({"type":"unavailable", "reason":"providerUnavailable"})) + ); + } + assert_eq!(h.provider.calls.load(Ordering::SeqCst), 0); + h.shutdown().await; + } + Ok(()) +} + +#[tokio::test] +async fn user_verification_rpc_dropped_caller_cancels_native_work() -> Result<()> { + let h = Harness::new(ConnectionOrigin::InProcess, || true).await?; + let (entered, waiting) = oneshot::channel(); + *h.provider.entered.lock().unwrap() = Some(entered); + let future = h.service.handle( + Operation::Verify(rpc::UserVerificationVerifyParams { + challenge: "AQ".into(), + title: "Approve".into(), + description: String::new(), + }), + Arc::clone(&h.session.rpc_gate), + CancellationToken::new(), + ConnectionOrigin::InProcess, + ); + { + let future = std::pin::pin!(future); + tokio::select! { result = future => panic!("native operation should block: {}", result.is_ok()), entered = waiting => entered? } + } + let permit = timeout(Duration::from_secs(/*secs*/ 5), h.service.worker.acquire()).await??; + drop(permit); + assert_eq!(h.provider.calls.load(Ordering::SeqCst), 1); + h.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn user_verification_rechecks_auth_after_response_queue_wait() -> Result<()> { + let mut h = Harness::new(ConnectionOrigin::Stdio, || true).await?; + h.provider.released.store(/*val*/ true, Ordering::Release); + let response = h + .service + .handle( + Operation::Verify(rpc::UserVerificationVerifyParams { + challenge: "AQ".into(), + title: "Approve".into(), + description: String::new(), + }), + Arc::clone(&h.session.rpc_gate), + CancellationToken::new(), + ConnectionOrigin::Stdio, + ) + .await + .map_err(|error| anyhow::anyhow!("{error:?}"))?; + for id in 0..16 { + h.outgoing + .send_response( + crate::outgoing_message::ConnectionRequestId { + connection_id: ConnectionId(1), + request_id: rpc::RequestId::Integer(id), + }, + rpc::UserVerificationDeleteResponse {}, + ) + .await; + } + let outgoing = Arc::clone(&h.outgoing); + let (payload, check) = response.into_parts(); + let sending = outgoing.send_response_as_checked( + crate::outgoing_message::ConnectionRequestId { + connection_id: ConnectionId(1), + request_id: rpc::RequestId::Integer(16), + }, + payload, + check, + ); + let mut sending = std::pin::pin!(sending); + assert!(matches!( + futures::poll!(&mut sending), + std::task::Poll::Pending + )); + write_auth(h.home.path(), "second")?; + h.auth.reload().await; + h.response().await; + sending.await; + for _ in 1..16 { + h.response().await; + } + let OutgoingMessage::Error(error) = h.response().await else { + panic!("queued response must recheck auth") + }; + assert_eq!( + error.error.data, + Some(json!({"type":"cancelled", "reason":"interrupted"})) + ); + h.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn user_verification_rejects_concurrent_native_workers() -> Result<()> { + let h = Harness::new(ConnectionOrigin::Stdio, || true).await?; + let permit = h.service.worker.clone().acquire_owned().await?; + let result = h + .service + .handle( + Operation::Status, + Arc::clone(&h.session.rpc_gate), + CancellationToken::new(), + ConnectionOrigin::Stdio, + ) + .await; + assert_eq!( + result.err().unwrap().message, + "Another user verification operation is still running." + ); + assert_eq!(h.provider.calls.load(Ordering::SeqCst), 0); + drop(permit); + h.shutdown().await; + Ok(()) +} diff --git a/codex-rs/app-server/src/user_verification_test_support.rs b/codex-rs/app-server/src/user_verification_test_support.rs new file mode 100644 index 0000000000..413be31d79 --- /dev/null +++ b/codex-rs/app-server/src/user_verification_test_support.rs @@ -0,0 +1,270 @@ +//! Dispatcher harness with a controllable native provider for lifecycle tests. +//! Public stdio and WebSocket framing are exercised in the integration suite. + +use super::*; +use crate::config_manager::ConfigManager; +use crate::message_processor::ConnectionSessionState; +use crate::message_processor::MessageProcessor; +use crate::message_processor::MessageProcessorArgs; +use crate::outgoing_message::ConnectionId; +use crate::outgoing_message::OutgoingEnvelope; +use crate::outgoing_message::OutgoingMessage; +use crate::outgoing_message::OutgoingMessageSender; +use crate::transport::AppServerTransport; +use anyhow::Result; +use app_test_support::ChatGptAuthFixture; +use app_test_support::write_chatgpt_auth; +use base64::Engine as _; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use codex_analytics::AnalyticsEventsClient; +use codex_analytics::AppServerRpcTransport; +use codex_arg0::Arg0DispatchPaths; +use codex_config::CloudConfigBundleLoader; +use codex_config::LoaderOverrides; +use codex_config::types::AuthCredentialsStoreMode; +use codex_core::config::ConfigBuilder; +use codex_exec_server::EnvironmentManager; +use codex_feedback::CodexFeedback; +use codex_protocol::protocol::SessionSource; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use tokio::sync::mpsc; +use tokio::time::timeout; + +#[derive(Default)] +pub(super) struct BlockingProvider { + pub(super) entered: Mutex>>, + pub(super) released: AtomicBool, + pub(super) calls: AtomicUsize, +} + +impl native::UserVerificationProvider for BlockingProvider { + fn status( + &self, + guard: &native::UserVerificationRequestGuard, + ) -> Result { + guard.check()?; + Ok(native::UserVerificationStatus { + credential: None, + unavailable_reason: Some(native::UserVerificationUnavailableReason::CredentialMissing), + unavailable_message: None, + }) + } + fn ensure_key( + &self, + _guard: &native::UserVerificationRequestGuard, + ) -> Result { + unreachable!("this test must not create keys") + } + fn delete( + &self, + _guard: &native::UserVerificationRequestGuard, + ) -> Result { + unreachable!("this test must not delete keys") + } + fn verify( + &self, + _request: &native::UserVerificationRequest, + guard: &native::UserVerificationRequestGuard, + ) -> Result { + self.calls.fetch_add(/*val*/ 1, Ordering::SeqCst); + if let Some(entered) = self.entered.lock().unwrap().take() { + let _ = entered.send(()); + } + let deadline = std::time::Instant::now() + Duration::from_secs(/*secs*/ 10); + while !self.released.load(Ordering::Acquire) { + guard.check()?; + assert!( + std::time::Instant::now() < deadline, + "test failed to release native worker" + ); + std::thread::sleep(Duration::from_millis(/*millis*/ 5)); + } + Ok(native::UserVerificationProof { + credential_id: "credential".into(), + signature: "signature".into(), + }) + } +} + +pub(super) struct Harness { + pub(super) processor: Arc, + pub(super) session: Arc, + pub(super) auth: Arc, + pub(super) service: Arc, + pub(super) provider: Arc, + pub(super) outgoing: Arc, + pub(super) messages: mpsc::Receiver, + pub(super) home: tempfile::TempDir, + transport: AppServerTransport, +} + +impl Harness { + pub(super) async fn new(origin: ConnectionOrigin, supported: fn() -> bool) -> Result { + let home = tempfile::tempdir()?; + write_auth(home.path(), "first")?; + let config = Arc::new( + ConfigBuilder::default() + .codex_home(home.path().into()) + .build() + .await?, + ); + let auth = AuthManager::shared_from_config( + config.as_ref(), + /*enable_codex_api_key_env*/ false, + ) + .await?; + let provider = Arc::new(BlockingProvider::default()); + let factory_provider = Arc::clone(&provider); + let service = Arc::new(Service { + auth_manager: Arc::clone(&auth), + provider: Arc::new(move |_| factory_provider.clone()), + platform_supported: true, + device_supported: supported, + worker: Arc::new(Semaphore::new(/*permits*/ 1)), + }); + let config_manager = ConfigManager::new( + home.path().into(), + Vec::new(), + LoaderOverrides::default(), + /*strict_config*/ false, + CloudConfigBundleLoader::default(), + Arg0DispatchPaths::default(), + Arc::new(codex_config::NoopThreadConfigLoader), + ); + let (sender, messages) = mpsc::channel(/*buffer*/ 16); + let outgoing = Arc::new(OutgoingMessageSender::new( + sender, + AnalyticsEventsClient::disabled(), + )); + let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs { + outgoing: Arc::clone(&outgoing), + analytics_events_client: AnalyticsEventsClient::disabled(), + arg0_paths: Arg0DispatchPaths::default(), + config, + config_manager, + environment_manager: Arc::new(EnvironmentManager::default_for_tests()), + feedback: CodexFeedback::new(), + log_db: None, + state_db: None, + config_warnings: Vec::new(), + session_source: SessionSource::Cli, + auth_manager: Arc::clone(&auth), + user_verification: Arc::clone(&service), + installation_id: "11111111-1111-4111-8111-111111111111".into(), + code_mode_session_provider: None, + rpc_transport: match origin { + ConnectionOrigin::InProcess => AppServerRpcTransport::InProcess, + ConnectionOrigin::Stdio | ConnectionOrigin::RemoteControl => { + AppServerRpcTransport::Stdio + } + ConnectionOrigin::WebSocket => AppServerRpcTransport::Websocket, + }, + remote_control_handle: None, + plugin_startup_tasks: None, + })); + Ok(Self { + processor, + session: Arc::new(ConnectionSessionState::new(origin)), + auth, + service, + provider, + outgoing, + messages, + home, + transport: match origin { + ConnectionOrigin::WebSocket => AppServerTransport::WebSocket { + bind_address: "127.0.0.1:0".parse()?, + }, + // Remote control can share a process whose primary listener is stdio. + // Keeping that case tests that authorization uses the connection's origin. + ConnectionOrigin::Stdio + | ConnectionOrigin::InProcess + | ConnectionOrigin::RemoteControl => AppServerTransport::Stdio, + }, + }) + } + + pub(super) async fn send(&self, id: i64, method: &str, params: serde_json::Value) { + self.processor + .process_request( + ConnectionId(1), + rpc::JSONRPCRequest { + id: rpc::RequestId::Integer(id), + method: method.into(), + params: Some(params), + trace: None, + }, + &self.transport, + Arc::clone(&self.session), + ) + .await; + } + + pub(super) async fn response(&mut self) -> OutgoingMessage { + let envelope = timeout(Duration::from_secs(/*secs*/ 10), self.messages.recv()) + .await + .expect("response deadline") + .expect("response channel"); + let OutgoingEnvelope::ToConnection { + connection_id, + message, + .. + } = envelope + else { + panic!("unexpected broadcast") + }; + assert_eq!(connection_id, ConnectionId(1)); + message + } + + pub(super) async fn initialize(&mut self, name: &str, opt_in: bool) { + self.send(/*id*/ 0, "initialize", json!({"clientInfo": {"name": name, "version": "1"}, "capabilities": {"experimentalApi": opt_in}})).await; + assert!(matches!( + self.response().await, + OutgoingMessage::Response(_) + )); + } + + pub(super) async fn start_verify(&self) -> Result<()> { + let (entered, waiting) = oneshot::channel(); + *self.provider.entered.lock().unwrap() = Some(entered); + self.send( + /*id*/ 1, + "userVerification/verify", + json!({"challenge": "AQID", "title": "Approve", "description": "Action"}), + ) + .await; + timeout(Duration::from_secs(/*secs*/ 10), waiting).await??; + Ok(()) + } + + pub(super) async fn shutdown(self) { + self.provider + .released + .store(/*val*/ true, Ordering::Release); + self.processor + .connection_closed(ConnectionId(1), &self.session) + .await; + self.processor.clear_runtime_references(); + self.processor.clear_all_thread_listeners().await; + self.processor.drain_background_tasks().await; + self.processor.shutdown_threads().await; + } +} + +pub(super) fn write_auth(home: &std::path::Path, token: &str) -> Result<()> { + let payload = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&json!({"jti": token, "https://api.openai.com/auth": {"chatgpt_account_id": "account-1", "chatgpt_account_user_id": "membership-1"}}))?); + write_chatgpt_auth( + home, + ChatGptAuthFixture::new(format!("header.{payload}.signature")) + .chatgpt_user_id("user-1") + .account_id("account-1") + .chatgpt_account_id("account-1"), + AuthCredentialsStoreMode::File, + ) +} diff --git a/codex-rs/app-server/src/user_verification_tests.rs b/codex-rs/app-server/src/user_verification_tests.rs new file mode 100644 index 0000000000..65970c750d --- /dev/null +++ b/codex-rs/app-server/src/user_verification_tests.rs @@ -0,0 +1,192 @@ +//! Verifies local key lifecycle, signing, and cancellation without transport concerns. + +use super::*; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +#[derive(Default)] +struct Provider { + key: Mutex>, + signed: Mutex>>, + deactivate_on_sign: Option>, +} + +impl native::UserVerificationProvider for Provider { + fn status( + &self, + _guard: &native::UserVerificationRequestGuard, + ) -> Result { + let credential = self.key.lock().unwrap().clone(); + Ok(native::UserVerificationStatus { + unavailable_reason: credential + .is_none() + .then_some(native::UserVerificationUnavailableReason::CredentialMissing), + unavailable_message: credential.is_none().then(|| "No key".into()), + credential, + }) + } + fn ensure_key( + &self, + guard: &native::UserVerificationRequestGuard, + ) -> Result { + guard.check()?; + let mut key = self.key.lock().unwrap(); + let created = key.is_none(); + let credential = key + .get_or_insert_with(|| native::UserVerificationKeyInfo { + credential_id: "credential".into(), + algorithm: "ecdsaP256Sha256X962".into(), + public_key: "public-key".into(), + }) + .clone(); + Ok(native::UserVerificationKeyCreation { + created, + credential, + }) + } + fn delete( + &self, + guard: &native::UserVerificationRequestGuard, + ) -> Result { + guard.check()?; + Ok(native::UserVerificationKeyDeletion { + deleted_credential_id: self.key.lock().unwrap().take().map(|key| key.credential_id), + }) + } + fn verify( + &self, + request: &native::UserVerificationRequest, + guard: &native::UserVerificationRequestGuard, + ) -> Result { + guard.check()?; + self.signed.lock().unwrap().push(request.challenge.clone()); + if let Some(active) = &self.deactivate_on_sign { + active.store(/*val*/ false, Ordering::Release); + } + Ok(native::UserVerificationProof { + credential_id: "credential".into(), + signature: "signature".into(), + }) + } +} + +#[tokio::test] +async fn local_enrollment_reuses_key_and_status_and_delete_have_no_signing_effects() { + let provider = Arc::new(Provider::default()); + for _ in 0..2 { + let response = run( + NativeOperation::Enroll, + provider.clone(), + native::UserVerificationRequestGuard::default(), + Arc::new(Semaphore::new(/*permits*/ 1)), + ) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(response.payload).unwrap(), + json!({"credentialId": "credential"}) + ); + } + let response = run( + NativeOperation::Status, + provider.clone(), + native::UserVerificationRequestGuard::default(), + Arc::new(Semaphore::new(/*permits*/ 1)), + ) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(response.payload).unwrap(), + json!({"credentialId": "credential", "unavailableReason": null, "unavailableMessage": null}) + ); + for _ in 0..2 { + let response = run( + NativeOperation::Delete, + provider.clone(), + native::UserVerificationRequestGuard::default(), + Arc::new(Semaphore::new(/*permits*/ 1)), + ) + .await + .unwrap(); + assert_eq!(serde_json::to_value(response.payload).unwrap(), json!({})); + } + assert_eq!(*provider.key.lock().unwrap(), None); + assert!(provider.signed.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn signing_uses_decoded_bytes_without_a_pending_elicitation() { + let provider = Arc::new(Provider::default()); + let request = validate(rpc::UserVerificationVerifyParams { + challenge: "AQID".into(), + title: "Approve".into(), + description: "Action".into(), + }) + .unwrap(); + let response = run( + NativeOperation::Verify(request), + provider.clone(), + native::UserVerificationRequestGuard::default(), + Arc::new(Semaphore::new(/*permits*/ 1)), + ) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(response.payload).unwrap(), + json!({"proof": {"credentialId": "credential", "signature": "signature"}}) + ); + assert_eq!(*provider.signed.lock().unwrap(), vec![vec![1, 2, 3]]); +} + +#[tokio::test] +async fn identity_change_during_native_prompt_discards_proof() { + let active = Arc::new(AtomicBool::new(/*v*/ true)); + let provider = Arc::new(Provider { + deactivate_on_sign: Some(active.clone()), + ..Default::default() + }); + let guard = native::UserVerificationRequestGuard::with_activity_check(move || { + active.load(Ordering::Acquire) + }); + let response = run( + NativeOperation::Verify(native::UserVerificationRequest { + challenge: vec![1], + title: "Approve".into(), + description: String::new(), + }), + provider, + guard, + Arc::new(Semaphore::new(/*permits*/ 1)), + ) + .await + .err() + .unwrap(); + assert_eq!( + response.data, + Some(json!({"type": "cancelled", "reason": "interrupted"})) + ); +} + +#[tokio::test] +async fn cancelled_operation_cannot_create_a_key() { + let provider = Arc::new(Provider::default()); + let guard = native::UserVerificationRequestGuard::default(); + guard.cancel(); + let response = run( + NativeOperation::Enroll, + provider.clone(), + guard, + Arc::new(Semaphore::new(/*permits*/ 1)), + ) + .await + .err() + .unwrap(); + assert_eq!( + response.data, + Some(json!({"type": "cancelled", "reason": "interrupted"})) + ); + assert_eq!(*provider.key.lock().unwrap(), None); +} diff --git a/codex-rs/app-server/tests/suite/v2/user_verification.rs b/codex-rs/app-server/tests/suite/v2/user_verification.rs index 1da2477f05..20a54a7fd8 100644 --- a/codex-rs/app-server/tests/suite/v2/user_verification.rs +++ b/codex-rs/app-server/tests/suite/v2/user_verification.rs @@ -1,3 +1,8 @@ +use super::connection_handling_websocket::connect_websocket; +use super::connection_handling_websocket::read_error_for_id; +use super::connection_handling_websocket::read_response_for_id; +use super::connection_handling_websocket::send_request; +use super::connection_handling_websocket::spawn_websocket_server; use anyhow::Result; use app_test_support::DEFAULT_CLIENT_NAME; use app_test_support::TestAppServer; @@ -50,6 +55,26 @@ async fn user_verification_methods_require_experimental_opt_in() -> Result<()> { Ok(()) } +#[tokio::test] +async fn user_verification_status_without_account_returns_local_readiness() -> Result<()> { + let mut app_server = TestAppServer::builder().build_initialized().await?; + let id = app_server + .send_raw_request("userVerification/status", Some(json!({}))) + .await?; + let response = app_server + .read_stream_until_response_message(RequestId::Integer(id)) + .await?; + assert_eq!( + response.result, + json!({ + "credentialId": null, + "unavailableReason": "providerUnavailable", + "unavailableMessage": "User verification is not available in this build or account." + }) + ); + Ok(()) +} + #[tokio::test] async fn user_verification_without_provider_returns_typed_unavailability() -> Result<()> { let mut app_server = TestAppServer::builder().build_initialized().await?; @@ -74,3 +99,96 @@ async fn user_verification_without_provider_returns_typed_unavailability() -> Re ); Ok(()) } + +#[tokio::test] +async fn user_verification_stdio_requires_initialization_then_validates_challenges() -> Result<()> { + let mut app_server = TestAppServer::builder().build().await?; + let params = json!({"challenge": "!", "title": "Approve", "description": ""}); + let id = app_server + .send_raw_request("userVerification/verify", Some(params.clone())) + .await?; + let response = app_server + .read_stream_until_error_message(RequestId::Integer(id)) + .await?; + assert_eq!( + response.error, + JSONRPCErrorError { + code: -32600, + message: "Not initialized".into(), + data: None, + } + ); + app_server.initialize().await?; + let id = app_server + .send_raw_request("userVerification/verify", Some(params)) + .await?; + let response = app_server + .read_stream_until_error_message(RequestId::Integer(id)) + .await?; + assert_eq!( + response.error, + JSONRPCErrorError { + code: -32602, + message: "Invalid verification challenge or display text.".into(), + data: Some(json!({"type": "invalidRequest", "reason": "invalidParams"})), + } + ); + Ok(()) +} + +#[tokio::test] +async fn user_verification_websocket_blocks_native_operations_before_validation() -> Result<()> { + // TestAppServer speaks stdio; use the existing transport fixture for real WebSocket frames. + let home = tempfile::tempdir()?; + let (mut process, address) = spawn_websocket_server(home.path()).await?; + let mut websocket = connect_websocket(address).await?; + send_request( + &mut websocket, + "initialize", + /*id*/ 0, + Some(json!({ + "clientInfo": {"name": "codex-tui", "version": "1"}, + "capabilities": {"experimentalApi": true}, + })), + ) + .await?; + read_response_for_id(&mut websocket, /*id*/ 0).await?; + send_request( + &mut websocket, + "userVerification/status", + /*id*/ 1, + Some(json!({})), + ) + .await?; + assert_eq!( + read_response_for_id(&mut websocket, /*id*/ 1).await?.result, + json!({ + "credentialId": null, + "unavailableReason": "providerUnavailable", + "unavailableMessage": "User verification is not available in this build or account." + }) + ); + + // This challenge produces invalidRequest over stdio. A network peer must be + // rejected before validation, regardless of account or biometric hardware. + for (method, params) in [ + ("userVerification/enroll", json!({})), + ("userVerification/delete", json!({})), + ( + "userVerification/verify", + json!({"challenge": "!", "title": "Approve", "description": ""}), + ), + ] { + send_request(&mut websocket, method, /*id*/ 2, Some(params)).await?; + assert_eq!( + read_error_for_id(&mut websocket, /*id*/ 2).await?.error, + JSONRPCErrorError { + code: -32603, + message: "User verification is not available in this build or account.".into(), + data: Some(json!({"type": "unavailable", "reason": "providerUnavailable"})), + } + ); + } + process.kill().await?; + Ok(()) +}