Support FIDO2 credential storage && write private key in factory mode. (

#221) * Add command to write private key in factory mode. * Support FIDO2 credential storage.
OneKeyHQ · Dec 30, 2024 · 78a1be9 · 78a1be9
1 parent 5bb965e
commit 78a1be9
Show file tree

Hide file tree

Showing 27 changed files with 376 additions and 39 deletions.
diff --git a/common/protob/messages-management.proto b/common/protob/messages-management.proto
@@ -643,6 +643,14 @@ message DeviceInfo {
     optional string pre_firmware = 6;
 }
 
+/**
+ * Request: Read SE private key
+ * @end
+ */
+message WriteSEPrivateKey {
+    required bytes private_key = 1;
+}
+
 /**
  * Request: Read SE public key
  * @end

diff --git a/common/protob/messages.proto b/common/protob/messages.proto
@@ -593,4 +593,5 @@ enum MessageType {
     MessageType_FileInfoList = 10024 [(wire_out) = true];
     MessageType_OnekeyGetFeatures = 10025 [(wire_in) = true];
     MessageType_OnekeyFeatures = 10026 [(wire_out) = true];
+    MessageType_WriteSEPrivateKey = 10027 [(wire_in) = true, (wire_bootloader) = true];
 }
diff --git a/core/embed/bootloader/main.c b/core/embed/bootloader/main.c
@@ -527,6 +527,9 @@ secbool bootloader_usb_loop_factory(const vendor_header* const vhdr,
       case MSG_NAME_TO_ID(GetDeviceInfo):  // GetDeviceInfo
         process_msg_GetDeviceInfo(USB_IFACE_NUM, msg_size, buf);
         break;
+      case MSG_NAME_TO_ID(WriteSEPrivateKey):  // WriteSEPrivateKey
+        process_msg_WriteSEPrivateKey(USB_IFACE_NUM, msg_size, buf);
+        break;
       case MSG_NAME_TO_ID(ReadSEPublicKey):  // ReadSEPublicKey
         process_msg_ReadSEPublicKey(USB_IFACE_NUM, msg_size, buf);
         break;

diff --git a/core/embed/bootloader/messages.c b/core/embed/bootloader/messages.c
@@ -339,6 +339,7 @@ static void send_msg_features(uint8_t iface_num,
     init_state |= device_serial_set() ? 1 : 0;
     init_state |= se_has_cerrificate() ? (1 << 2) : 0;
     MSG_SEND_ASSIGN_VALUE(initstates, init_state);
+    MSG_SEND_ASSIGN_VALUE(onekey_device_type, OneKeyDeviceType_PRO);
 
   } else {
     MSG_SEND_ASSIGN_STRING(vendor, "onekey.so");
@@ -671,6 +672,25 @@ void process_msg_GetDeviceInfo(uint8_t iface_num, uint32_t msg_size,
   MSG_SEND(DeviceInfo);
 }
 
+void process_msg_WriteSEPrivateKey(uint8_t iface_num, uint32_t msg_size,
+                                   uint8_t *buf) {
+  MSG_RECV_INIT(WriteSEPrivateKey);
+  MSG_RECV(WriteSEPrivateKey);
+
+  if (msg_recv.private_key.size != 32) {
+    send_failure(iface_num, FailureType_Failure_ProcessError,
+                 "Private key size invalid");
+    return;
+  }
+
+  if (se_set_private_key_extern(msg_recv.private_key.bytes)) {
+    send_success(iface_num, "Write private key success");
+  } else {
+    send_failure(iface_num, FailureType_Failure_ProcessError,
+                 "Write private key failed");
+  }
+}
+
 void process_msg_ReadSEPublicKey(uint8_t iface_num, uint32_t msg_size,
                                  uint8_t *buf) {
   uint8_t pubkey[64] = {0};
@@ -728,6 +748,13 @@ void process_msg_SESignMessage(uint8_t iface_num, uint32_t msg_size,
 
   MSG_SEND_INIT(SEMessageSignature);
 
+  if (se_sign_message_with_write_key((uint8_t *)msg_recv.message.bytes,
+                                     msg_recv.message.size, sign)) {
+    MSG_SEND_ASSIGN_REQUIRED_BYTES(signature, sign, 64);
+    MSG_SEND(SEMessageSignature);
+    return;
+  }
+
   if (se_sign_message((uint8_t *)msg_recv.message.bytes, msg_recv.message.size,
                       sign)) {
     MSG_SEND_ASSIGN_REQUIRED_BYTES(signature, sign, 64);

diff --git a/core/embed/bootloader/messages.h b/core/embed/bootloader/messages.h
@@ -68,6 +68,8 @@ void process_msg_DeviceInfoSettings(uint8_t iface_num, uint32_t msg_size,
                                     uint8_t *buf);
 void process_msg_GetDeviceInfo(uint8_t iface_num, uint32_t msg_size,
                                uint8_t *buf);
+void process_msg_WriteSEPrivateKey(uint8_t iface_num, uint32_t msg_size,
+                                   uint8_t *buf);
 void process_msg_ReadSEPublicKey(uint8_t iface_num, uint32_t msg_size,
                                  uint8_t *buf);
 void process_msg_WriteSEPublicCert(uint8_t iface_num, uint32_t msg_size,

diff --git a/core/embed/bootloader/protob/messages.options b/core/embed/bootloader/protob/messages.options
@@ -107,6 +107,7 @@ DeviceInfo.NFT_voucher                  max_size:32
 DeviceInfo.cpu_info                     max_size:16
 DeviceInfo.pre_firmware                 max_size:16
 
+WriteSEPrivateKey.private_key           max_size:32
 SEPublicKey.public_key                  max_size:64
 WriteSEPublicCert.public_cert           max_size:416
 SEPublicCert.public_cert                max_size:416

diff --git a/core/embed/bootloader/protob/messages.pb.c b/core/embed/bootloader/protob/messages.pb.c
@@ -99,6 +99,9 @@ PB_BIND(GetDeviceInfo, GetDeviceInfo, AUTO)
 PB_BIND(DeviceInfo, DeviceInfo, AUTO)
 
 
+PB_BIND(WriteSEPrivateKey, WriteSEPrivateKey, AUTO)
+
+
 PB_BIND(ReadSEPublicKey, ReadSEPublicKey, AUTO)
 
 

diff --git a/core/embed/bootloader/protob/messages.pb.h b/core/embed/bootloader/protob/messages.pb.h
@@ -36,6 +36,7 @@ typedef enum _MessageType {
     MessageType_MessageType_SEMessageSignature = 10013, 
     MessageType_MessageType_OnekeyGetFeatures = 10025, 
     MessageType_MessageType_OnekeyFeatures = 10026, 
+    MessageType_MessageType_WriteSEPrivateKey = 10027, 
     MessageType_MessageType_Reboot = 30000, 
     MessageType_MessageType_FirmwareUpdateEmmc = 30001, 
     MessageType_MessageType_EmmcFixPermission = 30100, 
@@ -487,6 +488,11 @@ typedef struct _Success {
     char message[256]; 
 } Success;
 
+typedef PB_BYTES_ARRAY_T(32) WriteSEPrivateKey_private_key_t;
+typedef struct _WriteSEPrivateKey { 
+    WriteSEPrivateKey_private_key_t private_key; 
+} WriteSEPrivateKey;
+
 typedef PB_BYTES_ARRAY_T(416) WriteSEPublicCert_public_cert_t;
 typedef struct _WriteSEPublicCert { 
     WriteSEPublicCert_public_cert_t public_cert; 
@@ -573,6 +579,7 @@ extern "C" {
 #define DeviceInfoSettings_init_default          {false, "", false, "", false, ""}
 #define GetDeviceInfo_init_default               {0}
 #define DeviceInfo_init_default                  {false, "", false, "", false, "", false, {0, {0}}, false, "", false, ""}
+#define WriteSEPrivateKey_init_default           {{0, {0}}}
 #define ReadSEPublicKey_init_default             {0}
 #define SEPublicKey_init_default                 {{0, {0}}}
 #define WriteSEPublicCert_init_default           {{0, {0}}}
@@ -611,6 +618,7 @@ extern "C" {
 #define DeviceInfoSettings_init_zero             {false, "", false, "", false, ""}
 #define GetDeviceInfo_init_zero                  {0}
 #define DeviceInfo_init_zero                     {false, "", false, "", false, "", false, {0, {0}}, false, "", false, ""}
+#define WriteSEPrivateKey_init_zero              {{0, {0}}}
 #define ReadSEPublicKey_init_zero                {0}
 #define SEPublicKey_init_zero                    {{0, {0}}}
 #define WriteSEPublicCert_init_zero              {{0, {0}}}
@@ -769,6 +777,7 @@ extern "C" {
 #define SEPublicKey_public_key_tag               1
 #define SESignMessage_message_tag                1
 #define Success_message_tag                      1
+#define WriteSEPrivateKey_private_key_tag        1
 #define WriteSEPublicCert_public_cert_tag        1
 #define EmmcFileRead_file_tag                    1
 #define EmmcFileRead_ui_percentage_tag           2
@@ -1062,6 +1071,11 @@ X(a, STATIC,   OPTIONAL, STRING,   pre_firmware,      6)
 #define DeviceInfo_CALLBACK NULL
 #define DeviceInfo_DEFAULT NULL
 
+#define WriteSEPrivateKey_FIELDLIST(X, a) \
+X(a, STATIC,   REQUIRED, BYTES,    private_key,       1)
+#define WriteSEPrivateKey_CALLBACK NULL
+#define WriteSEPrivateKey_DEFAULT NULL
+
 #define ReadSEPublicKey_FIELDLIST(X, a) \
 
 #define ReadSEPublicKey_CALLBACK NULL
@@ -1128,6 +1142,7 @@ extern const pb_msgdesc_t EmmcDirRemove_msg;
 extern const pb_msgdesc_t DeviceInfoSettings_msg;
 extern const pb_msgdesc_t GetDeviceInfo_msg;
 extern const pb_msgdesc_t DeviceInfo_msg;
+extern const pb_msgdesc_t WriteSEPrivateKey_msg;
 extern const pb_msgdesc_t ReadSEPublicKey_msg;
 extern const pb_msgdesc_t SEPublicKey_msg;
 extern const pb_msgdesc_t WriteSEPublicCert_msg;
@@ -1168,6 +1183,7 @@ extern const pb_msgdesc_t SEMessageSignature_msg;
 #define DeviceInfoSettings_fields &DeviceInfoSettings_msg
 #define GetDeviceInfo_fields &GetDeviceInfo_msg
 #define DeviceInfo_fields &DeviceInfo_msg
+#define WriteSEPrivateKey_fields &WriteSEPrivateKey_msg
 #define ReadSEPublicKey_fields &ReadSEPublicKey_msg
 #define SEPublicKey_fields &SEPublicKey_msg
 #define WriteSEPublicCert_fields &WriteSEPublicCert_msg
@@ -1214,6 +1230,7 @@ extern const pb_msgdesc_t SEMessageSignature_msg;
 #define SESignMessage_size                       1027
 #define Success_size                             258
 #define WipeDevice_size                          0
+#define WriteSEPrivateKey_size                   34
 #define WriteSEPublicCert_size                   419
 
 #ifdef __cplusplus

diff --git a/core/embed/bootloader/protob/messages.proto b/core/embed/bootloader/protob/messages.proto
@@ -31,6 +31,7 @@ enum MessageType {
     MessageType_SEMessageSignature = 10013;
     MessageType_OnekeyGetFeatures = 10025;
     MessageType_OnekeyFeatures = 10026;
+    MessageType_WriteSEPrivateKey = 10027;
 
     MessageType_Reboot = 30000;
     MessageType_FirmwareUpdateEmmc = 30001;
@@ -525,6 +526,14 @@ message DeviceInfo {
     optional string pre_firmware = 6;
 }
 
+/**
+ * Request: Read SE private key
+ * @end
+ */
+message WriteSEPrivateKey {
+    required bytes private_key = 1;
+}
+
 /**
  * Request: Read SE public key
  * @end

diff --git a/core/embed/extmod/modtrezorconfig/modtrezorconfig.c b/core/embed/extmod/modtrezorconfig/modtrezorconfig.c
@@ -393,6 +393,20 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_trezorconfig_get_val_len_obj, 2,
 ///     value fails.
 ///     """
 STATIC mp_obj_t mod_trezorconfig_get(size_t n_args, const mp_obj_t *args) {
+  uint8_t app = trezor_obj_get_uint8(args[0]);
+  // webauthn resident credentials, FIDO2
+  if (app == 4) {
+    uint32_t index = trezor_obj_get_uint(args[1]);
+    uint16_t len = sizeof(CTAP_credential_id_storage) -
+                   FIDO2_RESIDENT_CREDENTIALS_HEADER_LEN;
+    CTAP_credential_id_storage cred_id = {0};
+
+    if (!se_get_fido2_resident_credentials(index, cred_id.rp_id_hash, &len)) {
+      return mp_const_none;
+    }
+    return mp_obj_new_bytes(cred_id.rp_id_hash, len);
+  }
+
   uint32_t key = trezor_obj_get_uint(args[1]);
 
   bool is_private = key & (1 << 31);
@@ -435,6 +449,23 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_trezorconfig_get_obj, 2, 3,
 ///     Sets a value of given key for given app.
 ///     """
 STATIC mp_obj_t mod_trezorconfig_set(size_t n_args, const mp_obj_t *args) {
+  uint8_t app = trezor_obj_get_uint8(args[0]);
+  // webauthn resident credentials, FIDO2
+  if (app == 4) {
+    uint32_t index = trezor_obj_get_uint(args[1]);
+
+    mp_buffer_info_t cred_id;
+    mp_get_buffer_raise(args[2], &cred_id, MP_BUFFER_READ);
+    if (cred_id.len > sizeof(CTAP_credential_id_storage) -
+                          FIDO2_RESIDENT_CREDENTIALS_HEADER_LEN) {
+      mp_raise_msg(&mp_type_RuntimeError, "Credential ID too long");
+    }
+    if (!se_set_fido2_resident_credentials(index, cred_id.buf, cred_id.len)) {
+      mp_raise_msg(&mp_type_RuntimeError, "Could not save value");
+    }
+    return mp_const_none;
+  }
+
   uint32_t key = trezor_obj_get_uint(args[1]);
   bool is_private = key & (1 << 31);
   secbool (*writer)(uint16_t, const void *, uint16_t) =
@@ -468,6 +499,16 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_trezorconfig_set_obj, 3, 4,
 ///     Deletes the given key of the given app.
 ///     """
 STATIC mp_obj_t mod_trezorconfig_delete(size_t n_args, const mp_obj_t *args) {
+  uint8_t app = trezor_obj_get_uint8(args[0]);
+  // webauthn resident credentials, FIDO2
+  if (app == 4) {
+    uint32_t index = trezor_obj_get_uint(args[1]);
+    if (!se_delete_fido2_resident_credentials(index)) {
+      mp_raise_msg(&mp_type_RuntimeError, "Could not delete value");
+    }
+    return mp_const_true;
+  }
+
   uint32_t key = trezor_obj_get_uint(args[1]);
   bool is_private = key & (1 << 31);
   secbool (*writer)(uint16_t, const void *, uint16_t) =

diff --git a/core/embed/extmod/modtrezorcrypto/modtrezorcrypto-nist256p1.h b/core/embed/extmod/modtrezorcrypto/modtrezorcrypto-nist256p1.h
@@ -203,16 +203,18 @@ STATIC mp_obj_t mod_trezorcrypto_nist256p1_verify_recover(mp_obj_t signature,
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_trezorcrypto_nist256p1_verify_recover_obj,
                                  mod_trezorcrypto_nist256p1_verify_recover);
 
-/// def multiply(secret_key: bytes, public_key: bytes) -> bytes:
+/// def multiply(secret_key: bytes, public_key: bytes, inner_secret_key: bool =
+/// False) -> bytes:
 ///     """
 ///     Multiplies point defined by public_key with scalar defined by
 ///     secret_key. Useful for ECDH.
 ///     """
-STATIC mp_obj_t mod_trezorcrypto_nist256p1_multiply(mp_obj_t secret_key,
-                                                    mp_obj_t public_key) {
+STATIC mp_obj_t mod_trezorcrypto_nist256p1_multiply(size_t n_args,
+                                                    const mp_obj_t *args) {
   mp_buffer_info_t sk = {0}, pk = {0};
-  mp_get_buffer_raise(secret_key, &sk, MP_BUFFER_READ);
-  mp_get_buffer_raise(public_key, &pk, MP_BUFFER_READ);
+  mp_get_buffer_raise(args[0], &sk, MP_BUFFER_READ);
+  mp_get_buffer_raise(args[1], &pk, MP_BUFFER_READ);
+  bool inner_secret_key = (n_args == 3 && args[2] == mp_const_true);
   if (sk.len != 32) {
     mp_raise_ValueError("Invalid length of secret key");
   }
@@ -222,19 +224,28 @@ STATIC mp_obj_t mod_trezorcrypto_nist256p1_multiply(mp_obj_t secret_key,
   vstr_t out = {0};
   vstr_init_len(&out, 65);
 #if USE_THD89
-  uint8_t pubkey[65] = {0};
-  if (pk.len == 33) {
-    if (!ecdsa_uncompress_pubkey(&nist256p1, pk.buf, pubkey)) {
-      mp_raise_ValueError("Invalid public key");
+  if (inner_secret_key) {
+    if (0 != ecdh_multiply(&nist256p1, (const uint8_t *)sk.buf,
+                           (const uint8_t *)pk.buf, (uint8_t *)out.buf)) {
+      vstr_clear(&out);
+      mp_raise_ValueError("Multiply failed");
     }
   } else {
-    memcpy(pubkey, pk.buf, 65);
-  }
-  if (0 != se_get_shared_key(NIST256P1_NAME, (const uint8_t *)pubkey,
-                             (uint8_t *)out.buf)) {
-    vstr_clear(&out);
-    mp_raise_ValueError("Multiply failed");
+    uint8_t pubkey[65] = {0};
+    if (pk.len == 33) {
+      if (!ecdsa_uncompress_pubkey(&nist256p1, pk.buf, pubkey)) {
+        mp_raise_ValueError("Invalid public key");
+      }
+    } else {
+      memcpy(pubkey, pk.buf, 65);
+    }
+    if (0 != se_get_shared_key(NIST256P1_NAME, (const uint8_t *)pubkey,
+                               (uint8_t *)out.buf)) {
+      vstr_clear(&out);
+      mp_raise_ValueError("Multiply failed");
+    }
   }
+
 #else
   if (0 != ecdh_multiply(&nist256p1, (const uint8_t *)sk.buf,
                          (const uint8_t *)pk.buf, (uint8_t *)out.buf)) {
@@ -244,8 +255,9 @@ STATIC mp_obj_t mod_trezorcrypto_nist256p1_multiply(mp_obj_t secret_key,
 #endif
   return mp_obj_new_str_from_vstr(&mp_type_bytes, &out);
 }
-STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_trezorcrypto_nist256p1_multiply_obj,
-                                 mod_trezorcrypto_nist256p1_multiply);
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(
+    mod_trezorcrypto_nist256p1_multiply_obj, 2, 3,
+    mod_trezorcrypto_nist256p1_multiply);
 
 STATIC const mp_rom_map_elem_t mod_trezorcrypto_nist256p1_globals_table[] = {
     {MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_nist256p1)},

diff --git a/core/embed/extmod/modtrezorcrypto/modtrezorcrypto-se-thd89.h b/core/embed/extmod/modtrezorcrypto/modtrezorcrypto-se-thd89.h
@@ -926,6 +926,20 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_1(
     mod_trezorcrypto_se_thd89_fido_att_sign_digest_obj,
     mod_trezorcrypto_se_thd89_fido_att_sign_digest);
 
+/// def fido_delete_all_credentials() -> None:
+///     """
+///     Delete all FIDO2 credentials.
+///     """
+STATIC mp_obj_t mod_trezorcrypto_se_thd89_fido_delete_all_credentials(void) {
+  se_delete_all_fido2_credentials();
+  return mp_const_none;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(
+    mod_trezorcrypto_se_thd89_fido_delete_all_credentials_obj,
+    mod_trezorcrypto_se_thd89_fido_delete_all_credentials);
+
+/// FIDO2_CRED_COUNT_MAX: int
+
 STATIC const mp_rom_map_elem_t mod_trezorcrypto_se_thd89_globals_table[] = {
     {MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_se_thd89)},
     {MP_ROM_QSTR(MP_QSTR_check),
@@ -994,6 +1008,10 @@ STATIC const mp_rom_map_elem_t mod_trezorcrypto_se_thd89_globals_table[] = {
      MP_ROM_PTR(&mod_trezorcrypto_se_thd89_fido_sign_digest_obj)},
     {MP_ROM_QSTR(MP_QSTR_fido_att_sign_digest),
      MP_ROM_PTR(&mod_trezorcrypto_se_thd89_fido_att_sign_digest_obj)},
+    {MP_ROM_QSTR(MP_QSTR_fido_delete_all_credentials),
+     MP_ROM_PTR(&mod_trezorcrypto_se_thd89_fido_delete_all_credentials_obj)},
+    {MP_ROM_QSTR(MP_QSTR_FIDO2_CRED_COUNT_MAX),
+     MP_ROM_INT(FIDO2_RESIDENT_CREDENTIALS_COUNT)},
 };
 STATIC MP_DEFINE_CONST_DICT(mod_trezorcrypto_se_thd89_globals,
                             mod_trezorcrypto_se_thd89_globals_table);
Original file line number	Diff line number	Diff line change
Expand Up		@@ -99,6 +99,9 @@ PB_BIND(GetDeviceInfo, GetDeviceInfo, AUTO)
		PB_BIND(DeviceInfo, DeviceInfo, AUTO)


		PB_BIND(WriteSEPrivateKey, WriteSEPrivateKey, AUTO)


		PB_BIND(ReadSEPublicKey, ReadSEPublicKey, AUTO)


Expand Down