From 82f4f6cfb6f8811bf2aac399a1b3866d80629fa0 Mon Sep 17 00:00:00 2001
From: Luca8991 <luca.berte98@hotmail.it>
Date: Mon, 4 Dec 2023 10:33:15 +0100
Subject: [PATCH] fix: remove client if it was connected on ws_open

---
 src/State.mo | 18 ++++++++++--------
 src/lib.mo   | 11 +++++++++++
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/src/State.mo b/src/State.mo
index 55daff6..a815c66 100644
--- a/src/State.mo
+++ b/src/State.mo
@@ -252,25 +252,27 @@ module {
       increment_gateway_clients_count(new_client.gateway_principal);
     };
 
+    /// Removes a client from the internal state
+    /// and call the on_close callback,
+    /// if the client was registered in the state.
     public func remove_client(client_key : ClientKey, handlers : WsHandlers) : async () {
       CLIENTS_WAITING_FOR_KEEP_ALIVE := TrieSet.delete(CLIENTS_WAITING_FOR_KEEP_ALIVE, client_key, Types.hashClientKey(client_key), Types.areClientKeysEqual);
       CURRENT_CLIENT_KEY_MAP.delete(client_key.client_principal);
       OUTGOING_MESSAGE_TO_CLIENT_NUM_MAP.delete(client_key);
       INCOMING_MESSAGE_FROM_CLIENT_NUM_MAP.delete(client_key);
 
-      let registered_client = REGISTERED_CLIENTS.remove(client_key);
-      switch (registered_client) {
+      switch (REGISTERED_CLIENTS.remove(client_key)) {
         case (?registered_client) {
           decrement_gateway_clients_count(registered_client.gateway_principal);
+
+          await handlers.call_on_close({
+            client_principal = client_key.client_principal;
+          });
         };
         case (null) {
-          Prelude.unreachable();
+          // Do nothing
         };
       };
-
-      await handlers.call_on_close({
-        client_principal = client_key.client_principal;
-      });
     };
 
     public func format_message_for_gateway_key(gateway_principal : Principal, nonce : Nat64) : Text {
@@ -413,7 +415,7 @@ module {
       };
     };
 
-    /// Deletes the an amount of [MESSAGES_TO_DELETE] messages from the queue
+    /// Deletes the an amount of [MESSAGES_TO_DELETE_COUNT] messages from the queue
     /// that are older than the ack interval.
     func delete_old_messages_for_gateway(gateway_principal : GatewayPrincipal) : Result<(), Text> {
       let ack_interval_ms = init_params.send_ack_interval_ms;
diff --git a/src/lib.mo b/src/lib.mo
index 70e302a..7fd4a81 100644
--- a/src/lib.mo
+++ b/src/lib.mo
@@ -98,6 +98,17 @@ module {
 				};
 			};
 
+			// check if there's a client already registered with the same principal
+			// and remove it if there is
+			switch (WS_STATE.get_client_key_from_principal(client_key.client_principal)) {
+				case (#Err(err)) {
+					// Do nothing
+				};
+				case (#Ok(old_client_key)) {
+					await WS_STATE.remove_client(old_client_key, handlers);
+				};
+			};
+
 			// initialize client maps
 			let new_client = Types.RegisteredClient(args.gateway_principal);
 			WS_STATE.add_client(client_key, new_client);