From 57cce79a46f6ab7e990a626251ca8f9d59cd095a Mon Sep 17 00:00:00 2001
From: Angus Gratton <angus@redyak.com.au>
Date: Wed, 8 Nov 2023 17:35:28 +1100
Subject: [PATCH] esp32/modsocket: Implement socket connect timeout.

LWIP doesn't implement a timeout for blocking connect(), and such a timeout
is not required by POSIX.  However, CPython will use the socket timeout for
blocking connect on most platforms.  The "principle of least surprise"
suggests we should support it on ESP32 as well (not to mention it's
useful!).

This provides the additional improvement that external exceptions (like
KeyboardInterrupt) are now handled immediately if they happen during
connect().  Previously Ctrl-C would not terminate a blocking connect until
connect() returned, but now it will.

Fixes issue #8326.

This work was funded through GitHub Sponsors.

Signed-off-by: Angus Gratton <angus@redyak.com.au>
---
 ports/esp32/modsocket.c | 94 ++++++++++++++++++++++++++++++++++++++---
 1 file changed, 88 insertions(+), 6 deletions(-)

diff --git a/ports/esp32/modsocket.c b/ports/esp32/modsocket.c
index fe76b5f76..cacc52dec 100644
--- a/ports/esp32/modsocket.c
+++ b/ports/esp32/modsocket.c
@@ -82,6 +82,8 @@ typedef struct _socket_obj_t {
     #endif
 } socket_obj_t;
 
+STATIC const char *TAG = "modsocket";
+
 void _socket_settimeout(socket_obj_t *sock, uint64_t timeout_ms);
 
 #if MICROPY_PY_SOCKET_EVENTS
@@ -364,16 +366,96 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_1(socket_accept_obj, socket_accept);
 STATIC mp_obj_t socket_connect(const mp_obj_t arg0, const mp_obj_t arg1) {
     socket_obj_t *self = MP_OBJ_TO_PTR(arg0);
     struct addrinfo *res;
+    bool blocking = false;
+    int flags;
+    int raise_err = 0;
+
     _socket_getaddrinfo(arg1, &res);
     MP_THREAD_GIL_EXIT();
     self->state = SOCKET_STATE_CONNECTED;
-    int r = lwip_connect(self->fd, res->ai_addr, res->ai_addrlen);
-    MP_THREAD_GIL_ENTER();
-    lwip_freeaddrinfo(res);
-    if (r != 0) {
-        mp_raise_OSError(errno);
+
+    flags = fcntl(self->fd, F_GETFL);
+
+    blocking = (flags & O_NONBLOCK) == 0;
+
+    if (blocking) {
+        // For blocking sockets, make the socket temporarily non-blocking and emulate
+        // blocking using select.
+        //
+        // This has two benefits:
+        //
+        // - Allows handling external exceptions while waiting for connect.
+        //
+        // - Allows emulating a connect timeout, which is not supported by LWIP or
+        //   required by POSIX but is normal behaviour for CPython.
+        if (fcntl(self->fd, F_SETFL, flags | O_NONBLOCK) < 0) {
+            ESP_LOGE(TAG, "fcntl set failed %d", errno); // Unexpected internal failure
+            raise_err = errno;
+        }
     }
 
+    if (raise_err == 0) {
+        // Try performing the actual connect. Expected to always return immediately.
+        int r = lwip_connect(self->fd, res->ai_addr, res->ai_addrlen);
+        if (r != 0) {
+            raise_err = errno;
+        }
+    }
+
+    if (blocking) {
+        // Set the socket back to blocking. We can still pass it to select() in this state.
+        int r = fcntl(self->fd, F_SETFL, flags);
+        if (r != 0 && (raise_err == 0 || raise_err == EINPROGRESS)) {
+            ESP_LOGE(TAG, "fcntl restore failed %d", errno); // Unexpected internal failure
+            raise_err = errno;
+        }
+    }
+
+    lwip_freeaddrinfo(res);
+
+    if (blocking && raise_err == EINPROGRESS) {
+        // Keep calling select() until the socket is marked writable (i.e. connected),
+        // or an error or a timeout occurs
+
+        // Note: _socket_settimeout() always sets self->retries != 0 on blocking sockets.
+
+        for (unsigned int i = 0; i <= self->retries; i++) {
+            struct timeval timeout = {
+                .tv_sec = 0,
+                .tv_usec = SOCKET_POLL_US,
+            };
+            fd_set wfds;
+            FD_ZERO(&wfds);
+            FD_SET(self->fd, &wfds);
+
+            int r = select(self->fd + 1, NULL, &wfds, NULL, &timeout);
+            if (r < 0) {
+                // Error condition
+                raise_err = errno;
+                break;
+            } else if (r > 0) {
+                // Select indicated the socket is writable. Check for any error.
+                socklen_t socklen = sizeof(raise_err);
+                r = getsockopt(self->fd, SOL_SOCKET, SO_ERROR, &raise_err, &socklen);
+                if (r < 0) {
+                    raise_err = errno;
+                }
+                break;
+            } else {
+                // Select timed out
+                raise_err = ETIMEDOUT;
+
+                MP_THREAD_GIL_ENTER();
+                check_for_exceptions();
+                MP_THREAD_GIL_EXIT();
+            }
+        }
+    }
+
+    MP_THREAD_GIL_ENTER();
+    if (raise_err) {
+        mp_raise_OSError(raise_err);
+    }
     return mp_const_none;
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(socket_connect_obj, socket_connect);
@@ -847,7 +929,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(esp_socket_getaddrinfo_obj, 2, 6, esp
 STATIC mp_obj_t esp_socket_initialize() {
     static int initialized = 0;
     if (!initialized) {
-        ESP_LOGI("modsocket", "Initializing");
+        ESP_LOGI(TAG, "Initializing");
         esp_netif_init();
         initialized = 1;
     }