fix: close realtime connection on logout to prevent ghost sessions (#1158)

packages/api/package.json

@@ -7,7 +7,7 @@
   "typings": "dist/index.d.ts",
   "type": "module",
   "scripts": {
-    "test": "echo \"Error: no test specified\" && exit 1",
+    "test": "yarn build && yarn esbuild src/EmbeddedChatApi.ts --bundle --platform=node --format=esm --main-fields=module,main --outfile=dist/EmbeddedChatApi.test.mjs && node --test test/*.test.mjs",
     "build": "rollup -c",
     "dev": "yarn parcel playground/index.html",
     "format": "prettier --write 'src/'",
@@ -16,6 +16,7 @@
   "author": "",
   "license": "ISC",
   "devDependencies": {
+    "esbuild": "^0.17.19",
     "parcel": "^2.10.3",
     "rollup": "^3.23.0",
     "rollup-plugin-dts": "^6.0.1",

packages/api/src/EmbeddedChatApi.ts

@@ -9,6 +9,14 @@ import {
 
 // mutliple typing status can come at the same time they should be processed in order.
 let typingHandlerLock = 0;
+
+class ConnectionCancelledError extends Error {
+  constructor() {
+    super("Connection attempt was cancelled");
+    this.name = "ConnectionCancelledError";
+  }
+}
+
 export default class EmbeddedChatApi {
   host: string;
   rid: string;
@@ -21,7 +29,11 @@ export default class EmbeddedChatApi {
   typingUsers: string[];
   auth: RocketChatAuth;
   private _connectPromise: Promise<void> | null = null;
+  private _cancelConnect: (() => void) | null = null;
   private _activeSubscriptions: { stop: () => void }[] = [];
+  // Incremented by close() so that a connect() which is still awaiting
+  // can detect that it has been superseded or cancelled.
+  private _connectionGeneration = 0;
   private _authListener: ((user: any) => void) | null = null;
 
   constructor(
@@ -231,6 +243,16 @@ export default class EmbeddedChatApi {
     } catch (err: any) {
       console.error(err instanceof Error ? err.message : err);
       return err;
+    } finally {
+      // Tear down the realtime connection even if the logout request
+      // failed, otherwise the socket stays subscribed to the room and
+      // keeps receiving messages for a session that is no longer
+      // authenticated.
+      try {
+        await this.close();
+      } catch (err: any) {
+        console.error(err instanceof Error ? err.message : err);
+      }
     }
   }
 
@@ -243,10 +265,30 @@ export default class EmbeddedChatApi {
     if (this._connectPromise) {
       return this._connectPromise;
     }
-    this._connectPromise = this._doConnect(isChannelPrivate).finally(() => {
-      this._connectPromise = null;
+
+    // Cancel previous subscriptions and snapshot the new connection generation
+    // before registering this attempt's cancellation callback.
+    this.close();
+    const generation = this._connectionGeneration;
+    let cancelConnect: () => void;
+    const cancelled = new Promise<never>((_, reject) => {
+      cancelConnect = () => reject(new ConnectionCancelledError());
     });
-    return this._connectPromise;
+    this._cancelConnect = cancelConnect!;
+
+    const connectPromise = Promise.race([
+      this._doConnect(generation, isChannelPrivate),
+      cancelled,
+    ]).finally(() => {
+      // A cancelled handshake may finish after a newer connect() has started.
+      // Do not let its cleanup clear the newer attempt.
+      if (this._connectPromise === connectPromise) {
+        this._connectPromise = null;
+        this._cancelConnect = null;
+      }
+    });
+    this._connectPromise = connectPromise;
+    return connectPromise;
   }
 
   private _normalizeMessage(data: any) {
@@ -272,17 +314,24 @@ export default class EmbeddedChatApi {
     }
   }
 
-  private async _doConnect(isChannelPrivate = false) {
+  private async _doConnect(generation: number, isChannelPrivate = false) {
+    // Every close() invalidates this generation, so a connect() interrupted
+    // while awaiting can stop before subscribing. The cancellation promise in
+    // connect() separately releases callers if the DDP handshake never settles.
+    const isStale = () => generation !== this._connectionGeneration;
     try {
-      this.close(); // before connection, all previous subscriptions should be cancelled
       await this.sdk.connection.connect();
+      if (isStale()) return;
 
       const currentUser = (await this.auth.getCurrentUser()) as any;
+      if (isStale()) return;
       const token = currentUser?.authToken || currentUser?.data?.authToken;
       if (token) {
         this._applyCredentials(currentUser);
         await this.sdk.account.loginWithToken(token);
+        if (isStale()) return;
         await this.joinRoom(isChannelPrivate);
+        if (isStale()) return;
       }
 
       // Subscribe to room messages
@@ -359,6 +408,12 @@ export default class EmbeddedChatApi {
         );
         this._activeSubscriptions.push(notifyUserUiSub);
       }
+
+      // A close() may have happened while the subscriptions above were
+      // being registered, so drop them if this connect is no longer current.
+      if (isStale()) {
+        this._stopSubscriptions();
+      }
     } catch (err) {
       this.close();
     }
@@ -551,9 +606,21 @@ export default class EmbeddedChatApi {
     }
   }
 
-  async close() {
+  private _stopSubscriptions() {
     this._activeSubscriptions.forEach((sub) => sub.stop());
     this._activeSubscriptions = [];
+  }
+
+  async close() {
+    // Invalidate and release any connect() currently waiting on a DDP
+    // handshake. The underlying promise may remain pending in the DDP client,
+    // but it can no longer block a subsequent connect().
+    this._connectionGeneration += 1;
+    this._stopSubscriptions();
+    const cancelConnect = this._cancelConnect;
+    this._cancelConnect = null;
+    this._connectPromise = null;
+    cancelConnect?.();
     this.sdk.connection.close();
   }
 

packages/api/test/connection-lifecycle.test.mjs

@@ -0,0 +1,97 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import EmbeddedChatApi from "../dist/EmbeddedChatApi.test.mjs";
+
+const withTimeout = (promise, message) =>
+  Promise.race([
+    promise,
+    new Promise((_, reject) =>
+      setTimeout(() => reject(new Error(message)), 250)
+    ),
+  ]);
+
+const createApi = () => {
+  const api = Object.create(EmbeddedChatApi.prototype);
+  let connectCalls = 0;
+  let loginCalls = 0;
+  let currentUser = { authToken: "old-token", userId: "old-user" };
+  let resolveFirstHandshake;
+  const firstHandshake = new Promise((resolve) => {
+    resolveFirstHandshake = resolve;
+  });
+
+  Object.assign(api, {
+    rid: "room-id",
+    _connectPromise: null,
+    _cancelConnect: null,
+    _activeSubscriptions: [],
+    _connectionGeneration: 0,
+    onMessageCallbacks: [],
+    onMessageDeleteCallbacks: [],
+    onTypingStatusCallbacks: [],
+    onActionTriggeredCallbacks: [],
+    onUiInteractionCallbacks: [],
+    typingUsers: [],
+    auth: {
+      logout: async () => {
+        currentUser = null;
+      },
+      getCurrentUser: async () => currentUser,
+    },
+    joinRoom: async () => undefined,
+    sdk: {
+      connection: {
+        connect: () => {
+          connectCalls += 1;
+          return connectCalls === 1 ? firstHandshake : Promise.resolve();
+        },
+        close: () => undefined,
+      },
+      account: {
+        uid: null,
+        loginWithToken: async () => {
+          loginCalls += 1;
+        },
+      },
+      rest: { setCredentials: () => undefined },
+      stream: () => ({ stop: () => undefined }),
+    },
+  });
+
+  return {
+    api,
+    getConnectCalls: () => connectCalls,
+    getLoginCalls: () => loginCalls,
+    login: () => {
+      currentUser = { authToken: "new-token", userId: "new-user" };
+    },
+    resolveFirstHandshake,
+  };
+};
+
+test("logout during handshake allows the next login to reconnect", async () => {
+  const { api, getConnectCalls, getLoginCalls, login, resolveFirstHandshake } =
+    createApi();
+
+  const interruptedConnect = api.connect();
+  assert.equal(getConnectCalls(), 1);
+
+  await api.logout();
+  await assert.rejects(
+    withTimeout(interruptedConnect, "the interrupted connection did not settle"),
+    { name: "ConnectionCancelledError" }
+  );
+
+  login();
+  await withTimeout(api.connect(), "the next connection did not settle");
+  assert.equal(getConnectCalls(), 2);
+  assert.equal(getLoginCalls(), 1);
+
+  // Let the abandoned DDP promise settle and verify its finalizer does not
+  // interfere with later connection attempts.
+  resolveFirstHandshake();
+  await Promise.resolve();
+  await withTimeout(api.connect(), "a later connection did not settle");
+  assert.equal(getConnectCalls(), 3);
+});

yarn.lock

@@ -2456,6 +2456,7 @@ __metadata:
     "@embeddedchat/auth": "workspace:^"
     "@rocket.chat/ddp-client": 1.0.9-rc.2
     "@rocket.chat/emitter": ^0.31.25
+    esbuild: ^0.17.19
     parcel: ^2.10.3
     rollup: ^3.23.0
     rollup-plugin-dts: ^6.0.1