trezor.msg: remove select, add poll

TODO: decide if we want to continue polling on interface numbers, or rather on usb/touch file objects

micropython/extmod/modtrezormsg/modtrezormsg-stm32.h

@@ -6,22 +6,4 @@
  */
 
 #include "usb.h"
-
-void msg_init(void)
-{
-}
-
-ssize_t msg_recv(uint8_t *iface, uint8_t *buf, size_t len)
-{
-    int i = usb_hid_read_select(0);
-    if (i < 0) {
-        return 0;
-    }
-    *iface = i;
-    return usb_hid_read(i, buf, len);
-}
-
-ssize_t msg_send(uint8_t iface, const uint8_t *buf, size_t len)
-{
-    return usb_hid_write_blocking(iface, buf, len, 1000); // 1s timeout
-}
+#include "touch.h"

micropython/extmod/modtrezormsg/modtrezormsg-unix.h

@@ -7,24 +7,29 @@
 
 #include <arpa/inet.h>
 #include <sys/socket.h>
+#include <sys/poll.h>
 #include <fcntl.h>
 #include <assert.h>
 #include <stdlib.h>
 
-#include "unix-usb-mock.h"
+#include "../../trezorhal/usb.h"
+#include "../../trezorhal/touch.h"
 
 #define TREZOR_UDP_PORT 21324
 
-static int s;
+static int sock;
 static struct sockaddr_in si_me, si_other;
 static socklen_t slen = 0;
 
-void msg_init(void)
-{
-    s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
-    assert(s != -1);
+int usb_init(const usb_dev_info_t *dev_info) {
+    (void)dev_info;
 
-    fcntl(s, F_SETFL, O_NONBLOCK);
+    sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
+    if (sock < 0) {
+        return -1;
+    }
+
+    fcntl(sock, F_SETFL, O_NONBLOCK);
 
     si_me.sin_family = AF_INET;
     const char *ip = getenv("TREZOR_UDP_IP");
@@ -40,38 +45,86 @@ void msg_init(void)
         si_me.sin_port = htons(TREZOR_UDP_PORT);
     }
 
-    int b;
-    b = bind(s, (struct sockaddr*)&si_me, sizeof(si_me));
-    assert(b != -1);
+    int b = bind(sock, (struct sockaddr*)&si_me, sizeof(si_me));
+    if (b < 0) {
+        return -1;
+    }
+
+    return 0;
+}
+
+int usb_deinit(void) {
+    return 0;
+}
+
+int usb_start(void) {
+    return 0;
+}
+
+int usb_stop(void) {
+    return 0;
 }
 
-ssize_t msg_send(uint8_t iface, const uint8_t *buf, size_t len);
+int usb_hid_add(const usb_hid_info_t *info) {
+    return 0;
+}
 
