#6656 Partner refresh tokens

.secrets.baseline

@@ -253,16 +253,16 @@
         "filename": "src/testUtils/factories/authFactories.ts",
         "hashed_secret": "d15e5a27160ace913d22871497c05a7e5bbbe2ef",
         "is_verified": false,
-        "line_number": 157
+        "line_number": 158
       },
       {
         "type": "Hex High Entropy String",
         "filename": "src/testUtils/factories/authFactories.ts",
         "hashed_secret": "ff998abc1ce6d8f01a675fa197368e44c8916e9c",
         "is_verified": false,
-        "line_number": 172
+        "line_number": 173
       }
     ]
   },
-  "generated_at": "2024-06-11T00:08:01Z"
+  "generated_at": "2024-06-20T16:25:50Z"
 }

eslint-local-rules/persistBackgroundData.txt

@@ -220,7 +220,6 @@
 ./src/starterBricks/factory.ts
 ./src/starterBricks/helpers.ts
 ./src/starterBricks/menuItemExtension.ts
-./src/starterBricks/panelExtension.ts
 ./src/starterBricks/quickBarExtension.tsx
 ./src/starterBricks/quickBarProviderExtension.tsx
 ./src/starterBricks/quickbarQueryReader.ts

package.json

@@ -59,6 +59,7 @@
     "ace-builds": "^1.35.0",
     "autocompleter": "^9.2.1",
     "axios": "^0.28.1",
+    "axios-auth-refresh": "^3.3.6",
     "batched-function": "^2.0.1",
     "bootstrap": "^4.6.0",
     "bootstrap-icons": "^1.11.3",

src/__mocks__/@/background/messenger/api.ts

@@ -54,3 +54,5 @@ export const performConfiguredRequestInBackground = getMethod(
 ) => Promise<RemoteResponse<TData>>;
 
 export const clearServiceCache = jest.fn();
+
+export const removeOAuth2Token = jest.fn();

src/auth/authStorage.ts

@@ -25,18 +25,17 @@ import {
 } from "./authTypes";
 import { isExtensionContext } from "webext-detect-page";
 import { expectContext } from "@/utils/expectContext";
-import { isEmpty, omit } from "lodash";
+import { omit } from "lodash";
 import { syncRemotePackages } from "@/registry/memoryRegistry";
 import { StorageItem } from "webext-storage";
 import { SimpleEventTarget } from "@/utils/SimpleEventTarget";
 import { ReusableAbortController } from "abort-utils";
+import { removeOAuth2Token } from "@/background/messenger/api";
 
 const extensionKeyStorage = new StorageItem("extensionKey", {
   defaultValue: {} as Partial<TokenAuthData>,
 });
-const partnerTokenStorage = new StorageItem("partnerToken", {
-  defaultValue: {} as Partial<PartnerAuthData>,
-});
+const partnerTokenStorage = new StorageItem<PartnerAuthData>("partnerToken");
 
 type AuthListener = (auth: Partial<TokenAuthData | PartnerAuthData>) => void;
 
@@ -48,11 +47,11 @@ const authChanges = new SimpleEventTarget<
 >();
 
 // Use listeners to allow inversion of control and avoid circular dependency with error reporter.
