From 84674a746333c2c92a2b18f0796479e323188933 Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Thu, 8 Aug 2019 15:37:05 +0200
Subject: [PATCH] core/usb: Add write_blocking() to trezor.io.HID.

---
 .../extmod/modtrezorio/modtrezorio-hid.h      | 19 +++++++++++++++++++
 core/embed/unix/usb.c                         | 13 +++++++++++++
 core/mocks/generated/trezorio.pyi             |  5 +++++
 3 files changed, 37 insertions(+)

diff --git a/core/embed/extmod/modtrezorio/modtrezorio-hid.h b/core/embed/extmod/modtrezorio/modtrezorio-hid.h
index a8f2c5e8e..28761e40e 100644
--- a/core/embed/extmod/modtrezorio/modtrezorio-hid.h
+++ b/core/embed/extmod/modtrezorio/modtrezorio-hid.h
@@ -128,10 +128,29 @@ 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 write_blocking(self, msg: bytes, timeout_ms: int) -> int:
+///     """
+///     Sends message using USB HID (device) or UDP (emulator).
+///     """
+STATIC mp_obj_t mod_trezorio_HID_write_blocking(mp_obj_t self, mp_obj_t msg,
+                                                mp_obj_t timeout_ms) {
+  mp_obj_HID_t *o = MP_OBJ_TO_PTR(self);
+  mp_buffer_info_t buf;
+  mp_get_buffer_raise(msg, &buf, MP_BUFFER_READ);
+  uint32_t timeout = trezor_obj_get_uint(timeout_ms);
+  ssize_t r =
+      usb_hid_write_blocking(o->info.iface_num, buf.buf, buf.len, timeout);
+  return MP_OBJ_NEW_SMALL_INT(r);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_3(mod_trezorio_HID_write_blocking_obj,
+                                 mod_trezorio_HID_write_blocking);
+
 STATIC const mp_rom_map_elem_t mod_trezorio_HID_locals_dict_table[] = {
     {MP_ROM_QSTR(MP_QSTR_iface_num),
      MP_ROM_PTR(&mod_trezorio_HID_iface_num_obj)},
     {MP_ROM_QSTR(MP_QSTR_write), MP_ROM_PTR(&mod_trezorio_HID_write_obj)},
+    {MP_ROM_QSTR(MP_QSTR_write_blocking),
+     MP_ROM_PTR(&mod_trezorio_HID_write_blocking_obj)},
 };
 STATIC MP_DEFINE_CONST_DICT(mod_trezorio_HID_locals_dict,
                             mod_trezorio_HID_locals_dict_table);
diff --git a/core/embed/unix/usb.c b/core/embed/unix/usb.c
index 360631bd8..1a439a594 100644
--- a/core/embed/unix/usb.c
+++ b/core/embed/unix/usb.c
@@ -24,6 +24,7 @@
 #include <string.h>
 #include <sys/poll.h>
 #include <sys/socket.h>
+#include <time.h>
 
 #include "touch.h"
 #include "usb.h"
@@ -227,6 +228,18 @@ int usb_hid_write(uint8_t iface_num, const uint8_t *buf, uint32_t len) {
   return usb_emulated_write(iface_num, buf, len);
 }
 
+int usb_hid_write_blocking(uint8_t iface_num, const uint8_t *buf, uint32_t len,
+                           int timeout) {
+  const uint32_t start = clock();
+  while (sectrue != usb_hid_can_write(iface_num)) {
+    if (timeout >= 0 &&
+        (1000 * (clock() - start)) / CLOCKS_PER_SEC >= timeout) {
+      return 0;  // Timeout
+    }
+  }
+  return usb_hid_write(iface_num, buf, len);
+}
+
 int usb_webusb_write(uint8_t iface_num, const uint8_t *buf, uint32_t len) {
   if (iface_num >= USBD_MAX_NUM_INTERFACES ||
       usb_ifaces[iface_num].type != USB_IFACE_TYPE_WEBUSB) {
diff --git a/core/mocks/generated/trezorio.pyi b/core/mocks/generated/trezorio.pyi
index 185be8b17..d3b3cb4ea 100644
--- a/core/mocks/generated/trezorio.pyi
+++ b/core/mocks/generated/trezorio.pyi
@@ -171,6 +171,11 @@ class HID:
         Sends message using USB HID (device) or UDP (emulator).
         """
 
+    def write_blocking(self, msg: bytes, timeout_ms: int) -> int:
+        """
+        Sends message using USB HID (device) or UDP (emulator).
+        """
+
 
 # extmod/modtrezorio/modtrezorio-poll.h
 def poll(ifaces: Iterable[int], list_ref: List, timeout_us: int) -> bool: