WIP - U2F counter

core/embed/rust/build.rs

@@ -251,7 +251,9 @@ fn generate_storagedevice_bindings() {
         // storage
         .allowlist_function("storage_set")
         .allowlist_function("storage_delete")
-        .allowlist_function("storage_get");
+        .allowlist_function("storage_get")
+        .allowlist_function("storage_set_counter")
+        .allowlist_function("storage_next_counter");
 
     // Write the bindings to a file in the OUT_DIR.
     bindings

core/embed/rust/librust_qstr.h

@@ -55,6 +55,8 @@ static void _librust_qstrs(void) {
   MP_QSTR_set_safety_check_level;
   MP_QSTR_get_sd_salt_auth_key;
   MP_QSTR_set_sd_salt_auth_key;
+  MP_QSTR_get_next_u2f_counter;
+  MP_QSTR_set_u2f_counter;
 
   MP_QSTR_set_timer_fn;
   MP_QSTR_touch_event;

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

@@ -15,6 +15,7 @@ use core::convert::{TryFrom, TryInto};
 // - raising custom exceptions into python
 
 const FLAG_PUBLIC: u8 = 0x80;
+const FLAGS_WRITE: u8 = 0xC0;
 
 const APP_DEVICE: u8 = 0x01;
 
@@ -93,7 +94,8 @@ extern "C" {
     // fn storage_set(key: u16, val: *const cty::c_void, len: u16) ->
     // secbool::Secbool;
     fn storage_set(key: u16, val: *const u8, len: u16) -> secbool::Secbool;
-
+    fn storage_next_counter(key: u16, count: *mut u32) -> secbool::Secbool;
+    fn storage_set_counter(key: u16, count: u32) -> secbool::Secbool;
 }
 
 extern "C" fn storagedevice_is_version_stored() -> Obj {
@@ -524,6 +526,24 @@ extern "C" fn storagedevice_set_sd_salt_auth_key(auth_key: Obj) -> Obj {
     unsafe { util::try_or_raise(block) }
 }
 
+extern "C" fn storagedevice_get_next_u2f_counter() -> Obj {
+    let block = || {
+        let key = _get_appkey_u2f(U2F_COUNTER, true);
+        storagedevice_storage_get_next_counter(key).try_into()
+    };
+    unsafe { util::try_or_raise(block) }
+}
+
+extern "C" fn storagedevice_set_u2f_counter(count: Obj) -> Obj {
+    let block = || {
+        let count = u32::try_from(count)?;
+
+        let key = _get_appkey_u2f(U2F_COUNTER, true);
+        Ok(storagedevice_storage_set_counter(key, count).into())
+    };
+    unsafe { util::try_or_raise(block) }
+}
+
 pub fn storagedevice_storage_get(key: u16) -> Vec<u8, MAX_LEN> {
     let mut buf: [u8; MAX_LEN] = [0; MAX_LEN];
     let mut len: u16 = 0;
@@ -567,6 +587,12 @@ pub fn storagedevice_storage_get_homescreen(key: u16) -> Vec<u8, HOMESCREEN_MAXS
     vector_result
 }
 
+pub fn storagedevice_storage_get_next_counter(key: u16) -> u32 {
+    let mut count: u32 = 0;
+    unsafe { storage_next_counter(key, &mut count as *mut _) };
+    count
+}
+
 pub fn storagedevice_storage_get_bool(key: u16) -> bool {
     let result = storagedevice_storage_get(key);
     result.len() == 1 && result[0] == _TRUE_BYTE
@@ -606,6 +632,10 @@ pub fn storagedevice_storage_set(key: u16, val: *const u8, len: u16) -> bool {
     matches!(unsafe { storage_set(key, val, len) }, secbool::TRUE)
 }
 
+pub fn storagedevice_storage_set_counter(key: u16, count: u32) -> bool {
+    matches!(unsafe { storage_set_counter(key, count) }, secbool::TRUE)
+}
+
 pub fn storagedevice_storage_set_bool(key: u16, val: bool) -> bool {
     let val = if val { [_TRUE_BYTE] } else { [_FALSE_BYTE] };
     storagedevice_storage_set(key, &val as *const _, 1)
@@ -650,6 +680,15 @@ fn _get_appkey(key: u8, is_public: bool) -> u16 {
     ((app as u16) << 8) | key as u16
 }
 
+fn _get_appkey_u2f(key: u8, writable_locked: bool) -> u16 {
+    let app = if writable_locked {
+        APP_DEVICE | FLAGS_WRITE
+    } else {
+        APP_DEVICE | FLAG_PUBLIC
+    };
+    ((app as u16) << 8) | key as u16
+}
+
 fn _normalize_autolock_delay(delay_ms: u32) -> u32 {
     if delay_ms < AUTOLOCK_DELAY_MINIMUM {
         AUTOLOCK_DELAY_MINIMUM
@@ -827,6 +866,14 @@ pub static mp_module_trezorstoragedevice: Module = obj_module! {
     /// def set_sd_salt_auth_key(auth_key: bytes | None) -> bool:
     ///     """The key used to check the authenticity of the SD card salt."""
     Qstr::MP_QSTR_set_sd_salt_auth_key => obj_fn_1!(storagedevice_set_sd_salt_auth_key).as_obj(),
+
+    /// def get_next_u2f_counter() -> int:
+    ///     """Get next U2F counter."""
+    Qstr::MP_QSTR_get_next_u2f_counter => obj_fn_0!(storagedevice_get_next_u2f_counter).as_obj(),
+
+    /// def set_u2f_counter(count: int) -> bool:
+    ///     """Set U2F counter."""
+    Qstr::MP_QSTR_set_u2f_counter => obj_fn_1!(storagedevice_set_u2f_counter).as_obj(),
 };
 
 #[cfg(test)]
@@ -845,6 +892,18 @@ mod tests {
         assert_eq!(result, 0x8111);
     }
 
+    #[test]
+    fn get_appkey_u2f_writable_locked() {
+        let result = _get_appkey_u2f(0x09, true);
+        assert_eq!(result, 0xC109);
+    }
+
+    #[test]
+    fn get_appkey_u2f_not_writable_locked() {
+        let result = _get_appkey_u2f(0x09, false);
+        assert_eq!(result, 0x8109);
+    }
+
     #[test]
     fn normalize_autolock_delay_small() {
         let result = _normalize_autolock_delay(123);

core/mocks/generated/trezorstoragedevice.pyi

@@ -193,3 +193,13 @@ def get_sd_salt_auth_key() -> bytes | None:
 # rust/src/storagedevice/storage_device.rs
 def set_sd_salt_auth_key(auth_key: bytes | None) -> bool:
     """The key used to check the authenticity of the SD card salt."""
+
+
+# rust/src/storagedevice/storage_device.rs
+def get_next_u2f_counter() -> int:
+    """Get next U2F counter."""
+
+
+# rust/src/storagedevice/storage_device.rs
+def set_u2f_counter(count: int) -> bool:
+    """Set U2F counter."""

core/src/apps/management/get_next_u2f_counter.py

@@ -1,4 +1,3 @@
-import storage.device
 from trezor import storagedevice, ui, wire
 from trezor.enums import ButtonRequestType
 from trezor.messages import GetNextU2FCounter, NextU2FCounter
@@ -20,4 +19,4 @@ async def get_next_u2f_counter(
         br_code=ButtonRequestType.ProtectCall,
     )
 
-    return NextU2FCounter(u2f_counter=storage.device.next_u2f_counter())
+    return NextU2FCounter(u2f_counter=storagedevice.get_next_u2f_counter())

core/src/apps/management/recovery_device/__init__.py

@@ -58,7 +58,7 @@ async def recovery_device(ctx: wire.Context, msg: RecoveryDevice) -> Success:
 
         storagedevice.set_passphrase_enabled(bool(msg.passphrase_protection))
         if msg.u2f_counter is not None:
-            storage.device.set_u2f_counter(msg.u2f_counter)
+            storagedevice.set_u2f_counter(msg.u2f_counter)
         if msg.label is not None:
             storage.device.set_label(msg.label)
 

core/src/apps/management/set_u2f_counter.py

@@ -1,4 +1,3 @@
-import storage.device
 from trezor import storagedevice, ui, wire
 from trezor.enums import ButtonRequestType
 from trezor.messages import SetU2FCounter, Success
@@ -21,6 +20,6 @@ async def set_u2f_counter(ctx: wire.Context, msg: SetU2FCounter) -> Success:
         br_code=ButtonRequestType.ProtectCall,
     )
 
-    storage.device.set_u2f_counter(msg.u2f_counter)
+    storagedevice.set_u2f_counter(msg.u2f_counter)
 
     return Success(message="U2F counter set")

core/src/apps/webauthn/credential.py

@@ -3,8 +3,7 @@ from micropython import const
 from typing import Iterable
 from ubinascii import hexlify
 
-import storage.device
-from trezor import log, utils
+from trezor import log, storagedevice, utils
 from trezor.crypto import bip32, chacha20poly1305, der, hashlib, hmac, random
 from trezor.crypto.curve import ed25519, nist256p1
 
@@ -92,7 +91,7 @@ class Credential:
         return None
 
     def next_signature_counter(self) -> int:
-        return storage.device.next_u2f_counter() or 0
+        return storagedevice.get_next_u2f_counter() or 0
 
     @staticmethod
     def from_bytes(data: bytes, rp_id_hash: bytes) -> "Credential":
@@ -128,7 +127,7 @@ class Fido2Credential(Credential):
         return True
 
     def generate_id(self) -> None:
-        self.creation_time = storage.device.next_u2f_counter() or 0
+        self.creation_time = self.next_signature_counter()
 
         if not self.check_required_fields():
             raise AssertionError

core/src/storage/__init__.py

@@ -37,7 +37,7 @@ def _migrate_from_version_01() -> None:
     # U2F counter wasn't public, so we are intentionally not using storage.device module.
     counter = common.get(common.APP_DEVICE, device.U2F_COUNTER)
     if counter is not None:
-        device.set_u2f_counter(int.from_bytes(counter, "big"))
+        storagedevice.set_u2f_counter(int.from_bytes(counter, "big"))
         # Delete the old, non-public U2F_COUNTER.
         common.delete(common.APP_DEVICE, device.U2F_COUNTER)
     set_current_version()

core/src/storage/device.py

@@ -247,12 +247,12 @@ def set_label(label: str) -> None:
 #     common.set(_NAMESPACE, _AUTOLOCK_DELAY_MS, delay_ms.to_bytes(4, "big"))
 
 
-def next_u2f_counter() -> int:
-    return common.next_counter(_NAMESPACE, U2F_COUNTER, writable_locked=True)
+# def next_u2f_counter() -> int:
+#     return common.next_counter(_NAMESPACE, U2F_COUNTER, writable_locked=True)
 
 
-def set_u2f_counter(count: int) -> None:
-    common.set_counter(_NAMESPACE, U2F_COUNTER, count, writable_locked=True)
+# def set_u2f_counter(count: int) -> None:
+#     common.set_counter(_NAMESPACE, U2F_COUNTER, count, writable_locked=True)
 
 
 # def set_slip39_identifier(identifier: int) -> None: