From ed79d35de9b4f2dc47f8654d260e76541d37507c Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Mon, 5 Aug 2024 14:17:41 +0200
Subject: [PATCH] feat(python): make webusb transport more resilient

* convert more USB errors into TransportExceptions
* add a timeout + infinite loop for read/write operations,
  so that they are interruptible by Python, instead of leaving
  the interface in a bad state when hard-killed
* (also ctrl+c now works if the process is waiting for usb)
---
 python/.changelog.d/4089.changed         |  1 +
 python/.changelog.d/4089.fixed           |  1 +
 python/src/trezorlib/transport/webusb.py | 48 +++++++++++++++++++-----
 3 files changed, 40 insertions(+), 10 deletions(-)
 create mode 100644 python/.changelog.d/4089.changed
 create mode 100644 python/.changelog.d/4089.fixed

diff --git a/python/.changelog.d/4089.changed b/python/.changelog.d/4089.changed
new file mode 100644
index 0000000000..8a8b3b80b5
--- /dev/null
+++ b/python/.changelog.d/4089.changed
@@ -0,0 +1 @@
+Most USB level errors are now converted to `TransportException`.
diff --git a/python/.changelog.d/4089.fixed b/python/.changelog.d/4089.fixed
new file mode 100644
index 0000000000..1626297f51
--- /dev/null
+++ b/python/.changelog.d/4089.fixed
@@ -0,0 +1 @@
+It is now possible to interrupt USB communication (via Ctrl+C, or a signal, or any other way).
diff --git a/python/src/trezorlib/transport/webusb.py b/python/src/trezorlib/transport/webusb.py
index 1b60df61bd..8e2d08147a 100644
--- a/python/src/trezorlib/transport/webusb.py
+++ b/python/src/trezorlib/transport/webusb.py
@@ -40,6 +40,9 @@ ENDPOINT = 1
 DEBUG_INTERFACE = 1
 DEBUG_ENDPOINT = 2
 
+USB_COMM_TIMEOUT_MS = 300
+WEBUSB_CHUNK_SIZE = 64
+
 
 class WebUsbHandle:
     def __init__(self, device: "usb1.USBDevice", debug: bool = False) -> None:
@@ -64,28 +67,53 @@ class WebUsbHandle:
 
     def close(self) -> None:
         if self.handle is not None:
-            self.handle.releaseInterface(self.interface)
-            self.handle.close()
+            try:
+                self.handle.releaseInterface(self.interface)
+                self.handle.close()
+            except Exception as e:
+                raise TransportException(f"USB close failed: {e}") from e
         self.handle = None
 
     def write_chunk(self, chunk: bytes) -> None:
         assert self.handle is not None
-        if len(chunk) != 64:
+        if len(chunk) != WEBUSB_CHUNK_SIZE:
             raise TransportException(f"Unexpected chunk size: {len(chunk)}")
         LOG.log(DUMP_PACKETS, f"writing packet: {chunk.hex()}")
-        self.handle.interruptWrite(self.endpoint, chunk)
+        while True:
+            try:
+                bytes_written = self.handle.interruptWrite(
+                    self.endpoint, chunk, USB_COMM_TIMEOUT_MS
+                )
+            except usb1.USBErrorTimeout as e:
+                bytes_written = e.transferred
+            except Exception as e:
+                raise TransportException(f"USB write failed: {e}") from e
+            if bytes_written == 0:
+                continue
+            if bytes_written != len(chunk):
+                raise TransportException(
+                    f"USB partial write: {bytes_written} out of {WEBUSB_CHUNK_SIZE}"
+                )
+            return
 
     def read_chunk(self) -> bytes:
         assert self.handle is not None
         endpoint = 0x80 | self.endpoint
         while True:
-            chunk = self.handle.interruptRead(endpoint, 64)
-            if chunk:
-                break
-            else:
-                time.sleep(0.001)
+            try:
+                chunk = self.handle.interruptRead(
+                    endpoint, WEBUSB_CHUNK_SIZE, USB_COMM_TIMEOUT_MS
+                )
+                if chunk:
+                    break
+                else:
+                    time.sleep(0.001)
+            except usb1.USBErrorTimeout:
+                pass
+            except Exception as e:
+                raise TransportException(f"USB read failed: {e}") from e
         LOG.log(DUMP_PACKETS, f"read packet: {chunk.hex()}")
-        if len(chunk) != 64:
+        if len(chunk) != WEBUSB_CHUNK_SIZE:
             raise TransportException(f"Unexpected chunk size: {len(chunk)}")
         return chunk