fixup! refactor(core): split polling can_read and reading from USB

wip
trezor · Dec 6, 2024 · 7d9a53c · 7d9a53c
1 parent c9188fb
commit 7d9a53c
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 46 deletions.
diff --git a/core/embed/io/usb/inc/io/usb.h b/core/embed/io/usb/inc/io/usb.h
@@ -26,6 +26,8 @@
 #include <io/usb_vcp.h>
 #include <io/usb_webusb.h>
 
+#define USB_PACKET_LEN 64
+
 // clang-format off
 //
 // USB stack high-level state machine

diff --git a/core/embed/upymod/modtrezorio/modtrezorio-hid.h b/core/embed/upymod/modtrezorio/modtrezorio-hid.h
@@ -141,7 +141,7 @@ STATIC mp_obj_t mod_trezorio_HID_write(mp_obj_t self, mp_obj_t msg) {
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_trezorio_HID_write_obj,
                                  mod_trezorio_HID_write);
 
-/// def read(self, buf: bytes, offset: int = 0, limit: int | None = None) -> int
+/// def read(self, buf: bytes, offset: int = 0) -> int
 ///     """
 ///     Reads message using USB HID (device) or UDP (emulator).
 ///     """
@@ -155,15 +155,23 @@ STATIC mp_obj_t mod_trezorio_HID_read(size_t n_args, const mp_obj_t *args) {
     offset = mp_obj_get_int(args[2]);
   }
 
