#6656 Partner refresh tokens

package.json

@@ -59,6 +59,7 @@
     "ace-builds": "^1.34.2",
     "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/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 chromeP from "webext-polyfill-kinda";
 
 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,32 @@ 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
+> {
+  let isError = false;
+  let storageValue: PartnerAuthData | undefined;
+  try {
+    storageValue = await partnerTokenStorage.get();
+  } catch {
+    isError = true;
+  }
+
+  // Backwards compatibility with old, looser type
+  if (isError || !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 +126,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> {
+  // // 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 chromeP.identity.removeCachedAuthToken({
+      token: partnerAuthData.token,
+    });
+  }
+
+  await partnerTokenStorage.remove();
 }
 
 /**
@@ -139,7 +164,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 +178,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";

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

@@ -0,0 +1,28 @@
+/*
+ * 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 { registry } from "@/background/messenger/api";
+import oauth2IntegrationDefinition from "@contrib/integrations/automation-anywhere-oauth2.yaml";
+
+jest.mocked(registry.find).mockResolvedValue({
+  id: (oauth2IntegrationDefinition!.metadata as any).id,
+  config: oauth2IntegrationDefinition,
+} as any);
+
+describe("launchAuthIntegration", () => {
+  it("throws error if no local auths are found", async () => {});
+});