exchange: add timeout in ticks

LedgerHQ · Jan 18, 2024 · eddc2fe · eddc2fe
1 parent e262a0c
commit eddc2fe
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 11 deletions.
diff --git a/speculos/api/apdu.py b/speculos/api/apdu.py
@@ -18,17 +18,23 @@ def __init__(self, seph: SeProxyHal):
         self._seph.apdu_callbacks.append(self.seph_apdu_callback)
         self.response: Optional[bytes]
 
-    def exchange(self, data: bytes) -> Generator[bytes, None, None]:
+    def exchange(self, data: bytes, tick_timeout: int = 0) -> Generator[bytes, None, None]:
         # force headers to be sent
         yield b""
 
+        tick_count_before_exchange = self._seph.socket_helper.ticks_count
+
         with self.endpoint_lock:  # Lock for a command/response for one client
             with self.response_condition:
                 self.response = None
             self._seph.to_app(data)
             with self.response_condition:
                 while self.response is None:
-                    self.response_condition.wait()
+                    self.response_condition.wait(0.1)
+                    exchange_tick_count = self._seph.socket_helper.ticks_count - tick_count_before_exchange
+
+                    if tick_timeout != 0 and exchange_tick_count > tick_timeout:
+                        yield
             yield json.dumps({"data": self.response.hex()}).encode()
 
     def seph_apdu_callback(self, data: bytes) -> None:
@@ -55,5 +61,6 @@ def post(self):
         except jsonschema.exceptions.ValidationError as e:
             return {"error": f"{e}"}, 400
 
-        data = bytes.fromhex(args.get("data"))
-        return Response(stream_with_context(self._bridge.exchange(data)), content_type="application/json")
+        tick_timeout = args.get("data")["tick_timeout"]
+        data = bytes.fromhex(args.get("data")["apdu"])
+        return Response(stream_with_context(self._bridge.exchange(data, tick_timeout)), content_type="application/json")
diff --git a/speculos/api/resources/apdu.schema b/speculos/api/resources/apdu.schema
@@ -3,7 +3,10 @@
 
     "type": "object",
     "properties": {
-        "data": { "type": "string", "pattern": "^([a-fA-F0-9]{2})+$" }
+        "data": {
+            "apdu": { "type": "string", "pattern": "^([a-fA-F0-9]{2})+$" },
+            "tick_timeout": { "type": "int"}
+        }
     },
     "required": [ "data" ],
     "additionalProperties": false

diff --git a/speculos/client.py b/speculos/client.py
@@ -141,10 +141,14 @@ def get_screenshot(self) -> bytes:
             check_status_code(response, "/screenshot")
             return response.content
 
-    def _apdu_exchange(self, data: bytes) -> bytes:
-        with self.session.post(f"{self.api_url}/apdu", json={"data": data.hex()}) as response:
-            apdu_response = ApduResponse(response)
-            return apdu_response.receive()
+    def _apdu_exchange(self, data: bytes, tick_timeout: int = 0) -> bytes:
+        try:
+            data_payload = {"apdu": data.hex(), "tick_timeout": tick_timeout}
+            with self.session.post(f"{self.api_url}/apdu", json={"data": data_payload}) as response:
+                apdu_response = ApduResponse(response)
+                return apdu_response.receive()
+        except:
+            raise TimeoutError()
 
     def _apdu_exchange_nowait(self, data: bytes) -> requests.Response:
         return self.session.post(f"{self.api_url}/apdu", json={"data": data.hex()}, stream=True)
@@ -239,9 +243,9 @@ def __exit__(
     ) -> None:
         self.stop()
 
-    def apdu_exchange(self, cla: int, ins: int, data: bytes = b"", p1: int = 0, p2: int = 0) -> bytes:
+    def apdu_exchange(self, cla: int, ins: int, data: bytes = b"", p1: int = 0, p2: int = 0, tick_timeout:int = 0) -> bytes:
         apdu = bytes([cla, ins, p1, p2, len(data)]) + data
-        return Api._apdu_exchange(self, apdu)
+        return Api._apdu_exchange(self, apdu, tick_timeout)
 
     @contextmanager
     def apdu_exchange_nowait(