From 16a5f2a1752135cf35c51c9d8031e8d05c623898 Mon Sep 17 00:00:00 2001
From: Louis Lam <louislam@users.noreply.github.com>
Date: Sat, 9 Dec 2023 17:02:14 +0800
Subject: [PATCH] Fix logout

---
 backend/dockge-server.ts                      | 30 +++++++++--
 .../socket-handlers/main-socket-handler.ts    |  2 +
 backend/util-common.ts                        |  5 ++
 extra/reset-password.ts                       | 51 +++++++++++++++++--
 frontend/src/mixins/socket.ts                 |  4 ++
 5 files changed, 86 insertions(+), 6 deletions(-)

diff --git a/backend/dockge-server.ts b/backend/dockge-server.ts
index c361d2f..b774bf6 100644
--- a/backend/dockge-server.ts
+++ b/backend/dockge-server.ts
@@ -605,9 +605,33 @@ export class DockgeServer {
     }
 
     /**
-     * Final function called before application exits
+     * Force connected sockets of a user to refresh and disconnect.
+     * Used for resetting password.
+     * @param {string} userID
+     * @param {string?} currentSocketID
      */
-    finalFunction() {
-        log.info("server", "Graceful shutdown successful!");
+    disconnectAllSocketClient(userID: number, currentSocketID? : string) {
+        for (const rawSocket of this.io.sockets.sockets.values()) {
+            let socket = rawSocket as DockgeSocket;
+            if (socket.userID === userID && socket.id !== currentSocketID) {
+                try {
+                    socket.emit("refresh");
+                    socket.disconnect();
+                } catch (e) {
+
+                }
+            }
+        }
     }
+
+    isSSL() {
+        return this.config.sslKey && this.config.sslCert;
+    }
+
+    getLocalWebSocketURL() {
+        const protocol = this.isSSL() ? "wss" : "ws";
+        const host = this.config.hostname || "localhost";
+        return `${protocol}://${host}:${this.config.port}`;
+    }
+
 }
diff --git a/backend/socket-handlers/main-socket-handler.ts b/backend/socket-handlers/main-socket-handler.ts
index bfdb45d..e82b072 100644
--- a/backend/socket-handlers/main-socket-handler.ts
+++ b/backend/socket-handlers/main-socket-handler.ts
@@ -211,6 +211,8 @@ export class MainSocketHandler extends SocketHandler {
                 let user = await doubleCheckPassword(socket, password.currentPassword);
                 await user.resetPassword(password.newPassword);
 
+                server.disconnectAllSocketClient(user.id, socket.id);
+
                 callback({
                     ok: true,
                     msg: "Password has been updated successfully.",
diff --git a/backend/util-common.ts b/backend/util-common.ts
index 71fd4d8..7c1544d 100644
--- a/backend/util-common.ts
+++ b/backend/util-common.ts
@@ -17,6 +17,11 @@ export interface LooseObject {
     [key: string]: any
 }
 
+export interface BaseRes {
+    ok: boolean;
+    msg?: string;
+}
+
 let randomBytes : (numBytes: number) => Uint8Array;
 initRandomBytes();
 
diff --git a/extra/reset-password.ts b/extra/reset-password.ts
index 8c2a2df..3d38fe1 100644
--- a/extra/reset-password.ts
+++ b/extra/reset-password.ts
@@ -4,6 +4,8 @@ import readline from "readline";
 import { User } from "../backend/models/user";
 import { DockgeServer } from "../backend/dockge-server";
 import { log } from "../backend/log";
+import { io } from "socket.io-client";
+import { BaseRes } from "../backend/util-common";
 
 console.log("== Dockge Reset Password Tool ==");
 
@@ -12,11 +14,10 @@ const rl = readline.createInterface({
     output: process.stdout
 });
 
+const server = new DockgeServer();
+
 export const main = async () => {
-    const server = new DockgeServer();
-
     // Check if
-
     console.log("Connecting the database");
     try {
         await Database.init(server);
@@ -47,6 +48,9 @@ export const main = async () => {
                     // Reset all sessions by reset jwt secret
                     await server.initJWTSecret();
 
+                    // Disconnect all other socket clients of the user
+                    await disconnectAllSocketClients(user.username, password);
+
                     break;
                 } else {
                     console.log("Passwords do not match, please try again.");
@@ -79,6 +83,47 @@ function question(question : string) : Promise<string> {
     });
 }
 
+function disconnectAllSocketClients(username : string, password : string) : Promise<void> {
+    return new Promise((resolve) => {
+        const url = server.getLocalWebSocketURL();
+
+        console.log("Connecting to " + url + " to disconnect all other socket clients");
+
+        // Disconnect all socket connections
+        const socket = io(url, {
+            transports: [ "websocket" ],
+            reconnection: false,
+            timeout: 5000,
+        });
+        socket.on("connect", () => {
+            socket.emit("login", {
+                username,
+                password,
+            }, (res : BaseRes) => {
+                if (res.ok) {
+                    console.log("Logged in.");
+                    socket.emit("disconnectOtherSocketClients");
+                } else {
+                    console.warn("Login failed.");
+                    console.warn("Please restart the server to disconnect all sessions.");
+                }
+                socket.close();
+            });
+        });
+
+        socket.on("connect_error", function () {
+            // The localWebSocketURL is not guaranteed to be working for some complicated Uptime Kuma setup
+            // Ask the user to restart the server manually
+            console.warn("Failed to connect to " + url);
+            console.warn("Please restart the server to disconnect all sessions manually.");
+            resolve();
+        });
+        socket.on("disconnect", () => {
+            resolve();
+        });
+    });
+}
+
 if (!process.env.TEST_BACKEND) {
     main();
 }
diff --git a/frontend/src/mixins/socket.ts b/frontend/src/mixins/socket.ts
index 30c74e7..2ba8464 100644
--- a/frontend/src/mixins/socket.ts
+++ b/frontend/src/mixins/socket.ts
@@ -202,6 +202,10 @@ export default defineComponent({
                     }
                 }
             });
+
+            socket.on("refresh", () => {
+                location.reload();
+            });
         },
 
         /**