OIDC: retrieve refreshToken from storage (#11250)

* test persistCredentials without a pickle key

* test setLoggedIn with pickle key

* lint

* type error

* extract token persisting code into function, persist refresh token

* store has_refresh_token too

* pass refreshToken from oidcAuthGrant into credentials

* rest restore session with pickle key

* retreive stored refresh token and add to credentials

* extract token decryption into function

* remove TODO

* comments

* prettier

* comment pedantry

* fix code smell - nullish coalesce instead of ||

* more comments

src/Lifecycle.ts

@@ -501,11 +501,41 @@ export interface IStoredSession {
     isUrl: string;
     hasAccessToken: boolean;
     accessToken: string | IEncryptedPayload;
+    hasRefreshToken: boolean;
+    refreshToken?: string | IEncryptedPayload;
     userId: string;
     deviceId: string;
     isGuest: boolean;
 }
 
+/**
+ * Retrieve a token, as stored by `persistCredentials`
+ * Attempts to migrate token from localStorage to idb
+ * @param storageKey key used to store the token, eg ACCESS_TOKEN_STORAGE_KEY
+ * @returns Promise that resolves to token or undefined
+ */
+async function getStoredToken(storageKey: string): Promise<string | undefined> {
+    let token: string | undefined;
+    try {
+        token = await StorageManager.idbLoad("account", storageKey);
+    } catch (e) {
+        logger.error(`StorageManager.idbLoad failed for account:${storageKey}`, e);
+    }
+    if (!token) {
+        token = localStorage.getItem(storageKey) ?? undefined;
+        if (token) {
+            try {
+                // try to migrate access token to IndexedDB if we can
+                await StorageManager.idbSave("account", storageKey, token);
+                localStorage.removeItem(storageKey);
+            } catch (e) {
+                logger.error(`migration of token ${storageKey} to IndexedDB failed`, e);
+            }
+        }
+    }
+    return token;
+}
+
 /**
  * Retrieves information about the stored session from the browser's storage. The session
  * may not be valid, as it is not tested for consistency here.
@@ -514,27 +544,14 @@ export interface IStoredSession {
 export async function getStoredSessionVars(): Promise<Partial<IStoredSession>> {
     const hsUrl = localStorage.getItem(HOMESERVER_URL_KEY) ?? undefined;
     const isUrl = localStorage.getItem(ID_SERVER_URL_KEY) ?? undefined;
-    let accessToken: string | undefined;
+
-    try {
+    const accessToken = await getStoredToken(ACCESS_TOKEN_STORAGE_KEY);
-        accessToken = await StorageManager.idbLoad("account", ACCESS_TOKEN_STORAGE_KEY);
+    const refreshToken = await getStoredToken(REFRESH_TOKEN_STORAGE_KEY);
-    } catch (e) {
+
-        logger.error("StorageManager.idbLoad failed for account:mx_access_token", e);
-    }
-    if (!accessToken) {
-        accessToken = localStorage.getItem(ACCESS_TOKEN_STORAGE_KEY) ?? undefined;
-        if (accessToken) {
-            try {
-                // try to migrate access token to IndexedDB if we can
-                await StorageManager.idbSave("account", ACCESS_TOKEN_STORAGE_KEY, accessToken);
-                localStorage.removeItem(ACCESS_TOKEN_STORAGE_KEY);
-            } catch (e) {
-                logger.error("migration of access token to IndexedDB failed", e);
-            }
-        }
-    }
     // if we pre-date storing "mx_has_access_token", but we retrieved an access
     // token, then we should say we have an access token
     const hasAccessToken = localStorage.getItem(HAS_ACCESS_TOKEN_STORAGE_KEY) === "true" || !!accessToken;
+    const hasRefreshToken = localStorage.getItem(HAS_REFRESH_TOKEN_STORAGE_KEY) === "true" || !!refreshToken;
     const userId = localStorage.getItem("mx_user_id") ?? undefined;
     const deviceId = localStorage.getItem("mx_device_id") ?? undefined;
 
@@ -546,7 +563,7 @@ export async function getStoredSessionVars(): Promise<Partial<IStoredSession>> {
         isGuest = localStorage.getItem("matrix-is-guest") === "true";
     }
 
-    return { hsUrl, isUrl, hasAccessToken, accessToken, userId, deviceId, isGuest };
+    return { hsUrl, isUrl, hasAccessToken, accessToken, refreshToken, hasRefreshToken, userId, deviceId, isGuest };
 }
 
 // The pickle key is a string of unspecified length and format.  For AES, we
@@ -585,6 +602,36 @@ async function abortLogin(): Promise<void> {
     }
 }
 
+const isEncryptedPayload = (token?: IEncryptedPayload | string | undefined): token is IEncryptedPayload => {
+    return !!token && typeof token !== "string";
+};
+/**
+ * Try to decrypt a token retrieved from storage
+ * Where token is not encrypted (plain text) returns the plain text token
+ * Where token is encrypted, attempts decryption. Returns successfully decrypted token, else undefined.
+ * @param pickleKey pickle key used during encryption of token, or undefined
+ * @param token
+ * @param tokenIv initialization vector used during encryption of token eg ACCESS_TOKEN_IV
+ * @returns the decrypted token, or the plain text token. Returns undefined when token cannot be decrypted
+ */
+async function tryDecryptToken(
+    pickleKey: string | undefined,
+    token: IEncryptedPayload | string | undefined,
+    tokenIv: string,
+): Promise<string | undefined> {
+    if (pickleKey && isEncryptedPayload(token)) {
+        const encrKey = await pickleKeyToAesKey(pickleKey);
+        const decryptedToken = await decryptAES(token, encrKey, tokenIv);
+        encrKey.fill(0);
+        return decryptedToken;
+    }
+    // if the token wasn't encrypted (plain string) just return it back
+    if (typeof token === "string") {
+        return token;
+    }
+    // otherwise return undefined
+}
+
 // returns a promise which resolves to true if a session is found in
 // localstorage
 //
@@ -602,7 +649,8 @@ export async function restoreFromLocalStorage(opts?: { ignoreGuest?: boolean }):
         return false;
     }
 