-  int limit;
-  if (n_args >= 3) {
-    limit = mp_obj_get_int(args[3]);
-  } else {
-    limit = buf.len - offset;
+  if (offset < 0) {
+    mp_raise_ValueError("Negative offset not allowed");
+  }
+
+  uint32_t buffer_space = buf.len - offset;
+
+  if (buffer_space < USB_PACKET_LEN) {
+    mp_raise_ValueError("Buffer too small");
+  }
+
+  ssize_t r = usb_hid_read(o->info.iface_num, &((uint8_t *)buf.buf)[offset],
+                           USB_PACKET_LEN);
+
+  if (r != USB_PACKET_LEN) {
+    mp_raise_msg(&mp_type_RuntimeError, "Unexpected read length");
   }
 
-  ssize_t r =
-      usb_hid_read(o->info.iface_num, &((uint8_t *)buf.buf)[offset], limit);
   return MP_OBJ_NEW_SMALL_INT(r);
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_trezorio_HID_read_obj, 3, 4,

diff --git a/core/embed/upymod/modtrezorio/modtrezorio-poll.h b/core/embed/upymod/modtrezorio/modtrezorio-poll.h
@@ -169,15 +169,12 @@ STATIC mp_obj_t mod_trezorio_poll(mp_obj_t ifaces, mp_obj_t list_ref,
         if ((sectrue == usb_hid_can_read(iface)) ||
             (sectrue == usb_webusb_can_read(iface))) {
           ret->items[0] = MP_OBJ_NEW_SMALL_INT(i);
-          ret->items[1] = MP_OBJ_NEW_SMALL_INT(64);
+          ret->items[1] = MP_OBJ_NEW_SMALL_INT(USB_PACKET_LEN);
           return mp_const_true;
         }
       } else if (mode == POLL_WRITE) {
-        if (sectrue == usb_hid_can_write(iface)) {
-          ret->items[0] = MP_OBJ_NEW_SMALL_INT(i);
-          ret->items[1] = mp_const_none;
-          return mp_const_true;
-        } else if (sectrue == usb_webusb_can_write(iface)) {
+        if ((sectrue == usb_hid_can_write(iface)) ||
+            (sectrue == usb_webusb_can_write(iface))) {
           ret->items[0] = MP_OBJ_NEW_SMALL_INT(i);
           ret->items[1] = mp_const_none;
           return mp_const_true;

diff --git a/core/embed/upymod/modtrezorio/modtrezorio-webusb.h b/core/embed/upymod/modtrezorio/modtrezorio-webusb.h
@@ -127,7 +127,7 @@ STATIC mp_obj_t mod_trezorio_WebUSB_write(mp_obj_t self, mp_obj_t msg) {
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_trezorio_WebUSB_write_obj,
                                  mod_trezorio_WebUSB_write);
 
-/// def read(self, buf: bytes, offset: int = 0, limit: int | None = None) -> int
+/// def read(self, buf: bytes, offset: int = 0) -> int
 ///     """
 ///     Reads message using USB WebUSB (device) or UDP (emulator).
 ///     """
@@ -141,18 +141,26 @@ STATIC mp_obj_t mod_trezorio_WebUSB_read(size_t n_args, const mp_obj_t *args) {
     offset = mp_obj_get_int(args[2]);
   }
 
-  int limit;
-  if (n_args >= 3) {
-    limit = mp_obj_get_int(args[3]);
-  } else {
-    limit = buf.len - offset;
+  if (offset < 0) {
+    mp_raise_ValueError("Negative offset not allowed");
+  }
+
+  uint32_t buffer_space = buf.len - offset;
+
+  if (buffer_space < USB_PACKET_LEN) {
+    mp_raise_ValueError("Buffer too small");
+  }
+
+  ssize_t r = usb_webusb_read(o->info.iface_num, &((uint8_t *)buf.buf)[offset],
+                              USB_PACKET_LEN);
+
+  if (r != USB_PACKET_LEN) {
+    mp_raise_msg(&mp_type_RuntimeError, "Unexpected read length");
   }
 
-  ssize_t r =
-      usb_webusb_read(o->info.iface_num, &((uint8_t *)buf.buf)[offset], limit);
   return MP_OBJ_NEW_SMALL_INT(r);
 }
-STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_trezorio_WebUSB_read_obj, 2, 4,
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_trezorio_WebUSB_read_obj, 2, 3,
                                            mod_trezorio_WebUSB_read);
 
 STATIC const mp_rom_map_elem_t mod_trezorio_WebUSB_locals_dict_table[] = {

diff --git a/core/mocks/generated/trezorio/__init__.pyi b/core/mocks/generated/trezorio/__init__.pyi
@@ -32,7 +32,7 @@ class HID:
         Sends message using USB HID (device) or UDP (emulator).
         """
 
-    def read(self, buf: bytes, offset: int = 0, limit: int | None = None) -> int
+    def read(self, buf: bytes, offset: int = 0) -> int
         """
         Reads message using USB HID (device) or UDP (emulator).
         """
@@ -154,7 +154,7 @@ class WebUSB:
         Sends message using USB WebUSB (device) or UDP (emulator).
         """
 
-    def read(self, buf: bytes, offset: int = 0, limit: int | None = None) -> int
+    def read(self, buf: bytes, offset: int = 0) -> int
         """
         Reads message using USB WebUSB (device) or UDP (emulator).
         """

diff --git a/core/src/apps/webauthn/fido2.py b/core/src/apps/webauthn/fido2.py
@@ -378,9 +378,7 @@ async def _read_cmd(iface: HID) -> Cmd | None:
     # wait for incoming command indefinitely
     msg_len = await read
     buf = bytearray(msg_len)
-    read_len = iface.read(buf, 0, msg_len)
-    if read_len != msg_len:
-        raise ValueError("Invalid length")
+    iface.read(buf, 0)
     while True:
         ifrm = overlay_struct(bytearray(buf), desc_init)
         bcnt = ifrm.bcnt
@@ -421,9 +419,7 @@ async def _read_cmd(iface: HID) -> Cmd | None:
             try:
                 msg_len = await read
                 buf = bytearray(msg_len)
-                read_len = iface.read(buf, 0, msg_len)
-                if read_len != msg_len:
-                    raise ValueError("Invalid length")
+                iface.read(buf, 0)
             except loop.Timeout:
                 if __debug__:
                     warning(__name__, "_ERR_MSG_TIMEOUT")

diff --git a/core/src/trezor/wire/codec/codec_v1.py b/core/src/trezor/wire/codec/codec_v1.py
@@ -27,10 +27,7 @@ async def read_message(iface: WireInterface, buffer: utils.BufferType) -> Messag
     # wait for initial report
     msg_len = await read
     report = bytearray(msg_len)
-    read_len = iface.read(report, 0, msg_len)
-    if read_len != msg_len:
-        print("read_len", read_len, "msg_len", msg_len)
-        raise CodecError("Invalid length")
+    iface.read(report, 0)
     if report[0] != _REP_MARKER:
         raise CodecError("Invalid magic")
     _, magic1, magic2, mtype, msize = ustruct.unpack(_REP_INIT, report)
@@ -57,9 +54,7 @@ async def read_message(iface: WireInterface, buffer: utils.BufferType) -> Messag
         # wait for continuation report
         msg_len = await read
         report = bytearray(msg_len)
-        read_len = iface.read(report, 0, msg_len)
-        if read_len != msg_len:
-            raise CodecError("Invalid length")
+        iface.read(report, 0)
         if report[0] != _REP_MARKER:
             raise CodecError("Invalid magic")
 

diff --git a/core/tests/test_trezor.wire.codec.codec_v1.py b/core/tests/test_trezor.wire.codec.codec_v1.py
@@ -25,17 +25,14 @@ def mock_read(self, packet, gen):
         self.packet = packet
         return gen.send(len(packet))
 
-    def read(self, buffer, offset=0, limit=None):
+    def read(self, buffer, offset=0):
         if self.packet is None:
             raise Exception("No packet to read")
-        if limit is None:
-            limit = len(buffer) - offset
 
-        if len(self.packet) > limit:
-            end = offset + limit
-            buffer[offset:end] = self.packet[:limit]
-            self.packet = None
-            return limit
+        buffer_space = len(buffer) - offset
+
+        if len(self.packet) > buffer_space:
+            raise Exception("Buffer too small")
         else:
             end = offset + len(self.packet)
             buffer[offset:end] = self.packet