-export function addListener(handler: AuthListener): void {
+export function addAuthListener(handler: AuthListener): void {
   authChanges.add(handler, { signal: controller.signal });
 }
 
-export function removeListener(handler: AuthListener): void {
+export function removeAuthListener(handler: AuthListener): void {
   authChanges.remove(handler);
 }
 
@@ -93,17 +92,29 @@ export async function getExtensionToken(): Promise<string | undefined> {
   return token;
 }
 
-export async function getPartnerAuthData(): Promise<Partial<PartnerAuthData>> {
-  return partnerTokenStorage.get();
+export async function getPartnerAuthData(): Promise<
+  PartnerAuthData | undefined
+> {
+  const storageValue = await partnerTokenStorage.get();
+  if (storageValue == null) {
+    return undefined;
+  }
+
+  // Backwards compatibility with old, looser type -- just clear the bad data and let the user log in again
+  if (!storageValue?.authId || !storageValue?.token) {
+    await clearPartnerAuthData();
+    return undefined;
+  }
+
+  return storageValue;
 }
 
 /**
  * Set authentication data when using the partner JWT to authenticate.
- *
- * @see clearPartnerAuth
  */
 export async function setPartnerAuthData(data: PartnerAuthData): Promise<void> {
-  if (!isEmpty(data.authId) && isEmpty(data.token)) {
+  // Backwards compatibility with old, looser type
+  if (data.token == null) {
     // Should use clearPartnerAuth for clearing the partner integration JWT
     throw new Error("Received null/blank token for partner integration");
   }
@@ -112,12 +123,23 @@ export async function setPartnerAuthData(data: PartnerAuthData): Promise<void> {
 }
 
 /**
- * Clear authentication data when using the partner JWT to authenticate.
- *
- * @see setPartnerAuthData
+ * Clear all partner OAuth2 tokens and reset api query caches
  */
-export async function clearPartnerAuth(): Promise<void> {
-  return partnerTokenStorage.set({});
+export async function clearPartnerAuthData(): Promise<void> {
+  // Old code that clears all cached auth tokens
+  // There is an issue with this api in Edge: https://github.com/w3c/webextensions/issues/648
+  // See: https://developer.chrome.com/docs/extensions/reference/identity/#method-clearAllCachedAuthTokens
+  // await chromeP.identity.clearAllCachedAuthTokens();
+
+  const partnerAuthData = await partnerTokenStorage.get();
+  if (partnerAuthData?.token) {
+    console.debug(
+      "Clearing partner auth for authId: " + partnerAuthData.authId,
+    );
+    await removeOAuth2Token(partnerAuthData.token);
+  }
+
+  await partnerTokenStorage.remove();
 }
 
 /**
@@ -139,7 +161,7 @@ export async function getAuthHeaders(): Promise<UnknownObject | null> {
     };
   }
 
-  if (partnerAuth?.token) {
+  if (partnerAuth) {
     return {
       ...partnerAuth.extraHeaders,
       // Put Authorization second to avoid overriding Authorization header. (Is defensive for now, currently
@@ -153,12 +175,6 @@ export async function getAuthHeaders(): Promise<UnknownObject | null> {
 
 /**
  * Return `true` if the extension is linked to the API. I.e., that the user is "logged in".
- *
- * NOTE: do not use this as a check before making an authenticated API call. Instead, use `maybeGetLinkedApiClient`
- * which avoids a race condition between the time the check is made and underlying `getExtensionToken` call to get
- * the token.
- *
- * @see maybeGetLinkedApiClient
  */
 export async function isLinked(): Promise<boolean> {
   return (await getAuthHeaders()) != null;

src/auth/authTypes.ts

@@ -117,21 +117,35 @@ export type PartnerAuthData = {
   /**
    * The service auth configuration for authenticating with the PixieBrix API.
    */
-  authId: UUID | null;
+  authId: UUID;
   /**
    * The JWT bearer token corresponding to the authId.
    */
-  token: string | null;
+  token: string;
   /**
    * The refresh token, if `offline_access` was included in scope.
    * @since 1.7.15
    */
   refreshToken: string | null;
-
   /**
    * Extra HTTP headers to send with every request.
    */
   extraHeaders: Record<string, string> | null;
+
+  /**
+   * The refresh request URL
+   * @since 2.0.3
+   */
+  refreshUrl: string | null;
+  /**
+   * The refresh request param payload
+   * @since 2.0.3
+   */
+  refreshParamPayload: Record<string, string> | null;
+  /**
+   * The refresh request extra headers
+   */
+  refreshExtraHeaders: Record<string, string> | null;
 };
 
 export type OrganizationAuthState = {

src/auth/featureFlagStorage.ts

@@ -19,7 +19,7 @@ import { CachedFunction } from "webext-storage-cache";
 import { expectContext } from "@/utils/expectContext";
 import { fetchFeatureFlagsInBackground } from "@/background/messenger/api";
 import { getMe } from "@/data/service/backgroundApi";
-import { addListener as addAuthStorageListener } from "@/auth/authStorage";
+import { addAuthListener as addAuthStorageListener } from "@/auth/authStorage";
 
 /**
  * Fetch the latest feature flags from the server.

src/auth/isAuthenticationAxiosError.ts

@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2024 PixieBrix, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+import type { AxiosError } from "axios";
+import { isObject } from "@/utils/objectUtils";
+
+const UNAUTHORIZED_STATUS_CODES = new Set([401, 403]);
+
+export function isAuthenticationAxiosError(
+  error: Pick<AxiosError, "response">,
+): boolean {
+  // Response should be an object
+  if (!isObject(error.response)) {
+    return false;
+  }
+
+  // Technically 403 is an authorization error and re-authenticating as the same user won't help. However, there is
+  // a case where the user just needs an updated JWT that contains the most up-to-date entitlements
+  if (UNAUTHORIZED_STATUS_CODES.has(error.response.status)) {
+    return true;
+  }
+
+  // Handle Automation Anywhere's Control Room expired JWT response. They'll return this from any endpoint instead
+  // of a proper error code.
+  if (
+    error.response.status === 400 &&
+    isObject(error.response.data) &&
+    error.response.data.message === "Access Token has expired"
+  ) {
+    return true;
+  }
+
+  return false;
+}

src/auth/useLinkState.ts

@@ -16,8 +16,8 @@
  */
 
 import {
-  addListener as addAuthListener,
-  removeListener as removeAuthListener,
+  addAuthListener,
+  removeAuthListener,
   isLinked,
 } from "@/auth/authStorage";
 import useAsyncExternalStore from "@/hooks/useAsyncExternalStore";

src/auth/usePartnerAuthData.ts

@@ -16,9 +16,9 @@
  */
 
 import {
-  addListener as addAuthListener,
+  addAuthListener,
   getPartnerAuthData,
-  removeListener as removeAuthListener,
+  removeAuthListener,
 } from "@/auth/authStorage";
 import useAsyncExternalStore from "@/hooks/useAsyncExternalStore";
 import type { AsyncState } from "@/types/sliceTypes";
@@ -38,7 +38,7 @@ const subscribe = (callback: () => void) => {
  * Use the current partner auth data. Automatically listens for changes and updates the state.
  * @see getPartnerAuthData
  */
-function usePartnerAuthData(): AsyncState<Partial<PartnerAuthData>> {
+function usePartnerAuthData(): AsyncState<PartnerAuthData | undefined> {
   return useAsyncExternalStore(subscribe, getPartnerAuthData);
 }
 

src/auth/useRequiredPartnerAuth.test.tsx

@@ -26,9 +26,10 @@ import {
   mockAuthenticatedMeApiResponse,
 } from "@/testUtils/userMock";
 import {
-  meWithPartnerApiResponseFactory,
   meApiResponseFactory,
   meOrganizationApiResponseFactory,
+  meWithPartnerApiResponseFactory,
+  partnerAuthDataFactory,
 } from "@/testUtils/factories/authFactories";
 import { renderHook } from "@/pageEditor/testHelpers";
 import { integrationConfigFactory } from "@/testUtils/factories/integrationFactories";
@@ -188,9 +189,7 @@ describe("useRequiredPartnerAuth", () => {
 
   test("has required partner integration", async () => {
     usePartnerAuthDataMock.mockReturnValue(
-      valueToAsyncState({
-        token: "NOTAREALTOKEN",
-      }),
+      valueToAsyncState(partnerAuthDataFactory()),
     );
 
     await mockAuthenticatedMeApiResponse(

src/background/auth/authStorage.ts

@@ -19,6 +19,7 @@ import { type UUID } from "@/types/stringTypes";
 import { expectContext } from "@/utils/expectContext";
 import { type AuthData } from "@/integrations/integrationTypes";
 import { oauth2Storage } from "@/auth/authConstants";
+import chromeP from "webext-polyfill-kinda";
 
 /**
  * Set the cached auth data for the given integration configId
@@ -94,3 +95,7 @@ export async function deleteCachedAuthData(configId: UUID): Promise<void> {
     );
   }
 }
+
+export async function removeOAuth2Token(token: string) {
+  await chromeP.identity.removeCachedAuthToken({ token });
+}

src/background/auth/partnerIntegrations/getPartnerPrincipals.test.ts

@@ -29,9 +29,9 @@ import {
   CONTROL_ROOM_OAUTH_INTEGRATION_ID,
   CONTROL_ROOM_TOKEN_INTEGRATION_ID,
 } from "@/integrations/constants";
-import { readRawConfigurations } from "@/integrations/registry";
 import { registry } from "@/background/messenger/api";
 import { type RegistryId } from "@/types/registryTypes";
+import { readRawConfigurations } from "@/integrations/util/readRawConfigurations";
 
 jest.mock("@/integrations/registry", () => {
   const actual = jest.requireActual("@/integrations/registry");
@@ -57,6 +57,7 @@ jest.mocked(registry.find).mockImplementation(async (id: RegistryId) => {
   } as any;
 });
 
+jest.mock("@/integrations/util/readRawConfigurations");
 const readRawConfigurationsMock = jest.mocked(readRawConfigurations);
 
 describe("getPartnerPrincipals", () => {