fix(sockets) fix connecting sockets not reporting close when context …

…or the socket it self is closed manually (oven-sh#13249)
luqmanoop · Aug 12, 2024 · dfa3a9a · dfa3a9a
1 parent 4447668
commit dfa3a9a
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 15 deletions.
diff --git a/packages/bun-usockets/src/context.c b/packages/bun-usockets/src/context.c
@@ -59,7 +59,14 @@ void us_listen_socket_close(int ssl, struct us_listen_socket_t *ls) {
 }
 
 void us_socket_context_close(int ssl, struct us_socket_context_t *context) {
-    /* Begin by closing all listen sockets */
+    /* First start closing pending connecting sockets*/
+    struct us_connecting_socket_t *c = context->head_connecting_sockets;
+    while (c) {
+        struct us_connecting_socket_t *nextC = c->next_pending;
+        us_connecting_socket_close(ssl, c);
+        c = nextC;
+    }
+    /* After this by closing all listen sockets */
     struct us_listen_socket_t *ls = context->head_listen_sockets;
     while (ls) {
         struct us_listen_socket_t *nextLS = (struct us_listen_socket_t *) ls->s.next;
@@ -118,6 +125,21 @@ void us_internal_socket_context_unlink_socket(int ssl, struct us_socket_context_
     }
     us_socket_context_unref(ssl, context);
 }
+void us_internal_socket_context_unlink_connecting_socket(int ssl, struct us_socket_context_t *context, struct us_connecting_socket_t *c) {
+    if (c->prev_pending == c->next_pending) {
+        context->head_connecting_sockets = 0;
+    } else {
+        if (c->prev_pending) {
+            c->prev_pending->next_pending = c->next_pending;
+        } else {
+            context->head_connecting_sockets = c->next_pending;
+        }
+        if (c->next_pending) {
+            c->next_pending->prev_pending = c->prev_pending;
+        }
+    }
+    us_socket_context_unref(ssl, context);
+}
 
 /* We always add in the top, so we don't modify any s.next */
 void us_internal_socket_context_link_listen_socket(struct us_socket_context_t *context, struct us_listen_socket_t *ls) {
@@ -131,6 +153,18 @@ void us_internal_socket_context_link_listen_socket(struct us_socket_context_t *c
     us_socket_context_ref(0, context);
 }
 
+void us_internal_socket_context_link_connecting_socket(int ssl, struct us_socket_context_t *context, struct us_connecting_socket_t *c) {
+    c->context = context;
+    c->next_pending = context->head_connecting_sockets;
+    c->prev_pending = 0;
+    if (context->head_connecting_sockets) {
+        context->head_connecting_sockets->prev_pending = c;
+    }
+    context->head_connecting_sockets = c;
+    us_socket_context_ref(ssl, context);
+}
+
+
 /* We always add in the top, so we don't modify any s.next */
 void us_internal_socket_context_link_socket(struct us_socket_context_t *context, struct us_socket_t *s) {
     s->context = context;
@@ -479,15 +513,14 @@ void *us_socket_context_connect(int ssl, struct us_socket_context_t *context, co
     }
 
     struct us_connecting_socket_t *c = us_calloc(1, sizeof(struct us_connecting_socket_t) + socket_ext_size);
-    c->socket_ext_size = socket_ext_size;
-    c->context = context;
-    us_socket_context_ref(ssl, context);     
+    c->socket_ext_size = socket_ext_size;  
     c->options = options;
     c->ssl = ssl > 0;
     c->timeout = 255;
     c->long_timeout = 255;
     c->pending_resolve_callback = 1;
     c->port = port;
+    us_internal_socket_context_link_connecting_socket(ssl, context, c);
 
 #ifdef _WIN32
     loop->uv_loop->active_handles++;
@@ -554,9 +587,6 @@ void us_internal_socket_after_resolve(struct us_connecting_socket_t *c) {
     }
     struct addrinfo_result *result = Bun__addrinfo_getRequestResult(c->addrinfo_req);
     if (result->error) {
-        c->error = result->error;
-        c->context->on_connect_error(c, result->error);
-        Bun__addrinfo_freeRequest(c->addrinfo_req, 0);
         us_connecting_socket_close(c->ssl, c);
         return;
     }
@@ -565,9 +595,6 @@ void us_internal_socket_after_resolve(struct us_connecting_socket_t *c) {
 
     int opened = start_connections(c, CONCURRENT_CONNECTIONS);
     if (opened == 0) {
-        c->error = ECONNREFUSED;
-        c->context->on_connect_error(c, ECONNREFUSED);
-        Bun__addrinfo_freeRequest(c->addrinfo_req, 1);
         us_connecting_socket_close(c->ssl, c);
         return;
     }
@@ -636,9 +663,6 @@ void us_internal_socket_after_open(struct us_socket_t *s, int error) {
                 // we have run out of addresses to attempt, signal the connection error
                 // but only if there are no other sockets in the list
                 if (opened == 0 && c->connecting_head == NULL) {
-                    c->error = ECONNREFUSED;
-                    c->context->on_connect_error(c, error);
-                    Bun__addrinfo_freeRequest(c->addrinfo_req, ECONNREFUSED);
                     us_connecting_socket_close(c->ssl, c);
                 }
             }
@@ -746,7 +770,10 @@ struct us_socket_t *us_socket_context_adopt_socket(int ssl, struct us_socket_con
         new_s = (struct us_socket_t *) us_poll_resize(&s->p, s->context->loop, sizeof(struct us_socket_t) + ext_size);
         if (c) {
             c->connecting_head = new_s;
+            struct us_socket_context_t *old_context = s->context;
             c->context = context;
+            us_internal_socket_context_link_connecting_socket(ssl, context, c);
+            us_internal_socket_context_unlink_connecting_socket(ssl, old_context, c);
         }
     }
     new_s->timeout = 255;

diff --git a/packages/bun-usockets/src/internal/internal.h b/packages/bun-usockets/src/internal/internal.h
@@ -176,6 +176,7 @@ struct us_socket_t {
 struct us_connecting_socket_t {
     alignas(LIBUS_EXT_ALIGNMENT) struct addrinfo_request *addrinfo_req;
     struct us_socket_context_t *context;
+    // this is used to track all dns resolutions in this connection
     struct us_connecting_socket_t *next;
     struct us_socket_t *connecting_head;
     int options;
@@ -186,6 +187,9 @@ struct us_connecting_socket_t {
     uint16_t port;
     int error;
     struct addrinfo *addrinfo_head;
+    // this is used to track pending connecting sockets in the context
+    struct us_connecting_socket_t* next_pending;
+    struct us_connecting_socket_t* prev_pending;
 };
 
 struct us_wrapped_socket_context_t {
@@ -264,6 +268,7 @@ struct us_socket_context_t {
   unsigned char long_timestamp;
   struct us_socket_t *head_sockets;
   struct us_listen_socket_t *head_listen_sockets;
+  struct us_connecting_socket_t *head_connecting_sockets;
   struct us_socket_t *iterator;
   struct us_socket_context_t *prev, *next;
 
@@ -434,6 +439,9 @@ us_internal_ssl_socket_open(us_internal_ssl_socket_r s, int is_client,
                             char *ip, int ip_length);
 
 int us_raw_root_certs(struct us_cert_string_t **out);
+
+void us_internal_socket_context_unlink_connecting_socket(int ssl, struct us_socket_context_t *context, struct us_connecting_socket_t *c);
+void us_internal_socket_context_link_connecting_socket(int ssl, struct us_socket_context_t *context, struct us_connecting_socket_t *c);
 #endif
 
 #endif // INTERNAL_H
diff --git a/packages/bun-usockets/src/socket.c b/packages/bun-usockets/src/socket.c
@@ -22,6 +22,7 @@
 #include <string.h>
 #include <stdint.h>
 #include <stdio.h>
+#include <errno.h>
 
 #ifndef WIN32
 #include <fcntl.h>
@@ -132,9 +133,10 @@ int us_socket_is_established(int ssl, struct us_socket_t *s) {
 void us_connecting_socket_free(int ssl, struct us_connecting_socket_t *c) {
     // we can't just free c immediately, as it may be enqueued in the dns_ready_head list
     // instead, we move it to a close list and free it after the iteration
+    us_internal_socket_context_unlink_connecting_socket(ssl, c->context, c);
+
     c->next = c->context->loop->data.closed_connecting_head;
     c->context->loop->data.closed_connecting_head = c;
-    us_socket_context_unref(ssl, c->context);
 }
 
 void us_connecting_socket_close(int ssl, struct us_connecting_socket_t *c) {
@@ -153,7 +155,15 @@ void us_connecting_socket_close(int ssl, struct us_connecting_socket_t *c) {
         /* Any socket with prev = context is marked as closed */
         s->prev = (struct us_socket_t *) s->context;
     }
-
+    if(!c->error) {
+        // if we have no error, we have to set that we were aborted aka we called close
+        c->error = ECONNABORTED;
+    }
+    c->context->on_connect_error(c, c->error);
+    if(c->addrinfo_req) {
+        Bun__addrinfo_freeRequest(c->addrinfo_req, c->error == ECONNREFUSED);
+        c->addrinfo_req = 0;
+    }
     // we can only schedule the socket to be freed if there is no pending callback
     // otherwise, the callback will see that the socket is closed and will free it
     if (!c->pending_resolve_callback) {