Make socket reads interruptable

tannewt · tannewt · commit 701e80a0252d · 2020-08-21T11:00:02.000-07:00
diff --git a/ports/esp32s2/common-hal/socketpool/Socket.c b/ports/esp32s2/common-hal/socketpool/Socket.c
@@ -55,9 +55,13 @@ bool common_hal_socketpool_socket_connect(socketpool_socket_obj_t* self, const c
         int flags;
         esp_err_t err = esp_tls_get_and_clear_last_error(self->tcp->error_handle, &esp_tls_code, &flags);
 
-        // mp_raise_espidf_MemoryError
-        mp_raise_OSError_msg_varg(translate("Unhandled ESP TLS error %d %d %x"), esp_tls_code, flags, err);
+        if (err == ESP_ERR_MBEDTLS_SSL_SETUP_FAILED) {
+            mp_raise_espidf_MemoryError();
+        } else {
+            mp_raise_OSError_msg_varg(translate("Unhandled ESP TLS error %d %d %x"), esp_tls_code, flags, err);
+        }
     }
+
     return self->connected;
 }
 
@@ -78,23 +82,35 @@ mp_uint_t common_hal_socketpool_socket_recv_into(socketpool_socket_obj_t* self,
     size_t received = 0;
     ssize_t last_read = 1;
     uint64_t start_ticks = supervisor_ticks_ms64();
+    int sockfd;
+    esp_err_t err = esp_tls_get_conn_sockfd(self->tcp, &sockfd);
+    if (err != ESP_OK) {
+        mp_raise_OSError(MP_EBADF);
+    }
     while (received < len &&
            last_read > 0 &&
            (self->timeout_ms == 0 || supervisor_ticks_ms64() - start_ticks <= self->timeout_ms) &&
            !mp_hal_is_interrupted()) {
         RUN_BACKGROUND_TASKS;
         size_t available = esp_tls_get_bytes_avail(self->tcp);
-        ESP_EARLY_LOGW(TAG, "available %d", available);
+        if (available == 0) {
+            // This reads the raw socket buffer and is used for non-TLS connections
+            // and between encrypted TLS blocks.
+            int status = lwip_ioctl(sockfd, FIONREAD, &available);
+            if (status < 0) {
+                // ESP_EARLY_LOGW(TAG, "ioctl fail. socket %d status %d errno %d available %d", sockfd, status, errno, available);
+                last_read = status;
+                break;
+            }
+        }
+        // ESP_EARLY_LOGW(TAG, "available %d", available);
         size_t remaining = len - received;
         if (available > remaining) {
             available = remaining;
         }
-        if (true || available > 0) {
-            if (available == 0) {
-                available = len - received;
-            }
+        if (available > 0) {
             last_read = esp_tls_conn_read(self->tcp, (void*) buf + received, available);
-            ESP_EARLY_LOGW(TAG, "read %d out of %d", last_read, available);
+            // ESP_EARLY_LOGW(TAG, "read %d out of %d", last_read, available);
             received += last_read;
         }
     }
diff --git a/ports/esp32s2/common-hal/wifi/Radio.c b/ports/esp32s2/common-hal/wifi/Radio.c
@@ -105,7 +105,7 @@ void common_hal_wifi_radio_stop_scanning_networks(wifi_radio_obj_t *self) {
     self->current_scan = NULL;
 }
 
-bool common_hal_wifi_radio_connect(wifi_radio_obj_t *self, uint8_t* ssid, size_t ssid_len, uint8_t* password, size_t password_len, uint8_t channel, mp_float_t timeout) {
+wifi_radio_error_t common_hal_wifi_radio_connect(wifi_radio_obj_t *self, uint8_t* ssid, size_t ssid_len, uint8_t* password, size_t password_len, uint8_t channel, mp_float_t timeout) {
     // check enabled
     wifi_config_t* config = &self->sta_config;
     memcpy(&config->sta.ssid, ssid, ssid_len);
@@ -116,6 +116,8 @@ bool common_hal_wifi_radio_connect(wifi_radio_obj_t *self, uint8_t* ssid, size_t
     config->sta.password[password_len] = 0;
     config->sta.channel = channel;
     esp_wifi_set_config(ESP_IF_WIFI_STA, config);
+    self->starting_retries = 5;
+    self->retries_left = 5;
     esp_wifi_connect();
 
     EventBits_t bits;
@@ -128,9 +130,14 @@ bool common_hal_wifi_radio_connect(wifi_radio_obj_t *self, uint8_t* ssid, size_t
             0);
     } while ((bits & (WIFI_CONNECTED_BIT | WIFI_DISCONNECTED_BIT)) == 0 && !mp_hal_is_interrupted());
     if ((bits & WIFI_DISCONNECTED_BIT) != 0) {
-        return false;
+        if (self->last_disconnect_reason == WIFI_REASON_AUTH_FAIL) {
+            return WIFI_RADIO_ERROR_AUTH;
+        } else if (self->last_disconnect_reason == WIFI_REASON_NO_AP_FOUND) {
+            return WIFI_RADIO_ERROR_NO_AP_FOUND;
+        }
+        return WIFI_RADIO_ERROR_UNKNOWN;
     }
-    return true;
+    return WIFI_RADIO_ERROR_NONE;
 }
 
 mp_obj_t common_hal_wifi_radio_get_ipv4_address(wifi_radio_obj_t *self) {
diff --git a/ports/esp32s2/common-hal/wifi/Radio.h b/ports/esp32s2/common-hal/wifi/Radio.h
@@ -50,6 +50,9 @@ typedef struct {
     bool started;
     bool ap_mode;
     bool sta_mode;
+    uint8_t retries_left;
+    uint8_t starting_retries;
+    uint8_t last_disconnect_reason;
 } wifi_radio_obj_t;
 
 #endif // MICROPY_INCLUDED_ESP32S2_COMMON_HAL_WIFI_RADIO_H
diff --git a/ports/esp32s2/common-hal/wifi/__init__.c b/ports/esp32s2/common-hal/wifi/__init__.c
@@ -54,11 +54,24 @@ static void event_handler(void* arg, esp_event_base_t event_base,
             ESP_EARLY_LOGW(TAG, "disconnected");
             wifi_event_sta_disconnected_t* d = (wifi_event_sta_disconnected_t*) event_data;
             uint8_t reason = d->reason;
-            if (reason != WIFI_REASON_ASSOC_LEAVE) {
-                // reconnect
-            }
             ESP_EARLY_LOGW(TAG, "reason %d 0x%02x", reason, reason);
+            if (radio->retries_left > 0 &&
+                    (reason == WIFI_REASON_AUTH_EXPIRE ||
+                     reason == WIFI_REASON_ASSOC_EXPIRE ||
+                     reason == WIFI_REASON_CONNECTION_FAIL ||
+                     reason == WIFI_REASON_4WAY_HANDSHAKE_TIMEOUT)) {
+                radio->retries_left--;
+                ESP_EARLY_LOGI(TAG, "Retrying connect. %d retries remaining", radio->retries_left);
+                esp_wifi_connect();
+                return;
+            }
+
+            radio->last_disconnect_reason = reason;
             xEventGroupSetBits(radio->event_group_handle, WIFI_DISCONNECTED_BIT);
+
+            // if (reason != WIFI_REASON_ASSOC_LEAVE) {
+            //     // reconnect
+            // }
         } else if (event_id == WIFI_EVENT_STA_AUTHMODE_CHANGE) {
         }
     }
@@ -78,6 +91,7 @@ static void event_handler(void* arg, esp_event_base_t event_base,
     // } else
     if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
         ESP_EARLY_LOGW(TAG, "got ip");
+        radio->retries_left = radio->starting_retries;
         xEventGroupSetBits(radio->event_group_handle, WIFI_CONNECTED_BIT);
     }
 }
diff --git a/ports/esp32s2/sdkconfig.defaults b/ports/esp32s2/sdkconfig.defaults
@@ -452,7 +452,7 @@ CONFIG_LWIP_MAX_SOCKETS=10
 # CONFIG_LWIP_SO_LINGER is not set
 CONFIG_LWIP_SO_REUSE=y
 CONFIG_LWIP_SO_REUSE_RXTOALL=y
-# CONFIG_LWIP_SO_RCVBUF is not set
+CONFIG_LWIP_SO_RCVBUF=y
 # CONFIG_LWIP_NETBUF_RECVINFO is not set
 CONFIG_LWIP_IP4_FRAG=y
 CONFIG_LWIP_IP6_FRAG=y
diff --git a/py/runtime.c b/py/runtime.c
@@ -1602,6 +1602,10 @@ NORETURN void mp_raise_OSError_msg_varg(const compressed_string_t *fmt, ...) {
     va_end(argptr);
 }
 
+NORETURN void mp_raise_ConnectionError(const compressed_string_t *msg) {
+    mp_raise_msg(&mp_type_ConnectionError, msg);
+}
+
 NORETURN void mp_raise_BrokenPipeError(void) {
     nlr_raise(mp_obj_new_exception_arg1(&mp_type_BrokenPipeError, MP_OBJ_NEW_SMALL_INT(MP_EPIPE)));
 }
diff --git a/py/runtime.h b/py/runtime.h
@@ -166,6 +166,7 @@ NORETURN void mp_raise_OSError(int errno_);
 NORETURN void mp_raise_OSError_errno_str(int errno_, mp_obj_t str);
 NORETURN void mp_raise_OSError_msg(const compressed_string_t *msg);
 NORETURN void mp_raise_OSError_msg_varg(const compressed_string_t *fmt, ...);
+NORETURN void mp_raise_ConnectionError(const compressed_string_t *msg);
 NORETURN void mp_raise_BrokenPipeError(void);
 NORETURN void mp_raise_NotImplementedError(const compressed_string_t *msg);
 NORETURN void mp_raise_NotImplementedError_varg(const compressed_string_t *fmt, ...);
diff --git a/shared-bindings/wifi/Radio.c b/shared-bindings/wifi/Radio.c
@@ -136,9 +136,21 @@ STATIC mp_obj_t wifi_radio_connect(size_t n_args, const mp_obj_t *pos_args, mp_m
     password.len = 0;
     if (args[ARG_password].u_obj != MP_OBJ_NULL) {
         mp_get_buffer_raise(args[ARG_password].u_obj, &password, MP_BUFFER_READ);
+        if (password.len > 0 && (password.len < 8 || password.len > 63)) {
+            mp_raise_ValueError(translate("WiFi password must be between 8 and 63 characters."));
+        }
     }
 
-    return mp_obj_new_bool(common_hal_wifi_radio_connect(self, ssid.buf, ssid.len, password.buf, password.len, args[ARG_channel].u_int, timeout));
+    wifi_radio_error_t error = common_hal_wifi_radio_connect(self, ssid.buf, ssid.len, password.buf, password.len, args[ARG_channel].u_int, timeout);
+    if (error == WIFI_RADIO_ERROR_AUTH) {
+        mp_raise_ConnectionError(translate("Authentication failure"));
+    } else if (error == WIFI_RADIO_ERROR_NO_AP_FOUND) {
+        mp_raise_ConnectionError(translate("No network with that ssid"));
+    } else if (error != WIFI_RADIO_ERROR_NONE) {
+        mp_raise_ConnectionError(translate("Unknown failure"));
+    }
+
+    return mp_const_none;
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_KW(wifi_radio_connect_obj, 1, wifi_radio_connect);
 
diff --git a/shared-bindings/wifi/Radio.h b/shared-bindings/wifi/Radio.h
@@ -35,6 +35,14 @@
 
 const mp_obj_type_t wifi_radio_type;
 
+
+typedef enum {
+    WIFI_RADIO_ERROR_NONE,
+    WIFI_RADIO_ERROR_UNKNOWN,
+    WIFI_RADIO_ERROR_AUTH,
+    WIFI_RADIO_ERROR_NO_AP_FOUND
+} wifi_radio_error_t;
+
 extern bool common_hal_wifi_radio_get_enabled(wifi_radio_obj_t *self);
 extern void common_hal_wifi_radio_set_enabled(wifi_radio_obj_t *self, bool enabled);
 
@@ -43,7 +51,7 @@ extern mp_obj_t common_hal_wifi_radio_get_mac_address(wifi_radio_obj_t *self);
 extern mp_obj_t common_hal_wifi_radio_start_scanning_networks(wifi_radio_obj_t *self);
 extern void common_hal_wifi_radio_stop_scanning_networks(wifi_radio_obj_t *self);
 
-extern bool common_hal_wifi_radio_connect(wifi_radio_obj_t *self, uint8_t* ssid, size_t ssid_len, uint8_t* password, size_t password_len, uint8_t channel, mp_float_t timeout);
+extern wifi_radio_error_t common_hal_wifi_radio_connect(wifi_radio_obj_t *self, uint8_t* ssid, size_t ssid_len, uint8_t* password, size_t password_len, uint8_t channel, mp_float_t timeout);
 
 extern mp_obj_t common_hal_wifi_radio_get_ipv4_address(wifi_radio_obj_t *self);
 

Original file line number	Diff line number	Diff line change
`@@ -1602,6 +1602,10 @@ NORETURN void mp_raise_OSError_msg_varg(const compressed_string_t *fmt, ...) {`
`1602`	`1602`	`va_end(argptr);`
`1603`	`1603`	`}`
`1604`	`1604`
	`1605`	`+NORETURN void mp_raise_ConnectionError(const compressed_string_t *msg) {`
	`1606`	`+ mp_raise_msg(&mp_type_ConnectionError, msg);`
	`1607`	`+}`
	`1608`	`+`
`1605`	`1609`	`NORETURN void mp_raise_BrokenPipeError(void) {`
`1606`	`1610`	`nlr_raise(mp_obj_new_exception_arg1(&mp_type_BrokenPipeError, MP_OBJ_NEW_SMALL_INT(MP_EPIPE)));`
`1607`	`1611`	`}`