diff --git a/Cargo.lock b/Cargo.lock index aa6e0644..9293ea4f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -821,6 +821,7 @@ dependencies = [ "serde_json", "tokio", "tokio-stream", + "tokio-util", "tracing", "tracing-subscriber", "zbus", @@ -4076,14 +4077,13 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.19" +version = "0.7.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52" +checksum = "d7fcaa8d55a2bdd6b83ace262b016eca0d79ee02818c5c1bcdf0305114081078" dependencies = [ "bytes", "futures-core", "futures-sink", - "libc", "pin-project-lite", "tokio", ] diff --git a/credentialsd/Cargo.toml b/credentialsd/Cargo.toml index 04f54f3a..347e7ee0 100644 --- a/credentialsd/Cargo.toml +++ b/credentialsd/Cargo.toml @@ -25,6 +25,7 @@ ring = "=0.17.14" serde_json = "=1.0.151" tokio = { version = "=1.53.1", features = ["rt-multi-thread"] } tokio-stream = "=0.1.19" +tokio-util = "=0.7.13" [dev-dependencies] gio = "=0.22.8" diff --git a/credentialsd/src/credential_service/hybrid.rs b/credentialsd/src/credential_service/hybrid.rs index 5f136346..9b31fb8f 100644 --- a/credentialsd/src/credential_service/hybrid.rs +++ b/credentialsd/src/credential_service/hybrid.rs @@ -7,6 +7,7 @@ use tokio::sync::{ broadcast, mpsc::{self, Sender}, }; +use tokio_util::sync::CancellationToken; use tracing::{debug, error}; use libwebauthn::transport::cable::channel::{CableUpdate, CableUxUpdate}; @@ -29,6 +30,7 @@ pub(crate) trait HybridHandler { fn start( &self, request: &CredentialRequest, + cancellation: CancellationToken, ) -> impl Stream + Unpin + Send + Sized + 'static; } @@ -44,6 +46,7 @@ impl HybridHandler for InternalHybridHandler { fn start( &self, request: &CredentialRequest, + cancellation: CancellationToken, ) -> impl Stream + Unpin + Send + Sized + 'static { tracing::debug!("Starting hybrid operation"); let request = request.clone(); @@ -95,60 +98,74 @@ impl HybridHandler for InternalHybridHandler { tracing::debug!("Polling hybrid channel for updates."); let response: Result = loop { - match &request { - CredentialRequest::CreatePublicKeyCredentialRequest(make_request) => { - match channel.webauthn_make_credential(make_request).await { - Ok(response) => break Ok(response.into()), - Err(WebAuthnError::Ctap(ctap_error)) => { - if ctap_error.is_retryable_user_error() { - tracing::debug!( - "Retrying credential creation operation because of CTAP error: {:?}", - ctap_error - ); - continue; - } else { - tracing::error!( - "Received CTAP unrecoverable CTAP error: {:?}", - ctap_error - ); - break Err(Error::AuthenticatorError); + tokio::select! { + result = async { + match &request { + CredentialRequest::CreatePublicKeyCredentialRequest(make_request) => { + match channel.webauthn_make_credential(make_request).await { + Ok(response) => Ok(response.into()), + Err(WebAuthnError::Ctap(ctap_error)) => { + if ctap_error.is_retryable_user_error() { + tracing::debug!( + "Retrying credential creation operation because of CTAP error: {:?}", + ctap_error + ); + Err(None) + } else { + tracing::error!( + "Received CTAP unrecoverable CTAP error: {:?}", + ctap_error + ); + Err(Some(Error::AuthenticatorError)) + } + } + Err(err) => { + tracing::error!( + "Received unrecoverable error from authenticator: {:?}", + err + ); + Err(Some(Error::AuthenticatorError)) + } } } - Err(err) => { - tracing::error!( - "Received unrecoverable error from authenticator: {:?}", - err - ); - break Err(Error::AuthenticatorError); - } - }; - } - CredentialRequest::GetPublicKeyCredentialRequest(get_request) => { - match channel.webauthn_get_assertion(get_request).await { - Ok(response) => break Ok(response.into()), - Err(WebAuthnError::Ctap(ctap_error)) => { - if ctap_error.is_retryable_user_error() { - tracing::debug!( - "Retrying assertion operation because of CTAP error: {:?}", - ctap_error - ); - continue; - } else { - tracing::error!( - "Received CTAP unrecoverable CTAP error: {:?}", - ctap_error - ); - break Err(Error::AuthenticatorError); + CredentialRequest::GetPublicKeyCredentialRequest(get_request) => { + match channel.webauthn_get_assertion(get_request).await { + Ok(response) => Ok(response.into()), + Err(WebAuthnError::Ctap(ctap_error)) => { + if ctap_error.is_retryable_user_error() { + tracing::debug!( + "Retrying assertion operation because of CTAP error: {:?}", + ctap_error + ); + Err(None) + } else { + tracing::error!( + "Received CTAP unrecoverable CTAP error: {:?}", + ctap_error + ); + Err(Some(Error::AuthenticatorError)) + } + } + Err(err) => { + tracing::error!( + "Received unrecoverable error from authenticator: {:?}", + err + ); + Err(Some(Error::AuthenticatorError)) + } } } - Err(err) => { - tracing::error!( - "Received unrecoverable error from authenticator: {:?}", - err - ); - break Err(Error::AuthenticatorError); - } - }; + } + } => { + match result { + Ok(response) => break Ok(response), + Err(Some(err)) => break Err(err), + Err(None) => continue, // Retryable error + } + } + _ = cancellation.cancelled() => { + tracing::debug!("Hybrid handler cancelled, stopping processing"); + break Err(Error::Internal("Request cancelled".to_string())); } } }; diff --git a/credentialsd/src/credential_service/mod.rs b/credentialsd/src/credential_service/mod.rs index 28188f03..a76ee65f 100644 --- a/credentialsd/src/credential_service/mod.rs +++ b/credentialsd/src/credential_service/mod.rs @@ -18,6 +18,7 @@ use libwebauthn::{ }; use nfc::{NfcEvent, NfcHandler, NfcState, NfcStateInternal}; use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; use credentialsd_common::model::{ BackgroundEvent, Device, Error as CredentialServiceError, Transport, @@ -38,6 +39,19 @@ pub use usb::UsbState; /// Identifier for a request to be used for cancellation. pub type RequestId = u32; +/// Helper function to sleep with cancellation support. +async fn cancellable_sleep( + duration: std::time::Duration, + cancellation: &CancellationToken, +) -> Result<(), CredentialServiceError> { + tokio::select! { + _ = tokio::time::sleep(duration) => Ok(()), + _ = cancellation.cancelled() => { + Err(CredentialServiceError::Internal("Request cancelled".to_string())) + } + } +} + /// Process-wide in-memory store so a security key's pinUvAuthToken is reused across ceremonies. fn persistent_token_store() -> Arc { static STORE: OnceLock> = OnceLock::new(); @@ -51,6 +65,7 @@ struct RequestContext { request: CredentialRequest, response_channel: oneshot::Sender>, request_id: RequestId, + cancellation: CancellationToken, } impl RequestContext { @@ -70,7 +85,7 @@ pub trait ManageDevice { &self, request: &CredentialRequest, tx: oneshot::Sender>, - ) -> Result; + ) -> Result<(RequestId, CancellationToken), CredentialServiceError>; async fn cancel_request(&self, request_id: RequestId); async fn get_available_public_key_devices(&self) -> Result, ()>; async fn start_discovery( @@ -109,8 +124,17 @@ impl &self, ) -> Pin + Send + 'static>> { let guard = self.ctx.lock().unwrap(); - if let Some(RequestContext { ref request, .. }) = *guard { - let stream = self.hybrid_handler.lock().unwrap().start(request); + if let Some(RequestContext { + ref request, + ref cancellation, + .. + }) = *guard + { + let stream = self + .hybrid_handler + .lock() + .unwrap() + .start(request, cancellation.clone()); let ctx = self.ctx.clone(); Box::pin(HybridStateStream { inner: stream, ctx }) } else { @@ -123,8 +147,17 @@ impl async fn get_usb_credential(&self) -> Pin + Send + 'static>> { let guard = self.ctx.lock().unwrap(); - if let Some(RequestContext { ref request, .. }) = *guard { - let stream = self.usb_handler.lock().unwrap().start(request); + if let Some(RequestContext { + ref request, + ref cancellation, + .. + }) = *guard + { + let stream = self + .usb_handler + .lock() + .unwrap() + .start(request, cancellation.clone()); let ctx = self.ctx.clone(); Box::pin(UsbStateStream { inner: stream, ctx }) } else { @@ -137,8 +170,17 @@ impl async fn _get_nfc_credential(&self) -> Pin + Send + 'static>> { let guard = self.ctx.lock().unwrap(); - if let Some(RequestContext { ref request, .. }) = *guard { - let stream = self._nfc_handler.lock().unwrap().start(request); + if let Some(RequestContext { + ref request, + ref cancellation, + .. + }) = *guard + { + let stream = self + ._nfc_handler + .lock() + .unwrap() + .start(request, cancellation.clone()); let ctx = self.ctx.clone(); Box::pin(NfcStateStream { inner: stream, ctx }) } else { @@ -158,23 +200,26 @@ impl Manage &self, request: &CredentialRequest, tx: oneshot::Sender>, - ) -> Result { + ) -> Result<(RequestId, CancellationToken), CredentialServiceError> { let mut cred_request = self.ctx.lock().unwrap(); if cred_request.is_some() { Err(CredentialServiceError::Internal( "Already a request in progress.".to_string(), )) } else { - let request_id: RequestId = rand::random(); + // Generate non-zero request ID + let request_id: RequestId = rand::random_range(1..=u32::MAX); + let cancellation = CancellationToken::new(); // TODO: Spawn a task here that will listen to the signals from ui_control_client. // Move the get_*_credential(), etc. from gateway to here. let ctx = RequestContext { request: request.clone(), response_channel: tx, request_id, + cancellation: cancellation.clone(), }; _ = cred_request.insert(ctx); - Ok(request_id) + Ok((request_id, cancellation)) } } @@ -184,7 +229,7 @@ impl Manage && request_id == ctx.request_id { tracing::debug!("Cancelling request {request_id}"); - // TODO: cancel sub-tasks: hybrid and USB streams. + ctx.cancellation.cancel(); // It's fine if the requestor is no longer listening for the response. // TODO: create Cancelled variant @@ -268,27 +313,33 @@ where match Box::pin(Box::pin(self).as_mut().inner.next()).poll(cx) { Poll::Pending => Poll::Pending, Poll::Ready(Some(HybridEvent { state })) => { - if let HybridStateInternal::Completed(hybrid_response) = &state { - let response = match &**hybrid_response { - AuthenticatorResponse::CredentialCreated(make_credential_response) => { - CredentialResponse::from_make_credential( - make_credential_response, - &["hybrid"], - "cross-platform", - ) - } - AuthenticatorResponse::CredentialsAsserted(get_assertion_response) => { - CredentialResponse::from_get_assertion( - // When doing hybrid, the authenticator is capable of displaying it's own UI. - // So we assume here, it only ever returns one assertion. - // In case this doesn't hold true, we have to implement credential selection here, - // as is done for USB. - &get_assertion_response.assertions[0], - "cross-platform", - ) - } - }; - complete_request(ctx, response.clone()); + match &state { + HybridStateInternal::Completed(hybrid_response) => { + let response = match &**hybrid_response { + AuthenticatorResponse::CredentialCreated(make_credential_response) => { + CredentialResponse::from_make_credential( + make_credential_response, + &["hybrid"], + "cross-platform", + ) + } + AuthenticatorResponse::CredentialsAsserted(get_assertion_response) => { + CredentialResponse::from_get_assertion( + // When doing hybrid, the authenticator is capable of displaying it's own UI. + // So we assume here, it only ever returns one assertion. + // In case this doesn't hold true, we have to implement credential selection here, + // as is done for USB. + &get_assertion_response.assertions[0], + "cross-platform", + ) + } + }; + complete_request(ctx, Ok(response.clone())); + } + HybridStateInternal::Failed => { + complete_request(ctx, Err(CredentialServiceError::AuthenticatorError)); + } + _ => {} } Poll::Ready(Some(state.into())) } @@ -316,8 +367,14 @@ where match Box::pin(Box::pin(self).as_mut().inner.next()).poll(cx) { Poll::Pending => Poll::Pending, Poll::Ready(Some(UsbEvent { state })) => { - if let UsbStateInternal::Completed(response) = &state { - complete_request(ctx, response.clone()); + match &state { + UsbStateInternal::Completed(response) => { + complete_request(ctx, Ok(response.clone())); + } + UsbStateInternal::Failed(error) => { + complete_request(ctx, Err(error.clone())); + } + _ => {} } Poll::Ready(Some(state.into())) } @@ -346,8 +403,14 @@ where match Box::pin(Box::pin(self).as_mut().inner.next()).poll(cx) { Poll::Pending => Poll::Pending, Poll::Ready(Some(NfcEvent { state })) => { - if let NfcStateInternal::Completed(response) = &state { - complete_request(ctx, response.clone()); + match &state { + NfcStateInternal::Completed(response) => { + complete_request(ctx, Ok(response.clone())); + } + NfcStateInternal::Failed(error) => { + complete_request(ctx, Err(error.clone())); + } + _ => {} } Poll::Ready(Some(state.into())) } @@ -390,10 +453,14 @@ impl From for DeviceStateUpdate { } } -fn complete_request(ctx: &Mutex>, response: CredentialResponse) { +fn complete_request( + ctx: &Mutex>, + response: Result, +) { match ctx.lock().unwrap().take() { Some(ctx) => { - ctx.send_response(Ok(response)); + ctx.cancellation.cancel(); + ctx.send_response(response); } _ => { tracing::error!("Tried to consume context to respond to caller, but none was found.") @@ -418,3 +485,831 @@ impl From for AuthenticatorResponse { Self::CredentialsAsserted(value) } } + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + // Mock handlers for testing + #[derive(Debug)] + struct MockUsbHandler; + impl UsbHandler for MockUsbHandler { + fn start( + &self, + _request: &CredentialRequest, + _cancellation: CancellationToken, + ) -> impl Stream + Send + Sized + Unpin + 'static { + futures::stream::empty() + } + } + + #[derive(Debug)] + struct MockHybridHandler; + impl HybridHandler for MockHybridHandler { + fn start( + &self, + _request: &CredentialRequest, + _cancellation: CancellationToken, + ) -> impl Stream + Unpin + Send + Sized + 'static { + futures::stream::empty() + } + } + + #[derive(Debug)] + struct MockNfcHandler; + impl NfcHandler for MockNfcHandler { + fn start( + &self, + _request: &CredentialRequest, + _cancellation: CancellationToken, + ) -> impl Stream + Send + Sized + Unpin + 'static { + futures::stream::empty() + } + } + + fn create_test_credential_response() -> CredentialResponse { + use libwebauthn::ops::webauthn::GetAssertionResponse; + + // Create a minimal GetAssertion response for testing + let get_assertion_response = GetAssertionResponse { + assertions: vec![libwebauthn::ops::webauthn::Assertion { + credential_id: None, + authenticator_data: libwebauthn::fido::AuthenticatorData { + rp_id_hash: [0u8; 32], + flags: libwebauthn::fido::AuthenticatorDataFlags::empty(), + signature_count: 0, + attested_credential: None, + extensions: None, + raw: None, + }, + signature: vec![], + user: None, + credentials_count: None, + user_selected: None, + unsigned_extensions_output: None, + transport: None, + }], + }; + + CredentialResponse::from_get_assertion( + &get_assertion_response.assertions[0], + "cross-platform", + ) + } + + async fn create_test_request() -> CredentialRequest { + use libwebauthn::ops::webauthn::{ + MakeCredentialRequest, OriginValidation, RequestSettings, idl::origin::RequestOrigin, + }; + + let request_json = r#" + { + "rp": { + "id": "example.com", + "name": "Example Relying Party" + }, + "user": { + "id": "MTIzNDU2NzgxMjM0NTY3ODEyMzQ1Njc4MTIzNDU2Nzg", + "name": "test@example.com", + "displayName": "Test User" + }, + "challenge": "MTIzNDU2NzgxMjM0NTY3ODEyMzQ1Njc4MTIzNDU2Nzg", + "pubKeyCredParams": [ + {"type": "public-key", "alg": -7} + ], + "timeout": 60000, + "excludeCredentials": [], + "authenticatorSelection": { + "residentKey": "discouraged", + "userVerification": "preferred" + }, + "attestation": "none" + } + "#; + + let request_origin: RequestOrigin = + "https://example.com".try_into().expect("Invalid origin"); + + let settings = RequestSettings { + origin: OriginValidation::Trust, + }; + + let make_credentials_request = + MakeCredentialRequest::prepare(&request_origin, request_json, &settings) + .await + .expect("Failed to parse request JSON"); + + CredentialRequest::CreatePublicKeyCredentialRequest(make_credentials_request) + } + + #[tokio::test] + async fn test_init_request_returns_token_and_id() { + let service = CredentialService::new(MockHybridHandler, MockNfcHandler, MockUsbHandler); + let (tx, _rx) = oneshot::channel(); + let request = create_test_request().await; + + let result = service.init_request(&request, tx).await; + + assert!(result.is_ok()); + let (request_id, cancellation_token) = result.unwrap(); + assert!(request_id > 0); + assert!(!cancellation_token.is_cancelled()); + } + + #[tokio::test] + async fn test_cancel_request_triggers_cancellation() { + let service = CredentialService::new(MockHybridHandler, MockNfcHandler, MockUsbHandler); + let (tx, _rx) = oneshot::channel(); + let request = create_test_request().await; + + let (request_id, cancellation_token) = service.init_request(&request, tx).await.unwrap(); + assert!(!cancellation_token.is_cancelled()); + + service.cancel_request(request_id).await; + assert!(cancellation_token.is_cancelled()); + } + + #[tokio::test] + async fn test_cancellable_sleep_completes_normally() { + let token = CancellationToken::new(); + let start = tokio::time::Instant::now(); + + let result = cancellable_sleep(Duration::from_millis(50), &token).await; + + assert!(result.is_ok()); + assert!(start.elapsed() >= Duration::from_millis(50)); + } + + #[tokio::test] + async fn test_cancellable_sleep_respects_cancellation() { + let token = CancellationToken::new(); + token.cancel(); // Pre-cancel the token + + let start = tokio::time::Instant::now(); + let result = cancellable_sleep(Duration::from_secs(5), &token).await; + + assert!(result.is_err()); + // Should return immediately, not after 5 seconds + assert!(start.elapsed() < Duration::from_millis(100)); + } + + #[tokio::test] + async fn test_init_request_rejects_concurrent() { + let service = CredentialService::new(MockHybridHandler, MockNfcHandler, MockUsbHandler); + let (tx1, _rx1) = oneshot::channel(); + let (tx2, _rx2) = oneshot::channel(); + let request = create_test_request().await; + + // First request should succeed + let result1 = service.init_request(&request, tx1).await; + assert!(result1.is_ok()); + + // Second concurrent request should fail + let result2 = service.init_request(&request, tx2).await; + assert!(result2.is_err()); + assert!( + result2 + .unwrap_err() + .to_string() + .contains("Already a request in progress") + ); + } + + // Mock handlers that track cancellation and emit configurable states + use std::sync::atomic::{AtomicBool, Ordering}; + + #[derive(Debug, Clone)] + struct CancellationTrackingUsbHandler { + cancelled: Arc, + emit_states: Arc>, + delay_ms: u64, + } + + impl CancellationTrackingUsbHandler { + fn new(emit_states: Vec, delay_ms: u64) -> Self { + Self { + cancelled: Arc::new(AtomicBool::new(false)), + emit_states: Arc::new(emit_states), + delay_ms, + } + } + + fn was_cancelled(&self) -> bool { + self.cancelled.load(Ordering::SeqCst) + } + } + + impl UsbHandler for CancellationTrackingUsbHandler { + fn start( + &self, + _request: &CredentialRequest, + cancellation: CancellationToken, + ) -> impl Stream + Send + Sized + Unpin + 'static { + let cancelled = self.cancelled.clone(); + let states = self.emit_states.clone(); + let delay_ms = self.delay_ms; + Box::pin(async_stream::stream! { + for state in states.iter() { + tokio::select! { + _ = tokio::time::sleep(Duration::from_millis(delay_ms)) => { + yield UsbEvent { state: state.clone() }; + } + _ = cancellation.cancelled() => { + cancelled.store(true, Ordering::SeqCst); + break; + } + } + } + }) + } + } + + #[derive(Debug, Clone)] + struct CancellationTrackingHybridHandler { + cancelled: Arc, + emit_states: Arc>, + delay_ms: u64, + } + + impl CancellationTrackingHybridHandler { + fn new(emit_states: Vec, delay_ms: u64) -> Self { + Self { + cancelled: Arc::new(AtomicBool::new(false)), + emit_states: Arc::new(emit_states), + delay_ms, + } + } + + fn was_cancelled(&self) -> bool { + self.cancelled.load(Ordering::SeqCst) + } + } + + impl HybridHandler for CancellationTrackingHybridHandler { + fn start( + &self, + _request: &CredentialRequest, + cancellation: CancellationToken, + ) -> impl Stream + Unpin + Send + Sized + 'static { + let cancelled = self.cancelled.clone(); + let states = self.emit_states.clone(); + let delay_ms = self.delay_ms; + Box::pin(async_stream::stream! { + for state in states.iter() { + tokio::select! { + _ = tokio::time::sleep(Duration::from_millis(delay_ms)) => { + yield HybridEvent { state: state.clone() }; + } + _ = cancellation.cancelled() => { + cancelled.store(true, Ordering::SeqCst); + break; + } + } + } + }) + } + } + + #[tokio::test] + async fn test_handler_respects_cancellation_during_polling() { + // Create a handler that would emit many states if not cancelled + let usb_handler = CancellationTrackingUsbHandler::new( + vec![ + UsbStateInternal::Waiting, + UsbStateInternal::Waiting, + UsbStateInternal::Waiting, + ], + 100, // 100ms between each state + ); + + let request = create_test_request().await; + let cancellation = CancellationToken::new(); + + let mut stream = usb_handler.start(&request, cancellation.clone()); + + // Collect first state + let first = stream.next().await; + assert!(first.is_some()); + + // Cancel immediately + cancellation.cancel(); + + // Stream should stop immediately - no more states should be emitted + let remaining: Vec<_> = stream.collect().await; + assert!( + remaining.is_empty(), + "Handler should not emit any more states after cancellation" + ); + + // Handler should have detected cancellation during collect + assert!(usb_handler.was_cancelled()); + } + + #[tokio::test] + async fn test_handler_stops_on_external_cancellation() { + // Test that handlers respect cancellation from external source + let hybrid_handler = CancellationTrackingHybridHandler::new( + vec![ + HybridStateInternal::Init("qr_code".to_string()), + HybridStateInternal::Connecting, + HybridStateInternal::Connected, + // Would continue forever if not cancelled + ], + 100, + ); + + let request = create_test_request().await; + let cancellation = CancellationToken::new(); + + let mut stream = hybrid_handler.start(&request, cancellation.clone()); + + // Collect first couple states + let _first = stream.next().await; + let _second = stream.next().await; + + // Simulate external cancellation (e.g., another transport completed) + cancellation.cancel(); + + // Stream should stop immediately - no more states should be emitted + let remaining: Vec<_> = stream.collect().await; + assert!( + remaining.is_empty(), + "Handler should not emit any more states after cancellation" + ); + + // Handler should have detected cancellation during collect + assert!(hybrid_handler.was_cancelled()); + } + + #[tokio::test] + async fn test_all_states_delivered_without_cancellation() { + // Test that all states are emitted when not cancelled + let hybrid_handler = CancellationTrackingHybridHandler::new( + vec![ + HybridStateInternal::Init("qr_code".to_string()), + HybridStateInternal::Connecting, + HybridStateInternal::Connected, + ], + 10, + ); + + let request = create_test_request().await; + let cancellation = CancellationToken::new(); + + let stream = hybrid_handler.start(&request, cancellation); + let states: Vec<_> = stream.collect().await; + + // Should have all 3 states + assert_eq!(states.len(), 3, "Should emit all states when not cancelled"); + assert!(matches!(states[0].state, HybridStateInternal::Init(_))); + assert!(matches!(states[1].state, HybridStateInternal::Connecting)); + assert!(matches!(states[2].state, HybridStateInternal::Connected)); + } + + #[tokio::test] + async fn test_cancel_request_by_id() { + let service = CredentialService::new(MockHybridHandler, MockNfcHandler, MockUsbHandler); + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + + let (request_id, cancellation_token) = service.init_request(&request, tx).await.unwrap(); + + // Token should not be cancelled initially + assert!(!cancellation_token.is_cancelled()); + + // Cancel by ID + service.cancel_request(request_id).await; + + // Token should now be cancelled + assert!(cancellation_token.is_cancelled()); + } + + #[tokio::test] + async fn test_multiple_handlers_all_cancelled() { + // All three handlers emit states slowly + let usb_handler = + CancellationTrackingUsbHandler::new(vec![UsbStateInternal::Waiting; 5], 100); + let hybrid_handler = CancellationTrackingHybridHandler::new( + vec![ + HybridStateInternal::Init("qr".to_string()), + HybridStateInternal::Connecting, + HybridStateInternal::Connected, + ], + 100, + ); + let nfc_handler = MockNfcHandler; + + let service = CredentialService::new(hybrid_handler, nfc_handler, usb_handler); + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + + let (request_id, cancellation_token) = service.init_request(&request, tx).await.unwrap(); + + // Start all handlers + let mut usb_stream = service.get_usb_credential().await; + let mut hybrid_stream = service.get_hybrid_credential().await; + + // Let them emit a couple states + let _ = usb_stream.next().await; + let _ = hybrid_stream.next().await; + + // Cancel the request + service.cancel_request(request_id).await; + + // Small delay for propagation + tokio::time::sleep(Duration::from_millis(50)).await; + + // Token should be cancelled + assert!(cancellation_token.is_cancelled()); + + // Streams should stop immediately - no more states should be emitted + let usb_remaining: Vec<_> = usb_stream.collect().await; + let hybrid_remaining: Vec<_> = hybrid_stream.collect().await; + + assert!( + usb_remaining.is_empty(), + "USB should not emit any more states after cancellation" + ); + assert!( + hybrid_remaining.is_empty(), + "Hybrid should not emit any more states after cancellation" + ); + } + + #[tokio::test] + async fn test_cancellation_cleans_up_request_context() { + let service = CredentialService::new(MockHybridHandler, MockNfcHandler, MockUsbHandler); + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + + let (request_id, _token) = service.init_request(&request, tx).await.unwrap(); + + // Cancel the request + service.cancel_request(request_id).await; + + // Should be able to start a new request now (context cleaned up) + let (tx2, _rx2) = oneshot::channel(); + let result = service.init_request(&request, tx2).await; + assert!( + result.is_ok(), + "Should be able to init new request after cancel" + ); + } + + #[tokio::test] + async fn test_cancel_with_unknown_id_is_noop() { + let service = CredentialService::new(MockHybridHandler, MockNfcHandler, MockUsbHandler); + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + + let (request_id, cancellation_token) = service.init_request(&request, tx).await.unwrap(); + + // Cancel with a different ID (should be a no-op) + let wrong_id = request_id.wrapping_add(1); + service.cancel_request(wrong_id).await; + + // Original request should still be active + assert!( + !cancellation_token.is_cancelled(), + "Token should not be cancelled with wrong ID" + ); + + // Now cancel with correct ID + service.cancel_request(request_id).await; + assert!( + cancellation_token.is_cancelled(), + "Token should be cancelled with correct ID" + ); + } + + #[tokio::test] + async fn test_cancel_with_no_active_request_is_noop() { + let service = CredentialService::new(MockHybridHandler, MockNfcHandler, MockUsbHandler); + + // Cancel when no request is active (should not crash or panic) + service.cancel_request(12345).await; + + // Should still be able to start a new request + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + let result = service.init_request(&request, tx).await; + assert!( + result.is_ok(), + "Should be able to init request after no-op cancel" + ); + } + + #[tokio::test] + async fn test_request_id_matches_on_init() { + let service = CredentialService::new(MockHybridHandler, MockNfcHandler, MockUsbHandler); + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + + let (request_id_1, _token_1) = service.init_request(&request, tx).await.unwrap(); + + // Cancel to free up the service + service.cancel_request(request_id_1).await; + + // Start a new request + let (tx2, _rx2) = oneshot::channel(); + let (request_id_2, _token_2) = service.init_request(&request, tx2).await.unwrap(); + + // IDs should be different (random) + assert_ne!( + request_id_1, request_id_2, + "Sequential requests should (almost certainly) have different IDs" + ); + } + + #[tokio::test] + async fn test_explicit_cancel_stops_all_transports() { + // USB would keep polling + let usb_handler = CancellationTrackingUsbHandler::new( + vec![ + UsbStateInternal::Waiting, + UsbStateInternal::Waiting, + UsbStateInternal::Waiting, + ], + 100, + ); + + // Hybrid would keep connecting + let hybrid_handler = CancellationTrackingHybridHandler::new( + vec![ + HybridStateInternal::Init("qr".to_string()), + HybridStateInternal::Connecting, + HybridStateInternal::Connected, + ], + 100, + ); + + // Clone handlers to verify cancellation later + let usb_handler_ref = usb_handler.clone(); + let hybrid_handler_ref = hybrid_handler.clone(); + assert!( + !usb_handler_ref.was_cancelled(), + "USB handler should not have detected cancellation" + ); + assert!( + !hybrid_handler_ref.was_cancelled(), + "Hybrid handler should not have detected cancellation" + ); + + let service = CredentialService::new(hybrid_handler, MockNfcHandler, usb_handler); + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + + let (request_id, cancellation_token) = service.init_request(&request, tx).await.unwrap(); + + // Start both streams + let mut usb_stream = service.get_usb_credential().await; + let mut hybrid_stream = service.get_hybrid_credential().await; + + // Let them emit a state each + let _ = usb_stream.next().await; + let _ = hybrid_stream.next().await; + + // Explicitly cancel the request + service.cancel_request(request_id).await; + + // Cancellation token should be triggered + assert!( + cancellation_token.is_cancelled(), + "Cancellation token should be triggered after cancel_request" + ); + + // Both streams should stop emitting states + let usb_remaining: Vec<_> = usb_stream.collect().await; + let hybrid_remaining: Vec<_> = hybrid_stream.collect().await; + + assert!( + usb_remaining.is_empty(), + "USB should not emit any more states after cancellation" + ); + assert!( + hybrid_remaining.is_empty(), + "Hybrid should not emit any more states after cancellation" + ); + + // Verify both handlers detected cancellation + assert!( + usb_handler_ref.was_cancelled(), + "USB handler should have detected cancellation" + ); + assert!( + hybrid_handler_ref.was_cancelled(), + "Hybrid handler should have detected cancellation" + ); + } + + #[tokio::test] + async fn test_failed_request_triggers_cancellation() { + use credentialsd_common::model::Error; + + // Handler that emits a Failed state + let usb_handler = CancellationTrackingUsbHandler::new( + vec![ + UsbStateInternal::Waiting, + UsbStateInternal::Failed(Error::Internal("test failure".to_string())), + ], + 10, + ); + + let service = CredentialService::new(MockHybridHandler, MockNfcHandler, usb_handler); + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + + let (_request_id, cancellation_token) = service.init_request(&request, tx).await.unwrap(); + + // Start the USB stream + let mut usb_stream = service.get_usb_credential().await; + + // Token should not be cancelled initially + assert!(!cancellation_token.is_cancelled()); + + // Consume states until we hit the Failed state + while let Some(state) = usb_stream.next().await { + if matches!(state, UsbState::Failed(_)) { + break; + } + } + + // The Failed state should have triggered cancellation + assert!( + cancellation_token.is_cancelled(), + "Cancellation token should be triggered when request fails" + ); + } + + #[tokio::test] + async fn test_failed_request_cancels_other_transports() { + use credentialsd_common::model::Error; + + // USB will fail quickly + let usb_handler = CancellationTrackingUsbHandler::new( + vec![ + UsbStateInternal::Waiting, + UsbStateInternal::Failed(Error::Internal("test".to_string())), + ], + 10, + ); + + // Hybrid would keep going if not cancelled + let hybrid_handler = CancellationTrackingHybridHandler::new( + vec![ + HybridStateInternal::Init("qr".to_string()), + HybridStateInternal::Connecting, + HybridStateInternal::Connected, + ], + 100, + ); + + // Clone handler to verify cancellation later + let hybrid_handler_ref = hybrid_handler.clone(); + + let service = CredentialService::new(hybrid_handler, MockNfcHandler, usb_handler); + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + + let (_request_id, cancellation_token) = service.init_request(&request, tx).await.unwrap(); + + // Start both streams + let mut usb_stream = service.get_usb_credential().await; + let mut hybrid_stream = service.get_hybrid_credential().await; + + // Collect one state from hybrid + let _ = hybrid_stream.next().await; + + // Consume USB until it fails + while let Some(state) = usb_stream.next().await { + if matches!(state, UsbState::Failed(_)) { + break; + } + } + + // Cancellation token should be triggered by USB failure + assert!( + cancellation_token.is_cancelled(), + "Cancellation token should be triggered when USB fails" + ); + + // Hybrid should stop emitting states (cancelled by USB failure) + let hybrid_remaining: Vec<_> = hybrid_stream.collect().await; + assert!( + hybrid_remaining.is_empty(), + "Hybrid should not emit any more states after USB fails" + ); + + // Verify the handler actually detected cancellation + assert!( + hybrid_handler_ref.was_cancelled(), + "Hybrid handler should have detected cancellation when USB failed" + ); + } + + #[tokio::test] + async fn test_completed_request_triggers_cancellation() { + let credential_response = create_test_credential_response(); + + // Handler that emits a Completed state + let usb_handler = CancellationTrackingUsbHandler::new( + vec![ + UsbStateInternal::Waiting, + UsbStateInternal::Completed(credential_response), + ], + 10, + ); + + let service = CredentialService::new(MockHybridHandler, MockNfcHandler, usb_handler); + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + + let (_request_id, cancellation_token) = service.init_request(&request, tx).await.unwrap(); + + // Start the USB stream + let mut usb_stream = service.get_usb_credential().await; + + // Token should not be cancelled initially + assert!(!cancellation_token.is_cancelled()); + + // Consume states until we hit the Completed state + while let Some(state) = usb_stream.next().await { + if matches!(state, UsbState::Completed) { + break; + } + } + + // The Completed state should have triggered cancellation + assert!( + cancellation_token.is_cancelled(), + "Cancellation token should be triggered when request completes successfully" + ); + } + + #[tokio::test] + async fn test_completed_request_cancels_other_transports() { + let credential_response = create_test_credential_response(); + + // USB will complete quickly + let usb_handler = CancellationTrackingUsbHandler::new( + vec![ + UsbStateInternal::Waiting, + UsbStateInternal::Completed(credential_response), + ], + 10, + ); + + // Hybrid would keep going if not cancelled + let hybrid_handler = CancellationTrackingHybridHandler::new( + vec![ + HybridStateInternal::Init("qr".to_string()), + HybridStateInternal::Connecting, + HybridStateInternal::Connected, + ], + 100, + ); + + // Clone handler to verify cancellation later + let hybrid_handler_ref = hybrid_handler.clone(); + + let service = CredentialService::new(hybrid_handler, MockNfcHandler, usb_handler); + let request = create_test_request().await; + let (tx, _rx) = oneshot::channel(); + + let (_request_id, cancellation_token) = service.init_request(&request, tx).await.unwrap(); + + // Start both streams + let mut usb_stream = service.get_usb_credential().await; + let mut hybrid_stream = service.get_hybrid_credential().await; + + // Collect one state from hybrid + let _ = hybrid_stream.next().await; + + // Consume USB until it completes + while let Some(state) = usb_stream.next().await { + if matches!(state, UsbState::Completed) { + break; + } + } + + // Cancellation token should be triggered by USB completion + assert!( + cancellation_token.is_cancelled(), + "Cancellation token should be triggered when USB completes" + ); + + // Hybrid should stop emitting states (cancelled by USB completion) + let hybrid_remaining: Vec<_> = hybrid_stream.collect().await; + assert!( + hybrid_remaining.is_empty(), + "Hybrid should not emit any more states after USB completes" + ); + + // Verify the handler actually detected cancellation + assert!( + hybrid_handler_ref.was_cancelled(), + "Hybrid handler should have detected cancellation when USB completed" + ); + } +} diff --git a/credentialsd/src/credential_service/nfc.rs b/credentialsd/src/credential_service/nfc.rs index 96b71021..ecdb5dfb 100644 --- a/credentialsd/src/credential_service/nfc.rs +++ b/credentialsd/src/credential_service/nfc.rs @@ -13,6 +13,7 @@ use libwebauthn::{ }; use tokio::sync::broadcast; use tokio::sync::mpsc::{self, Receiver, Sender, WeakSender}; +use tokio_util::sync::CancellationToken; use tracing::{debug, warn}; use credentialsd_common::model::{BackgroundEvent, Credential, Error, PinNotSetError}; @@ -26,6 +27,7 @@ pub(crate) trait NfcHandler { fn start( &self, request: &CredentialRequest, + cancellation: CancellationToken, ) -> impl Stream + Send + Sized + Unpin + 'static; } @@ -36,29 +38,38 @@ impl InProcessNfcHandler { async fn process_idle_waiting( failures: &mut usize, prev_nfc_state: &NfcStateInternal, + cancellation: &CancellationToken, ) -> Result { - match libwebauthn::transport::nfc::get_nfc_device().await { - Ok(Some(nfc_device)) => Ok(NfcStateInternal::Connected(nfc_device)), - Ok(None) => { - let state = NfcStateInternal::Waiting; - Ok(state) - } - Err(err) => { - *failures += 1; - if *failures == 5 { - Err(Error::Internal(format!( - "Failed to list NFC authenticators: {:?}. Cancelling NFC state updates.", - err - ))) - } else { - tracing::warn!( - "Failed to list NFC authenticators: {:?}. Throttling NFC state updates", - err - ); - tokio::time::sleep(Duration::from_secs(1)).await; - Ok(prev_nfc_state.clone()) + tokio::select! { + result = libwebauthn::transport::nfc::get_nfc_device() => { + match result { + Ok(Some(nfc_device)) => Ok(NfcStateInternal::Connected(nfc_device)), + Ok(None) => { + let state = NfcStateInternal::Waiting; + Ok(state) + } + Err(err) => { + *failures += 1; + if *failures == 5 { + Err(Error::Internal(format!( + "Failed to list NFC authenticators: {:?}. Cancelling NFC state updates.", + err + ))) + } else { + tracing::warn!( + "Failed to list NFC authenticators: {:?}. Throttling NFC state updates", + err + ); + super::cancellable_sleep(Duration::from_secs(1), cancellation).await?; + Ok(prev_nfc_state.clone()) + } + } } } + _ = cancellation.cancelled() => { + tracing::debug!("NFC idle polling cancelled"); + Err(Error::Internal("Request cancelled".to_string())) + } } } @@ -158,6 +169,7 @@ impl InProcessNfcHandler { async fn process( tx: Sender, cred_request: CredentialRequest, + cancellation: CancellationToken, ) -> Result<(), Error> { let mut state = NfcStateInternal::Idle; let (signal_tx, mut signal_rx) = mpsc::channel(256); @@ -169,34 +181,54 @@ impl InProcessNfcHandler { loop { tracing::debug!("current nfc state: {:?}", state); let prev_nfc_state = state; - let next_nfc_state = match prev_nfc_state { - NfcStateInternal::Idle | NfcStateInternal::Waiting => { - Self::process_idle_waiting(&mut failures, &prev_nfc_state).await - } - NfcStateInternal::Connected(device) => { - let signal_tx2 = signal_tx.clone(); - let cred_request = cred_request.clone(); - tokio::spawn(async move { - handle_events(&cred_request, device, &signal_tx2).await; - }); - Self::process_user_interaction(&mut signal_rx, &cred_tx).await + + tokio::select! { + next_nfc_state = async { + match prev_nfc_state { + NfcStateInternal::Idle | NfcStateInternal::Waiting => { + Self::process_idle_waiting(&mut failures, &prev_nfc_state, &cancellation).await + } + NfcStateInternal::Connected(device) => { + let signal_tx2 = signal_tx.clone(); + let cred_request = cred_request.clone(); + tokio::spawn(async move { + handle_events(&cred_request, device, &signal_tx2).await; + }); + Self::process_user_interaction(&mut signal_rx, &cred_tx).await + } + NfcStateInternal::NeedsPin { .. } + | NfcStateInternal::PinNotSet { .. } + | NfcStateInternal::NeedsUserVerification { .. } => { + Self::process_user_interaction(&mut signal_rx, &cred_tx).await + } + NfcStateInternal::SelectCredential { + response, + cred_tx: _, + } => Self::process_select_credential(response, &mut cred_rx).await, + // Terminal states - preserve state unchanged, will break loop after sending + NfcStateInternal::Completed(_) | NfcStateInternal::Failed(_) => { + Ok(prev_nfc_state.clone()) + } + } + } => { + state = next_nfc_state.unwrap_or_else(NfcStateInternal::Failed); + + tx.send(state.clone()).await.map_err(|_| { + Error::Internal("NFC state channel receiver closed prematurely".to_string()) + })?; + + // Check for terminal states AFTER sending + match state { + NfcStateInternal::Completed(_) => break Ok(()), + NfcStateInternal::Failed(err) => break Err(err), + _ => {} + } } - NfcStateInternal::NeedsPin { .. } - | NfcStateInternal::PinNotSet { .. } - | NfcStateInternal::NeedsUserVerification { .. } => { - Self::process_user_interaction(&mut signal_rx, &cred_tx).await + _ = cancellation.cancelled() => { + tracing::debug!("NFC handler cancelled, stopping processing"); + break Err(Error::Internal("Request cancelled".to_string())); } - NfcStateInternal::SelectCredential { - response, - cred_tx: _, - } => Self::process_select_credential(response, &mut cred_rx).await, - NfcStateInternal::Completed(_) => break Ok(()), - NfcStateInternal::Failed(err) => break Err(err), - }; - state = next_nfc_state.unwrap_or_else(NfcStateInternal::Failed); - tx.send(state.clone()).await.map_err(|_| { - Error::Internal("NFC state channel receiver closed prematurely".to_string()) - })?; + } } } } @@ -284,13 +316,14 @@ impl NfcHandler for InProcessNfcHandler { fn start( &self, request: &CredentialRequest, + cancellation: CancellationToken, ) -> impl Stream + Send + Sized + Unpin + 'static { let request = request.clone(); let (tx, mut rx) = mpsc::channel(32); tokio::spawn(async move { // TODO: instead of logging error here, push the errors into the // stream so credential service can handle/forward them to the UI - if let Err(err) = InProcessNfcHandler::process(tx, request).await { + if let Err(err) = InProcessNfcHandler::process(tx, request, cancellation).await { tracing::error!("Error getting credential from NFC: {:?}", err); } }); diff --git a/credentialsd/src/credential_service/usb.rs b/credentialsd/src/credential_service/usb.rs index 3ea66f07..754febe5 100644 --- a/credentialsd/src/credential_service/usb.rs +++ b/credentialsd/src/credential_service/usb.rs @@ -18,6 +18,7 @@ use tokio::sync::{ Mutex as AsyncMutex, broadcast, mpsc::{self, Receiver, Sender, WeakSender}, }; +use tokio_util::sync::CancellationToken; use tracing::{debug, warn}; use credentialsd_common::model::{BackgroundEvent, Credential, Error, PinNotSetError}; @@ -30,6 +31,7 @@ pub(crate) trait UsbHandler { fn start( &self, request: &CredentialRequest, + cancellation: CancellationToken, ) -> impl Stream + Send + Sized + Unpin + 'static; } @@ -40,38 +42,47 @@ impl InProcessUsbHandler { async fn process_idle_waiting( failures: &mut usize, prev_usb_state: &UsbStateInternal, + cancellation: &CancellationToken, ) -> Result { - match libwebauthn::transport::hid::list_devices().await { - Ok(hid_devices) => { - if hid_devices.is_empty() { - tokio::time::sleep(Duration::from_millis(50)).await; - let state = UsbStateInternal::Waiting; - Ok(state) - } else { - Ok(UsbStateInternal::SelectingDevice(hid_devices)) + tokio::select! { + result = libwebauthn::transport::hid::list_devices() => { + match result { + Ok(hid_devices) => { + if hid_devices.is_empty() { + super::cancellable_sleep(Duration::from_millis(50), cancellation).await?; + Ok(UsbStateInternal::Waiting) + } else { + Ok(UsbStateInternal::SelectingDevice(hid_devices)) + } + } + Err(err) => { + *failures += 1; + if *failures == 5 { + Err(Error::Internal(format!( + "Failed to list USB authenticators: {:?}. Cancelling USB state updates.", + err + ))) + } else { + tracing::warn!( + "Failed to list USB authenticators: {:?}. Throttling USB state updates", + err + ); + super::cancellable_sleep(Duration::from_secs(1), cancellation).await?; + Ok(prev_usb_state.clone()) + } + } } } - Err(err) => { - *failures += 1; - if *failures == 5 { - Err(Error::Internal(format!( - "Failed to list USB authenticators: {:?}. Cancelling USB state updates.", - err - ))) - } else { - tracing::warn!( - "Failed to list USB authenticators: {:?}. Throttling USB state updates", - err - ); - tokio::time::sleep(Duration::from_secs(1)).await; - Ok(prev_usb_state.clone()) - } + _ = cancellation.cancelled() => { + tracing::debug!("USB idle polling cancelled"); + Err(Error::Internal("Request cancelled".to_string())) } } } async fn process_selecting_device( hid_devices: &[HidDevice], + cancellation: &CancellationToken, ) -> Result { let expected_answers = hid_devices.len(); let (blinking_tx, mut blinking_rx) = @@ -125,20 +136,40 @@ impl InProcessUsbHandler { tracing::info!("Waiting for user interaction"); drop(blinking_tx); let mut state = UsbStateInternal::Idle; - while let Some(msg) = blinking_rx.recv().await { - match msg { - Some(idx) => { - let (device, _handle) = channel_map.remove(&idx).unwrap(); - tracing::info!("User selected device {device:?}."); + loop { + tokio::select! { + maybe_msg = blinking_rx.recv() => { + let Some(msg) = maybe_msg else { + // All blink tasks finished without a selection. + break; + }; + match msg { + Some(idx) => { + let (device, _handle) = channel_map.remove(&idx).unwrap(); + tracing::info!("User selected device {device:?}."); + for (_key, (device, handle)) in channel_map.into_iter() { + tracing::info!("Cancelling device {device:?}."); + handle.cancel_ongoing_operation().await; + } + state = UsbStateInternal::Connected(Arc::new(AsyncMutex::new(device))); + break; + } + None => { + continue; + } + } + } + _ = cancellation.cancelled() => { + // The request was cancelled (e.g. another transport completed, or + // the user cancelled). Stop all blinking devices. This interrupts + // the blocking HID read within the transport (≤100ms) and sends a + // CTAP CANCEL frame to each device. + tracing::debug!("USB device selection cancelled"); for (_key, (device, handle)) in channel_map.into_iter() { - tracing::info!("Cancelling device {device:?}."); + tracing::info!("Cancelling blinking device {device:?}."); handle.cancel_ongoing_operation().await; } - state = UsbStateInternal::Connected(Arc::new(AsyncMutex::new(device))); - break; - } - None => { - continue; + return Err(Error::Internal("Request cancelled".to_string())); } } } @@ -242,6 +273,7 @@ impl InProcessUsbHandler { async fn process( tx: Sender, cred_request: CredentialRequest, + cancellation: CancellationToken, ) -> Result<(), Error> { let mut state = UsbStateInternal::Idle; let (signal_tx, mut signal_rx) = mpsc::channel(256); @@ -253,50 +285,71 @@ impl InProcessUsbHandler { loop { tracing::trace!("current usb state: {:?}", state); let prev_usb_state = state; - let next_usb_state = match prev_usb_state { - UsbStateInternal::Idle | UsbStateInternal::Waiting => { - Self::process_idle_waiting(&mut failures, &prev_usb_state).await - } - UsbStateInternal::SelectingDevice(ref hid_devices) => { - Self::process_selecting_device(hid_devices.as_slice()).await - } - UsbStateInternal::Connected(ref device) => { - let device = std::sync::Arc::clone(device); - let signal_tx2 = signal_tx.clone(); - let cred_request = cred_request.clone(); - tokio::spawn(async move { - handle_events(&cred_request, device.clone(), &signal_tx2).await; - }); - Self::process_user_interaction(&mut signal_rx, &cred_tx).await - } - UsbStateInternal::NeedsPin { .. } - | UsbStateInternal::PinNotSet { .. } - | UsbStateInternal::NeedsUserVerification { .. } - | UsbStateInternal::NeedsUserPresence => { - Self::process_user_interaction(&mut signal_rx, &cred_tx).await + + tokio::select! { + next_usb_state = async { + match prev_usb_state { + UsbStateInternal::Idle | UsbStateInternal::Waiting => { + Self::process_idle_waiting(&mut failures, &prev_usb_state, &cancellation).await + } + UsbStateInternal::SelectingDevice(ref hid_devices) => { + Self::process_selecting_device(hid_devices.as_slice(), &cancellation).await + } + UsbStateInternal::Connected(ref device) => { + let device = std::sync::Arc::clone(device); + let signal_tx2 = signal_tx.clone(); + let cred_request = cred_request.clone(); + let cancellation = cancellation.clone(); + tokio::spawn(async move { + handle_events(&cred_request, device.clone(), &signal_tx2, cancellation).await; + }); + Self::process_user_interaction(&mut signal_rx, &cred_tx).await + } + UsbStateInternal::NeedsPin { .. } + | UsbStateInternal::PinNotSet { .. } + | UsbStateInternal::NeedsUserVerification { .. } + | UsbStateInternal::NeedsUserPresence => { + Self::process_user_interaction(&mut signal_rx, &cred_tx).await + } + UsbStateInternal::SelectCredential { + ref response, + cred_tx: _, + } => Self::process_select_credential(response, &mut cred_rx).await, + // Terminal states - preserve state unchanged, will break loop after sending + UsbStateInternal::Completed(_) | UsbStateInternal::Failed(_) => { + Ok(prev_usb_state.clone()) + } + } + } => { + state = next_usb_state.unwrap_or_else(UsbStateInternal::Failed); + // Usually, comparing the Discrimimant is enough, but PinNotSet/NeedsPin + // can be repeated multiple times with different or the same error reasons + // (PIN wrong, PIN too short, PIN too long, etc.) + let state_changed = match (&state, &prev_usb_state) { + (UsbStateInternal::PinNotSet { .. }, UsbStateInternal::PinNotSet { .. }) => true, + (UsbStateInternal::NeedsPin { .. }, UsbStateInternal::NeedsPin { .. }) => true, + (new_state, old_state) => { + std::mem::discriminant(new_state) != std::mem::discriminant(old_state) + } + }; + if state_changed { + tracing::debug!("USB current state: {state:?}"); + tx.send(state.clone()).await.map_err(|_| { + Error::Internal("USB state channel receiver closed prematurely".to_string()) + })?; + } + + // Check for terminal states AFTER sending + match state { + UsbStateInternal::Completed(_) => break Ok(()), + UsbStateInternal::Failed(err) => break Err(err), + _ => {} + } } - UsbStateInternal::SelectCredential { - ref response, - cred_tx: _, - } => Self::process_select_credential(response, &mut cred_rx).await, - UsbStateInternal::Completed(_) => break Ok(()), - UsbStateInternal::Failed(err) => break Err(err), - }; - state = next_usb_state.unwrap_or_else(UsbStateInternal::Failed); - // Usually, comparing the Discrimimant is enough, but PinNotSet can be - // repeated multiple times with different or the same error reasons - // (PIN too short, PIN too long, etc.) - let state_changed = match (&state, &prev_usb_state) { - (UsbStateInternal::PinNotSet { .. }, UsbStateInternal::PinNotSet { .. }) => true, - (new_state, old_state) => { - std::mem::discriminant(new_state) != std::mem::discriminant(old_state) + _ = cancellation.cancelled() => { + tracing::debug!("USB handler cancelled, stopping processing"); + break Err(Error::Internal("Request cancelled".to_string())); } - }; - if state_changed { - tracing::debug!("USB current state: {state:?}"); - tx.send(state.clone()).await.map_err(|_| { - Error::Internal("USB state channel receiver closed prematurely".to_string()) - })?; } } } @@ -306,6 +359,7 @@ async fn handle_events( cred_request: &CredentialRequest, device: Arc>, signal_tx: &Sender>, + cancellation: CancellationToken, ) { let mut device = device.lock().await; let device_debug = device.to_string(); @@ -322,6 +376,7 @@ async fn handle_events( ); } Ok(mut channel) => { + let cancel_handle = channel.get_handle(); let signal_tx2 = signal_tx.clone().downgrade(); let ux_updates_rx = channel.get_ux_update_receiver(); tokio::spawn(async move { @@ -332,49 +387,58 @@ async fn handle_events( "Polling for credential from USB authenticator {}", &device_debug ); - let response: Result = loop { - let response = match cred_request { - CredentialRequest::CreatePublicKeyCredentialRequest(make_cred_request) => { - channel - .webauthn_make_credential(make_cred_request) - .await - .map(|response| { - UsbUvMessage::ReceivedCredentials(Box::new(response.into())) - }) - } - CredentialRequest::GetPublicKeyCredentialRequest(get_cred_request) => channel - .webauthn_get_assertion(get_cred_request) - .await - .map(|response| { - UsbUvMessage::ReceivedCredentials(Box::new(response.into())) - }), - }; - match response { - Ok(response) => { - tracing::debug!("Received credential from USB authenticator"); - break Ok(response); - } - Err(WebAuthnError::Ctap(ctap_error)) - if ctap_error.is_retryable_user_error() => - { - warn!("Retrying WebAuthn credential operation"); - continue; - } - Err(err) => { - tracing::warn!( - "Failed to make/get credential with USB authenticator: {:?}", - err - ); - break Err(err); + let response: Result = tokio::select! { + res = async { + loop { + let response = match cred_request { + CredentialRequest::CreatePublicKeyCredentialRequest(make_cred_request) => { + channel + .webauthn_make_credential(make_cred_request) + .await + .map(|response| { + UsbUvMessage::ReceivedCredentials(Box::new(response.into())) + }) + } + CredentialRequest::GetPublicKeyCredentialRequest(get_cred_request) => channel + .webauthn_get_assertion(get_cred_request) + .await + .map(|response| { + UsbUvMessage::ReceivedCredentials(Box::new(response.into())) + }), + }; + match response { + Ok(response) => { + tracing::debug!("Received credential from USB authenticator"); + break Ok(response); + } + Err(WebAuthnError::Ctap(ctap_error)) + if ctap_error.is_retryable_user_error() => + { + warn!("Retrying WebAuthn credential operation"); + continue; + } + Err(err) => { + tracing::warn!( + "Failed to make/get credential with USB authenticator: {:?}", + err + ); + break Err(err); + } + } } + .map_err(|err| match err { + WebAuthnError::Ctap(CtapError::PINAuthBlocked) => Error::PinAttemptsExhausted, + WebAuthnError::Ctap(CtapError::NoCredentials) => Error::NoCredentials, + WebAuthnError::Ctap(CtapError::CredentialExcluded) => Error::CredentialExcluded, + _ => Error::AuthenticatorError, + }) + } => res, + _ = cancellation.cancelled() => { + tracing::debug!("USB ceremony cancelled, interrupting authenticator operation"); + cancel_handle.cancel_ongoing_operation().await; + Err(Error::Internal("Request cancelled".to_string())) } - } - .map_err(|err| match err { - WebAuthnError::Ctap(CtapError::PINAuthBlocked) => Error::PinAttemptsExhausted, - WebAuthnError::Ctap(CtapError::NoCredentials) => Error::NoCredentials, - WebAuthnError::Ctap(CtapError::CredentialExcluded) => Error::CredentialExcluded, - _ => Error::AuthenticatorError, - }); + }; if let Err(err) = signal_tx.send(response).await { tracing::error!("Failed to notify that ceremony completed: {:?}", err); } @@ -386,13 +450,14 @@ impl UsbHandler for InProcessUsbHandler { fn start( &self, request: &CredentialRequest, + cancellation: CancellationToken, ) -> impl Stream + Send + Sized + Unpin + 'static { let request = request.clone(); let (tx, mut rx) = mpsc::channel(32); tokio::spawn(async move { // TODO: instead of logging error here, push the errors into the // stream so credential service can handle/forward them to the UI - if let Err(err) = InProcessUsbHandler::process(tx, request).await { + if let Err(err) = InProcessUsbHandler::process(tx, request, cancellation).await { tracing::error!("Error getting credential from USB: {:?}", err); } }); diff --git a/credentialsd/src/dbus/flow_control.rs b/credentialsd/src/dbus/flow_control.rs index f196a811..e2d41db9 100644 --- a/credentialsd/src/dbus/flow_control.rs +++ b/credentialsd/src/dbus/flow_control.rs @@ -95,7 +95,7 @@ async fn handle, ) -> Result { let (request_tx, request_rx) = oneshot::channel(); - let request_id = svc.lock().await.init_request(&msg, request_tx).await?; + let (request_id, cancellation_token) = svc.lock().await.init_request(&msg, request_tx).await?; let operation = msg.operation(); let rp_id = msg.relying_party_id().to_string(); @@ -146,7 +146,13 @@ async fn handle>>> = Arc::new(Mutex::new(None)); let set_pin_tx: Arc>>> = Arc::new(Mutex::new(None)); let cred_selector_tx = Arc::new(Mutex::new(None)); - while let Some(ui_request) = flow.receive_ui_event().await { + loop { + tokio::select! { + ui_request = flow.receive_ui_event() => { + let Some(ui_request) = ui_request else { + tracing::debug!("UI event stream closed"); + break; + }; match ui_request { UserInteractedEvent::DiscoveryRequested => { let client_pin_tx = client_pin_tx.clone(); @@ -260,9 +266,16 @@ async fn handle { - tracing::debug!(%request_id, "Cancelling request"); - svc.lock().await.cancel_request(request_id).await; + UserInteractedEvent::RequestCancelled => { + tracing::debug!(%request_id, "Cancelling request"); + svc.lock().await.cancel_request(request_id).await; + break; + } + } + } + _ = cancellation_token.cancelled() => { + tracing::debug!("Request cancelled, stopping UI event handler"); + break; } } } @@ -286,6 +299,7 @@ fn forward_background_event_stream( break; } } + tracing::debug!("Background event stream ended"); }); }