ui: port tx fee confirm to async

diff --git a/src/rust/bitbox02-rust/src/workflow/transaction.rs b/src/rust/bitbox02-rust/src/workflow/transaction.rs
index 22e0e3e..0c4d722 100644
--- a/src/rust/bitbox02-rust/src/workflow/transaction.rs
+++ b/src/rust/bitbox02-rust/src/workflow/transaction.rs
@@ -2,10 +2,6 @@
 
 use crate::hal::Ui;
 
-use crate::bb02_async::option_no_screensaver;
-use core::cell::RefCell;
-
-use alloc::boxed::Box;
 use alloc::string::String;
 
 pub struct UserAbort;
@@ -23,18 +19,10 @@ fn format_percentage(p: f64) -> String {
 }
 
 pub async fn verify_total_fee(total: &str, fee: &str, longtouch: bool) -> Result<(), UserAbort> {
-    let result = RefCell::new(None as Option<Result<(), UserAbort>>);
-
-    let mut component = bitbox02::ui::confirm_transaction_fee_create(
-        total,
-        fee,
-        longtouch,
-        Box::new(|ok| {
-            *result.borrow_mut() = Some(if ok { Ok(()) } else { Err(UserAbort) });
-        }),
-    );
-    component.screen_stack_push();
-    option_no_screensaver(&result).await
+    match bitbox02::ui::confirm_transaction_fee(total, fee, longtouch).await {
+        bitbox02::ui::ConfirmResponse::Approved => Ok(()),
+        bitbox02::ui::ConfirmResponse::Cancelled => Err(UserAbort),
+    }
 }
 
 pub async fn verify_total_fee_maybe_warn(
diff --git a/src/rust/bitbox02/src/ui/ui.rs b/src/rust/bitbox02/src/ui/ui.rs
index c836b4b..5321510 100644
--- a/src/rust/bitbox02/src/ui/ui.rs
+++ b/src/rust/bitbox02/src/ui/ui.rs
@@ -21,7 +21,6 @@ use core::marker::PhantomData;
 pub struct Component<'a> {
     component: *mut bitbox02_sys::component_t,
     is_pushed: bool,
-    on_drop: Option<Box<dyn FnMut()>>,
     // This is used to have the result callbacks outlive the component.
     _p: PhantomData<&'a ()>,
 }
@@ -46,9 +45,6 @@ impl Drop for Component<'_> {
         unsafe {
             bitbox02_sys::ui_screen_stack_pop();
         }
-        if let Some(ref mut on_drop) = self.on_drop {
-            (*on_drop)();
-        }
     }
 }
 
