diff --git a/include/ssh_server.h b/include/ssh_server.h index ed5092c..0f8419e 100644 --- a/include/ssh_server.h +++ b/include/ssh_server.h @@ -23,6 +23,8 @@ typedef struct client { char command_output[2048]; pthread_t thread; bool connected; + int ref_count; /* Reference count for safe cleanup */ + pthread_mutex_t ref_lock; /* Lock for ref_count */ } client_t; /* Initialize SSH server */ diff --git a/src/chat_room.c b/src/chat_room.c index 0033a98..65b4c5f 100644 --- a/src/chat_room.c +++ b/src/chat_room.c @@ -81,11 +81,18 @@ void room_broadcast(chat_room_t *room, const message_t *msg) { /* Add to history */ room_add_message(room, msg); - /* Get copy of client list for iteration */ + /* Get copy of client list and increment ref counts */ client_t **clients_copy = calloc(room->client_count, sizeof(client_t*)); int count = room->client_count; memcpy(clients_copy, room->clients, count * sizeof(client_t*)); + /* Increment reference count for each client */ + for (int i = 0; i < count; i++) { + pthread_mutex_lock(&clients_copy[i]->ref_lock); + clients_copy[i]->ref_count++; + pthread_mutex_unlock(&clients_copy[i]->ref_lock); + } + pthread_rwlock_unlock(&room->lock); /* Render to each client (outside of lock) */ @@ -95,6 +102,26 @@ void room_broadcast(chat_room_t *room, const message_t *msg) { client->command_output[0] == '\0') { tui_render_screen(client); } + + /* Decrement reference count and free if needed */ + pthread_mutex_lock(&client->ref_lock); + client->ref_count--; + int ref = client->ref_count; + pthread_mutex_unlock(&client->ref_lock); + + if (ref == 0) { + /* Safe to free now */ + if (client->channel) { + ssh_channel_close(client->channel); + ssh_channel_free(client->channel); + } + if (client->session) { + ssh_disconnect(client->session); + ssh_free(client->session); + } + pthread_mutex_destroy(&client->ref_lock); + free(client); + } } free(clients_copy); diff --git a/src/tui.c b/src/tui.c index 58cdc15..7df9317 100644 --- a/src/tui.c +++ b/src/tui.c @@ -169,8 +169,12 @@ void tui_render_command_output(client_t *client) { } pos += snprintf(buffer + pos, sizeof(buffer) - pos, ANSI_RESET "\r\n"); - /* Command output */ - char *line = strtok(client->command_output, "\n"); + /* Command output - use a copy to avoid strtok corruption */ + char output_copy[2048]; + strncpy(output_copy, client->command_output, sizeof(output_copy) - 1); + output_copy[sizeof(output_copy) - 1] = '\0'; + + char *line = strtok(output_copy, "\n"); int line_count = 0; int max_lines = client->height - 2;