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

4 files changed, 14 insertions, 7 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();
diff --git a/shared-bindings/ssl/SSLContext.c b/shared-bindings/ssl/SSLContext.c
index 9d4df7261..44e9e6bbf 100644
--- a/shared-bindings/ssl/SSLContext.c
+++ b/shared-bindings/ssl/SSLContext.c
@@ -51,7 +51,7 @@ STATIC mp_obj_t ssl_sslcontext_make_new(const mp_obj_type_t *type, size_t n_args
     return MP_OBJ_FROM_PTR(s);
 }
 
-//| def wrap_socket(sock: socketpool.Socket, *, server_side: bool = False, server_hostname: Optional[str] = None) -> socketpool.Socket:
+//| def wrap_socket(sock: socketpool.Socket, *, server_side: bool = False, server_hostname: Optional[str] = None) -> ssl.SSLSocket:
 //|     """Wraps the socket into a socket-compatible class that handles SSL negotiation.
 //|        The socket must be of type SOCK_STREAM."""
 //|     ...
diff --git a/shared-bindings/ssl/SSLSocket.c b/shared-bindings/ssl/SSLSocket.c
index 154d3d1d4..cd2daeb3e 100644
--- a/shared-bindings/ssl/SSLSocket.c
+++ b/shared-bindings/ssl/SSLSocket.c
@@ -45,7 +45,7 @@
 //|        recv that do not allocate bytes objects."""
 //|
 
-//|     def __enter__(self) -> Socket:
+//|     def __enter__(self) -> SSLSocket:
 //|         """No-op used by Context Managers."""
 //|         ...
 //|
@@ -63,7 +63,7 @@ STATIC mp_obj_t ssl_sslsocket___exit__(size_t n_args, const mp_obj_t *args) {
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(ssl_sslsocket___exit___obj, 4, 4, ssl_sslsocket___exit__);
 
-//|     def accept(self) -> Tuple[Socket, Tuple[str, int]]:
+//|     def accept(self) -> Tuple[SSLSocket, Tuple[str, int]]:
 //|         """Accept a connection on a listening socket of type SOCK_STREAM,
 //|         creating a new socket of type SOCK_STREAM.
 //|         Returns a tuple of (new_socket, remote_address)"""