@@ -148,7 +144,6 @@ pub async fn trinary_input_string(
     let mut component = Component {
         component,
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     };
     component.screen_stack_push();
@@ -232,7 +227,6 @@ pub async fn confirm(params: &ConfirmParams<'_>) -> ConfirmResponse {
     let mut component = Component {
         component,
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     };
     component.screen_stack_push();
@@ -270,7 +264,6 @@ pub fn status_create<'a>(text: &str, status_success: bool) -> Component<'a> {
     Component {
         component,
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     }
 }
@@ -315,7 +308,6 @@ pub async fn sdcard() -> SdcardResponse {
     let mut component = Component {
         component,
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     };
     component.screen_stack_push();
@@ -437,7 +429,6 @@ pub async fn menu(params: MenuParams<'_>) -> MenuResponse {
     let mut component = Component {
         component,
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     };
     component.screen_stack_push();
@@ -540,7 +531,6 @@ pub async fn trinary_choice(
     let mut component = Component {
         component,
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     };
     component.screen_stack_push();
@@ -603,7 +593,6 @@ pub async fn confirm_transaction_address(amount: &str, address: &str) -> Confirm
     let mut component = Component {
         component,
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     };
     component.screen_stack_push();
@@ -625,36 +614,67 @@ pub async fn confirm_transaction_address(amount: &str, address: &str) -> Confirm
     .await
 }
 
-pub fn confirm_transaction_fee_create<'a, 'b>(
-    amount: &'a str,
-    fee: &'a str,
-    longtouch: bool,
-    callback: AcceptRejectCb<'b>,
-) -> Component<'b> {
-    unsafe extern "C" fn c_callback(result: bool, user_data: *mut c_void) {
-        let callback = user_data as *mut AcceptRejectCb;
-        unsafe { (*callback)(result) };
+pub async fn confirm_transaction_fee(amount: &str, fee: &str, longtouch: bool) -> ConfirmResponse {
+    let _no_screensaver = crate::screen_saver::ScreensaverInhibitor::new();
+
+    // Shared between the async context and the c callback
+    struct SharedState {
+        waker: Option<Waker>,
+        result: Option<ConfirmResponse>,
+    }
+    let shared_state = Box::new(RefCell::new(SharedState {
+        waker: None,
+        result: None,
+    }));
+    let shared_state_ptr = shared_state.as_ref() as *const RefCell<SharedState> as *mut c_void;
+
+    unsafe extern "C" fn callback(result: bool, user_data: *mut c_void) {
+        let shared_state = unsafe { &*(user_data as *mut RefCell<SharedState>) };
+        let mut shared_state = shared_state.borrow_mut();
+        if shared_state.result.is_none() {
+            shared_state.result = Some(if result {
+                ConfirmResponse::Approved
+            } else {
+                ConfirmResponse::Cancelled
+            });
+            if let Some(waker) = shared_state.waker.as_ref() {
+                waker.wake_by_ref();
+            }
+        }
     }
 
-    let user_data = Box::into_raw(Box::new(callback)) as *mut c_void;
     let component = unsafe {
         bitbox02_sys::confirm_transaction_fee_create(
             util::strings::str_to_cstr_vec(amount).unwrap().as_ptr(), // copied in C
             util::strings::str_to_cstr_vec(fee).unwrap().as_ptr(),    // copied in C
             longtouch,
-            Some(c_callback as _),
-            user_data,
+            Some(callback),
+            shared_state_ptr, // passed to callback as `user_data`.
         )
     };
-    Component {
+
+    let mut component = Component {
         component,
         is_pushed: false,
-        on_drop: Some(Box::new(move || unsafe {
-            // Drop all callbacks.
-            drop(Box::from_raw(user_data as *mut AcceptRejectCb));
-        })),
         _p: PhantomData,
-    }
+    };
+    component.screen_stack_push();
+
+    core::future::poll_fn({
+        let shared_state = &shared_state;
+        move |cx| {
+            let mut shared_state = shared_state.borrow_mut();
+
+            if let Some(result) = shared_state.result.take() {
+                Poll::Ready(result)
+            } else {
+                // Store the waker so the callback can wake up this task
+                shared_state.waker = Some(cx.waker().clone());
+                Poll::Pending
+            }
+        }
+    })
+    .await
 }
 
 pub fn screen_stack_pop_all() {
@@ -673,7 +693,6 @@ pub fn progress_create<'a>(title: &str) -> Component<'a> {
     Component {
         component,
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     }
 }
@@ -686,7 +705,6 @@ pub fn empty_create<'a>() -> Component<'a> {
     Component {
         component: unsafe { bitbox02_sys::empty_create() },
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     }
 }
@@ -714,7 +732,6 @@ where
     Component {
         component,
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     }
 }
@@ -750,7 +767,6 @@ pub async fn choose_orientation() -> bool {
     let mut component = Component {
         component,
         is_pushed: false,
-        on_drop: None,
         _p: PhantomData,
     };
     component.screen_stack_push();
diff --git a/src/rust/bitbox02/src/ui/ui_stub.rs b/src/rust/bitbox02/src/ui/ui_stub.rs
index e1aa402..8aa474e 100644
--- a/src/rust/bitbox02/src/ui/ui_stub.rs
+++ b/src/rust/bitbox02/src/ui/ui_stub.rs
@@ -78,12 +78,11 @@ pub async fn confirm_transaction_address(_amount: &str, _address: &str) -> Confi
     panic!("not used");
 }
 
-pub fn confirm_transaction_fee_create<'a, 'b>(
-    _amount: &'a str,
-    _fee: &'a str,
+pub async fn confirm_transaction_fee(
+    _amount: &str,
+    _fee: &str,
     _longtouch: bool,
-    _callback: AcceptRejectCb<'b>,
-) -> Component<'b> {
+) -> ConfirmResponse {
     panic!("not used");
 }
 
diff --git a/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs b/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs
index 2c615af..42ede4b 100644
--- a/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs
+++ b/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs
@@ -96,21 +96,16 @@ pub async fn confirm_transaction_address(_amount: &str, _address: &str) -> Confi
     ConfirmResponse::Approved
 }
 
-pub fn confirm_transaction_fee_create<'a, 'b>(
-    _amount: &'a str,
-    _fee: &'a str,
+pub async fn confirm_transaction_fee(
+    _amount: &str,
+    _fee: &str,
     _longtouch: bool,
-    mut callback: AcceptRejectCb<'b>,
-) -> Component<'b> {
+) -> ConfirmResponse {
     crate::print_stdout(&format!(
         "CONFIRM TRANSACTION FEE SCREEN START\nAMOUNT: {}\nFEE: {}\nCONFIRM TRANSACTION FEE SCREEN END\n",
         _amount, _fee
     ));
-    callback(true);
-    Component {
-        is_pushed: false,
-        _p: PhantomData,
-    }
+    ConfirmResponse::Approved
 }
 
 pub fn screen_stack_pop_all() {}