#6284: Generic PKCE refresh token function

src/background/refreshGoogleToken.test.ts → src/background/refreshPKCEToken.test.ts

@@ -1,4 +1,4 @@
 /*
  * Copyright (C) 2023 PixieBrix, Inc.
  *
  * This program is free software: you can redistribute it and/or modify
@@ -15,9 +15,8 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import refreshGoogleToken from "@/background/refreshGoogleToken";
+import refreshPKCEToken from "@/background/refreshPKCEToken";
 import { getCachedAuthData, setCachedAuthData } from "@/background/auth";
-import { GOOGLE_OAUTH_PKCE_INTEGRATION_ID } from "@/services/constants";
 import { appApiMock } from "@/testUtils/appApiMock";
 import { sanitizedIntegrationConfigFactory } from "@/testUtils/factories/integrationFactories";
 import { type IntegrationConfig } from "@/types/integrationTypes";
@@ -26,6 +25,8 @@
 import googleDefinition from "@contrib/integrations/google-oauth2-pkce.yaml";
 import { locator } from "@/background/locator";
 
+const PKCE_INTEGRATION_ID = (googleDefinition.metadata as any).id;
+
 jest.mock("@/background/auth", () => ({
   getCachedAuthData: jest.fn().mockRejectedValue(new Error("Not mocked")),
   setCachedAuthData: jest.fn(),
@@ -45,7 +46,7 @@
 jest.mocked(registry.find).mockImplementation(
   async () =>
     ({
-      id: (googleDefinition.metadata as any).id,
+      id: PKCE_INTEGRATION_ID,
       config: googleDefinition,
     } as any)
 );
@@ -54,7 +55,7 @@
 const setCachedAuthDataMock = jest.mocked(setCachedAuthData);
 const readRawConfigurationsMock = jest.mocked(readRawConfigurations);
 
-describe("refresh partner token", () => {
+describe("refresh PKCE token", () => {
   beforeEach(() => {
     getCachedAuthDataMock.mockReset();
     setCachedAuthDataMock.mockReset();
@@ -63,42 +64,42 @@
     appApiMock.resetHistory();
   });
 
-  it("throws if integration configuration is not google pkce", async () => {
+  it("throws if integration configuration is not pkce", async () => {
     const integrationConfig = sanitizedIntegrationConfigFactory();
 
-    await expect(refreshGoogleToken(integrationConfig)).rejects.toThrow(Error);
+    await expect(refreshPKCEToken(integrationConfig)).rejects.toThrow(Error);
   });
 
   it("nop if no cached auth data", async () => {
     const integrationConfig = sanitizedIntegrationConfigFactory({
-      serviceId: GOOGLE_OAUTH_PKCE_INTEGRATION_ID,
+      serviceId: PKCE_INTEGRATION_ID,
     });
 
-    const isTokenRefreshed = await refreshGoogleToken(integrationConfig);
+    const isTokenRefreshed = await refreshPKCEToken(integrationConfig);
 
     expect(isTokenRefreshed).toBe(false);
     expect(getCachedAuthDataMock).toHaveBeenCalledOnce();
   });
 
   it("nop if no refresh token", async () => {
     const integrationConfig = sanitizedIntegrationConfigFactory({
-      serviceId: GOOGLE_OAUTH_PKCE_INTEGRATION_ID,
+      serviceId: PKCE_INTEGRATION_ID,
     });
 
     getCachedAuthDataMock.mockResolvedValue({
       access_token: "notatoken",
       _oauthBrand: undefined,
     });
 
-    const isTokenRefreshed = await refreshGoogleToken(integrationConfig);
+    const isTokenRefreshed = await refreshPKCEToken(integrationConfig);
 
     expect(isTokenRefreshed).toBe(false);
     expect(appApiMock.history.post).toHaveLength(0);
   });
 
-  it("refreshes token", async () => {
+  it("refreshes token with refresh token in response", async () => {
     const integrationConfig = sanitizedIntegrationConfigFactory({
-      serviceId: GOOGLE_OAUTH_PKCE_INTEGRATION_ID,
+      serviceId: PKCE_INTEGRATION_ID,
     });
 
     getCachedAuthDataMock.mockResolvedValue({
@@ -121,7 +122,7 @@
       refresh_token: "notarefreshtoken2",
     });
 
-    const isTokenRefreshed = await refreshGoogleToken(integrationConfig);
+    const isTokenRefreshed = await refreshPKCEToken(integrationConfig);
 
     expect(isTokenRefreshed).toBe(true);
     expect(appApiMock.history.post).toHaveLength(1);
@@ -131,9 +132,43 @@
     });
   });
 
+  it("refreshes token without refresh token in response", async () => {
+    const integrationConfig = sanitizedIntegrationConfigFactory({
+      serviceId: PKCE_INTEGRATION_ID,
+    });
+
+    getCachedAuthDataMock.mockResolvedValue({
+      access_token: "notatoken",
+      refresh_token: "notarefreshtoken",
+      _oauthBrand: undefined,
+    });
+
+    readRawConfigurationsMock.mockResolvedValue([
+      {
+        id: integrationConfig.id,
+        config: {},
+      } as IntegrationConfig,
+    ]);
+    await locator.refreshLocal();
+
+    appApiMock.onGet("/api/services/shared/").reply(200, []);
+    appApiMock.onPost().reply(200, {
+      access_token: "notatoken2",
+    });
+
+    const isTokenRefreshed = await refreshPKCEToken(integrationConfig);
+
+    expect(isTokenRefreshed).toBe(true);
+    expect(appApiMock.history.post).toHaveLength(1);
+    expect(setCachedAuthDataMock).toHaveBeenCalledWith(integrationConfig.id, {
+      access_token: "notatoken2",
+      refresh_token: "notarefreshtoken",
+    });
+  });
+
   it("throws on authorization error", async () => {
     const integrationConfig = sanitizedIntegrationConfigFactory({
-      serviceId: GOOGLE_OAUTH_PKCE_INTEGRATION_ID,
+      serviceId: PKCE_INTEGRATION_ID,
     });
 
     getCachedAuthDataMock.mockResolvedValue({
@@ -153,7 +188,7 @@
     appApiMock.onGet("/api/services/shared/").reply(200, []);
     appApiMock.onPost().reply(401);
 
-    await expect(refreshGoogleToken(integrationConfig)).rejects.toThrow(
+    await expect(refreshPKCEToken(integrationConfig)).rejects.toThrow(
       "Request failed with status code 401"
     );
   });

src/background/refreshGoogleToken.ts → src/background/refreshPKCEToken.ts

@@ -1,4 +1,4 @@
 /*
  * Copyright (C) 2023 PixieBrix, Inc.
  *
  * This program is free software: you can redistribute it and/or modify
@@ -19,54 +19,64 @@
 import { expectContext } from "@/utils/expectContext";
 import serviceRegistry from "@/services/registry";
 import { type SanitizedIntegrationConfig } from "@/types/integrationTypes";
-import { GOOGLE_OAUTH_PKCE_INTEGRATION_ID } from "@/services/constants";
+import { OAUTH_PKCE_INTEGRATION_IDS } from "@/services/constants";
 import { getCachedAuthData, setCachedAuthData } from "@/background/auth";
 import { locator as serviceLocator } from "@/background/locator";
 
 /**
- * Refresh a Google OAuth2 PKCE token. NOOP if a refresh token is not available.
+ * Refresh an OAuth2 PKCE token. NOOP if a refresh token is not available.
  * @returns True if the token was successfully refreshed. False if the token refresh was not attempted.
- * @throws AxiosError if the token refresh failed or Error if the integration is not a Google OAuth2 PKCE integration.
+ * @throws AxiosError if the token refresh failed or Error if the integration is not an OAuth2 PKCE integration.
  */
