Fix stubs error, out of sockets error, invalid TLS leak

2 files changed, 11 insertions, 4 deletions

diff --git a/ports/esp32s2/common-hal/socketpool/Socket.c b/ports/esp32s2/common-hal/socketpool/Socket.c
index 4cbf4cff2..0022c49c6 100644
--- a/ports/esp32s2/common-hal/socketpool/Socket.c
+++ b/ports/esp32s2/common-hal/socketpool/Socket.c
@@ -44,8 +44,8 @@ void socket_reset(void) {
     for (size_t i = 0; i < MP_ARRAY_SIZE(open_socket_handles); i++) {
         if (open_socket_handles[i]) {
             if (open_socket_handles[i]->num > 0) {
+                // Close automatically clears socket handle
                 common_hal_socketpool_socket_close(open_socket_handles[i]);
-                open_socket_handles[i] = NULL;
             } else {
                 open_socket_handles[i] = NULL;
             }
@@ -136,6 +136,12 @@ void common_hal_socketpool_socket_close(socketpool_socket_obj_t* self) {
         lwip_close(self->num);
         self->num = -1;
     }
+    // Remove socket record
+    for (size_t i = 0; i < MP_ARRAY_SIZE(open_socket_handles); i++) {
+        if (open_socket_handles[i] == self) {
+            open_socket_handles[i] = NULL;
+        }
+    }
 }
 
 bool common_hal_socketpool_socket_connect(socketpool_socket_obj_t* self,
diff --git a/ports/esp32s2/common-hal/ssl/SSLContext.c b/ports/esp32s2/common-hal/ssl/SSLContext.c
index c0179399d..6b05905aa 100644
--- a/ports/esp32s2/common-hal/ssl/SSLContext.c
+++ b/ports/esp32s2/common-hal/ssl/SSLContext.c
@@ -38,14 +38,15 @@ void common_hal_ssl_sslcontext_construct(ssl_sslcontext_obj_t* self) {
 ssl_sslsocket_obj_t* common_hal_ssl_sslcontext_wrap_socket(ssl_sslcontext_obj_t* self,
     socketpool_socket_obj_t* socket, bool server_side, const char* server_hostname) {
 
+    if (socket->type != SOCK_STREAM || socket->num != -1) {
+        mp_raise_RuntimeError(translate("Invalid socket for TLS"));
+    }
+
     ssl_sslsocket_obj_t *sock = m_new_obj_with_finaliser(ssl_sslsocket_obj_t);
     sock->base.type = &ssl_sslsocket_type;
     sock->ssl_context = self;
     sock->sock = socket;
 
-    if (socket->type != SOCK_STREAM || socket->num != -1) {
-        mp_raise_RuntimeError(translate("Invalid socket for TLS"));
-    }
     esp_tls_t* tls_handle = esp_tls_init();
     if (tls_handle == NULL) {
         mp_raise_espidf_MemoryError();