feat: enable refresh token rotation for salesforce

previously, we were not using the new refresh token if one was provided

Author: paustint

apps/api/src/app/routes/route.middleware.ts

@@ -357,14 +357,14 @@ export async function getOrgForRequest(
   }
 
   // Handle org refresh - then remove event listener if refreshed
-  const handleRefresh = async (accessToken: string, refreshToken: string) => {
+  const handleRefresh = async (newAccessToken: string, newRefreshToken: string) => {
     // Refresh event will be fired when renewed access token
     // to store it in your storage for next request
     try {
-      if (!refreshToken) {
+      if (!newRefreshToken) {
         return;
       }
-      await salesforceOrgsDb.updateAccessToken_UNSAFE({ accessToken, refreshToken, org, userId: user.id });
+      await salesforceOrgsDb.updateAccessToken_UNSAFE({ accessToken: newAccessToken, refreshToken: newRefreshToken, org, userId: user.id });
     } catch (ex) {
       logger.error({ requestId, ...getExceptionLog(ex) }, '[ORG][REFRESH] Error saving refresh token');
     }
@@ -378,6 +378,25 @@ export async function getOrgForRequest(
     }
   };
 
+  // Re-reads current tokens from DB so concurrent workers that lose the refresh token rotation race
+  // can retry with the tokens written by the worker that won.
+  const getFreshTokens = async () => {
+    try {
+      const freshOrg = await salesforceOrgsDb.findByUniqueId_UNSAFE(user.id, uniqueId);
+      if (!freshOrg) {
+        return null;
+      }
+      const [freshAccessToken, freshRefreshToken] = await sfdcEncService.decryptAccessToken({
+        encryptedAccessToken: freshOrg.accessToken,
+        userId: user.id,
+      });
+      return { accessToken: freshAccessToken, refreshToken: freshRefreshToken };
+    } catch (ex) {
+      logger.error({ requestId, ...getExceptionLog(ex) }, '[ORG][REFRESH] Error fetching fresh tokens for race condition check');
+      return null;
+    }
+  };
+
   const jetstreamConn = new ApiConnection(
     {
       apiRequestAdapter: getApiRequestFactoryFn(fetch),
@@ -392,6 +411,7 @@ export async function getOrgForRequest(
       logger,
       sfdcClientId: ENV.SFDC_CONSUMER_KEY,
       sfdcClientSecret: ENV.SFDC_CONSUMER_SECRET,
+      getFreshTokens,
     },
     handleRefresh,
     handleConnectionError,
@@ -574,7 +594,6 @@ export function setPermissionPolicy(_req: express.Request, res: express.Response
   next();
 }
 
-
 export function setCacheControlForApiRoutes(_req: express.Request, res: express.Response, next: express.NextFunction) {
   res.setHeader('Cache-Control', 'no-store, no-cache, must-revalidate, proxy-revalidate, max-age=0');
   next();

apps/jetstream-desktop/src/utils/route.utils.ts

@@ -200,14 +200,14 @@ export function initApiConnection(
   const accessToken = plaintext.slice(0, spaceIndex);
   const refreshToken = plaintext.slice(spaceIndex + 1);
 
-  const handleRefresh = async (accessToken: string, refreshToken: string) => {
+  const handleRefresh = async (newAccessToken: string, newRefreshToken: string) => {
     // Refresh event will be fired when renewed access token
     // to store it in your storage for next request
     try {
-      if (!refreshToken) {
+      if (!newRefreshToken) {
         return;
       }
-      await updateAccessTokens(org.uniqueId, { accessToken, refreshToken });
+      await updateAccessTokens(org.uniqueId, { accessToken: newAccessToken, refreshToken: newRefreshToken });
     } catch (ex) {
       logger.error('[ORG][REFRESH] Error saving refresh token', getErrorMessage(ex));
     }
@@ -221,6 +221,26 @@ export function initApiConnection(
     }
   };
 
+  // Re-reads current tokens from the in-memory store so concurrent requests that lose the
+  // refresh token rotation race can retry with the tokens written by the request that won.
+  const getFreshTokens = async () => {
+    try {
+      const freshOrg = getSalesforceOrgById(org.uniqueId);
+      if (!freshOrg) {
+        return null;
+      }
+      const plaintext = decryptTokenPortable(freshOrg.accessToken);
+      const spaceIndex = plaintext.indexOf(' ');
+      if (spaceIndex === -1) {
+        return null;
+      }
+      return { accessToken: plaintext.slice(0, spaceIndex), refreshToken: plaintext.slice(spaceIndex + 1) };
+    } catch (ex) {
+      logger.error('[ORG][REFRESH] Error fetching fresh tokens for race condition check', getErrorMessage(ex));
+      return null;
+    }
+  };
+
   const jetstreamConn = new ApiConnection(
     {
       apiRequestAdapter: getApiRequestFactoryFn(fetch),
@@ -234,6 +254,7 @@ export function initApiConnection(
       logger: logger as any,
       enableLogging: false,
       sfdcClientId: ENV.DESKTOP_SFDC_CLIENT_ID,
+      getFreshTokens,
     },
     handleRefresh,
     handleConnectionError,

libs/salesforce-api/src/lib/callout-adapter.ts

@@ -1,7 +1,17 @@
 import { ERROR_MESSAGES, HTTP } from '@jetstream/shared/constants';
+import { getErrorMessageAndStackObj } from '@jetstream/shared/utils';
 import { parse } from '@jetstreamapp/simple-xml';
 import isObject from 'lodash/isObject';
-import { ApiRequestOptions, ApiRequestOutputType, BulkXmlErrorResponse, FetchFn, FetchResponse, Logger, SoapErrorResponse } from './types';
+import {
+  ApiRequestOptions,
+  ApiRequestOutputType,
+  BulkXmlErrorResponse,
+  FetchFn,
+  FetchResponse,
+  Logger,
+  SessionInfo,
+  SoapErrorResponse,
+} from './types';
 
 const SOAP_API_AUTH_ERROR_REGEX = /<faultcode>[a-zA-Z]+:INVALID_SESSION_ID<\/faultcode>/;
 // Shows up for certain API requests, such as Identity
@@ -46,9 +56,14 @@ function parseXml(value: string) {
 export function getApiRequestFactoryFn(fetch: FetchFn) {
   return (
     logger: Logger,
-    onRefresh?: (accessToken: string) => void,
+    onRefresh?: (accessToken: string, refreshToken?: string) => void,
     onConnectionError?: (accessToken: string) => void,
+    /**
+     * Enable logging only applies to request/response data
+     * other logging for refresh flow and logic errors will still be logged
+     */
     enableLogging?: boolean,
+    getFreshTokens?: () => Promise<Pick<SessionInfo, 'accessToken' | 'refreshToken'> | null>,
   ) => {
     const apiRequest = async <Response = unknown>(options: ApiRequestOptions, attemptRefresh = true): Promise<Response> => {
       // eslint-disable-next-line prefer-const
@@ -125,18 +140,37 @@ export function getApiRequestFactoryFn(fetch: FetchFn) {
             sessionInfo.refreshToken
           ) {
             try {
-              // if 401 and we have a refresh token, then attempt to refresh the token
-              const { access_token: newAccessToken } = await exchangeRefreshToken(fetch, sessionInfo);
-              onRefresh?.(newAccessToken);
+              logger.debug({ url, method, status: response.status }, '[TOKEN REFRESH] Attempting token refresh');
+              const { access_token: newAccessToken, refresh_token: newRefreshToken } = await exchangeRefreshToken(fetch, sessionInfo);
+              logger.debug({ url, method, tokenRotated: !!newRefreshToken }, '[TOKEN REFRESH] Token refresh successful');
+              onRefresh?.(newAccessToken, newRefreshToken);
               // replace token in body
               if (typeof options.body === 'string' && options.body.includes(accessToken)) {
                 // if the response is soap, we need to return the response as is
                 options.body = options.body.replace(accessToken, newAccessToken);
               }
 
               return apiRequest({ ...options, sessionInfo: { ...sessionInfo, accessToken: newAccessToken } }, false);
-            } catch {
-              logger.warn('Unable to refresh accessToken');
+            } catch (ex) {
+              logger.warn({ url, method, ...getErrorMessageAndStackObj(ex) }, '[TOKEN REFRESH] Unable to refresh accessToken');
+
+              // Check if another worker already refreshed (race condition on token rotation).
+              // If the DB has a different access token, a concurrent request won the race — retry with fresh tokens.
+              if (getFreshTokens) {
+                try {
+                  const freshTokens = await getFreshTokens();
+                  if (freshTokens && freshTokens.accessToken !== accessToken) {
+                    logger.info({ url, method }, '[TOKEN REFRESH] Concurrent refresh detected — retrying with tokens from another worker');
+                    return apiRequest({ ...options, sessionInfo: { ...sessionInfo, ...freshTokens } }, false);
+                  }
+                } catch (freshEx) {
+                  logger.warn(
+                    { url, method, ...getErrorMessageAndStackObj(freshEx) },
+                    '[TOKEN REFRESH] Failed to retrieve fresh tokens for race condition check',
+                  );
+                }
+              }
+
               responseText = ERROR_MESSAGES.SFDC_EXPIRED_TOKEN;
               onConnectionError?.(ERROR_MESSAGES.SFDC_EXPIRED_TOKEN);
             }
@@ -192,7 +226,10 @@ function handleSalesforceApiError(outputType: ApiRequestOutputType, responseText
   return output;
 }
 
-function exchangeRefreshToken(fetch: FetchFn, sessionInfo: ApiRequestOptions['sessionInfo']): Promise<{ access_token: string }> {
+function exchangeRefreshToken(
+  fetch: FetchFn,
+  sessionInfo: ApiRequestOptions['sessionInfo'],
+): Promise<{ access_token: string; refresh_token?: string }> {
   const searchParams = new URLSearchParams({
     grant_type: 'refresh_token',
   });

libs/salesforce-api/src/lib/connection.ts

@@ -23,6 +23,8 @@ export interface ApiConnectionOptions {
   sfdcClientId?: string;
   sfdcClientSecret?: string;
   logger: Logger;
+  /** Re-reads current tokens from the source of truth (e.g. DB) to handle concurrent refresh token rotation across workers */
+  getFreshTokens?: () => Promise<Pick<SessionInfo, 'accessToken' | 'refreshToken'> | null>;
 }
 
 export class ApiConnection {
@@ -56,12 +58,19 @@ export class ApiConnection {
       sfdcClientId,
       sfdcClientSecret,
       logger,
+      getFreshTokens,
     }: ApiConnectionOptions,
     refreshCallback?: (accessToken: string, refreshToken: string) => void,
     onConnectionError?: (error: string) => void,
   ) {
     this.logger = logger;
-    this.apiRequest = apiRequestAdapter(logger, this.handleRefresh.bind(this), this.handleConnectionError.bind(this), enableLogging);
+    this.apiRequest = apiRequestAdapter(
+      logger,
+      this.handleRefresh.bind(this),
+      this.handleConnectionError.bind(this),
+      enableLogging,
+      getFreshTokens,
+    );
     this.refreshCallback = refreshCallback;
     this.onConnectionError = onConnectionError;
     this.sessionInfo = {
@@ -122,8 +131,11 @@ export class ApiConnection {
     this.sessionInfo.userId = userId ?? this.sessionInfo.userId;
   }
 
-  public handleRefresh(accessToken: string) {
+  public handleRefresh(accessToken: string, newRefreshToken?: string) {
     this.sessionInfo.accessToken = accessToken;
+    if (newRefreshToken) {
+      this.sessionInfo.refreshToken = newRefreshToken;
+    }
     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
     this.refreshCallback?.(accessToken, this.sessionInfo.refreshToken!);
   }