From 8f0778b2092a4af1b7109b997a3c9c14766e31e8 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Tue, 17 Mar 2020 16:43:19 +1100
Subject: [PATCH] extmod/modlwip: Properly handle non-blocking and timeout on
 UDP recv.

Fixes UDP non-blocking recv so it returns EAGAIN instead of ETIMEDOUT.
Timeout waiting for incoming data is also improved by replacing 100ms delay
with poll_sockets(), as is done in other parts of this module.

Fixes issue #5759.
---
 extmod/modlwip.c                     | 23 +++++++++++------------
 tests/extmod/usocket_udp_nonblock.py | 20 ++++++++++++++++++++
 2 files changed, 31 insertions(+), 12 deletions(-)
 create mode 100644 tests/extmod/usocket_udp_nonblock.py

diff --git a/extmod/modlwip.c b/extmod/modlwip.c
index af0d4cbcdf..321d3de8ac 100644
--- a/extmod/modlwip.c
+++ b/extmod/modlwip.c
@@ -598,21 +598,20 @@ STATIC mp_uint_t lwip_raw_udp_send(lwip_socket_obj_t *socket, const byte *buf, m
 STATIC mp_uint_t lwip_raw_udp_receive(lwip_socket_obj_t *socket, byte *buf, mp_uint_t len, byte *ip, mp_uint_t *port, int *_errno) {
 
     if (socket->incoming.pbuf == NULL) {
-        if (socket->timeout != -1) {
-            for (mp_uint_t retries = socket->timeout / 100; retries--;) {
-                mp_hal_delay_ms(100);
-                if (socket->incoming.pbuf != NULL) {
-                    break;
-                }
-            }
-            if (socket->incoming.pbuf == NULL) {
+        if (socket->timeout == 0) {
+            // Non-blocking socket.
+            *_errno = MP_EAGAIN;
+            return -1;
+        }
+
+        // Wait for data to arrive on UDP socket.
+        mp_uint_t start = mp_hal_ticks_ms();
+        while (socket->incoming.pbuf == NULL) {
+            if (socket->timeout != -1 && mp_hal_ticks_ms() - start > socket->timeout) {
                 *_errno = MP_ETIMEDOUT;
                 return -1;
             }
-        } else {
-            while (socket->incoming.pbuf == NULL) {
-                poll_sockets();
-            }
+            poll_sockets();
         }
     }
 
diff --git a/tests/extmod/usocket_udp_nonblock.py b/tests/extmod/usocket_udp_nonblock.py
new file mode 100644
index 0000000000..63197584da
--- /dev/null
+++ b/tests/extmod/usocket_udp_nonblock.py
@@ -0,0 +1,20 @@
+# test non-blocking UDP sockets
+
+try:
+    import usocket as socket, uerrno as errno
+except ImportError:
+    try:
+        import socket, errno
+    except ImportError:
+        print("SKIP")
+        raise SystemExit
+
+
+s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+s.bind(socket.getaddrinfo('127.0.0.1', 8000)[0][-1])
+s.settimeout(0)
+
+try:
+    s.recv(1)
+except OSError as er:
+    print('EAGAIN:', er.args[0] == errno.EAGAIN)