Merge pull request #477 from LedgerHQ/xch/fix-screenshot-race-condition

Fix race screenshot race condition and ticker pause behavior

CHANGELOG.md

@@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.8.6] - 2024-04-11
+
+### Fixed
+
+- Fix race condition on screenshot generation
+- Fix ticker pause not waiting for end of tick processing
+
 ## [0.8.5] - 2024-04-11
 
 ### Fixed

speculos/mcu/display.py

@@ -10,6 +10,7 @@
     cache = lru_cache(maxsize=None)
 from PIL import Image
 from socket import socket
+from threading import Lock
 from typing import Any, Dict, IO, List, Optional, Tuple, Union
 
 from speculos.observer import TextEvent
@@ -80,6 +81,7 @@ class FrameBuffer:
     def __init__(self, model: str):
         self.pixels: PixelColorMapping = {}
         self.screenshot_pixels: PixelColorMapping = {}
+        self.screenshot_pixels_lock = Lock()
         self.default_color = 0
         self.draw_default_color = False
         self.reset_screeshot_pixels = False
@@ -123,10 +125,13 @@ def draw_rect(self, x0: int, y0: int, width: int, height: int, color: int) -> Li
         return []
 
     def _get_image(self) -> bytes:
-        data = bytearray(self.default_color.to_bytes(3, "big")) * self._width * self._height
-        for (x, y), color in self.screenshot_pixels.items():
-            pos = 3 * (y * self._width + x)
-            data[pos:pos + 3] = color.to_bytes(3, "big")
+        # This call is made from the Speculos API thread
+        # Protect screenshot_pixels for concurrent Write during this Read
+        with self.screenshot_pixels_lock:
+            data = bytearray(self.default_color.to_bytes(3, "big")) * self._width * self._height
+            for (x, y), color in self.screenshot_pixels.items():
+                pos = 3 * (y * self._width + x)
+                data[pos:pos + 3] = color.to_bytes(3, "big")
         return bytes(data)
 
     def _get_screenshot_iobytes_value(self) -> bytes:
@@ -142,10 +147,13 @@ def take_screenshot(self) -> Tuple[Tuple[int, int], bytes]:
         return self.current_screen_size, self._get_image()
 
     def update_screenshot(self) -> None:
-        if self.reset_screeshot_pixels:
-            self.screenshot_pixels = {}
-            self.reset_screeshot_pixels = False
-        self.screenshot_pixels.update(self.pixels)
+        # This call is made from the MCU/Seproxyhal thread
+        # Protect screenshot_pixels for concurrent Read during this Write
+        with self.screenshot_pixels_lock:
+            if self.reset_screeshot_pixels:
+                self.screenshot_pixels = {}
+                self.reset_screeshot_pixels = False
+            self.screenshot_pixels.update(self.pixels)
 
     def update_public_screenshot(self) -> None:
         # Stax/Flex only

speculos/mcu/seproxyhal.py

@@ -66,7 +66,7 @@ class SephTag(IntEnum):
 
 
 class TimeTickerDaemon(threading.Thread):
-    def __init__(self, add_tick: Callable, *args, **kwargs):
+    def __init__(self, add_tick: Callable, wait_until_tick_is_processed: Callable, *args, **kwargs):
         """
         Initializes the Ticker Daemon
         The goal of this daemon is to emulate the time spent in the application by sending it
@@ -81,6 +81,7 @@ def __init__(self, add_tick: Callable, *args, **kwargs):
         self._paused = False
         self._resume_cond = threading.Condition()
         self.add_tick = add_tick
+        self.wait_until_tick_is_processed = wait_until_tick_is_processed
 
     def pause(self):
         """
@@ -106,6 +107,7 @@ def _wait_if_time_paused(self):
         Internal function to handle the pause
         """
         while self.paused:
+            self.wait_until_tick_is_processed()
             self._paused = True
             with self._resume_cond:
                 self._resume_cond.wait()
@@ -199,18 +201,21 @@ def queue_packet(self, tag: SephTag, data: bytes = b'', priority: bool = False):
         with self.queue_condition:
             self.queue_condition.notify()
 
-    def add_tick(self, wait_fully_processed=False):
+    def wait_until_tick_is_processed(self):
+        # Wait until the app has finished processing the tick
+        while self.tick_requested or self.queue or not self.status_event.is_set():
+            self.status_event.wait()
+
+    def add_tick(self, wait_until_tick_is_processed=False):
         """Request sending of a ticker event to the app"""
         self.tick_requested = True
 
         # notify this thread that a new event is available
         with self.queue_condition:
             self.queue_condition.notify()
 
-        if wait_fully_processed:
-            # Wait until the app have finished processing the tick
-            while self.tick_requested or self.queue or not self.status_event.is_set():
-                self.status_event.wait()
+        if wait_until_tick_is_processed:
+            self.wait_until_tick_is_processed()
 
     def run(self):
         while not self.stop:
@@ -261,7 +266,8 @@ def __init__(self,
         self.socket_helper = SocketHelper(self._socket, self.status_event)
         self.socket_helper.start()
 
-        self.time_ticker_thread = TimeTickerDaemon(self.socket_helper.add_tick)
+        self.time_ticker_thread = TimeTickerDaemon(self.socket_helper.add_tick,
+                                                   self.socket_helper.wait_until_tick_is_processed)
         self.time_ticker_thread.start()
 
         self.usb = usb.USB(self.socket_helper.queue_packet, transport=transport)
@@ -473,7 +479,7 @@ def handle_ticker_request(self, action):
         elif action == "resume":
             self.time_ticker_thread.resume()
         elif action == "single-step":
-            self.time_ticker_thread.add_tick(wait_fully_processed=True)
+            self.time_ticker_thread.add_tick(wait_until_tick_is_processed=True)
 
     def handle_wait(self, delay: float):
         '''Wait for a specified delay, taking account real time seen by the app.'''
@@ -484,7 +490,7 @@ def handle_wait(self, delay: float):
                 time.sleep(TICKER_DELAY)
         else:
             for _ in range(expected_ticks):
-                self.time_ticker_thread.add_tick(wait_fully_processed=True)
+                self.time_ticker_thread.add_tick(wait_until_tick_is_processed=True)
 
     def to_app(self, packet: bytes):
         '''