-    const { hsUrl, isUrl, hasAccessToken, accessToken, userId, deviceId, isGuest } = await getStoredSessionVars();
+    const { hsUrl, isUrl, hasAccessToken, accessToken, refreshToken, userId, deviceId, isGuest } =
+        await getStoredSessionVars();
 
     if (hasAccessToken && !accessToken) {
         await abortLogin();
@@ -614,18 +662,14 @@ export async function restoreFromLocalStorage(opts?: { ignoreGuest?: boolean }):
             return false;
         }
 
-        let decryptedAccessToken = accessToken;
+        const pickleKey = (await PlatformPeg.get()?.getPickleKey(userId, deviceId ?? "")) ?? undefined;
-        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_IV);
-                encrKey.fill(0);
-            }
         } else {
             logger.log("No pickle key available");
         }
+        const decryptedAccessToken = await tryDecryptToken(pickleKey, accessToken, ACCESS_TOKEN_IV);
+        const decryptedRefreshToken = await tryDecryptToken(pickleKey, refreshToken, REFRESH_TOKEN_IV);
 
         const freshLogin = sessionStorage.getItem("mx_fresh_login") === "true";
         sessionStorage.removeItem("mx_fresh_login");
@@ -635,7 +679,8 @@ export async function restoreFromLocalStorage(opts?: { ignoreGuest?: boolean }):
             {
                 userId: userId,
                 deviceId: deviceId,
-                accessToken: decryptedAccessToken as string,
+                accessToken: decryptedAccessToken!,
+                refreshToken: decryptedRefreshToken,
                 homeserverUrl: hsUrl,
                 identityServerUrl: isUrl,
                 guest: isGuest,

test/Lifecycle-test.ts

@@ -161,6 +161,8 @@ describe("Lifecycle", () => {
         accessToken,
     };
 
+    const refreshToken = "test-refresh-token";
+
     const encryptedTokenShapedObject = {
         ciphertext: expect.any(String),
         iv: expect.any(String),
@@ -285,6 +287,45 @@ describe("Lifecycle", () => {
 
                     expect(MatrixClientPeg.start).toHaveBeenCalled();
                 });
+
+                describe("with a refresh token", () => {
+                    beforeEach(() => {
+                        initLocalStorageMock({
+                            ...localStorageSession,
+                            mx_refresh_token: refreshToken,
+                        });
+                        initIdbMock(idbStorageSession);
+                    });
+
+                    it("should persist credentials", async () => {
+                        expect(await restoreFromLocalStorage()).toEqual(true);
+
+                        // refresh token from storage is re-persisted
+                        expect(localStorage.setItem).toHaveBeenCalledWith("mx_has_refresh_token", "true");
+                        expect(StorageManager.idbSave).toHaveBeenCalledWith(
+                            "account",
+                            "mx_refresh_token",
+                            refreshToken,
+                        );
+                    });
+
+                    it("should create new matrix client with credentials", async () => {
+                        expect(await restoreFromLocalStorage()).toEqual(true);
+
+                        expect(MatrixClientPeg.replaceUsingCreds).toHaveBeenCalledWith({
+                            userId,
+                            accessToken,
+                            // refreshToken included in credentials
+                            refreshToken,
+                            homeserverUrl,
+                            identityServerUrl,
+                            deviceId,
+                            freshLogin: false,
+                            guest: false,
+                            pickleKey: undefined,
+                        });
+                    });
+                });
             });
 
             describe("with a pickle key", () => {
@@ -344,6 +385,47 @@ describe("Lifecycle", () => {
                         pickleKey: expect.any(String),
                     });
                 });
+
+                describe("with a refresh token", () => {
+                    beforeEach(async () => {
+                        initLocalStorageMock({});
+                        initIdbMock({});
+                        // setup storage with a session with encrypted token
+                        await setLoggedIn({
+                            ...credentials,
+                            refreshToken,
+                        });
+                    });
+
+                    it("should persist credentials", async () => {
+                        expect(await restoreFromLocalStorage()).toEqual(true);
+
+                        // refresh token from storage is re-persisted
+                        expect(localStorage.setItem).toHaveBeenCalledWith("mx_has_refresh_token", "true");
+                        expect(StorageManager.idbSave).toHaveBeenCalledWith(
+                            "account",
+                            "mx_refresh_token",
+                            encryptedTokenShapedObject,
+                        );
+                    });
+
+                    it("should create new matrix client with credentials", async () => {
+                        expect(await restoreFromLocalStorage()).toEqual(true);
+
+                        expect(MatrixClientPeg.replaceUsingCreds).toHaveBeenCalledWith({
+                            userId,
+                            accessToken,
+                            // refreshToken included in credentials
+                            refreshToken,
+                            homeserverUrl,
+                            identityServerUrl,
+                            deviceId,
+                            freshLogin: false,
+                            guest: false,
+                            pickleKey: expect.any(String),
+                        });
+                    });
+                });
             });
 
             it("should show a toast if the matrix server version is unsupported", async () => {