From bf02fe0ddfd84f616bb9b89cf50a71ba93c856f2 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <stick@gk2.sk>
Date: Mon, 23 May 2016 17:53:42 +0200
Subject: [PATCH] trezor.msg is now iface aware

---
 extmod/modtrezormsg/modtrezormsg-stmhal.h |  4 ++--
 extmod/modtrezormsg/modtrezormsg-unix.h   | 21 ++++++++++----------
 extmod/modtrezormsg/modtrezormsg.c        | 24 +++++++++++++----------
 src/trezor/msg.py                         |  8 ++++----
 4 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/extmod/modtrezormsg/modtrezormsg-stmhal.h b/extmod/modtrezormsg/modtrezormsg-stmhal.h
index 67ab6b47..60c9d7bf 100644
--- a/extmod/modtrezormsg/modtrezormsg-stmhal.h
+++ b/extmod/modtrezormsg/modtrezormsg-stmhal.h
@@ -4,12 +4,12 @@ void msg_init(void)
 {
 }
 
-const uint8_t *msg_recv(void)
+ssize_t msg_recv(uint8_t *iface, uint8_t *buf, size_t len)
 {
     return 0;
 }
 
-int msg_send(uint8_t *buf, size_t len)
+ssize_t msg_send(uint8_t iface, const uint8_t *buf, size_t len)
 {
     return -1;
 }
diff --git a/extmod/modtrezormsg/modtrezormsg-unix.h b/extmod/modtrezormsg/modtrezormsg-unix.h
index 95faecd3..2cc95133 100644
--- a/extmod/modtrezormsg/modtrezormsg-unix.h
+++ b/extmod/modtrezormsg/modtrezormsg-unix.h
@@ -36,26 +36,25 @@ void msg_init(void)
     assert(b != -1);
 }
 
-#define RECV_BUFLEN 64
-
-const uint8_t *msg_recv(void)
+ssize_t msg_recv(uint8_t *iface, uint8_t *buf, size_t len)
 {
-    static uint8_t buf[RECV_BUFLEN];
     struct sockaddr_in si;
     socklen_t sl = sizeof(si);
-    memset(buf, 0, sizeof(buf));
-    int len = recvfrom(s, buf, RECV_BUFLEN, MSG_DONTWAIT, (struct sockaddr *)&si, &sl);
-    if (len < 0) {
-        return 0;
+    memset(buf, 0, len);
+    iface = 0; // UDP uses always interface 0
+    size_t r = recvfrom(s, buf, len, MSG_DONTWAIT, (struct sockaddr *)&si, &sl);
+    if (r < 0) {
+        return r;
     }
     si_other = si;
     slen = sl;
-    return buf;
+    return r;
 }
 
-int msg_send(uint8_t *buf, size_t len)
+ssize_t msg_send(uint8_t iface, const uint8_t *buf, size_t len)
 {
-    int r = -1;
+    (void)iface; // ignore interface for UDP
+    ssize_t r = -1;
     if (slen) {
         r = sendto(s, buf, len, MSG_DONTWAIT, (const struct sockaddr *)&si_other, slen);
     }
diff --git a/extmod/modtrezormsg/modtrezormsg.c b/extmod/modtrezormsg/modtrezormsg.c
index 24b52137..55fea793 100644
--- a/extmod/modtrezormsg/modtrezormsg.c
+++ b/extmod/modtrezormsg/modtrezormsg.c
@@ -36,17 +36,18 @@ STATIC mp_obj_t mod_TrezorMsg_Msg_make_new(const mp_obj_type_t *type, size_t n_a
     return MP_OBJ_FROM_PTR(o);
 }
 
-/// def trezor.msg.send(message) -> int
+/// def trezor.msg.send(iface: int, message: bytes) -> int
 ///
 /// Sends message using USB HID (device) or UDP (emulator).
 ///
-STATIC mp_obj_t mod_TrezorMsg_Msg_send(mp_obj_t self, mp_obj_t message) {
+STATIC mp_obj_t mod_TrezorMsg_Msg_send(mp_obj_t self, mp_obj_t iface, mp_obj_t message) {
+    uint8_t iface_num = mp_obj_get_int(iface);
     mp_buffer_info_t msg;
     mp_get_buffer_raise(message, &msg, MP_BUFFER_READ);
-    int r = msg_send(msg.buf, msg.len);
+    ssize_t r = msg_send(iface_num, msg.buf, msg.len);
     return MP_OBJ_NEW_SMALL_INT(r);
 }
-STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_TrezorMsg_Msg_send_obj, mod_TrezorMsg_Msg_send);
+STATIC MP_DEFINE_CONST_FUN_OBJ_3(mod_TrezorMsg_Msg_send_obj, mod_TrezorMsg_Msg_send);
 
 #define TICK_RESOLUTION 1000
 
@@ -69,14 +70,17 @@ STATIC mp_obj_t mod_TrezorMsg_Msg_select(mp_obj_t self, mp_obj_t timeout_us) {
             tuple->items[2] = MP_OBJ_NEW_SMALL_INT((e & 0xFF));
             return MP_OBJ_FROM_PTR(tuple);
         }
-        const uint8_t *m = msg_recv();
-        if (m) {
+        uint8_t iface;
+        uint8_t recvbuf[64];
+        ssize_t l = msg_recv(&iface, recvbuf, 64);
+        if (l > 0) {
             vstr_t vstr;
-            vstr_init_len(&vstr, 64);
-            memcpy(vstr.buf, m, 64);
-            mp_obj_tuple_t *tuple = MP_OBJ_TO_PTR(mp_obj_new_tuple(2, NULL));
+            vstr_init_len(&vstr, l);
+            memcpy(vstr.buf, recvbuf, l);
+            mp_obj_tuple_t *tuple = MP_OBJ_TO_PTR(mp_obj_new_tuple(3, NULL));
             tuple->items[0] = MP_OBJ_NEW_SMALL_INT(8);
-            tuple->items[1] = mp_obj_new_str_from_vstr(&mp_type_bytes, &vstr);
+            tuple->items[1] = MP_OBJ_NEW_SMALL_INT(iface);
+            tuple->items[2] = mp_obj_new_str_from_vstr(&mp_type_bytes, &vstr);
             return MP_OBJ_FROM_PTR(tuple);
          }
         if (timeout <= 0) {
diff --git a/src/trezor/msg.py b/src/trezor/msg.py
index 8970554c..8cc8fa5b 100644
--- a/src/trezor/msg.py
+++ b/src/trezor/msg.py
@@ -10,8 +10,8 @@ def select(timeout_us):
     return _msg.select(timeout_us)
 
 
-def send(msg):
-    return _msg.send(msg)
+def send(iface, msg):
+    return _msg.send(iface, msg)
 
 
 REPORT_LEN = const(64)
@@ -20,7 +20,7 @@ HEADER_MAGIC = const(35)  # '#'
 
 
 def read():
-    _, rep = yield loop.Select(loop.HID_READ)
+    _, iface, rep = yield loop.Select(loop.HID_READ)
     assert rep[0] == REPORT_NUM
     return rep
 
@@ -69,7 +69,7 @@ def write_wire_msg(mtype, mbuf):
         while i < len(data):
             data[i] = 0
             i += 1
-        send(rep)
+        send(0, rep)
         mbuf = mbuf[n:]
         data = rep[1:]