-export default async function refreshGoogleToken(
+export default async function refreshPKCEToken(
   integrationConfig: SanitizedIntegrationConfig
 ): Promise<boolean> {
   expectContext("background");
 
-  if (integrationConfig.serviceId !== GOOGLE_OAUTH_PKCE_INTEGRATION_ID) {
+  const integrationId = integrationConfig.serviceId;
+
+  // Instead of hardcoding the list, we could check the integration definition for a "code_challenge_method" field.
+  // If it exists, that means it's a PKCE integration. See isOAuth2PKCE in the pixiebrix-app repo.
+  if (!OAUTH_PKCE_INTEGRATION_IDS.includes(integrationId)) {
     throw new Error(
-      `Expected integration to be ${GOOGLE_OAUTH_PKCE_INTEGRATION_ID}, but got ${integrationConfig.serviceId}`
+      `Expected OAuth2 PKCE integration, but got ${integrationConfig.serviceId}`
     );
   }
 
   const cachedAuthData = await getCachedAuthData(integrationConfig.id);
 
   if (integrationConfig.id && cachedAuthData?.refresh_token) {
-    console.debug("Refreshing google token");
+    console.debug("Refreshing PKCE token");
 
-    const integration = await serviceRegistry.lookup(
-      GOOGLE_OAUTH_PKCE_INTEGRATION_ID
-    );
+    const integration = await serviceRegistry.lookup(integrationId);
     const { config } = await serviceLocator.findIntegrationConfig(
       integrationConfig.id
     );
-    const context = integration.getOAuth2Context(config);
+    const { tokenUrl, client_id, client_secret } =
+      integration.getOAuth2Context(config);
 
     // https://axios-http.com/docs/urlencoded
     const params = new URLSearchParams();
+
     params.append("grant_type", "refresh_token");
     params.append("refresh_token", cachedAuthData.refresh_token as string);
-    params.append("client_id", context.client_id);
-    params.append("client_secret", context.client_secret);
+    params.append("client_id", client_id);
+
+    if (client_secret) {
+      params.append("client_secret", client_secret);
+    }
 
-    const { data } = await axios.post(context.tokenUrl, params);
+    const { data } = await axios.post(tokenUrl, params);
 
-    // The Google refresh token response doesn't include the refresh token. Let's re-add it, so it doesn't get removed.
+    // Add the cached refresh token to the response if it's missing, so it doesn't get removed.
+    // - The Google refresh token response doesn't include a refresh token.
+    // - The Azure refresh token response includes a new refresh token that we should replace the cached one with.
+    //   See https://learn.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-auth-code-flow#refresh-the-access-token
     data.refresh_token ??= cachedAuthData.refresh_token;
 
     await setCachedAuthData(integrationConfig.id, data);
 
-    console.debug("Successfully refreshed google token");
+    console.debug("Successfully refreshed PKCE token");
 
     return true;
   }

src/background/requests.ts

@@ -35,7 +35,7 @@ import {
 import { expectContext } from "@/utils/expectContext";
 import { absoluteApiUrl } from "@/services/apiClient";
 import {
-  GOOGLE_OAUTH_PKCE_INTEGRATION_ID,
+  OAUTH_PKCE_INTEGRATION_IDS,
   PIXIEBRIX_INTEGRATION_ID,
 } from "@/services/constants";
 import { type ProxyResponseData, type RemoteResponse } from "@/types/contract";
@@ -63,7 +63,7 @@ import {
   type SecretsConfig,
 } from "@/types/integrationTypes";
 import { type MessageContext } from "@/types/loggerTypes";
-import refreshGoogleToken from "@/background/refreshGoogleToken";
+import refreshPKCEToken from "@/background/refreshPKCEToken";
 import reportError from "@/telemetry/reportError";
 import { isAbsoluteUrl } from "@/utils/urlUtils";
 import { isObject } from "@/utils/objectUtils";
@@ -286,20 +286,17 @@ async function performConfiguredRequest(
     if (axiosError && isAuthenticationError(axiosError)) {
       const service = await serviceRegistry.lookup(serviceConfig.serviceId);
       if (service.isOAuth2 || service.isToken) {
-        if (service.id === GOOGLE_OAUTH_PKCE_INTEGRATION_ID) {
+        if (OAUTH_PKCE_INTEGRATION_IDS.includes(service.id)) {
           try {
-            const isTokenRefreshed = await refreshGoogleToken(serviceConfig);
+            const isTokenRefreshed = await refreshPKCEToken(serviceConfig);
 
             if (isTokenRefreshed) {
               return serializableAxiosRequest(
                 await authenticate(serviceConfig, requestConfig)
               );
             }
           } catch (error) {
-            console.warn(
-              `Failed to refresh ${GOOGLE_OAUTH_PKCE_INTEGRATION_ID} token:`,
-              error
-            );
+            console.warn(`Failed to refresh ${service.id} token:`, error);
 
             // An authentication error can occur if the refresh token was revoked. Besides that, there should be
             // no reason for the refresh to fail. Report the error if it's not an authentication error.

src/services/constants.ts

@@ -34,5 +34,14 @@ export const CONTROL_ROOM_OAUTH_INTEGRATION_ID: RegistryId = validateRegistryId(
   "automation-anywhere/oauth2"
 );
 
-export const GOOGLE_OAUTH_PKCE_INTEGRATION_ID: RegistryId =
+const GOOGLE_OAUTH_PKCE_INTEGRATION_ID: RegistryId =
   validateRegistryId("google/oauth2-pkce");
+
+const MICROSOFT_OAUTH_PKCE_INTEGRATION_ID: RegistryId = validateRegistryId(
+  "microsoft/oauth2-pkce"
+);
+
+export const OAUTH_PKCE_INTEGRATION_IDS = [
+  GOOGLE_OAUTH_PKCE_INTEGRATION_ID,
+  MICROSOFT_OAUTH_PKCE_INTEGRATION_ID,
+];