1 files changed, 14 insertions, 7 deletions

diff --git a/src/liana/server.c b/src/liana/server.c
index 5065425..791fe0e 100644
--- a/src/liana/server.c
+++ b/src/liana/server.c
@@ -124,7 +124,7 @@ static void free_connection(struct lia_node_connection *conn)
     struct lia_node *node = conn->node;
     al_assert(conn->handler);
     conn->handler->free(&conn->handler);
-    cch_entry_return_handle(conn->node->entry, &conn->handle);
+    cch_entry_return_handle(node->entry, &conn->handle);
     nn_packet_pool_free(&conn->pool);
     bool removed = al_array_remove(node->connections, conn);
     al_assert(removed);
@@ -272,31 +272,32 @@ static void signal_callback(void *userdata)
 
     nn_thread_join(&conn->thread);
     nn_signal_stop(&conn->signal);
-    al_array_remove(node->requests, conn);
 
     struct nn_packet_stream *stream = conn->stream;
     struct nn_packet *packet = conn->packet;
     conn->packet = NULL;
 
+    if (packet) {
+        al_array_remove(node->requests, conn);
+    }
+
     if (!packet || conn->errored) {
         conn->handler->free(&conn->handler);
         cch_entry_return_handle(node->entry, &conn->handle);
     }
 
-    if (!packet) {
-        // Connection was closed before init was done.
+    if (!packet) { // Connection was closed before init was done.
+        al_free(conn);
         if (should_free_node(node)) {
             free_node(node);
         }
         nn_packet_stream_free(stream);
         al_free(stream);
-        al_free(conn);
     } else if (conn->errored) {
         // We must return the packet before disconnecting.
         nn_packet_stream_return_packet(stream, packet);
-        al_free(conn);
-        conn = NULL;
         demote_and_disconnect_stream(server, stream);
+        al_free(conn);
     } else {
         conn->id = get_incremental_id(server);
         al_array_push(node->connections, conn);
@@ -337,10 +338,12 @@ static struct lia_node_connection *get_connection_from_id(struct lia_node *node,
 static void pre_init_connection_closed_callback(void *userdata, struct nn_packet_stream *stream)
 {
     struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
+    struct lia_node *node = conn->node;
     struct nn_packet *packet = conn->packet;
     // Checked in signal_callback and will signal to cleanup the connection.
     conn->packet = NULL;
     nn_packet_stream_return_packet(stream, packet);
+    al_array_remove(node->requests, conn);
 }
 
 static void packet_callback(void *userdata, struct nn_packet_stream *stream, struct nn_packet *packet)
@@ -439,6 +442,10 @@ static nn_thread_result NNWT_THREADCALL init_duration_thread(void *userdata)
 {
     struct lia_node *node = (struct lia_node *)userdata;
     nn_thread_set_name("liana_init_dur");
+#if 0 // This is really bad for large resources and is just for testing.
+    lia_prepare_visual_data(&node->handle, &node->visual);
+    cch_handle_seek(&node->handle, 0, SEEK_SET);
+#endif
     if (!node->handler->init(node->handler, &node->handle)) {
         node->errored = true;
     } else {