-ssize_t msg_recv(uint8_t *iface, uint8_t *buf, size_t len)
-{
+int usb_vcp_add(const usb_vcp_info_t *info) {
+    return 0;
+}
+
+int usb_hid_can_read(uint8_t iface_num) {
+    struct pollfd fds[] = {
+        { sock, POLLIN, 0 },
+    };
+    int r = poll(fds, 1, 0);
+    if (r > 0) {
+        return 1;
+    } else {
+        return 0;
+    }
+}
+
+int usb_hid_can_write(uint8_t iface_num) {
+    struct pollfd fds[] = {
+        { sock, POLLOUT, 0 },
+    };
+    int r = poll(fds, 1, 0);
+    if (r > 0) {
+        return 1;
+    } else {
+        return 0;
+    }
+}
+
+int usb_hid_read(uint8_t iface_num, uint8_t *buf, uint32_t len) {
     struct sockaddr_in si;
     socklen_t sl = sizeof(si);
-    memset(buf, 0, len);
-    *iface = 0; // TODO: return proper interface
-    ssize_t r = recvfrom(s, buf, len, MSG_DONTWAIT, (struct sockaddr *)&si, &sl);
+    ssize_t r = recvfrom(sock, buf, len, MSG_DONTWAIT, (struct sockaddr *)&si, &sl);
     if (r < 0) {
         return r;
     }
     si_other = si;
     slen = sl;
-    if (r == 8 && memcmp("PINGPING", buf, 8) == 0) {
-        msg_send(0, (const uint8_t *)"PONGPONG", 8);
+    static const char *ping_req = "PINGPING";
+    static const char *ping_resp = "PONGPONG";
+    if (r == strlen(ping_req) && memcmp(ping_req, buf, strlen(ping_req)) == 0) {
+        usb_hid_write(0, (const uint8_t *)ping_resp, strlen(ping_resp));
         return 0;
     }
     return r;
 }
 
-ssize_t msg_send(uint8_t iface, const uint8_t *buf, size_t len)
-{
-    (void)iface; // TODO: ignore interface for now
+int usb_hid_write(uint8_t iface_num, const uint8_t *buf, uint32_t len) {
     ssize_t r = len;
     if (slen > 0) {
-        r = sendto(s, buf, len, MSG_DONTWAIT, (const struct sockaddr *)&si_other, slen);
+        r = sendto(sock, buf, len, MSG_DONTWAIT, (const struct sockaddr *)&si_other, slen);
     }
     return r;
 }
+
+void pendsv_kbd_intr(void) {
+}
+
+void mp_hal_set_vcp_iface(int iface_num) {
+}

micropython/extmod/modtrezormsg/modtrezormsg.c

@@ -12,6 +12,7 @@
 #include "py/runtime.h"
 #include "py/mphal.h"
 #include "py/objstr.h"
+#include "py/smallint.h"
 
 #if MICROPY_PY_TREZORMSG
 
@@ -327,7 +328,6 @@ typedef struct _mp_obj_Msg_t {
 ///     '''
 STATIC mp_obj_t mod_trezormsg_Msg_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
     mp_arg_check_num(n_args, n_kw, 0, 0, false);
-    msg_init();
     mp_obj_Msg_t *o = m_new_obj(mp_obj_Msg_t);
     o->base.type = type;
     o->usb_info = mp_const_none;
@@ -404,7 +404,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_3(mod_trezormsg_Msg_init_usb_obj, mod_trezormsg_M
 
 /// def deinit_usb(self) -> None:
 ///     '''
-///     Cleans up the USB stack
+///     Cleans up the USB stack.
 ///     '''
 STATIC mp_obj_t mod_trezormsg_Msg_deinit_usb(mp_obj_t self) {
 
@@ -428,55 +428,11 @@ STATIC mp_obj_t mod_trezormsg_Msg_send(mp_obj_t self, mp_obj_t iface, mp_obj_t m
     uint8_t i = mp_obj_get_int(iface);
     mp_buffer_info_t msg;
     mp_get_buffer_raise(message, &msg, MP_BUFFER_READ);
-    ssize_t r = msg_send(i, msg.buf, msg.len);
+    ssize_t r = usb_hid_write(i, msg.buf, msg.len);
     return MP_OBJ_NEW_SMALL_INT(r);
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_3(mod_trezormsg_Msg_send_obj, mod_trezormsg_Msg_send);
 
-#define TICK_RESOLUTION 1000
-#define TOUCH_IFACE 255
-extern uint32_t touch_read(void); // defined in HAL
-
-/// def select(self, timeout_us: int) -> tuple:
-///     '''
-///     Polls the event queue and returns the event object.
-///     Function returns None if timeout specified in microseconds is reached.
-///     '''
-STATIC mp_obj_t mod_trezormsg_Msg_select(mp_obj_t self, mp_obj_t timeout_us) {
-    // mp_obj_Msg_t *o = MP_OBJ_TO_PTR(self);
-    int timeout = mp_obj_get_int(timeout_us);
-    if (timeout < 0) {
-        timeout = 0;
-    }
-    for (;;) {
-        uint32_t e = touch_read();
-        if (e) {
-            mp_obj_tuple_t *tuple = MP_OBJ_TO_PTR(mp_obj_new_tuple(4, NULL));
-            tuple->items[0] = MP_OBJ_NEW_SMALL_INT(TOUCH_IFACE);
-            tuple->items[1] = MP_OBJ_NEW_SMALL_INT((e & 0xFF0000) >> 16); // event type
-            tuple->items[2] = MP_OBJ_NEW_SMALL_INT((e & 0xFF00) >> 8); // x position
-            tuple->items[3] = MP_OBJ_NEW_SMALL_INT((e & 0xFF)); // y position
-            return MP_OBJ_FROM_PTR(tuple);
-        }
-        uint8_t iface;
-        uint8_t recvbuf[64];
-        ssize_t l = msg_recv(&iface, recvbuf, 64);
-        if (l > 0) {
-            mp_obj_tuple_t *tuple = MP_OBJ_TO_PTR(mp_obj_new_tuple(2, NULL));
-            tuple->items[0] = MP_OBJ_NEW_SMALL_INT(iface);
-            tuple->items[1] = mp_obj_new_str_of_type(&mp_type_bytes, recvbuf, l);
-            return MP_OBJ_FROM_PTR(tuple);
-        }
-        if (timeout <= 0) {
-            break;
-        }
-        mp_hal_delay_us(TICK_RESOLUTION);
-        timeout -= TICK_RESOLUTION;
-    }
-    return mp_const_none;
-}
-STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_trezormsg_Msg_select_obj, mod_trezormsg_Msg_select);
-
 STATIC mp_obj_t mod_trezormsg_Msg___del__(mp_obj_t self) {
     mp_obj_Msg_t *o = MP_OBJ_TO_PTR(self);
     if (o->usb_info != mp_const_none || o->usb_ifaces != mp_const_none) {
@@ -492,7 +448,6 @@ STATIC const mp_rom_map_elem_t mod_trezormsg_Msg_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_init_usb), MP_ROM_PTR(&mod_trezormsg_Msg_init_usb_obj) },
     { MP_ROM_QSTR(MP_QSTR_deinit_usb), MP_ROM_PTR(&mod_trezormsg_Msg_deinit_usb_obj) },
     { MP_ROM_QSTR(MP_QSTR_send), MP_ROM_PTR(&mod_trezormsg_Msg_send_obj) },
-    { MP_ROM_QSTR(MP_QSTR_select), MP_ROM_PTR(&mod_trezormsg_Msg_select_obj) },
 };
 STATIC MP_DEFINE_CONST_DICT(mod_trezormsg_Msg_locals_dict, mod_trezormsg_Msg_locals_dict_table);
 
@@ -503,12 +458,94 @@ STATIC const mp_obj_type_t mod_trezormsg_Msg_type = {
     .locals_dict = (void*)&mod_trezormsg_Msg_locals_dict,
 };
 
+STATIC mp_obj_t read_bytes_from_hid(uint8_t iface) {
+    uint8_t buf[64];
+    int l = usb_hid_read(iface, buf, sizeof(buf));
+    if (l > 0) {
+        return mp_obj_new_str_of_type(&mp_type_bytes, buf, l);
+    } else {
+        return mp_const_empty_bytes;
+    }
+}
+
+#define TOUCH_IFACE (255)
+#define POLL_READ  (0x0000)
+#define POLL_WRITE (0x0100)
+
+/// def poll(ifaces: Iterable[int], list_ref: List, timeout_us: int) -> bool:
+///     '''
+///     '''
+STATIC mp_obj_t mod_trezormsg_poll(mp_obj_t ifaces, mp_obj_t list_ref, mp_obj_t timeout_us) {
+    mp_obj_list_t *ret = MP_OBJ_TO_PTR(list_ref);
+    if (!MP_OBJ_IS_TYPE(list_ref, &mp_type_list) || ret->len < 2) {
+        mp_raise_TypeError("invalid list_ref");
+    }
+
+    const mp_uint_t timeout = mp_obj_get_int(timeout_us);
+    const mp_uint_t deadline = mp_hal_ticks_us() + timeout;
+    mp_obj_iter_buf_t iterbuf;
+
+    for (;;) {
+        mp_obj_t iter = mp_getiter(ifaces, &iterbuf);
+        mp_obj_t item;
+        while ((item = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
+            const mp_uint_t i = mp_obj_int_get_truncated(item);
+            const mp_uint_t iface = i & 0x00FF;
+            const mp_uint_t mode = i & 0xFF00;
+
+            if (iface == TOUCH_IFACE) {
+                uint32_t evt = touch_read();
+                if (evt) {
+                    mp_obj_tuple_t *tuple = MP_OBJ_TO_PTR(mp_obj_new_tuple(3, NULL));
+                    tuple->items[0] = MP_OBJ_NEW_SMALL_INT((evt & 0xFF0000) >> 16); // event type
+                    tuple->items[1] = MP_OBJ_NEW_SMALL_INT((evt & 0xFF00) >> 8); // x position
+                    tuple->items[2] = MP_OBJ_NEW_SMALL_INT((evt & 0xFF)); // y position
+                    ret->items[0] = MP_OBJ_NEW_SMALL_INT(i);
+                    ret->items[1] = MP_OBJ_FROM_PTR(tuple);
+                    return mp_const_true;
+                }
+            } else
+            if (mode == POLL_READ) {
+                if (usb_hid_can_read(iface)) {
+                    ret->items[0] = MP_OBJ_NEW_SMALL_INT(i);
+                    ret->items[1] = read_bytes_from_hid(iface);
+                    return mp_const_true;
+                }
+            } else
+            if (mode == POLL_WRITE) {
+                if (usb_hid_can_write(iface)) {
+                    ret->items[0] = MP_OBJ_NEW_SMALL_INT(i);
+                    ret->items[1] = mp_const_none;
+                    return mp_const_true;
+                }
+            }
+        }
+
+        if (mp_hal_ticks_us() >= deadline) {
+            break;
+        } else {
+            MICROPY_EVENT_POLL_HOOK
+        }
+    }
+
+    return mp_const_false;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_3(mod_trezormsg_poll_obj, mod_trezormsg_poll);
+
 STATIC const mp_rom_map_elem_t mp_module_trezormsg_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_trezormsg) },
     { MP_ROM_QSTR(MP_QSTR_USB), MP_ROM_PTR(&mod_trezormsg_USB_type) },
     { MP_ROM_QSTR(MP_QSTR_HID), MP_ROM_PTR(&mod_trezormsg_HID_type) },
     { MP_ROM_QSTR(MP_QSTR_VCP), MP_ROM_PTR(&mod_trezormsg_VCP_type) },
     { MP_ROM_QSTR(MP_QSTR_Msg), MP_ROM_PTR(&mod_trezormsg_Msg_type) },
+
+    { MP_ROM_QSTR(MP_QSTR_poll), MP_ROM_PTR(&mod_trezormsg_poll_obj) },
+    { MP_ROM_QSTR(MP_QSTR_TOUCH), MP_OBJ_NEW_SMALL_INT(TOUCH_IFACE) },
+    { MP_ROM_QSTR(MP_QSTR_TOUCH_START), MP_OBJ_NEW_SMALL_INT((TOUCH_START & 0xFF0000) >> 16) },
+    { MP_ROM_QSTR(MP_QSTR_TOUCH_MOVE), MP_OBJ_NEW_SMALL_INT((TOUCH_MOVE & 0xFF0000) >> 16) },
+    { MP_ROM_QSTR(MP_QSTR_TOUCH_END), MP_OBJ_NEW_SMALL_INT((TOUCH_END & 0xFF0000) >> 16) },
+    { MP_ROM_QSTR(MP_QSTR_POLL_READ), MP_OBJ_NEW_SMALL_INT(POLL_READ) },
+    { MP_ROM_QSTR(MP_QSTR_POLL_WRITE), MP_OBJ_NEW_SMALL_INT(POLL_WRITE) },
 };
 STATIC MP_DEFINE_CONST_DICT(mp_module_trezormsg_globals, mp_module_trezormsg_globals_table);
 

micropython/extmod/modtrezormsg/unix-usb-mock.h

@@ -1,31 +0,0 @@
-#include "../../trezorhal/usb.h"
-
-int usb_init(const usb_dev_info_t *dev_info) {
-    return 0;
-}
-
-int usb_deinit(void) {
-    return 0;
-}
-
-int usb_start(void) {
-    return 0;
-}
-
-int usb_stop(void) {
-    return 0;
-}
-
-int usb_hid_add(const usb_hid_info_t *info) {
-    return 0;
-}
-
-int usb_vcp_add(const usb_vcp_info_t *info) {
-    return 0;
-}
-
-void pendsv_kbd_intr(void) {
-}
-
-void mp_hal_set_vcp_iface(int iface_num) {
-}

micropython/unix/mpconfigport.h

@@ -288,6 +288,8 @@ typedef unsigned int mp_uint_t; // must be pointer size
 #endif
 #endif
 
+#define MICROPY_EVENT_POLL_HOOK  mp_hal_delay_ms(1);
+
 // Cannot include <sys/types.h>, as it may lead to symbol name clashes
 #if _FILE_OFFSET_BITS == 64 && !defined(__LP64__)
 typedef long long mp_off_t;

src/trezor/loop.py

@@ -14,10 +14,15 @@ from micropython import const
 from trezor import msg
 from trezor import log
 
-TOUCH = const(255)  # interface
-TOUCH_START = const(1)  # event
-TOUCH_MOVE = const(2)  # event
-TOUCH_END = const(4)  # event
+import trezormsg
+
+TOUCH = trezormsg.TOUCH
+TOUCH_START = trezormsg.TOUCH_START
+TOUCH_MOVE = trezormsg.TOUCH_MOVE
+TOUCH_END = trezormsg.TOUCH_END
+
+READ = trezormsg.POLL_READ
+WRITE = trezormsg.POLL_WRITE
 
 after_step_hook = None  # function, called after each task step
 
@@ -84,6 +89,7 @@ def run_forever():
         global log_delay_pos
 
     task_entry = [0, 0, 0]  # deadline, task, value
+    msg_entry = [0, 0]  # iface | flags, value
     while True:
         # compute the maximum amount of time we can wait for a message
         if _scheduled_tasks:
@@ -97,13 +103,11 @@ def run_forever():
             log_delay_rb[log_delay_pos] = delay
             log_delay_pos = (log_delay_pos + 1) % log_delay_rb_len
 
-        msg_entry = msg.select(delay)
-        if msg_entry is not None:
+        if trezormsg.poll(_paused_tasks, msg_entry, delay):
             # message received, run tasks paused on the interface
-            msg_iface, *msg_value = msg_entry
-            msg_tasks = _paused_tasks.pop(msg_iface, ())
+            msg_tasks = _paused_tasks.pop(msg_entry[0], ())
             for task in msg_tasks:
-                _step_task(task, msg_value)
+                _step_task(task, msg_entry[1])
         else:
             # timeout occurred, run the first scheduled task
             if _scheduled_tasks:

src/trezor/wire/__init__.py

@@ -160,13 +160,13 @@ def _write_report(report):
 
 
 def _dispatch_reports():
+    read = loop.select(_interface)
     while True:
-        report, = yield loop.Select(_interface)
-        report = memoryview(report)
+        report = yield read
         # if __debug__:
         #     log.debug(__name__, 'read report %s', ubinascii.hexlify(report))
         sessions.dispatch(
-            report, _session_open, _session_close, _session_unknown)
+            memoryview(report), _session_open, _session_close, _session_unknown)
 
 
 def _session_open(session_id=None):