From 48c60b644a9da2d77def5052c1fdc088289f899e Mon Sep 17 00:00:00 2001
From: tychovrahe <brunam@seznam.cz>
Date: Thu, 24 Aug 2023 19:16:09 +0200
Subject: [PATCH] feat(core): add install restricted screen on installation
 with locked bootloader

[no changelog]
---
 core/embed/bootloader/bootui.c   | 12 ++++++++++++
 core/embed/bootloader/bootui.h   |  1 +
 core/embed/bootloader/main.c     | 16 +++++++---------
 core/embed/bootloader/messages.c |  8 ++++----
 core/embed/bootloader/messages.h |  6 +++---
 5 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/core/embed/bootloader/bootui.c b/core/embed/bootloader/bootui.c
index 369b479af7..39a0a0ed29 100644
--- a/core/embed/bootloader/bootui.c
+++ b/core/embed/bootloader/bootui.c
@@ -264,6 +264,18 @@ void ui_screen_fail(void) { screen_install_fail(); }
 uint32_t ui_screen_unlock_bootloader_confirm(void) {
   return screen_unlock_bootloader_confirm();
 }
+
+void ui_screen_install_restricted(void) {
+  display_clear();
+  screen_fatal_error_rust(
+      "INSTALL RESTRICTED",
+      "Installation of custom firmware is currently restricted.",
+      "Please visit\ntrezor.io/bootloader");
+
+  display_refresh();
+}
+#else
+void ui_screen_install_restricted(void) { screen_install_fail(); }
 #endif
 
 // general functions
diff --git a/core/embed/bootloader/bootui.h b/core/embed/bootloader/bootui.h
index 36f7b2cb35..a73349d018 100644
--- a/core/embed/bootloader/bootui.h
+++ b/core/embed/bootloader/bootui.h
@@ -61,6 +61,7 @@ void ui_screen_wipe_progress(int pos, int len);
 void ui_screen_done(uint8_t restart_seconds, secbool full_redraw);
 
 void ui_screen_fail(void);
+void ui_screen_install_restricted(void);
 
 void ui_fadein(void);
 void ui_fadeout(void);
diff --git a/core/embed/bootloader/main.c b/core/embed/bootloader/main.c
index 605fae8286..9772841ccb 100644
--- a/core/embed/bootloader/main.c
+++ b/core/embed/bootloader/main.c
@@ -188,7 +188,11 @@ static usb_result_t bootloader_usb_loop(const vendor_header *const vhdr,
       case MessageType_MessageType_FirmwareUpload:
         r = process_msg_FirmwareUpload(USB_IFACE_NUM, msg_size, buf);
         if (r < 0 && r != UPLOAD_ERR_USER_ABORT) {  // error, but not user abort
-          ui_screen_fail();
+          if (r == UPLOAD_ERR_BOOTLOADER_LOCKED) {
+            ui_screen_install_restricted();
+          } else {
+            ui_screen_fail();
+          }
           usb_stop();
           usb_deinit();
           return SHUTDOWN;
@@ -225,7 +229,7 @@ static usb_result_t bootloader_usb_loop(const vendor_header *const vhdr,
           usb_deinit();
           return RETURN;
         }
-        process_msg_AttestationDelete(USB_IFACE_NUM, msg_size, buf);
+        process_msg_UnlockBootloader(USB_IFACE_NUM, msg_size, buf);
         screen_unlock_bootloader_success();
         hal_delay(100);
         usb_stop();
@@ -549,13 +553,7 @@ int bootloader_main(void) {
 
 #ifdef USE_OPTIGA
   if (((vhdr.vtrust & VTRUST_SECRET) != 0) && (sectrue != secret_wiped())) {
-    display_clear();
-    screen_fatal_error_rust(
-        "INSTALL RESTRICTED",
-        "Installation of custom firmware is currently restricted.",
-        "Please visit\ntrezor.io/bootloader");
-
-    display_refresh();
+    ui_screen_install_restricted();
     return 1;
   }
 #endif
diff --git a/core/embed/bootloader/messages.c b/core/embed/bootloader/messages.c
index 33c4a6a790..c6860b58bb 100644
--- a/core/embed/bootloader/messages.c
+++ b/core/embed/bootloader/messages.c
@@ -577,9 +577,9 @@ int process_msg_FirmwareUpload(uint8_t iface_num, uint32_t msg_size,
       if (sectrue != secret_wiped() && ((vhdr.vtrust & VTRUST_SECRET) != 0)) {
         MSG_SEND_INIT(Failure);
         MSG_SEND_ASSIGN_VALUE(code, FailureType_Failure_ProcessError);
-        MSG_SEND_ASSIGN_STRING(message, "Attestation present");
+        MSG_SEND_ASSIGN_STRING(message, "Install restricted");
         MSG_SEND(Failure);
-        return UPLOAD_ERR_ATTESTATION_PRESENT;
+        return UPLOAD_ERR_BOOTLOADER_LOCKED;
       }
 #endif
 
@@ -735,8 +735,8 @@ void process_msg_unknown(uint8_t iface_num, uint32_t msg_size, uint8_t *buf) {
 }
 
 #ifdef USE_OPTIGA
-void process_msg_AttestationDelete(uint8_t iface_num, uint32_t msg_size,
-                                   uint8_t *buf) {
+void process_msg_UnlockBootloader(uint8_t iface_num, uint32_t msg_size,
+                                  uint8_t *buf) {
   secret_erase();
   MSG_SEND_INIT(Success);
   MSG_SEND(Success);
diff --git a/core/embed/bootloader/messages.h b/core/embed/bootloader/messages.h
index be8952941d..45a128560b 100644
--- a/core/embed/bootloader/messages.h
+++ b/core/embed/bootloader/messages.h
@@ -41,7 +41,7 @@ enum {
   UPLOAD_ERR_USER_ABORT = -7,
   UPLOAD_ERR_FIRMWARE_TOO_BIG = -8,
   UPLOAD_ERR_INVALID_CHUNK_HASH = -9,
-  UPLOAD_ERR_ATTESTATION_PRESENT = -10,
+  UPLOAD_ERR_BOOTLOADER_LOCKED = -10,
 };
 
 enum {
@@ -69,8 +69,8 @@ int process_msg_WipeDevice(uint8_t iface_num, uint32_t msg_size, uint8_t *buf);
 void process_msg_unknown(uint8_t iface_num, uint32_t msg_size, uint8_t *buf);
 
 #ifdef USE_OPTIGA
-void process_msg_AttestationDelete(uint8_t iface_num, uint32_t msg_size,
-                                   uint8_t *buf);
+void process_msg_UnlockBootloader(uint8_t iface_num, uint32_t msg_size,
+                                  uint8_t *buf);
 #endif
 
 secbool bootloader_WipeDevice(void);