From 8a18cfe0d4e6859bff7acaf18113a53f05c46d67 Mon Sep 17 00:00:00 2001
From: tychovrahe <brunam@seznam.cz>
Date: Wed, 6 Mar 2024 15:53:50 +0100
Subject: [PATCH] feat(core): use U5 DHUK to encrypt optiga pairing secret in
 flash

[no changelog]
---
 core/embed/bootloader/main.c          |  4 --
 core/embed/firmware/main.c            | 10 ++--
 core/embed/prodtest/main.c            |  4 ++
 core/embed/prodtest/optiga_prodtest.c | 85 +++++++++++++++++----------
 core/embed/trezorhal/secret.h         |  8 ++-
 core/embed/trezorhal/stm32f4/secret.c | 11 +++-
 core/embed/trezorhal/stm32u5/secret.c | 37 +++++++++---
 7 files changed, 106 insertions(+), 53 deletions(-)

diff --git a/core/embed/bootloader/main.c b/core/embed/bootloader/main.c
index 49268983b1..1153cce6ef 100644
--- a/core/embed/bootloader/main.c
+++ b/core/embed/bootloader/main.c
@@ -442,10 +442,6 @@ int bootloader_main(void) {
 #endif
 #endif
 
-#ifdef STM32U5
-  ensure(secret_ensure_initialized(), "secret reinitialized");
-#endif
-
   ui_screen_boot_empty(false);
 
   mpu_config_bootloader();
diff --git a/core/embed/firmware/main.c b/core/embed/firmware/main.c
index db1d244eb9..61769b0001 100644
--- a/core/embed/firmware/main.c
+++ b/core/embed/firmware/main.c
@@ -141,9 +141,13 @@ int main(void) {
 
   unit_variant_init();
 
+#ifdef STM32U5
+  secure_aes_init();
+#endif
+
 #ifdef USE_OPTIGA
   uint8_t secret[SECRET_OPTIGA_KEY_LEN] = {0};
-  secbool secret_ok = secret_optiga_extract(secret);
+  secbool secret_ok = secret_optiga_get(secret);
 #endif
 
   mpu_config_firmware_initial();
@@ -166,10 +170,6 @@ int main(void) {
   set_core_clock(CLOCK_180_MHZ);
 #endif
 
-#ifdef STM32U5
-  secure_aes_init();
-#endif
-
 #ifdef USE_BUTTON
   button_init();
 #endif
diff --git a/core/embed/prodtest/main.c b/core/embed/prodtest/main.c
index c6a7144d45..7897fea989 100644
--- a/core/embed/prodtest/main.c
+++ b/core/embed/prodtest/main.c
@@ -59,6 +59,7 @@
 #include "memzero.h"
 
 #ifdef STM32U5
+#include "secure_aes.h"
 #include "stm32u5xx_ll_utils.h"
 #else
 #include "stm32f4xx_ll_utils.h"
@@ -584,6 +585,9 @@ int main(void) {
   display_reinit();
   display_orientation(0);
   random_delays_init();
+#ifdef STM32U5
+  secure_aes_init();
+#endif
 #ifdef USE_HASH_PROCESSOR
   hash_processor_init();
 #endif
diff --git a/core/embed/prodtest/optiga_prodtest.c b/core/embed/prodtest/optiga_prodtest.c
index 87e45317ef..05fbcb1289 100644
--- a/core/embed/prodtest/optiga_prodtest.c
+++ b/core/embed/prodtest/optiga_prodtest.c
@@ -33,12 +33,21 @@
 #include "secret.h"
 #include "sha2.h"
 
+#include TREZOR_BOARD
+
+#ifdef STM32U5
+#include "secure_aes.h"
+#endif
+
 typedef enum {
   OPTIGA_PAIRING_UNPAIRED = 0,
   OPTIGA_PAIRING_PAIRED,
   OPTIGA_PAIRING_ERR_RNG,
-  OPTIGA_PAIRING_ERR_READ,
-  OPTIGA_PAIRING_ERR_HANDSHAKE,
+  OPTIGA_PAIRING_ERR_READ_FLASH,
+  OPTIGA_PAIRING_ERR_WRITE_FLASH,
+  OPTIGA_PAIRING_ERR_WRITE_OPTIGA,
+  OPTIGA_PAIRING_ERR_HANDSHAKE1,
+  OPTIGA_PAIRING_ERR_HANDSHAKE2,
 } optiga_pairing;
 
 static optiga_pairing optiga_pairing_state = OPTIGA_PAIRING_UNPAIRED;
@@ -72,11 +81,20 @@ static bool optiga_paired(void) {
     case OPTIGA_PAIRING_ERR_RNG:
       details = "optiga_get_random error";
       break;
-    case OPTIGA_PAIRING_ERR_READ:
-      details = "failed to read pairing secret";
+    case OPTIGA_PAIRING_ERR_READ_FLASH:
+      details = "failed to read pairing secret from flash";
       break;
-    case OPTIGA_PAIRING_ERR_HANDSHAKE:
-      details = "optiga_sec_chan_handshake";
+    case OPTIGA_PAIRING_ERR_WRITE_FLASH:
+      details = "failed to write pairing secret to flash";
+      break;
+    case OPTIGA_PAIRING_ERR_WRITE_OPTIGA:
+      details = "failed to write pairing secret to Optiga";
+      break;
+    case OPTIGA_PAIRING_ERR_HANDSHAKE1:
+      details = "failed optiga_sec_chan_handshake 1";
+      break;
+    case OPTIGA_PAIRING_ERR_HANDSHAKE2:
+      details = "failed optiga_sec_chan_handshake 2";
       break;
     default:
       break;
@@ -122,50 +140,57 @@ static bool set_metadata(uint16_t oid, const optiga_metadata *metadata) {
 }
 
 void pair_optiga(void) {
-  // The pairing key may already be written and locked. The success of the
-  // pairing procedure is determined by optiga_sec_chan_handshake(). Therefore
-  // it is OK for some of the intermediate operations to fail.
-
   uint8_t secret[SECRET_OPTIGA_KEY_LEN] = {0};
-  optiga_result ret = OPTIGA_SUCCESS;
 
-  if (secret_optiga_extract(secret) != sectrue) {
+  if (secret_optiga_get(secret) != sectrue) {
+    // Generate the pairing secret.
+    if (OPTIGA_SUCCESS != optiga_get_random(secret, sizeof(secret))) {
+      optiga_pairing_state = OPTIGA_PAIRING_ERR_RNG;
+      return;
+    }
+    random_xor(secret, sizeof(secret));
+
     // Enable writing the pairing secret to OPTIGA.
     optiga_metadata metadata = {0};
     metadata.change = OPTIGA_META_ACCESS_ALWAYS;
     metadata.execute = OPTIGA_META_ACCESS_ALWAYS;
     metadata.data_type = TYPE_PTFBIND;
-    set_metadata(OID_KEY_PAIRING, &metadata);  // Ignore result.
+    (void)set_metadata(OID_KEY_PAIRING, &metadata);
 
-    // Generate pairing secret.
-    ret = optiga_get_random(secret, sizeof(secret));
-    if (OPTIGA_SUCCESS != ret) {
-      optiga_pairing_state = OPTIGA_PAIRING_ERR_RNG;
+    // Store the pairing secret in OPTIGA.
+    if (OPTIGA_SUCCESS != optiga_set_data_object(OID_KEY_PAIRING, false, secret,
+                                                 sizeof(secret))) {
+      optiga_pairing_state = OPTIGA_PAIRING_ERR_WRITE_OPTIGA;
       return;
     }
 
-    // Store pairing secret.
-    ret =
-        optiga_set_data_object(OID_KEY_PAIRING, false, secret, sizeof(secret));
-    if (OPTIGA_SUCCESS == ret) {
-      secret_erase();
-      secret_write_header();
-      secret_write(secret, SECRET_OPTIGA_KEY_OFFSET, SECRET_OPTIGA_KEY_LEN);
+    // Execute the handshake to verify that the secret was stored correctly in
+    // Optiga.
+    if (OPTIGA_SUCCESS != optiga_sec_chan_handshake(secret, sizeof(secret))) {
+      optiga_pairing_state = OPTIGA_PAIRING_ERR_HANDSHAKE1;
+      return;
     }
 
-    // Verify whether the secret was stored correctly in flash and OPTIGA.
+    // Store the pairing secret in the flash memory.
+    if (sectrue != secret_optiga_set(secret)) {
+      optiga_pairing_state = OPTIGA_PAIRING_ERR_WRITE_FLASH;
+      return;
+    }
+
+    // Reload the pairing secret from the flash memory.
     memzero(secret, sizeof(secret));
-    if (secret_read(secret, SECRET_OPTIGA_KEY_OFFSET, SECRET_OPTIGA_KEY_LEN) !=
-        sectrue) {
-      optiga_pairing_state = OPTIGA_PAIRING_ERR_READ;
+    if (sectrue != secret_optiga_get(secret)) {
+      optiga_pairing_state = OPTIGA_PAIRING_ERR_READ_FLASH;
       return;
     }
   }
 
-  ret = optiga_sec_chan_handshake(secret, sizeof(secret));
+  // Execute the handshake to verify that the secret is stored correctly in both
+  // Optiga and MCU flash.
+  optiga_result ret = optiga_sec_chan_handshake(secret, sizeof(secret));
   memzero(secret, sizeof(secret));
   if (OPTIGA_SUCCESS != ret) {
-    optiga_pairing_state = OPTIGA_PAIRING_ERR_HANDSHAKE;
+    optiga_pairing_state = OPTIGA_PAIRING_ERR_HANDSHAKE2;
     return;
   }
 
diff --git a/core/embed/trezorhal/secret.h b/core/embed/trezorhal/secret.h
index 202d5952cb..054a8b407e 100644
--- a/core/embed/trezorhal/secret.h
+++ b/core/embed/trezorhal/secret.h
@@ -15,7 +15,7 @@
 
 secbool secret_bootloader_locked(void);
 
-void secret_write(uint8_t* data, uint32_t offset, uint32_t len);
+void secret_write(const uint8_t* data, uint32_t offset, uint32_t len);
 
 secbool secret_read(uint8_t* data, uint32_t offset, uint32_t len);
 
@@ -31,12 +31,14 @@ void secret_hide(void);
 
 void secret_write_header(void);
 
+secbool secret_optiga_set(const uint8_t secret[SECRET_OPTIGA_KEY_LEN]);
+
+secbool secret_optiga_get(uint8_t dest[SECRET_OPTIGA_KEY_LEN]);
+
 void secret_optiga_backup(void);
 
 void secret_optiga_hide(void);
 
-secbool secret_optiga_extract(uint8_t* dest);
-
 void secret_bhk_lock(void);
 
 secbool secret_bhk_locked(void);
diff --git a/core/embed/trezorhal/stm32f4/secret.c b/core/embed/trezorhal/stm32f4/secret.c
index 3e274fe9df..e29fba097c 100644
--- a/core/embed/trezorhal/stm32f4/secret.c
+++ b/core/embed/trezorhal/stm32f4/secret.c
@@ -37,7 +37,7 @@ void secret_write_header(void) {
   secret_write(header, 0, SECRET_HEADER_LEN);
 }
 
-void secret_write(uint8_t* data, uint32_t offset, uint32_t len) {
+void secret_write(const uint8_t* data, uint32_t offset, uint32_t len) {
   ensure(flash_unlock_write(), "secret write");
   for (int i = 0; i < len; i++) {
     ensure(flash_area_write_byte(&SECRET_AREA, offset + i, data[i]),
@@ -76,6 +76,13 @@ void secret_erase(void) {
   ensure(flash_area_erase(&SECRET_AREA, NULL), "secret erase");
 }
 
-secbool secret_optiga_extract(uint8_t* dest) {
+secbool secret_optiga_set(const uint8_t secret[SECRET_OPTIGA_KEY_LEN]) {
+  secret_erase();
+  secret_write_header();
+  secret_write(secret, SECRET_OPTIGA_KEY_OFFSET, SECRET_OPTIGA_KEY_LEN);
+  return sectrue;
+}
+
+secbool secret_optiga_get(uint8_t dest[SECRET_OPTIGA_KEY_LEN]) {
   return secret_read(dest, SECRET_OPTIGA_KEY_OFFSET, SECRET_OPTIGA_KEY_LEN);
 }
diff --git a/core/embed/trezorhal/stm32u5/secret.c b/core/embed/trezorhal/stm32u5/secret.c
index 9f2efdebfc..2dbe2a0724 100644
--- a/core/embed/trezorhal/stm32u5/secret.c
+++ b/core/embed/trezorhal/stm32u5/secret.c
@@ -6,6 +6,7 @@
 #include "memzero.h"
 #include "model.h"
 #include "rng.h"
+#include "secure_aes.h"
 
 static secbool bootloader_locked = secfalse;
 
@@ -48,7 +49,7 @@ void secret_write_header(void) {
   secret_write(header, 0, SECRET_HEADER_LEN);
 }
 
-void secret_write(uint8_t *data, uint32_t offset, uint32_t len) {
+void secret_write(const uint8_t *data, uint32_t offset, uint32_t len) {
   ensure(flash_unlock_write(), "secret write");
   for (int i = 0; i < len / 16; i++) {
     ensure(flash_area_write_quadword(&SECRET_AREA, offset + (i * 16),
@@ -142,6 +143,18 @@ secbool secret_optiga_present(void) {
   return secret_present(SECRET_OPTIGA_KEY_OFFSET, SECRET_OPTIGA_KEY_LEN);
 }
 
+secbool secret_optiga_set(const uint8_t secret[SECRET_OPTIGA_KEY_LEN]) {
+  uint8_t secret_enc[SECRET_OPTIGA_KEY_LEN] = {0};
+  if (sectrue != secure_aes_ecb_encrypt_hw(secret, sizeof(secret_enc),
+                                           secret_enc, SECURE_AES_KEY_DHUK)) {
+    return secfalse;
+  }
+  secret_write(secret_enc, SECRET_OPTIGA_KEY_OFFSET, SECRET_OPTIGA_KEY_LEN);
+  memzero(secret_enc, sizeof(secret_enc));
+  secret_optiga_backup();
+  return sectrue;
+}
+
 void secret_optiga_backup(void) {
   uint32_t secret[SECRET_OPTIGA_KEY_LEN / sizeof(uint32_t)] = {0};
   secbool ok = secret_read((uint8_t *)secret, SECRET_OPTIGA_KEY_OFFSET,
@@ -162,22 +175,28 @@ void secret_optiga_backup(void) {
   memzero(secret, sizeof(secret));
 }
 
-secbool secret_optiga_extract(uint8_t *dest) {
+secbool secret_optiga_get(uint8_t dest[SECRET_OPTIGA_KEY_LEN]) {
+  uint32_t secret[SECRET_OPTIGA_KEY_LEN / sizeof(uint32_t)] = {0};
+
   bool all_zero = true;
   volatile uint32_t *reg1 = &TAMP->BKP8R;
   for (int i = 0; i < 8; i++) {
-    uint32_t val = *reg1++;
+    secret[i] = reg1[i];
 
-    if (val != 0) {
+    if (secret[i] != 0) {
       all_zero = false;
     }
-
-    for (int j = 0; j < 4; j++) {
-      dest[i * 4 + j] = (val >> (j * 8)) & 0xFF;
-    }
   }
 
-  return all_zero ? secfalse : sectrue;
+  if (all_zero) {
+    return secfalse;
+  }
+
+  secbool res = secure_aes_ecb_decrypt_hw(
+      (uint8_t *)secret, SECRET_OPTIGA_KEY_LEN, dest, SECURE_AES_KEY_DHUK);
+
+  memzero(secret, sizeof(secret));
+  return res;
 }
 
 void secret_optiga_hide(void) {