OIDC: refresh tokens (#11699)

* 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

* very messy poc

* comments

* prettier

* comment pedantry

* working refresh without persistence

* extract token persistence functions to utils

* add sugar

* implement TokenRefresher class with persistence

* tidying

* persist idTokenClaims

* persist idTokenClaims

* tests

* remove unused cde

* create token refresher during doSetLoggedIn

* tidying

* also tidying

* update Lifecycle test replaceUsingCreds calls

* tidy

* test tokenrefresher creation in login flow

* test token refresher

* Update src/utils/oidc/TokenRefresher.ts

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* use literal value for m.authentication

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* improve comments

---------

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

src/Lifecycle.ts

@@ -18,7 +18,7 @@ limitations under the License.
 */
 
 import { ReactNode } from "react";
-import { createClient, MatrixClient, SSOAction } from "matrix-js-sdk/src/matrix";
+import { createClient, MatrixClient, SSOAction, OidcTokenRefresher } from "matrix-js-sdk/src/matrix";
 import { InvalidStoreError } from "matrix-js-sdk/src/errors";
 import { IEncryptedPayload } from "matrix-js-sdk/src/crypto/aes";
 import { QueryDict } from "matrix-js-sdk/src/utils";
@@ -65,7 +65,12 @@ import { OverwriteLoginPayload } from "./dispatcher/payloads/OverwriteLoginPaylo
 import { SdkContextClass } from "./contexts/SDKContext";
 import { messageForLoginError } from "./utils/ErrorUtils";
 import { completeOidcLogin } from "./utils/oidc/authorize";
-import { persistOidcAuthenticatedSettings } from "./utils/oidc/persistOidcSettings";
+import {
+    getStoredOidcClientId,
+    getStoredOidcIdTokenClaims,
+    getStoredOidcTokenIssuer,
+    persistOidcAuthenticatedSettings,
+} from "./utils/oidc/persistOidcSettings";
 import GenericToast from "./components/views/toasts/GenericToast";
 import {
     ACCESS_TOKEN_IV,
@@ -78,6 +83,7 @@ import {
     REFRESH_TOKEN_STORAGE_KEY,
     tryDecryptToken,
 } from "./utils/tokens/tokens";
+import { TokenRefresher } from "./utils/oidc/TokenRefresher";
 
 const HOMESERVER_URL_KEY = "mx_hs_url";
 const ID_SERVER_URL_KEY = "mx_is_url";
@@ -746,6 +752,45 @@ export async function hydrateSession(credentials: IMatrixClientCreds): Promise<M
     return doSetLoggedIn(credentials, overwrite);
 }
 
+/**
+ * When we have a authenticated via OIDC-native flow and have a refresh token
+ * try to create a token refresher.
+ * @param credentials from current session
+ * @returns Promise that resolves to a TokenRefresher, or undefined
+ */
+async function createOidcTokenRefresher(credentials: IMatrixClientCreds): Promise<OidcTokenRefresher | undefined> {
+    if (!credentials.refreshToken) {
+        return;
+    }
+    // stored token issuer indicates we authenticated via OIDC-native flow
+    const tokenIssuer = getStoredOidcTokenIssuer();
+    if (!tokenIssuer) {
+        return;
+    }
+    try {
+        const clientId = getStoredOidcClientId();
+        const idTokenClaims = getStoredOidcIdTokenClaims();
+        const redirectUri = window.location.origin;
+        const deviceId = credentials.deviceId;
+        if (!deviceId) {
+            throw new Error("Expected deviceId in user credentials.");
+        }
+        const tokenRefresher = new TokenRefresher(
+            { issuer: tokenIssuer },
+            clientId,
+            redirectUri,
+            deviceId,
+            idTokenClaims!,
+            credentials.userId,
+        );
+        // wait for the OIDC client to initialise
+        await tokenRefresher.oidcClientReady;
+        return tokenRefresher;
+    } catch (error) {
+        logger.error("Failed to initialise OIDC token refresher", error);
+    }
+}
+
 /**
  * optionally clears localstorage, persists new credentials
  * to localstorage, starts the new client.
@@ -787,9 +832,11 @@ async function doSetLoggedIn(credentials: IMatrixClientCreds, clearStorageEnable
         await abortLogin();
     }
 
+    const tokenRefresher = await createOidcTokenRefresher(credentials);
+
     // check the session lock just before creating the new client
     checkSessionLock();
-    MatrixClientPeg.replaceUsingCreds(credentials);
+    MatrixClientPeg.replaceUsingCreds(credentials, tokenRefresher?.doRefreshAccessToken.bind(tokenRefresher));
     const client = MatrixClientPeg.safeGet();
 
     setSentryUser(credentials.userId);

src/MatrixClientPeg.ts

@@ -27,6 +27,7 @@ import {
     IStartClientOpts,
     MatrixClient,
     MemoryStore,
+    TokenRefreshFunction,
 } from "matrix-js-sdk/src/matrix";
 import * as utils from "matrix-js-sdk/src/utils";
 import { verificationMethods } from "matrix-js-sdk/src/crypto";
@@ -122,8 +123,10 @@ export interface IMatrixClientPeg {
      * homeserver / identity server URLs and active credentials
      *
      * @param {IMatrixClientCreds} creds The new credentials to use.
+     * @param {TokenRefreshFunction} tokenRefreshFunction OPTIONAL function used by MatrixClient to attempt token refresh
+     *          see {@link ICreateClientOpts.tokenRefreshFunction}
      */
-    replaceUsingCreds(creds: IMatrixClientCreds): void;
+    replaceUsingCreds(creds: IMatrixClientCreds, tokenRefreshFunction?: TokenRefreshFunction): void;
 }
 
 /**
@@ -196,8 +199,8 @@ class MatrixClientPegClass implements IMatrixClientPeg {
         }
     }
 
-    public replaceUsingCreds(creds: IMatrixClientCreds): void {
-        this.createClient(creds);
+    public replaceUsingCreds(creds: IMatrixClientCreds, tokenRefreshFunction?: TokenRefreshFunction): void {
+        this.createClient(creds, tokenRefreshFunction);
     }
 
     private onUnexpectedStoreClose = async (): Promise<void> => {
@@ -378,11 +381,13 @@ class MatrixClientPegClass implements IMatrixClientPeg {
         });
     }
 
-    private createClient(creds: IMatrixClientCreds): void {
+    private createClient(creds: IMatrixClientCreds, tokenRefreshFunction?: TokenRefreshFunction): void {
         const opts: ICreateClientOpts = {
             baseUrl: creds.homeserverUrl,
             idBaseUrl: creds.identityServerUrl,
             accessToken: creds.accessToken,
+            refreshToken: creds.refreshToken,
+            tokenRefreshFunction,
             userId: creds.userId,
             deviceId: creds.deviceId,
             pickleKey: creds.pickleKey,

src/utils/oidc/TokenRefresher.ts

@@ -0,0 +1,47 @@
+/*
+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 { IDelegatedAuthConfig, OidcTokenRefresher, AccessTokens } from "matrix-js-sdk/src/matrix";
+import { IdTokenClaims } from "oidc-client-ts";
+
+import PlatformPeg from "../../PlatformPeg";
+import { persistAccessTokenInStorage, persistRefreshTokenInStorage } from "../tokens/tokens";
+
+/**
+ * OidcTokenRefresher that implements token persistence.
+ * Stores tokens in the same way as login flow in Lifecycle.
+ */
+export class TokenRefresher extends OidcTokenRefresher {
+    private readonly deviceId!: string;
+
+    public constructor(
+        authConfig: IDelegatedAuthConfig,
+        clientId: string,
+        redirectUri: string,
+        deviceId: string,
+        idTokenClaims: IdTokenClaims,
+        private readonly userId: string,
+    ) {
+        super(authConfig, clientId, deviceId, redirectUri, idTokenClaims);
+        this.deviceId = deviceId;
+    }
+
+    public async persistTokens({ accessToken, refreshToken }: AccessTokens): Promise<void> {
+        const pickleKey = (await PlatformPeg.get()?.getPickleKey(this.userId, this.deviceId)) ?? undefined;
+        await persistAccessTokenInStorage(accessToken, pickleKey);
+        await persistRefreshTokenInStorage(refreshToken, pickleKey);
+    }
+}

src/utils/oidc/persistOidcSettings.ts

@@ -57,3 +57,15 @@ export const getStoredOidcClientId = (): string => {
     }
     return clientId;
 };
+
+/**
+ * Retrieve stored id token claims from session storage
+ * @returns idtokenclaims or undefined
+ */
+export const getStoredOidcIdTokenClaims = (): IdTokenClaims | undefined => {
+    const idTokenClaims = sessionStorage.getItem(idTokenClaimsStorageKey);
+    if (!idTokenClaims) {
+        return;
+    }
+    return JSON.parse(idTokenClaims) as IdTokenClaims;
+};