WIP - migrating storage/device.py into rust

core/SConscript.firmware

@@ -200,6 +200,7 @@ SOURCE_MOD += [
 # rust mods
 SOURCE_MOD += [
     'embed/extmod/rustmods/modtrezorproto.c',
+    'embed/extmod/rustmods/modtrezorstoragedevice.c',
 ]
 if UI2:
     SOURCE_MOD += [

core/SConscript.unix

@@ -197,6 +197,7 @@ SOURCE_MOD += [
 # rust mods
 SOURCE_MOD += [
     'embed/extmod/rustmods/modtrezorproto.c',
+    'embed/extmod/rustmods/modtrezorstoragedevice.c',
 ]
 if UI2:
     SOURCE_MOD += [

core/embed/extmod/rustmods/modtrezorstoragedevice.c

@@ -0,0 +1,65 @@
+/*
+ * This file is part of the Trezor project, https://trezor.io/
+ *
+ * Copyright (c) SatoshiLabs
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include "py/runtime.h"
+
+#if MICROPY_PY_TREZORSTORAGEDEVICE
+
+#include "librust.h"
+
+
+/// def is_version_stored() -> bool:
+///     """Whether version is in storage."""
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorutils_storagedevice_is_version_stored_obj,
+                                 storagedevice_is_version_stored);
+
+// /// def get_version() -> bool:
+// ///     """Get from storage."""
+// STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_trezorutils_storagedevice_get_version_obj,
+//                                  storagedevice_get_version);
+
+/// def set_version(version: bytes) -> bool:
+///     """Save to storage."""
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_trezorutils_storagedevice_set_version_obj,
+                                 storagedevice_set_version);
+
+
+STATIC const mp_rom_map_elem_t mp_module_trezorstoragedevice_globals_table[] = {
+    {MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_trezorstoragedevice)},
+
+    {MP_ROM_QSTR(MP_QSTR_is_version_stored),
+     MP_ROM_PTR(&mod_trezorutils_storagedevice_is_version_stored_obj)},
+    // {MP_ROM_QSTR(MP_QSTR_get_version),
+    //  MP_ROM_PTR(&mod_trezorutils_storagedevice_get_version_obj)},
+    {MP_ROM_QSTR(MP_QSTR_set_version),
+     MP_ROM_PTR(&mod_trezorutils_storagedevice_set_version_obj)},
+};
+
+STATIC MP_DEFINE_CONST_DICT(mp_module_trezorstoragedevice_globals,
+                            mp_module_trezorstoragedevice_globals_table);
+
+const mp_obj_module_t mp_module_trezorstoragedevice = {
+    .base = {&mp_type_module},
+    .globals = (mp_obj_dict_t *)&mp_module_trezorstoragedevice_globals,
+};
+
+MP_REGISTER_MODULE(MP_QSTR_trezorstoragedevice, mp_module_trezorstoragedevice,
+                   MICROPY_PY_TREZORSTORAGEDEVICE);
+
+#endif  // MICROPY_PY_TREZORSTORAGEDEVICE

core/embed/firmware/mpconfigport.h

@@ -155,6 +155,7 @@
 
 #define MICROPY_PY_TREZORCONFIG     (1)
 #define MICROPY_PY_TREZORCRYPTO     (1)
+#define MICROPY_PY_TREZORSTORAGEDEVICE (1)
 #define MICROPY_PY_TREZORIO         (1)
 #define MICROPY_PY_TREZORUI         (1)
 #define MICROPY_PY_TREZORUTILS      (1)

core/embed/rust/build.rs

@@ -6,6 +6,7 @@ fn main() {
     generate_qstr_bindings();
     generate_micropython_bindings();
     generate_trezorhal_bindings();
+    generate_storagedevice_bindings();
     #[cfg(feature = "test")]
     link_core_objects();
 }
@@ -234,6 +235,31 @@ fn generate_trezorhal_bindings() {
         .unwrap();
 }
 
+fn generate_storagedevice_bindings() {
+    // TODO: how to avoid defining `storage.h` in this directory
+    // (somehow specifying its location)
+    // core/vendor/trezor-storage/storage.h
+
+    let out_path = env::var("OUT_DIR").unwrap();
+
+    // Tell cargo to invalidate the built crate whenever the header changes.
+    println!("cargo:rerun-if-changed=storage.h");
+
+    let bindings = prepare_bindings()
+        .header("storage.h")
+        // storage
+        .allowlist_function("storage_set")
+        .allowlist_function("storage_delete")
+        .allowlist_function("storage_get");
+
+    // Write the bindings to a file in the OUT_DIR.
+    bindings
+        .generate()
+        .expect("Unable to generate bindings")
+        .write_to_file(PathBuf::from(out_path).join("storage.rs"))
+        .unwrap();
+}
+
 fn is_firmware() -> bool {
     let target = env::var("TARGET").unwrap();
     target.starts_with("thumbv7")

core/embed/rust/librust.h

@@ -9,6 +9,10 @@ mp_obj_t protobuf_decode(mp_obj_t buf, mp_obj_t def,
 mp_obj_t protobuf_len(mp_obj_t obj);
 mp_obj_t protobuf_encode(mp_obj_t buf, mp_obj_t obj);
 
+mp_obj_t storagedevice_is_version_stored();
+// mp_obj_t storagedevice_get_version();
+mp_obj_t storagedevice_set_version(mp_obj_t key);
+
 #ifdef TREZOR_EMULATOR
 mp_obj_t protobuf_debug_msg_type();
 mp_obj_t protobuf_debug_msg_def_type();
@@ -18,4 +22,4 @@ extern mp_obj_module_t mp_module_trezorui2;
 
 #ifdef TREZOR_EMULATOR
 mp_obj_t ui_debug_layout_type();
-#endif
+#endif

core/embed/rust/src/lib.rs

@@ -5,6 +5,7 @@
 #![allow(dead_code)]
 
 mod error;
+mod storagedevice;
 #[macro_use]
 mod micropython;
 mod protobuf;

core/embed/rust/src/storagedevice/mod.rs

@@ -0,0 +1 @@
+mod storagedevice;

core/embed/rust/src/storagedevice/storagedevice.rs

@@ -0,0 +1,114 @@
+use crate::{
+    micropython::{buffer::Buffer, obj::Obj},
+    trezorhal::secbool,
+    util,
+};
+
+use core::convert::{TryFrom, TryInto};
+
+// TODO: transfer this into a struct with field specifying data type and
+// max_length
+const APP_DEVICE: u16 = 0x01;
+
+const DEVICE_ID: u16 = 0x01;
+const _VERSION: u16 = 0x02;
+const _MNEMONIC_SECRET: u16 = 0x03;
+const _LANGUAGE: u16 = 0x04;
+const _LABEL: u16 = 0x05;
+const _USE_PASSPHRASE: u16 = 0x06;
+const _HOMESCREEN: u16 = 0x07;
+const _NEEDS_BACKUP: u16 = 0x08;
+const _FLAGS: u16 = 0x09;
+const U2F_COUNTER: u16 = 0x0A;
+const _PASSPHRASE_ALWAYS_ON_DEVICE: u16 = 0x0B;
+const _UNFINISHED_BACKUP: u16 = 0x0C;
+const _AUTOLOCK_DELAY_MS: u16 = 0x0D;
+const _NO_BACKUP: u16 = 0x0E;
+const _BACKUP_TYPE: u16 = 0x0F;
+const _ROTATION: u16 = 0x10;
+const _SLIP39_IDENTIFIER: u16 = 0x11;
+const _SLIP39_ITERATION_EXPONENT: u16 = 0x12;
+const _SD_SALT_AUTH_KEY: u16 = 0x13;
+const INITIALIZED: u16 = 0x14;
+const _SAFETY_CHECK_LEVEL: u16 = 0x15;
+const _EXPERIMENTAL_FEATURES: u16 = 0x16;
+
+extern "C" {
+    // storage.h
+    fn storage_has(key: u16) -> secbool::Secbool;
+    fn storage_delete(key: u16) -> secbool::Secbool;
+    fn storage_get(
+        key: u16,
+        // val: *mut cty::c_void,
+        // val: *mut BufferMut,
+        val: *mut u8,
+        max_len: u16,
+        len: *mut u16,
+    ) -> secbool::Secbool;
+    // fn storage_set(key: u16, val: *const cty::c_void, len: u16) ->
+    // secbool::Secbool;
+    fn storage_set(key: u16, val: Buffer, len: u16) -> secbool::Secbool;
+}
+
+#[no_mangle]
+pub extern "C" fn storagedevice_is_version_stored() -> Obj {
+    let block = || {
+        let key: u16 = 0x0102;
+        let result = storagedevice_storage_has(key);
+        Ok(result.into())
+    };
+    unsafe { util::try_or_raise(block) }
+}
+
+// #[no_mangle]
+// pub extern "C" fn storagedevice_get_version() -> Obj {
+//     let block = || {
+//         let key: u16 = 0x0102;
+//         let result = storagedevice_storage_get(key);
+//         Ok(result.into())
+//     };
+//     unsafe { util::try_or_raise(block) }
+// }
+
+#[no_mangle]
+pub extern "C" fn storagedevice_set_version(version: Obj) -> Obj {
+    let block = || {
+        let value = Buffer::try_from(version)?;
+
+        let key: u16 = 0x0102;
+        let result = storagedevice_storage_set(key, value);
+        Ok(result.into())
+    };
+    unsafe { util::try_or_raise(block) }
+}
+
+// pub fn storagedevice_storage_get(key: u16) -> BufferMut {
+//     const MAX_LEN: usize = 300;
+//     // let mut buf: [u8; MAX_LEN] = [0; MAX_LEN];
+//     let mut buf: BufferMut;
+//     let mut len: u16 = 0;
+//     unsafe { storage_get(key, &mut buf as *mut _, MAX_LEN as u16, &mut len as
+// *mut _) };     buf[..len as usize] as BufferMut
+// }
+
+pub fn storagedevice_storage_set(key: u16, value: Buffer) -> bool {
+    let len = value.len();
+    match unsafe { storage_set(key, value, len as u16) } {
+        secbool::TRUE => true,
+        _ => false,
+    }
+}
+
+pub fn storagedevice_storage_has(key: u16) -> bool {
+    match unsafe { storage_has(key) } {
+        secbool::TRUE => true,
+        _ => false,
+    }
+}
+
+pub fn storagedevice_storage_delete(key: u16) -> bool {
+    match unsafe { storage_delete(key) } {
+        secbool::TRUE => true,
+        _ => false,
+    }
+}

core/embed/rust/src/trezorhal/secbool.rs

@@ -1,2 +1,3 @@
 pub const TRUE: u32 = 0xAAAA_AAAA;
 pub const FALSE: u32 = 0x0000_0000;
+pub type Secbool = u32;

core/embed/rust/storage.h

@@ -0,0 +1,75 @@
+/*
+ * This file is part of the Trezor project, https://trezor.io/
+ *
+ * Copyright (c) SatoshiLabs
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#ifndef __STORAGE_H__
+#define __STORAGE_H__
+
+#include <stddef.h>
+#include <stdint.h>
+#include "secbool.h"
+
+// The length of the external salt in bytes.
+#define EXTERNAL_SALT_SIZE 32
+
+// If the top bit of APP is set, then the value is not encrypted.
+#define FLAG_PUBLIC 0x80
+
+// If the top two bits of APP are set, then the value is not encrypted and it
+// can be written even when the storage is locked.
+#define FLAGS_WRITE 0xC0
+
+// The maximum value of app_id which is the six least significant bits of APP.
+#define MAX_APPID 0x3F
+
+// The PIN value corresponding to an empty PIN.
+extern const uint8_t *PIN_EMPTY;
+#define PIN_EMPTY_LEN 0
+
+typedef secbool (*PIN_UI_WAIT_CALLBACK)(uint32_t wait, uint32_t progress,
+                                        const char *message);
+
+void storage_init(PIN_UI_WAIT_CALLBACK callback, const uint8_t *salt,
+                  const uint16_t salt_len);
+void storage_wipe(void);
+secbool storage_is_unlocked(void);
+void storage_lock(void);
+secbool storage_unlock(const uint8_t *pin, size_t pin_len,
+                       const uint8_t *ext_salt);
+secbool storage_has_pin(void);
+secbool storage_pin_fails_increase(void);
+uint32_t storage_get_pin_rem(void);
+secbool storage_change_pin(const uint8_t *oldpin, size_t oldpin_len,
+                           const uint8_t *newpin, size_t newpin_len,
+                           const uint8_t *old_ext_salt,
+                           const uint8_t *new_ext_salt);
+void storage_ensure_not_wipe_code(const uint8_t *pin, size_t pin_len);
+secbool storage_has_wipe_code(void);
+secbool storage_change_wipe_code(const uint8_t *pin, size_t pin_len,
+                                 const uint8_t *ext_salt,
+                                 const uint8_t *wipe_code,
+                                 size_t wipe_code_len);
+secbool storage_has(const uint16_t key);
+secbool storage_get(const uint16_t key, void *val, const uint16_t max_len,
+                    uint16_t *len);
+secbool storage_set(const uint16_t key, const void *val, const uint16_t len);
+secbool storage_delete(const uint16_t key);
+secbool storage_set_counter(const uint16_t key, const uint32_t count);
+secbool storage_next_counter(const uint16_t key, uint32_t *count);
+
+#endif

core/embed/unix/mpconfigport.h

@@ -198,6 +198,7 @@ extern const struct _mp_print_t mp_stderr_print;
 
 #define MICROPY_PY_TREZORCONFIG     (1)
 #define MICROPY_PY_TREZORCRYPTO     (1)
+#define MICROPY_PY_TREZORSTORAGEDEVICE (1)
 #define MICROPY_PY_TREZORIO         (1)
 #define MICROPY_PY_TREZORUI         (1)
 #define MICROPY_PY_TREZORUTILS      (1)

core/mocks/generated/trezorstoragedevice.pyi

@@ -0,0 +1,11 @@
+from typing import *
+
+
+# extmod/rustmods/modtrezorstoragedevice.c
+def is_version_stored() -> bool:
+    """Whether version is in storage."""
+
+
+# extmod/rustmods/modtrezorstoragedevice.c
+def set_version(version: bytes) -> bool:
+    """Save to storage."""

core/src/storage/__init__.py

@@ -1,9 +1,11 @@
 from storage import cache, common, device
 from trezor import config
 
+import trezorstoragedevice
+
 
 def set_current_version() -> None:
-    device.set_version(common.STORAGE_VERSION_CURRENT)
+    trezorstoragedevice.set_version(common.STORAGE_VERSION_CURRENT)
 
 
 def wipe() -> None:
@@ -19,7 +21,7 @@ def init_unlocked() -> None:
 
     # In FWs <= 2.3.1 'version' denoted whether the device is initialized or not.
     # In 2.3.2 we have introduced a new field 'initialized' for that.
-    if device.is_version_stored() and not device.is_initialized():
+    if trezorstoragedevice.is_version_stored() and not device.is_initialized():
         common.set_bool(common.APP_DEVICE, device.INITIALIZED, True, public=True)
 
 

core/src/storage/device.py

@@ -5,6 +5,8 @@ from ubinascii import hexlify
 import storage.cache
 from storage import common
 
+import trezorstoragedevice
+
 if TYPE_CHECKING:
     from trezor.enums import BackupType
     from typing_extensions import Literal
@@ -60,16 +62,16 @@ AUTOLOCK_DELAY_MAXIMUM = 0x2000_0000  # ~6 days
 SD_SALT_AUTH_KEY_LEN_BYTES = const(16)
 
 
-def is_version_stored() -> bool:
-    return bool(common.get(_NAMESPACE, _VERSION))
+# def is_version_stored() -> bool:
+#     return bool(common.get(_NAMESPACE, _VERSION))
 
 
 def get_version() -> bytes | None:
     return common.get(_NAMESPACE, _VERSION)
 
 
-def set_version(version: bytes) -> None:
-    common.set(_NAMESPACE, _VERSION, version)
+# def set_version(version: bytes) -> None:
+#     common.set(_NAMESPACE, _VERSION, version)
 
 
 def is_initialized() -> bool:
@@ -163,7 +165,7 @@ def store_mnemonic_secret(
     needs_backup: bool = False,
     no_backup: bool = False,
 ) -> None:
-    set_version(common.STORAGE_VERSION_CURRENT)
+    trezorstoragedevice.set_version(common.STORAGE_VERSION_CURRENT)
     common.set(_NAMESPACE, _MNEMONIC_SECRET, secret)
     common.set_uint8(_NAMESPACE, _BACKUP_TYPE, backup_type)
     common.set_true_or_delete(_NAMESPACE, _NO_BACKUP, no_backup)