OIDC: Token refresher class (#3769)

* add tokenRefresher class

* export generateScope

* export oidc from matrix

* mark experimental

* Apply suggestions from code review

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

* remove some vars in test

* make TokenRefresher un-abstract, comments and improvements

* remove invalid jsdoc

* Update src/oidc/tokenRefresher.ts

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

* Code review improvements

* document TokenRefreshFunction

---------

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

spec/unit/oidc/tokenRefresher.spec.ts

@@ -0,0 +1,270 @@
+/**
+ * @jest-environment jsdom
+ */
+
+/*
+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 fetchMock from "fetch-mock-jest";
+
+import { OidcTokenRefresher } from "../../../src";
+import { logger } from "../../../src/logger";
+import { makeDelegatedAuthConfig } from "../../test-utils/oidc";
+
+describe("OidcTokenRefresher", () => {
+    // OidcTokenRefresher props
+    // see class declaration for info
+    const authConfig = {
+        issuer: "https://issuer.org/",
+    };
+    const clientId = "test-client-id";
+    const redirectUri = "https://test.org";
+    const deviceId = "abc123";
+    const idTokenClaims = {
+        exp: Date.now() / 1000 + 100000,
+        aud: clientId,
+        iss: authConfig.issuer,
+        sub: "123",
+        iat: 123,
+    };
+    // used to mock a valid token response, as consumed by OidcClient library
+    const scope = `openid urn:matrix:org.matrix.msc2967.client:api:* urn:matrix:org.matrix.msc2967.client:device:${deviceId}`;
+
+    // auth config used in mocked calls to OP .well-known
+    const config = makeDelegatedAuthConfig(authConfig.issuer);
+
+    const makeTokenResponse = (accessToken: string, refreshToken?: string) => ({
+        access_token: accessToken,
+        refresh_token: refreshToken,
+        token_type: "Bearer",
+        expires_in: 300,
+        scope: scope,
+    });
+
+    beforeEach(() => {
+        fetchMock.get(`${config.metadata.issuer}.well-known/openid-configuration`, config.metadata);
+        fetchMock.get(`${config.metadata.issuer}jwks`, {
+            status: 200,
+            headers: {
+                "Content-Type": "application/json",
+            },
+            keys: [],
+        });
+
+        fetchMock.post(config.metadata.token_endpoint, {
+            status: 200,
+            headers: {
+                "Content-Type": "application/json",
+            },
+            ...makeTokenResponse("new-access-token", "new-refresh-token"),
+        });
+    });
+
+    afterEach(() => {
+        jest.restoreAllMocks();
+        fetchMock.resetBehavior();
+    });
+
+    it("throws when oidc client cannot be initialised", async () => {
+        jest.spyOn(logger, "error");
+        fetchMock.get(
+            `${config.metadata.issuer}.well-known/openid-configuration`,
+            {
+                ok: false,
+                status: 404,
+            },
+            { overwriteRoutes: true },
+        );
+        const refresher = new OidcTokenRefresher(authConfig, clientId, redirectUri, deviceId, idTokenClaims);
+        await expect(refresher.oidcClientReady).rejects.toThrow();
+        expect(logger.error).toHaveBeenCalledWith(
+            "Failed to initialise OIDC client.",
+            // error from OidcClient
+            expect.any(Error),
+        );
+    });
+
+    it("initialises oidc client", async () => {
+        const refresher = new OidcTokenRefresher(authConfig, clientId, redirectUri, deviceId, idTokenClaims);
+        await refresher.oidcClientReady;
+
+        // @ts-ignore peek at private property to see we initialised the client correctly
+        expect(refresher.oidcClient.settings).toEqual(
+            expect.objectContaining({
+                client_id: clientId,
+                redirect_uri: redirectUri,
+                authority: authConfig.issuer,
+                scope,
+            }),
+        );
+    });
+
+    describe("doRefreshAccessToken()", () => {
+        it("should throw when oidcClient has not been initialised", async () => {
+            const refresher = new OidcTokenRefresher(authConfig, clientId, redirectUri, deviceId, idTokenClaims);
+            await expect(refresher.doRefreshAccessToken("token")).rejects.toThrow(
+                "Cannot get new token before OIDC client is initialised.",
+            );
+        });
+
+        it("should refresh the tokens", async () => {
+            const refresher = new OidcTokenRefresher(authConfig, clientId, redirectUri, deviceId, idTokenClaims);
+            await refresher.oidcClientReady;
+
+            const result = await refresher.doRefreshAccessToken("refresh-token");
+
+            expect(fetchMock).toHaveFetched(config.metadata.token_endpoint, {
+                method: "POST",
+            });
+
+            expect(result).toEqual({
+                accessToken: "new-access-token",
+                refreshToken: "new-refresh-token",
+            });
+        });
+
+        it("should persist the new tokens", async () => {
+            const refresher = new OidcTokenRefresher(authConfig, clientId, redirectUri, deviceId, idTokenClaims);
+            await refresher.oidcClientReady;
+            // spy on our stub
+            jest.spyOn(refresher, "persistTokens");
+
+            await refresher.doRefreshAccessToken("refresh-token");
+
+            expect(refresher.persistTokens).toHaveBeenCalledWith({
+                accessToken: "new-access-token",
+                refreshToken: "new-refresh-token",
+            });
+        });
+
+        it("should only have one inflight refresh request at once", async () => {
+            fetchMock
+                .postOnce(
+                    config.metadata.token_endpoint,
+                    {
+                        status: 200,
+                        headers: {
+                            "Content-Type": "application/json",
+                        },
+                        ...makeTokenResponse("first-new-access-token", "first-new-refresh-token"),
+                    },
+                    { overwriteRoutes: true },
+                )
+                .postOnce(
+                    config.metadata.token_endpoint,
+                    {
+                        status: 200,
+                        headers: {
+                            "Content-Type": "application/json",
+                        },
+                        ...makeTokenResponse("second-new-access-token", "second-new-refresh-token"),
+                    },
+                    { overwriteRoutes: false },
+                );
+
+            const refresher = new OidcTokenRefresher(authConfig, clientId, redirectUri, deviceId, idTokenClaims);
+            await refresher.oidcClientReady;
+            // reset call counts
+            fetchMock.resetHistory();
+
+            const refreshToken = "refresh-token";
+            const first = refresher.doRefreshAccessToken(refreshToken);
+            const second = refresher.doRefreshAccessToken(refreshToken);
+
+            const result1 = await second;
+            const result2 = await first;
+
+            // only one call to token endpoint
+            expect(fetchMock).toHaveFetchedTimes(1, config.metadata.token_endpoint);
+            expect(result1).toEqual({
+                accessToken: "first-new-access-token",
+                refreshToken: "first-new-refresh-token",
+            });
+            // same response
+            expect(result1).toEqual(result2);
+
+            // call again after first request resolves
+            const third = await refresher.doRefreshAccessToken("first-new-refresh-token");
+
+            // called token endpoint, got new tokens
+            expect(third).toEqual({
+                accessToken: "second-new-access-token",
+                refreshToken: "second-new-refresh-token",
+            });
+        });
+
+        it("should log and rethrow when token refresh fails", async () => {
+            fetchMock.post(
+                config.metadata.token_endpoint,
+                {
+                    status: 503,
+                    headers: {
+                        "Content-Type": "application/json",
+                    },
+                },
+                { overwriteRoutes: true },
+            );
+
+            const refresher = new OidcTokenRefresher(authConfig, clientId, redirectUri, deviceId, idTokenClaims);
+            await refresher.oidcClientReady;
+
+            await expect(refresher.doRefreshAccessToken("refresh-token")).rejects.toThrow();
+        });
+
+        it("should make fresh request after a failed request", async () => {
+            // make sure inflight request is cleared after a failure
+            fetchMock
+                .postOnce(
+                    config.metadata.token_endpoint,
+                    {
+                        status: 503,
+                        headers: {
+                            "Content-Type": "application/json",
+                        },
+                    },
+                    { overwriteRoutes: true },
+                )
+                .postOnce(
+                    config.metadata.token_endpoint,
+                    {
+                        status: 200,
+                        headers: {
+                            "Content-Type": "application/json",
+                        },
+                        ...makeTokenResponse("second-new-access-token", "second-new-refresh-token"),
+                    },
+                    { overwriteRoutes: false },
+                );
+
+            const refresher = new OidcTokenRefresher(authConfig, clientId, redirectUri, deviceId, idTokenClaims);
+            await refresher.oidcClientReady;
+            // reset call counts
+            fetchMock.resetHistory();
+
+            // first call fails
+            await expect(refresher.doRefreshAccessToken("refresh-token")).rejects.toThrow();
+
+            // call again after first request resolves
+            const result = await refresher.doRefreshAccessToken("first-new-refresh-token");
+
+            // called token endpoint, got new tokens
+            expect(result).toEqual({
+                accessToken: "second-new-access-token",
+                refreshToken: "second-new-refresh-token",
+            });
+        });
+    });
+});

src/http-api/interface.ts

@@ -18,6 +18,21 @@ import { MatrixError } from "./errors";
 
 export type Body = Record<string, any> | BodyInit;
 
+/**
+ * @experimental
+ * Unencrypted access and (optional) refresh token
+ */
+export type AccessTokens = {
+    accessToken: string;
+    refreshToken?: string;
+};
+// @TODO(kerrya) add link to IHttpOpts and CreateClientOpts when token refresh is added there
+/**
+ * @experimental
+ * Function that performs token refresh using the given refreshToken.
+ * Returns a promise that resolves to the refreshed access and (optional) refresh tokens.
+ */
+export type TokenRefreshFunction = (refreshToken: string) => Promise<AccessTokens>;
 export interface IHttpOpts {
     fetchFn?: typeof global.fetch;
 

src/matrix.ts

@@ -42,6 +42,7 @@ export * from "./models/typed-event-emitter";
 export * from "./models/user";
 export * from "./models/device";
 export * from "./models/search-result";
+export * from "./oidc";
 export * from "./scheduler";
 export * from "./filter";
 export * from "./timeline-window";

src/oidc/authorize.ts

@@ -51,9 +51,9 @@ export type AuthorizationParams = {
  * Generate the scope used in authorization request with OIDC OP
  * @returns scope
  */
-const generateScope = (): string => {
-    const deviceId = randomString(10);
-    return `openid urn:matrix:org.matrix.msc2967.client:api:* urn:matrix:org.matrix.msc2967.client:device:${deviceId}`;
+export const generateScope = (deviceId?: string): string => {
+    const safeDeviceId = deviceId ?? randomString(10);
+    return `openid urn:matrix:org.matrix.msc2967.client:api:* urn:matrix:org.matrix.msc2967.client:device:${safeDeviceId}`;
 };
 
 // https://www.rfc-editor.org/rfc/rfc7636

src/oidc/index.ts

@@ -0,0 +1,17 @@
+/*
+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.
+*/
+
+export * from "./tokenRefresher";

src/oidc/tokenRefresher.ts

@@ -0,0 +1,150 @@
+/*
+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 { IdTokenClaims, OidcClient, WebStorageStateStore } from "oidc-client-ts";
+
+import { AccessTokens } from "../http-api";
+import { IDelegatedAuthConfig } from "../client";
+import { generateScope } from "./authorize";
+import { discoverAndValidateAuthenticationConfig } from "./discovery";
+import { logger } from "../logger";
+
+/**
+ * @experimental
+ * Class responsible for refreshing OIDC access tokens
+ *
+ * Client implementations will likely want to override {@link persistTokens} to persist tokens after successful refresh
+ *
+ */
+export class OidcTokenRefresher {
+    /**
+     * Promise which will complete once the OidcClient has been initialised
+     * and is ready to start refreshing tokens.
+     *
+     * Will reject if the client initialisation fails.
+     */
+    public readonly oidcClientReady!: Promise<void>;
+    private oidcClient!: OidcClient;
+    private inflightRefreshRequest?: Promise<AccessTokens>;
+
+    public constructor(
+        /**
+         * Delegated auth config as found in matrix client .well-known
+         */
+        authConfig: IDelegatedAuthConfig,
+        /**
+         * id of this client as registered with the OP
+         */
+        clientId: string,
+        /**
+         * redirectUri as registered with OP
+         */
+        redirectUri: string,
+        /**
+         * Device ID of current session
+         */
+        deviceId: string,
+        /**
+         * idTokenClaims as returned from authorization grant
+         * used to validate tokens
+         */
+        private readonly idTokenClaims: IdTokenClaims,
+    ) {
+        this.oidcClientReady = this.initialiseOidcClient(authConfig, clientId, deviceId, redirectUri);
+    }
+
+    private async initialiseOidcClient(
+        authConfig: IDelegatedAuthConfig,
+        clientId: string,
+        deviceId: string,
+        redirectUri: string,
+    ): Promise<void> {
+        try {
+            const config = await discoverAndValidateAuthenticationConfig(authConfig);
+
+            const scope = generateScope(deviceId);
+
+            this.oidcClient = new OidcClient({
+                ...config.metadata,
+                client_id: clientId,
+                scope,
+                redirect_uri: redirectUri,
+                authority: config.metadata.issuer,
+                stateStore: new WebStorageStateStore({ prefix: "mx_oidc_", store: window.sessionStorage }),
+            });
+        } catch (error) {
+            logger.error("Failed to initialise OIDC client.", error);
+            throw new Error("Failed to initialise OIDC client.");
+        }
+    }
+
+    /**
+     * Attempt token refresh using given refresh token
+     * @param refreshToken - refresh token to use in request with token issuer
+     * @returns tokens - Promise that resolves with new access and refresh tokens
+     * @throws when token refresh fails
+     */
+    public async doRefreshAccessToken(refreshToken: string): Promise<AccessTokens> {
+        if (!this.inflightRefreshRequest) {
+            this.inflightRefreshRequest = this.getNewTokens(refreshToken);
+        }
+        try {
+            const tokens = await this.inflightRefreshRequest;
+            return tokens;
+        } finally {
+            this.inflightRefreshRequest = undefined;
+        }
+    }
+
+    /**
+     * Persist the new tokens, called after tokens are successfully refreshed.
+     *
+     * This function is intended to be overriden by the consumer when persistence is necessary.
+     *
+     * @param accessToken - new access token
+     * @param refreshToken - OPTIONAL new refresh token
+     */
+    public async persistTokens(_tokens: { accessToken: string; refreshToken?: string }): Promise<void> {
+        // NOOP
+    }
+
+    private async getNewTokens(refreshToken: string): Promise<AccessTokens> {
+        if (!this.oidcClient) {
+            throw new Error("Cannot get new token before OIDC client is initialised.");
+        }
+
+        const refreshTokenState = {
+            refresh_token: refreshToken,
+            session_state: "test",
+            data: undefined,
+            profile: this.idTokenClaims,
+        };
+
+        const response = await this.oidcClient.useRefreshToken({
+            state: refreshTokenState,
+            timeoutInSeconds: 300,
+        });
+
+        const tokens = {
+            accessToken: response.access_token,
+            refreshToken: response.refresh_token,
+        };
+
+        await this.persistTokens(tokens);
+
+        return tokens;
+    }
+}