OIDC: extract token persistence functions to utils (#11690)

* extract token persistence functions to utils

* add sugar

src/Lifecycle.ts

@@ -20,7 +20,7 @@ limitations under the License.
 import { ReactNode } from "react";
 import { createClient, MatrixClient, SSOAction } from "matrix-js-sdk/src/matrix";
 import { InvalidStoreError } from "matrix-js-sdk/src/errors";
-import { decryptAES, encryptAES, IEncryptedPayload } from "matrix-js-sdk/src/crypto/aes";
+import { IEncryptedPayload } from "matrix-js-sdk/src/crypto/aes";
 import { QueryDict } from "matrix-js-sdk/src/utils";
 import { logger } from "matrix-js-sdk/src/logger";
 import { MINIMUM_MATRIX_VERSION } from "matrix-js-sdk/src/version-support";
@@ -67,27 +67,21 @@ import { messageForLoginError } from "./utils/ErrorUtils";
 import { completeOidcLogin } from "./utils/oidc/authorize";
 import { persistOidcAuthenticatedSettings } from "./utils/oidc/persistOidcSettings";
 import GenericToast from "./components/views/toasts/GenericToast";
+import {
+    ACCESS_TOKEN_IV,
+    ACCESS_TOKEN_STORAGE_KEY,
+    HAS_ACCESS_TOKEN_STORAGE_KEY,
+    HAS_REFRESH_TOKEN_STORAGE_KEY,
+    persistAccessTokenInStorage,
+    persistRefreshTokenInStorage,
+    REFRESH_TOKEN_IV,
+    REFRESH_TOKEN_STORAGE_KEY,
+    tryDecryptToken,
+} from "./utils/tokens/tokens";
 
 const HOMESERVER_URL_KEY = "mx_hs_url";
 const ID_SERVER_URL_KEY = "mx_is_url";
 
-/*
- * Keys used when storing the tokens in indexeddb or localstorage
- */
-const ACCESS_TOKEN_STORAGE_KEY = "mx_access_token";
-const REFRESH_TOKEN_STORAGE_KEY = "mx_refresh_token";
-/*
- * Used as initialization vector during encryption in persistTokenInStorage
- * And decryption in restoreFromLocalStorage
- */
-const ACCESS_TOKEN_IV = "access_token";
-const REFRESH_TOKEN_IV = "refresh_token";
-/*
- * Keys for localstorage items which indicate whether we expect a token in indexeddb.
- */
-const HAS_ACCESS_TOKEN_STORAGE_KEY = "mx_has_access_token";
-const HAS_REFRESH_TOKEN_STORAGE_KEY = "mx_has_refresh_token";
-
 dis.register((payload) => {
     if (payload.action === Action.TriggerLogout) {
         // noinspection JSIgnoredPromiseFromCall - we don't care if it fails
@@ -566,32 +560,6 @@ export async function getStoredSessionVars(): Promise<Partial<IStoredSession>> {
     return { hsUrl, isUrl, hasAccessToken, accessToken, refreshToken, hasRefreshToken, userId, deviceId, isGuest };
 }
 
-// The pickle key is a string of unspecified length and format.  For AES, we
-// need a 256-bit Uint8Array. So we HKDF the pickle key to generate the AES
-// key.  The AES key should be zeroed after it is used.
-async function pickleKeyToAesKey(pickleKey: string): Promise<Uint8Array> {
-    const pickleKeyBuffer = new Uint8Array(pickleKey.length);
-    for (let i = 0; i < pickleKey.length; i++) {
-        pickleKeyBuffer[i] = pickleKey.charCodeAt(i);
-    }
-    const hkdfKey = await window.crypto.subtle.importKey("raw", pickleKeyBuffer, "HKDF", false, ["deriveBits"]);
-    pickleKeyBuffer.fill(0);
-    return new Uint8Array(
-        await window.crypto.subtle.deriveBits(
-            {
-                name: "HKDF",
-                hash: "SHA-256",
-                // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-                // @ts-ignore: https://github.com/microsoft/TypeScript-DOM-lib-generator/pull/879
-                salt: new Uint8Array(32),
-                info: new Uint8Array(0),
-            },
-            hkdfKey,
-            256,
-        ),
-    );
-}
-
 async function abortLogin(): Promise<void> {
     const signOut = await showStorageEvictedDialog();
     if (signOut) {
@@ -602,36 +570,6 @@ 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
 //
@@ -901,73 +839,6 @@ async function showStorageEvictedDialog(): Promise<boolean> {
 // `instanceof`. Babel 7 supports this natively in their class handling.
 class AbortLoginAndRebuildStorage extends Error {}
 
-/**
- * Persist a token in storage
- * When pickle key is present, will attempt to encrypt the token
- * Stores in idb, falling back to localStorage
- *
- * @param storageKey key used to store the token
- * @param initializationVector Initialization vector for encrypting the token. Only used when `pickleKey` is present
- * @param token the token to store, when undefined any existing token at the storageKey is removed from storage
- * @param pickleKey optional pickle key used to encrypt token
- * @param hasTokenStorageKey Localstorage key for an item which stores whether we expect to have a token in indexeddb, eg "mx_has_access_token".
- */
-async function persistTokenInStorage(
-    storageKey: string,
-    initializationVector: string,
-    token: string | undefined,
-    pickleKey: string | undefined,
-    hasTokenStorageKey: string,
-): Promise<void> {
-    // store whether we expect to find a token, to detect the case
-    // where IndexedDB is blown away
-    if (token) {
-        localStorage.setItem(hasTokenStorageKey, "true");
-    } else {
-        localStorage.removeItem(hasTokenStorageKey);
-    }
-
-    if (pickleKey) {
-        let encryptedToken: IEncryptedPayload | undefined;
-        try {
-            if (!token) {
-                throw new Error("No token: not attempting encryption");
-            }
-            // try to encrypt the access token using the pickle key
-            const encrKey = await pickleKeyToAesKey(pickleKey);
-            encryptedToken = await encryptAES(token, encrKey, initializationVector);
-            encrKey.fill(0);
-        } catch (e) {
-            logger.warn("Could not encrypt access token", e);
-        }
-        try {
-            // save either the encrypted access token, or the plain access
-            // token if we were unable to encrypt (e.g. if the browser doesn't
-            // have WebCrypto).
-            await StorageManager.idbSave("account", storageKey, encryptedToken || token);
-        } catch (e) {
-            // if we couldn't save to indexedDB, fall back to localStorage.  We
-            // store the access token unencrypted since localStorage only saves
-            // strings.
-            if (!!token) {
-                localStorage.setItem(storageKey, token);
-            } else {
-                localStorage.removeItem(storageKey);
-            }
-        }
-    } else {
-        try {
-            await StorageManager.idbSave("account", storageKey, token);
-        } catch (e) {
-            if (!!token) {
-                localStorage.setItem(storageKey, token);
-            } else {
-                localStorage.removeItem(storageKey);
-            }
-        }
-    }
-}
-
 async function persistCredentials(credentials: IMatrixClientCreds): Promise<void> {
     localStorage.setItem(HOMESERVER_URL_KEY, credentials.homeserverUrl);
     if (credentials.identityServerUrl) {
@@ -976,20 +847,8 @@ async function persistCredentials(credentials: IMatrixClientCreds): Promise<void
     localStorage.setItem("mx_user_id", credentials.userId);
     localStorage.setItem("mx_is_guest", JSON.stringify(credentials.guest));
 
-    await persistTokenInStorage(
+    await persistAccessTokenInStorage(credentials.accessToken, credentials.pickleKey);
-        ACCESS_TOKEN_STORAGE_KEY,
+    await persistRefreshTokenInStorage(credentials.refreshToken, credentials.pickleKey);
-        ACCESS_TOKEN_IV,
-        credentials.accessToken,
-        credentials.pickleKey,
-        HAS_ACCESS_TOKEN_STORAGE_KEY,
-    );
-    await persistTokenInStorage(
-        REFRESH_TOKEN_STORAGE_KEY,
-        REFRESH_TOKEN_IV,
-        credentials.refreshToken,
-        credentials.pickleKey,
-        HAS_REFRESH_TOKEN_STORAGE_KEY,
-    );
 
     if (credentials.pickleKey) {
         localStorage.setItem("mx_has_pickle_key", String(true));

src/utils/tokens/tokens.ts

@@ -0,0 +1,202 @@
+/*
+Copyright 2023 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 { decryptAES, encryptAES, IEncryptedPayload } from "matrix-js-sdk/src/crypto/aes";
+import { logger } from "matrix-js-sdk/src/logger";
+
+import * as StorageManager from "../StorageManager";
+
+/**
+ * Utility functions related to the storage and retrieval of access tokens
+ */
+
+/*
+ * Keys used when storing the tokens in indexeddb or localstorage
+ */
+export const ACCESS_TOKEN_STORAGE_KEY = "mx_access_token";
+export const REFRESH_TOKEN_STORAGE_KEY = "mx_refresh_token";
+/*
+ * Used as initialization vector during encryption in persistTokenInStorage
+ * And decryption in restoreFromLocalStorage
+ */
+export const ACCESS_TOKEN_IV = "access_token";
+export const REFRESH_TOKEN_IV = "refresh_token";
+/*
+ * Keys for localstorage items which indicate whether we expect a token in indexeddb.
+ */
+export const HAS_ACCESS_TOKEN_STORAGE_KEY = "mx_has_access_token";
+export const HAS_REFRESH_TOKEN_STORAGE_KEY = "mx_has_refresh_token";
+
+/**
+ * The pickle key is a string of unspecified length and format.  For AES, we need a 256-bit Uint8Array. So we HKDF the pickle key to generate the AES key.  The AES key should be zeroed after it is used.
+ * @param pickleKey
+ * @returns AES key
+ */
+async function pickleKeyToAesKey(pickleKey: string): Promise<Uint8Array> {
+    const pickleKeyBuffer = new Uint8Array(pickleKey.length);
+    for (let i = 0; i < pickleKey.length; i++) {
+        pickleKeyBuffer[i] = pickleKey.charCodeAt(i);
+    }
+    const hkdfKey = await window.crypto.subtle.importKey("raw", pickleKeyBuffer, "HKDF", false, ["deriveBits"]);
+    pickleKeyBuffer.fill(0);
+    return new Uint8Array(
+        await window.crypto.subtle.deriveBits(
+            {
+                name: "HKDF",
+                hash: "SHA-256",
+                // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+                // @ts-ignore: https://github.com/microsoft/TypeScript-DOM-lib-generator/pull/879
+                salt: new Uint8Array(32),
+                info: new Uint8Array(0),
+            },
+            hkdfKey,
+            256,
+        ),
+    );
+}
+
+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
+ */
+export 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
+}
+
+/**
+ * Persist a token in storage
+ * When pickle key is present, will attempt to encrypt the token
+ * Stores in idb, falling back to localStorage
+ *
+ * @param storageKey key used to store the token
+ * @param initializationVector Initialization vector for encrypting the token. Only used when `pickleKey` is present
+ * @param token the token to store, when undefined any existing token at the storageKey is removed from storage
+ * @param pickleKey optional pickle key used to encrypt token
+ * @param hasTokenStorageKey Localstorage key for an item which stores whether we expect to have a token in indexeddb, eg "mx_has_access_token".
+ */
+export async function persistTokenInStorage(
+    storageKey: string,
+    initializationVector: string,
+    token: string | undefined,
+    pickleKey: string | undefined,
+    hasTokenStorageKey: string,
+): Promise<void> {
+    // store whether we expect to find a token, to detect the case
+    // where IndexedDB is blown away
+    if (token) {
+        localStorage.setItem(hasTokenStorageKey, "true");
+    } else {
+        localStorage.removeItem(hasTokenStorageKey);
+    }
+
+    if (pickleKey) {
+        let encryptedToken: IEncryptedPayload | undefined;
+        try {
+            if (!token) {
+                throw new Error("No token: not attempting encryption");
+            }
+            // try to encrypt the access token using the pickle key
+            const encrKey = await pickleKeyToAesKey(pickleKey);
+            encryptedToken = await encryptAES(token, encrKey, initializationVector);
+            encrKey.fill(0);
+        } catch (e) {
+            logger.warn("Could not encrypt access token", e);
+        }
+        try {
+            // save either the encrypted access token, or the plain access
+            // token if we were unable to encrypt (e.g. if the browser doesn't
+            // have WebCrypto).
+            await StorageManager.idbSave("account", storageKey, encryptedToken || token);
+        } catch (e) {
+            // if we couldn't save to indexedDB, fall back to localStorage.  We
+            // store the access token unencrypted since localStorage only saves
+            // strings.
+            if (!!token) {
+                localStorage.setItem(storageKey, token);
+            } else {
+                localStorage.removeItem(storageKey);
+            }
+        }
+    } else {
+        try {
+            await StorageManager.idbSave("account", storageKey, token);
+        } catch (e) {
+            if (!!token) {
+                localStorage.setItem(storageKey, token);
+            } else {
+                localStorage.removeItem(storageKey);
+            }
+        }
+    }
+}
+
+/**
+ * Wraps persistTokenInStorage with accessToken storage keys
+ * @param token the token to store, when undefined any existing accessToken is removed from storage
+ * @param pickleKey optional pickle key used to encrypt token
+ */
+export async function persistAccessTokenInStorage(
+    token: string | undefined,
+    pickleKey: string | undefined,
+): Promise<void> {
+    return persistTokenInStorage(
+        ACCESS_TOKEN_STORAGE_KEY,
+        ACCESS_TOKEN_IV,
+        token,
+        pickleKey,
+        HAS_ACCESS_TOKEN_STORAGE_KEY,
+    );
+}
+
+/**
+ * Wraps persistTokenInStorage with refreshToken storage keys
+ * @param token the token to store, when undefined any existing refreshToken is removed from storage
+ * @param pickleKey optional pickle key used to encrypt token
+ */
+export async function persistRefreshTokenInStorage(
+    token: string | undefined,
+    pickleKey: string | undefined,
+): Promise<void> {
+    return persistTokenInStorage(
+        REFRESH_TOKEN_STORAGE_KEY,
+        REFRESH_TOKEN_IV,
+        token,
+        pickleKey,
+        HAS_REFRESH_TOKEN_STORAGE_KEY,
+    );
+}