avoid flicker when changing from lockscreen to waiting screen

The BitBox shows the lockscreen ("See the BitBoxApp") first, and when
the noise connection is initiazed, switches to the BitBox logo, which
stays there as the waiting/default screen until reconnected.

The noise connection is initiated by the host after unlock finishes,
so in between, "See the BitBoxApP" flashes for a split second.

This commit solves this by

1) not pushing the lockscreen as an extra component to be popped, but
folding it into the top/default component

2) updating that screen when the app initiates the connection if the
device is not initialized

3) if initialized, updating the screen on OP_UNLOCK (see comment in code)

This prevents flashing/flickering if the device is initialized. There
is still a small flicker when the device is uninitialized, as the
screen is updated before the pairing confirmation appears. This was
the case before this commit too.

Author: benma

src/rust/bitbox02-rust/src/hww.rs

@@ -110,6 +110,13 @@ fn api_attestation(usb_in: &[u8]) -> Vec<u8> {
 }
 
 async fn _process_packet(hal: &mut impl crate::hal::Hal, usb_in: Vec<u8>) -> Vec<u8> {
+    // Update the waiting screen from "See the BitBoxApp" to the logo, now that the host is
+    // connected. When the device is initialized, we delay this until the unlock call, otherwise
+    // there would be a flicker where the logo would be shown before the host invokes unlock.
+    if !bitbox02::memory::is_initialized() || usb_in.as_slice() == &[OP_UNLOCK] {
+        bitbox02::ui::screen_process_host_connected();
+    }
+
     match usb_in.split_first() {
         Some((&OP_UNLOCK, b"")) => return api_unlock(hal).await,
         Some((&OP_ATTESTATION, rest)) => return api_attestation(rest),

src/rust/bitbox02-rust/src/hww/noise.rs

@@ -82,8 +82,6 @@ pub(crate) async fn process(
 ) -> Result<(), Error> {
     match usb_in.split_first() {
         Some((&OP_I_CAN_HAS_HANDSHAEK, b"")) => {
-            // The previous screen was "See the BitBoxApp".
-            // Since a handshake was requested, a client was connected, so we pop that screen.
             // Pairing is the start of a session, so we clean the screen stack in case
             // we started a new session in the middle of something.
             bitbox02::ui::screen_stack_pop_all();

src/rust/bitbox02-sys/build.rs

@@ -129,6 +129,7 @@ const ALLOWLIST_FNS: &[&str] = &[
     "reset_ble",
     "screen_print_debug",
     "screen_process",
+    "screen_process_host_connected",
     "screen_saver_disable",
     "screen_saver_enable",
     "sd_card_inserted",

src/rust/bitbox02/src/ui.rs

@@ -25,3 +25,7 @@ mod types;
 mod ui;
 
 pub use ui::*;
+
+pub fn screen_process_host_connected() {
+    unsafe { bitbox02_sys::screen_process_host_connected() }
+}

src/ui/components/waiting.c

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "waiting.h"
+#include "lockscreen.h"
 
 #include "image.h"
 #include "ui_images.h"
@@ -21,13 +22,12 @@
 #include <screen.h>
 #include <ui/ui_util.h>
 
+#include <stdbool.h>
 #include <string.h>
 
-static void _render(component_t* component)
-{
-    // TODO - add an interesting animation?
-    ui_util_component_render_subcomponents(component);
-}
+typedef struct {
+    bool host_connected;
+} data_t;
 
 /********************************** Component Functions **********************************/
 
@@ -36,7 +36,7 @@ static void _render(component_t* component)
  */
 static component_functions_t _component_functions = {
     .cleanup = ui_util_component_cleanup,
-    .render = _render,
+    .render = ui_util_component_render_subcomponents,
     .on_event = ui_util_on_event_noop,
 };
 
@@ -47,6 +47,12 @@ static component_functions_t _component_functions = {
  */
 component_t* waiting_create(void)
 {
+    data_t* data = malloc(sizeof(data_t));
+    if (!data) {
+        Abort("Error: malloc waiting data");
+    }
+    memset(data, 0, sizeof(data_t));
+
     component_t* waiting = malloc(sizeof(component_t));
     if (!waiting) {
         Abort("Error: malloc waiting");
@@ -57,14 +63,36 @@ component_t* waiting_create(void)
     waiting->dimension.height = SCREEN_HEIGHT;
     waiting->position.top = 0;
     waiting->position.left = 0;
+    waiting->data = data;
+
+    ui_util_add_sub_component(waiting, lockscreen_create());
+
+    return waiting;
+}
+
+void waiting_host_connected(component_t* component)
+{
+    data_t* data = (data_t*)component->data;
+    if (data->host_connected) {
+        return;
+    }
+    data->host_connected = true;
+
+    if (component->sub_components.amount != 1) {
+        // Sanity check to avoid memory bugs, should never happen.
+        return;
+    }
+
+    ui_util_component_cleanup(component->sub_components.sub_components[0]);
+
     image_logo_data_t logo = image_logo_data();
     component_t* bb2_logo = image_create(
         logo.buffer.data,
         logo.buffer.len,
         logo.dimensions.width,
         logo.dimensions.height,
         CENTER,
-        waiting);
-    ui_util_add_sub_component(waiting, bb2_logo);
-    return waiting;
+        component);
+
+    component->sub_components.sub_components[0] = bb2_logo;
 }

src/ui/components/waiting.h

@@ -18,8 +18,14 @@
 #include <ui/component.h>
 
 /**
- * Creates a waiting screen.
+ * Creates a waiting screen. It starts out with a lockscreen (see lockscreen.c). Once
+ * `waiting_host_connected()` is called, the waiting screen will switch to showing the logo image.
  */
 component_t* waiting_create(void);
 
+/**
+ * Switch the waiting screen to show the BitBox logo instead.
+ */
+void waiting_host_connected(component_t* component);
+
 #endif

src/ui/screen_process.c

@@ -32,16 +32,21 @@ void ui_screen_render_component(component_t* component)
     UG_SendBuffer();
 }
 
+static component_t* _waiting_screen = NULL;
 static component_t* _get_waiting_screen(void)
 {
-    static component_t* waiting_screen = NULL;
-    if (waiting_screen == NULL) {
-        waiting_screen = waiting_create();
-        if (waiting_screen == NULL) {
+    if (_waiting_screen == NULL) {
+        _waiting_screen = waiting_create();
+        if (_waiting_screen == NULL) {
             Abort("Could not create\nwaiting screen");
         }
     }
-    return waiting_screen;
+    return _waiting_screen;
+}
+
+void screen_process_host_connected(void)
+{
+    waiting_host_connected(_get_waiting_screen());
 }
 
 component_t* screen_process_get_top_component(void)

src/ui/screen_process.h

@@ -27,6 +27,11 @@ void ui_screen_render_component(component_t* component);
  */
 component_t* screen_process_get_top_component(void);
 
+/**
+ * Wraps `waiting_host_connected()` for the waiting screen.
+ */
+void screen_process_host_connected(void);
+
 /**
  * Runs the UI once.
  *

src/workflow/orientation_screen.c

@@ -73,7 +73,6 @@ static void _idle_timer_cb(const struct timer_task* const timer_task)
     }
 
     usb_start();
-    ui_screen_stack_push(lockscreen_create());
 }
 #endif