From fdecf124c4b2d38940877c73712453a2118b7ccd Mon Sep 17 00:00:00 2001 From: Niklas Dusenlund Date: Mon, 6 Apr 2026 23:00:49 +0200 Subject: [PATCH 1/3] util: Replace AsyncOption with futures::completion --- src/rust/bitbox02-rust/src/async_usb.rs | 59 +++++++++++-------------- src/rust/util/src/bb02_async.rs | 20 --------- 2 files changed, 27 insertions(+), 52 deletions(-) diff --git a/src/rust/bitbox02-rust/src/async_usb.rs b/src/rust/bitbox02-rust/src/async_usb.rs index 26d58cd106..902a16e5f9 100644 --- a/src/rust/bitbox02-rust/src/async_usb.rs +++ b/src/rust/bitbox02-rust/src/async_usb.rs @@ -7,7 +7,8 @@ use alloc::boxed::Box; use alloc::vec::Vec; use core::cell::RefCell; use core::task::Poll; -use util::bb02_async::{Task, option, spin as spin_task}; +use util::bb02_async::{Task, spin as spin_task}; +use util::futures::completion; type UsbOut = Vec; type UsbIn = Vec; @@ -24,22 +25,11 @@ enum WaitingForNextRequestState { /// Since we have a strict request<->response model, we always need to send a response before /// getting another request. In this state, we are ready to send a response and are waiting for /// the host to fetch it. - SendingResponse(UsbOut), + SendingResponse(UsbOut, completion::Responder), /// Host got the response, now we are waiting for the next request by the host. - AwaitingRequest, + AwaitingRequest(completion::Responder), } -/// A safer version of `Option`. RefCell so we cannot accidentally borrow illegally. -struct SafeNextRequest(RefCell>); - -/// Safety: this implements Sync even though it is not thread safe. This is okay, as we -/// run only in a single thread in the BitBox02. -unsafe impl Sync for SafeNextRequest {} - -/// An option resolving the `next_request()` future. It is `Some(...)` once a request we've been -/// waiting for arrives. See `next_request()` for more details. -static NEXT_REQUEST: SafeNextRequest = SafeNextRequest(RefCell::new(None)); - /// Describes the global state of an api query. The documentation of /// the variants apply to the HWW stack, but have analogous meaning in /// the U2F stack. @@ -101,7 +91,7 @@ where pub fn waiting_for_next_request() -> bool { matches!( *USB_TASK_STATE.0.borrow(), - UsbTaskState::Running(Some(_), WaitingForNextRequestState::AwaitingRequest) + UsbTaskState::Running(Some(_), WaitingForNextRequestState::AwaitingRequest(_)) ) } @@ -113,14 +103,15 @@ pub fn is_idle() -> bool { /// this, otherwise this function panics. pub fn on_next_request(usb_in: &[u8]) { let mut state = USB_TASK_STATE.0.borrow_mut(); - match *state { + match &mut *state { UsbTaskState::Running( Some(_), - ref mut next_request_state @ WaitingForNextRequestState::AwaitingRequest, + next_request_state @ WaitingForNextRequestState::AwaitingRequest(_), ) => { - // Resolve NEXT_REQUEST future. - *NEXT_REQUEST.0.borrow_mut() = Some(usb_in.to_vec()); - + let WaitingForNextRequestState::AwaitingRequest(responder) = next_request_state else { + unreachable!(); + }; + responder.resolve(usb_in.to_vec()); *next_request_state = WaitingForNextRequestState::Idle; } _ => panic!("on_next_request: wrong state"), @@ -188,12 +179,15 @@ pub fn take_response() -> Result { match &mut *state { UsbTaskState::Nothing => Err(CopyResponseErr::NotRunning), UsbTaskState::Running(Some(_), next_request_state) => { - if let WaitingForNextRequestState::SendingResponse(response) = next_request_state { - let response = core::mem::take(response); - *next_request_state = WaitingForNextRequestState::AwaitingRequest; - Ok(response) - } else { - Err(CopyResponseErr::NotReady) + match core::mem::replace(next_request_state, WaitingForNextRequestState::Idle) { + WaitingForNextRequestState::SendingResponse(response, responder) => { + *next_request_state = WaitingForNextRequestState::AwaitingRequest(responder); + Ok(response) + } + next_request_state_value => { + *next_request_state = next_request_state_value; + Err(CopyResponseErr::NotReady) + } } } UsbTaskState::Running(_, _) => Err(CopyResponseErr::NotReady), @@ -213,7 +207,6 @@ pub fn take_response() -> Result { /// able to read the result (e.g. when resetting the BLE chip as part of a task), so another task /// can spawn afterwards immediately instead of being blocked by stale executor state. pub fn cancel() { - let _ = NEXT_REQUEST.0.borrow_mut().take(); let mut state = USB_TASK_STATE.0.borrow_mut(); *state = UsbTaskState::Nothing; } @@ -223,17 +216,19 @@ pub fn cancel() { pub async fn next_request(response: UsbOut) -> UsbIn { // Scope so that `state` is dropped before `.await`, see // https://rust-lang.github.io/rust-clippy/master/index.html#await_holding_refcell_ref - { + let result = { let mut state = USB_TASK_STATE.0.borrow_mut(); + let (responder, result) = completion::completion(); match *state { UsbTaskState::Running(None, ref mut next_request_state) => { - *next_request_state = WaitingForNextRequestState::SendingResponse(response); + *next_request_state = + WaitingForNextRequestState::SendingResponse(response, responder); } _ => panic!("next_request() called in wrong state"), } - } - - option(&NEXT_REQUEST.0).await + result + }; + result.await } #[cfg(test)] diff --git a/src/rust/util/src/bb02_async.rs b/src/rust/util/src/bb02_async.rs index fd4129fe31..48ff756b12 100644 --- a/src/rust/util/src/bb02_async.rs +++ b/src/rust/util/src/bb02_async.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 use alloc::boxed::Box; -use core::cell::RefCell; use core::pin::Pin; use core::task::{Context, Poll}; @@ -19,25 +18,6 @@ pub fn spin(task: &mut Task) -> Poll { task.as_mut().poll(context) } -/// Implements the Option future, see `option()`. -pub struct AsyncOption<'a, O>(&'a RefCell>); - -impl core::future::Future for AsyncOption<'_, O> { - type Output = O; - fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { - match self.0.borrow_mut().take() { - None => Poll::Pending, - Some(output) => Poll::Ready(output), - } - } -} - -/// Waits for an option to contain a value and returns that value, leaving `None` in its place. -/// E.g. `assert_eq!(option(&Some(42)).await, 42)`. -pub fn option(option: &RefCell>) -> AsyncOption<'_, O> { - AsyncOption(option) -} - /// Polls a future until the result is available. #[cfg(feature = "testing")] pub fn block_on(task: impl core::future::Future) -> O { From 58400c7b40945ab4c1ee5e797ab178ad83dff07c Mon Sep 17 00:00:00 2001 From: Niklas Dusenlund Date: Mon, 6 Apr 2026 23:15:18 +0200 Subject: [PATCH 2/3] util: remove spin and waker_fn --- src/rust/bitbox02-rust/src/async_usb.rs | 7 ++-- src/rust/util/src/bb02_async.rs | 13 ++----- src/rust/util/src/lib.rs | 1 - src/rust/util/src/waker_fn.rs | 47 ------------------------- 4 files changed, 6 insertions(+), 62 deletions(-) delete mode 100644 src/rust/util/src/waker_fn.rs diff --git a/src/rust/bitbox02-rust/src/async_usb.rs b/src/rust/bitbox02-rust/src/async_usb.rs index 902a16e5f9..50908bc579 100644 --- a/src/rust/bitbox02-rust/src/async_usb.rs +++ b/src/rust/bitbox02-rust/src/async_usb.rs @@ -6,8 +6,8 @@ use alloc::boxed::Box; use alloc::vec::Vec; use core::cell::RefCell; -use core::task::Poll; -use util::bb02_async::{Task, spin as spin_task}; +use core::task::{Context, Poll, Waker}; +use util::bb02_async::Task; use util::futures::completion; type UsbOut = Vec; @@ -135,7 +135,8 @@ pub fn spin() { _ => None, }; if let Some(ref mut task) = popped_task { - let spin_result = spin_task(task); + let context = &mut Context::from_waker(Waker::noop()); + let spin_result = task.as_mut().poll(context); if matches!(*USB_TASK_STATE.0.borrow(), UsbTaskState::Nothing) { // The task was cancelled while it was running, so there is nothing to do with the // result. diff --git a/src/rust/util/src/bb02_async.rs b/src/rust/util/src/bb02_async.rs index 48ff756b12..4e0246a6ba 100644 --- a/src/rust/util/src/bb02_async.rs +++ b/src/rust/util/src/bb02_async.rs @@ -2,23 +2,14 @@ use alloc::boxed::Box; use core::pin::Pin; -use core::task::{Context, Poll}; /// Task is the top-level future which can be polled by an executor. /// Note that other futures awaited inside do not have to be pinned. -/// The 'a lifetime allows to spin a boxed/pinned future that is not +/// The 'a lifetime allows to poll a boxed/pinned future that is not /// 'static, or a future with non-'static input param references. pub type Task<'a, O> = Pin + 'a>>; -/// A primitive poll invocation for a task, with no waking functionality. -pub fn spin(task: &mut Task) -> Poll { - // TODO: statically allocate the context. - let waker = crate::waker_fn::waker_fn(|| {}); - let context = &mut Context::from_waker(&waker); - task.as_mut().poll(context) -} - -/// Polls a future until the result is available. +/// Busy wait until `task` resolves. Ignores Waker. #[cfg(feature = "testing")] pub fn block_on(task: impl core::future::Future) -> O { let mut task: crate::bb02_async::Task = alloc::boxed::Box::pin(task); diff --git a/src/rust/util/src/lib.rs b/src/rust/util/src/lib.rs index 2ba2a6dd85..3d75e62645 100644 --- a/src/rust/util/src/lib.rs +++ b/src/rust/util/src/lib.rs @@ -12,7 +12,6 @@ pub mod futures; pub mod log; pub mod name; pub mod strings; -mod waker_fn; #[cfg(feature = "p256")] mod p256; diff --git a/src/rust/util/src/waker_fn.rs b/src/rust/util/src/waker_fn.rs deleted file mode 100644 index 16a94b43c1..0000000000 --- a/src/rust/util/src/waker_fn.rs +++ /dev/null @@ -1,47 +0,0 @@ -// This file was taken from here: -// https://github.com/async-rs/async-task/blob/b7a249680490991f92cc2144d4eff65e4effb3b7/src/waker_fn.rs -// https://github.com/async-rs/async-task/blob/b7a249680490991f92cc2144d4eff65e4effb3b7/LICENSE-APACHE - -use alloc::sync::Arc; -use core::mem::{self, ManuallyDrop}; -use core::task::{RawWaker, RawWakerVTable, Waker}; - -/// Creates a waker from a wake function. -/// -/// The function gets called every time the waker is woken. -pub fn waker_fn(f: F) -> Waker { - let raw = Arc::into_raw(Arc::new(f)) as *const (); - let vtable = &Helper::::VTABLE; - unsafe { Waker::from_raw(RawWaker::new(raw, vtable)) } -} - -struct Helper(F); - -impl Helper { - const VTABLE: RawWakerVTable = RawWakerVTable::new( - Self::clone_waker, - Self::wake, - Self::wake_by_ref, - Self::drop_waker, - ); - - unsafe fn clone_waker(ptr: *const ()) -> RawWaker { - let arc = ManuallyDrop::new(unsafe { Arc::from_raw(ptr as *const F) }); - mem::forget(arc.clone()); - RawWaker::new(ptr, &Self::VTABLE) - } - - unsafe fn wake(ptr: *const ()) { - let arc = unsafe { Arc::from_raw(ptr as *const F) }; - (arc)(); - } - - unsafe fn wake_by_ref(ptr: *const ()) { - let arc = ManuallyDrop::new(unsafe { Arc::from_raw(ptr as *const F) }); - (arc)(); - } - - unsafe fn drop_waker(ptr: *const ()) { - drop(unsafe { Arc::from_raw(ptr as *const F) }); - } -} From 4a0aa3627eec8491caf0683f4c4487bd93a920d6 Mon Sep 17 00:00:00 2001 From: Niklas Dusenlund Date: Mon, 6 Apr 2026 23:18:42 +0200 Subject: [PATCH 3/3] util: Move Task to the only usage site --- src/rust/bitbox02-rust/src/async_usb.rs | 12 +++++++----- src/rust/util/src/bb02_async.rs | 9 --------- 2 files changed, 7 insertions(+), 14 deletions(-) diff --git a/src/rust/bitbox02-rust/src/async_usb.rs b/src/rust/bitbox02-rust/src/async_usb.rs index 50908bc579..26a82a905e 100644 --- a/src/rust/bitbox02-rust/src/async_usb.rs +++ b/src/rust/bitbox02-rust/src/async_usb.rs @@ -6,12 +6,14 @@ use alloc::boxed::Box; use alloc::vec::Vec; use core::cell::RefCell; +use core::future::Future; +use core::pin::Pin; use core::task::{Context, Poll, Waker}; -use util::bb02_async::Task; use util::futures::completion; type UsbOut = Vec; type UsbIn = Vec; +type UsbTask = Pin + 'static>>; /// If a task is running (see `UsbTaskState`), this state is active and manages waiting for another /// request from the host in an async fashion. This allows to have multi-request workflows in async. @@ -33,7 +35,7 @@ enum WaitingForNextRequestState { /// Describes the global state of an api query. The documentation of /// the variants apply to the HWW stack, but have analogous meaning in /// the U2F stack. -enum UsbTaskState<'a> { +enum UsbTaskState { /// Waiting for a new query, nothing to do. Nothing, /// A query came in which launched a task, which is now running (e.g. user is entering a @@ -46,7 +48,7 @@ enum UsbTaskState<'a> { /// /// The second element manages waiting for another request while processing a request, allowing /// multi-request workflows. - Running(Option>, WaitingForNextRequestState), + Running(Option, WaitingForNextRequestState), /// The task has finished and written the result, so the USB response is available. We are now /// waiting for the host to fetch it (HWW_REQ_RETRY). For short-circuited or non-async api /// calls, the result might be returned immediately in response to HWW_REQ_NEW. @@ -54,7 +56,7 @@ enum UsbTaskState<'a> { } /// A safer version of UsbTaskState. RefCell so we cannot accidentally borrow illegally. -struct SafeUsbTaskState(RefCell>); +struct SafeUsbTaskState(RefCell); /// Safety: this implements Sync even though it is not thread safe. This is okay, as we /// run only in a single thread in the BitBox02. @@ -75,7 +77,7 @@ where let mut state = USB_TASK_STATE.0.borrow_mut(); match *state { UsbTaskState::Nothing => { - let task: Task = Box::pin(workflow(usb_in.to_vec())); + let task: UsbTask = Box::pin(workflow(usb_in.to_vec())); *state = UsbTaskState::Running(Some(task), WaitingForNextRequestState::Idle); } diff --git a/src/rust/util/src/bb02_async.rs b/src/rust/util/src/bb02_async.rs index 4e0246a6ba..2af2216145 100644 --- a/src/rust/util/src/bb02_async.rs +++ b/src/rust/util/src/bb02_async.rs @@ -1,14 +1,5 @@ // SPDX-License-Identifier: Apache-2.0 -use alloc::boxed::Box; -use core::pin::Pin; - -/// Task is the top-level future which can be polled by an executor. -/// Note that other futures awaited inside do not have to be pinned. -/// The 'a lifetime allows to poll a boxed/pinned future that is not -/// 'static, or a future with non-'static input param references. -pub type Task<'a, O> = Pin + 'a>>; - /// Busy wait until `task` resolves. Ignores Waker. #[cfg(feature = "testing")] pub fn block_on(task: impl core::future::Future) -> O {