diff --git a/core/src/apps/management/recovery_device/keyboard_bip39.py b/core/src/apps/management/recovery_device/keyboard_bip39.py index 91a08284ec..9bf5e58813 100644 --- a/core/src/apps/management/recovery_device/keyboard_bip39.py +++ b/core/src/apps/management/recovery_device/keyboard_bip39.py @@ -202,7 +202,7 @@ class Bip39Keyboard(ui.Layout): race = race_touch result = await race - if touch in race.finished: + if isinstance(result, tuple): event, x, y = result self.dispatch(event, x, y) else: diff --git a/core/src/apps/management/recovery_device/keyboard_slip39.py b/core/src/apps/management/recovery_device/keyboard_slip39.py index 0d5a2d39c0..5c5902d963 100644 --- a/core/src/apps/management/recovery_device/keyboard_slip39.py +++ b/core/src/apps/management/recovery_device/keyboard_slip39.py @@ -212,7 +212,7 @@ class Slip39Keyboard(ui.Layout): race = race_touch result = await race - if touch in race.finished: + if isinstance(result, tuple): event, x, y = result self.dispatch(event, x, y) else: diff --git a/core/src/trezor/loop.py b/core/src/trezor/loop.py index a1c461044e..88a2dbf3af 100644 --- a/core/src/trezor/loop.py +++ b/core/src/trezor/loop.py @@ -46,6 +46,13 @@ if __debug__: synthetic_events = [] # type: List[Tuple[int, Any]] +class SyscallTimeout(Exception): + pass + + +_TIMEOUT_ERROR = SyscallTimeout() + + def schedule( task: Task, value: Any = None, *, deadline: int = None, finalizer: Finalizer = None ) -> None: @@ -176,11 +183,10 @@ def _step(task: Task, value: Any) -> None: else: if isinstance(result, Syscall): result.handle(task) - elif result is None: - schedule(task) else: if __debug__: log.error(__name__, "unknown syscall: %s", result) + raise RuntimeError if after_step_hook: after_step_hook() @@ -221,6 +227,8 @@ class sleep(Syscall): deadline = utime.ticks_add(utime.ticks_us(), self.delay_us) schedule(task, deadline, deadline=deadline) +YIELD = sleep(0) + class wait(Syscall): """ @@ -234,11 +242,68 @@ class wait(Syscall): >>> event, x, y = await loop.wait(io.TOUCH) # await touch event """ - def __init__(self, msg_iface: int) -> None: + # The wait class implements a coroutine interface, so that it can be scheduled + # as a regular task. When it is resumed, it can perform cleanup and then give + # control back to the callback task. + # By returning a Syscall instance that does nothing in its handle() method, it + # ensures that the wait() instance will never be automatically resumed. + _DO_NOT_RESCHEDULE = Syscall() + _TIMEOUT_INDICATOR = object() + + def __init__(self, msg_iface: int, timeout_us: int = None) -> None: self.msg_iface = msg_iface + self.timeout_us = timeout_us + self.callback = None # type: Optional[Task] def handle(self, task: Task) -> None: - pause(task, self.msg_iface) + # We pause (and optionally schedule) ourselves instead of the calling task. + # When resumed, send() or throw() will give control to the calling task, + # after performing cleanup. + pause(self, self.msg_iface) + if self.timeout_us is not None: + deadline = utime.ticks_add(utime.ticks_us(), self.timeout_us) + schedule(self, wait._TIMEOUT_ERROR, deadline=deadline) + + def send(self, value) -> Any: + assert self.callback is not None + if value is wait._TIMEOUT_INDICATOR: + # we were resumed from the timeout + # discard the i/o wait + _paused[self.msg_iface].discard(self) + # convert the value to an exception + value = _TIMEOUT_ERROR + elif self.timeout_us is not None: + # we were resumed from the i/o wait, AND timeout was specified + # discard the scheduled timeout + _queue.discard(self) + _step(self.callback, value) + return wait._DO_NOT_RESCHEDULE + + def throw(self, exception, _value=None, _traceback=None) -> None: + assert self.callback is not None + # An exception was thrown to us. + # This should not happen unless (a) the i/o sent it, or (b) we were closed + # externally. + # Discard both the timeout and the i/o wait, because we don't know which + # caused this, if any. + _queue.discard(self) + _paused[self.msg_iface].discard(self) + # resume the callback anyway + _step(self.callback, exception) + return wait._DO_NOT_RESCHEDULE + + def close(self) -> None: + pass + + def __iter__(self) -> Task: # type: ignore + try: + return (yield self) + except: # noqa: E722 + # exception was raised on the waiting task externally with + # close() or throw(), kill the children tasks and re-raise + _queue.discard(self) + _paused[self.msg_iface].discard(self) + raise _type_gen = type((lambda: (yield))()) diff --git a/core/src/trezor/ui/__init__.py b/core/src/trezor/ui/__init__.py index 4d8a672e69..f78d113573 100644 --- a/core/src/trezor/ui/__init__.py +++ b/core/src/trezor/ui/__init__.py @@ -380,4 +380,4 @@ class Layout(Component): def wait_until_layout_is_running() -> Awaitable[None]: # type: ignore while not layout_chan.takers: - yield + yield loop.YIELD diff --git a/core/src/trezor/ui/passphrase.py b/core/src/trezor/ui/passphrase.py index c2713fce1a..b996c205a3 100644 --- a/core/src/trezor/ui/passphrase.py +++ b/core/src/trezor/ui/passphrase.py @@ -208,21 +208,20 @@ class PassphraseKeyboard(ui.Layout): async def handle_input(self) -> None: touch = loop.wait(io.TOUCH) - timeout = loop.sleep(1000 * 1000 * 1) - race_touch = loop.race(touch) - race_timeout = loop.race(touch, timeout) + touch_timeout = loop.wait(io.TOUCH, timeout_us=1000 * 1000 * 1) + # timeout = loop.sleep(1000 * 1000 * 1) + # race_touch = loop.race(touch) + # race_timeout = loop.race(touch, timeout) while True: if self.pending_button is not None: - race = race_timeout + touch_source = touch_timeout else: - race = race_touch - result = await race + touch_source = touch - if touch in race.finished: - event, x, y = result - self.dispatch(event, x, y) - else: + try: + self.dispatch(*(await touch_source)) + except loop.SyscallTimeout: self.on_timeout() async def handle_paging(self) -> None: