Introduce WebIPC to ensure token refreshes happen exactly once

See included README.md for docs

src/Lifecycle.ts

@@ -59,7 +59,9 @@ import LazyLoadingDisabledDialog from "./components/views/dialogs/LazyLoadingDis
 import SessionRestoreErrorDialog from "./components/views/dialogs/SessionRestoreErrorDialog";
 import StorageEvictedDialog from "./components/views/dialogs/StorageEvictedDialog";
 import { setSentryUser } from "./sentry";
-import { TokenLifecycle } from "./TokenLifecycle";
+import { IRenewedMatrixClientCreds, TokenLifecycle } from "./TokenLifecycle";
+import { WebIPC } from "./ipc/WebIPC";
+import { Op } from "./ipc/types";
 
 const HOMESERVER_URL_KEY = "mx_hs_url";
 const ID_SERVER_URL_KEY = "mx_is_url";
@@ -313,7 +315,7 @@ export interface IStoredSession {
     deviceId: string;
     isGuest: boolean;
     accessTokenExpiryTs?: number; // set if the token expires
-    accessTokenRefreshToken?: string; // set if the token can be renewed
+    accessTokenRefreshToken?: string | IEncryptedPayload; // set if the token can be renewed
 }
 
 /**
@@ -444,6 +446,41 @@ async function abortLogin() {
     }
 }
 
+export async function getRenewedStoredSessionVars(): Promise<IRenewedMatrixClientCreds> {
+    const {
+        userId,
+        deviceId,
+        accessToken,
+        accessTokenExpiryTs,
+        accessTokenRefreshToken,
+    } = await getStoredSessionVars();
+
+    let decryptedAccessToken = accessToken;
+    let decryptedRefreshToken = accessTokenRefreshToken;
+    const pickleKey = await PlatformPeg.get().getPickleKey(userId, deviceId);
+    if (pickleKey) {
+        logger.log("Got pickle key");
+        if (typeof accessToken !== "string") {
+            const encrKey = await pickleKeyToAesKey(pickleKey);
+            decryptedAccessToken = await decryptAES(accessToken, encrKey, "access_token");
+            encrKey.fill(0);
+        }
+        if (accessTokenRefreshToken && typeof accessTokenRefreshToken !== "string") {
+            const encrKey = await pickleKeyToAesKey(pickleKey);
+            decryptedRefreshToken = await decryptAES(accessTokenRefreshToken, encrKey, "refresh_token");
+            encrKey.fill(0);
+        }
+    } else {
+        logger.log("No pickle key available");
+    }
+
+    return {
+        accessToken: decryptedAccessToken as string,
+        accessTokenExpiryTs: accessTokenExpiryTs,
+        accessTokenRefreshToken: decryptedRefreshToken as string,
+    };
+}
+
 // returns a promise which resolves to true if a session is found in
 // localstorage
 //
@@ -470,7 +507,6 @@ export async function restoreFromLocalStorage(opts?: { ignoreGuest?: boolean }):
         deviceId,
         isGuest,
         accessTokenExpiryTs,
-        accessTokenRefreshToken,
     } = await getStoredSessionVars();
 
     if (hasAccessToken && !accessToken) {
@@ -483,24 +519,11 @@ export async function restoreFromLocalStorage(opts?: { ignoreGuest?: boolean }):
             return false;
         }
 
-        let decryptedAccessToken = accessToken;
-        let decryptedRefreshToken = accessTokenRefreshToken;
         const pickleKey = await PlatformPeg.get().getPickleKey(userId, deviceId);
-        if (pickleKey) {
-            logger.log("Got pickle key");
-            if (typeof accessToken !== "string") {
-                const encrKey = await pickleKeyToAesKey(pickleKey);
-                decryptedAccessToken = await decryptAES(accessToken, encrKey, "access_token");
-                encrKey.fill(0);
-            }
-            if (accessTokenRefreshToken && typeof accessTokenRefreshToken !== "string") {
-                const encrKey = await pickleKeyToAesKey(pickleKey);
-                decryptedRefreshToken = await decryptAES(accessTokenRefreshToken, encrKey, "refresh_token");
-                encrKey.fill(0);
-            }
-        } else {
-            logger.log("No pickle key available");
-        }
+        const {
+            accessToken: decryptedAccessToken,
+            accessTokenRefreshToken: decryptedRefreshToken,
+        } = await getRenewedStoredSessionVars();
 
         const freshLogin = sessionStorage.getItem("mx_fresh_login") === "true";
         sessionStorage.removeItem("mx_fresh_login");
@@ -516,7 +539,7 @@ export async function restoreFromLocalStorage(opts?: { ignoreGuest?: boolean }):
             pickleKey: pickleKey,
             freshLogin: freshLogin,
             accessTokenExpiryTs: accessTokenExpiryTs,
-            accessTokenRefreshToken: decryptedRefreshToken,
+            accessTokenRefreshToken: decryptedRefreshToken as string,
         }, false);
         return true;
     } else {
@@ -631,6 +654,16 @@ export async function hydrateSessionInPlace(credentials: IMatrixClientCreds): Pr
     // reset the token timers
     TokenLifecycle.instance.startTimers(credentials);
 
+    // Now that everything is squared away, tell everyone we updated the token
+    if (WebIPC.instance.isCurrentlyLeader) {
+        WebIPC.instance.transport.send({
+            operation: Op.AccessTokenUpdated,
+            clientId: WebIPC.instance.clientId,
+            version: 1,
+            payload: {},
+        });
+    }
+
     return cli;
 }
 

src/TokenLifecycle.ts

@@ -18,7 +18,10 @@ import { logger } from "matrix-js-sdk/src/logger";
 import { MatrixClient } from "matrix-js-sdk/src";
 
 import { IMatrixClientCreds, MatrixClientPeg } from "./MatrixClientPeg";
-import { hydrateSessionInPlace } from "./Lifecycle";
+import { getRenewedStoredSessionVars, hydrateSessionInPlace } from "./Lifecycle";
+import { WebIPC } from "./ipc/WebIPC";
+import { Op, Operation } from "./ipc/types";
+import { TRANSPORT_EVENT } from "./ipc/transport/ITransport";
 
 export interface IRenewedMatrixClientCreds extends Pick<IMatrixClientCreds,
     "accessToken" | "accessTokenExpiryTs" | "accessTokenRefreshToken"> {}
@@ -31,6 +34,8 @@ export class TokenLifecycle {
     protected constructor() {
         // we only really want one of these floating around, so private-ish
         // constructor. Protected allows for unit tests.
+
+        this.registerIpcHandlers();
     }
 
     // noinspection JSMethodCanBeStatic
@@ -71,31 +76,86 @@ export class TokenLifecycle {
         }
     }
 
+    private registerIpcHandlers() {
+        WebIPC.instance.transport.on(TRANSPORT_EVENT, async (op: Operation) => {
+            if (op.operation !== Op.AccessTokenUpdated) return;
+
+            logger.info("TokenLifecyle#remoteUpdate: Leader triggered token refresh - updating client creds");
+            await hydrateSessionInPlace({
+                ...MatrixClientPeg.getCredentials(),
+                ...(await getRenewedStoredSessionVars()),
+            });
+        });
+
+        WebIPC.instance.transport.on(TRANSPORT_EVENT, async (op: Operation) => {
+            if (op.operation !== Op.RefreshToken) return;
+            if (!WebIPC.instance.isCurrentlyLeader) return; // not for us
+
+            logger.info("TokenLifecyle#remoteUpdate: Remote client wants a token refresh - honouring");
+            await this.forceTokenExchange();
+        });
+    }
+
     // noinspection JSMethodCanBeStatic
     private async doTokenRefresh(
         credentials: IMatrixClientCreds,
         client: MatrixClient,
     ): Promise<IRenewedMatrixClientCreds> {
-        // TODO: @@ TR - Lock so this+dependents only happens once per client
-        try {
-            const newCreds = await client.refreshToken(credentials.accessTokenRefreshToken);
-            return {
-                // We use the browser's local time to do two things:
-                // 1. Avoid having to write code that counts down and stores a "time left" variable
-                // 2. Work around any time drift weirdness by assuming the user's local machine will
-                //    drift consistently with itself.
-                // We additionally add our own safety buffer when renewing tokens to avoid cases where
-                // the time drift is accelerating.
-                accessTokenExpiryTs: Date.now() + newCreds.expires_in_ms,
-                accessToken: newCreds.access_token,
-                accessTokenRefreshToken: newCreds.refresh_token,
-            };
-        } catch (e) {
-            if (e.errcode === "M_UNKNOWN_TOKEN") {
-                // Emit the logout manually because the function inhibits it.
-                client.emit("Session.logged_out", e);
-            } else {
-                throw e; // we can't do anything with it, so re-throw
+        if (!WebIPC.instance.isCurrentlyLeader) {
+            logger.info("TokenLifecycle#doTokenRefresh: Asking for update from leader");
+            await new Promise<void>(resolve => {
+                // We set a timer in case the leader is suddenly non-responsive. In these cases we'll
+                // end up returning the currently stored credentials, which shouldn't be an issue.
+                const timerId = setTimeout(() => {
+                    logger.info("TokenLifecycle#doTokenRefresh: Leader failed to respond in time");
+                    fn({
+                        operation: Op.AccessTokenUpdated,
+                        clientId: "unknown", // not parsed, but needed for types
+                        version: 1,
+                        payload: {},
+                    });
+                }, 120000); // 2 minutes (120 seconds) to cover average timeout from server + buffer
+                const fn = (op: Operation) => {
+                    if (op.operation === Op.AccessTokenUpdated) {
+                        WebIPC.instance.transport.removeListener(TRANSPORT_EVENT, fn);
+                        clearTimeout(timerId);
+                        resolve();
+                    }
+                };
+                WebIPC.instance.transport.on(TRANSPORT_EVENT, fn);
+
+                // Ask the leader to refresh the token
+                WebIPC.instance.transport.send({
+                    operation: Op.RefreshToken,
+                    clientId: WebIPC.instance.clientId,
+                    version: 1,
+                    payload: {},
+                });
+            });
+            const { accessToken, accessTokenRefreshToken, accessTokenExpiryTs } = await getRenewedStoredSessionVars();
+            return { accessToken, accessTokenRefreshToken, accessTokenExpiryTs };
+        } else {
+            logger.info("TokenLifecycle#doTokenRefresh: Refreshing token as current leader");
+            try {
+                const newCreds = await client.refreshToken(credentials.accessTokenRefreshToken);
+                return {
+                    // We use the browser's local time to do two things:
+                    // 1. Avoid having to write code that counts down and stores a "time left" variable
+                    // 2. Work around any time drift weirdness by assuming the user's local machine will
+                    //    drift consistently with itself.
+                    // We additionally add our own safety buffer when renewing tokens to avoid cases where
+                    // the time drift is accelerating.
+                    accessTokenExpiryTs: Date.now() + newCreds.expires_in_ms,
+                    accessToken: newCreds.access_token,
+                    accessTokenRefreshToken: newCreds.refresh_token,
+                };
+            } catch (e) {
+                if (e.errcode === "M_UNKNOWN_TOKEN") {
+                    // Emit the logout manually because the function inhibits it.
+                    client.emit("Session.logged_out", e);
+                } else {
+                    throw e; // we can't do anything with it, so re-throw
+                }
             }
         }
     }

src/ipc/Election.ts

@@ -0,0 +1,62 @@
+/*
+Copyright 2022 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import { lexicographicCompare } from "matrix-js-sdk/src/utils";
+
+export const ELECTION_VALIDITY_MS = 10000; // 10 seconds
+
+export class Election {
+    public readonly votes = new Map<string, string>();
+
+    private _winner: string;
+
+    constructor(public readonly electionId: string, public readonly startTs: number) {
+    }
+
+    public get winner(): string {
+        return this._winner;
+    }
+
+    public get hasEnded(): boolean {
+        return Date.now() >= (this.startTs + ELECTION_VALIDITY_MS);
+    }
+
+    public addVote(clientId: string, vote: string) {
+        if (!this.votes.has(clientId)) {
+            this.votes.set(clientId, vote);
+        }
+    }
+
+    /**
+     * Finalizes the election, returning the elected leader
+     */
+    public finalize(): string {
+        if (this.winner) {
+            return this.winner;
+        }
+
+        const clientVoteTuples = Array.from(this.votes.entries());
+        this._winner = clientVoteTuples.sort(([aId, aVote], [bId, bVote]) => {
+            const voteCompare = lexicographicCompare(aVote, bVote);
+            if (voteCompare === 0) {
+                return lexicographicCompare(aId, bId);
+            }
+            return voteCompare;
+        })[0][0]; // first tuple, client ID from that tuple
+
+        return this.winner;
+    }
+}