feat: Implement support for exposing refreshHandler callback to handle token refresh.

src/auth/oauth2client.ts

@@ -298,6 +298,15 @@ export interface GenerateAuthUrlOpts {
   code_challenge?: string;
 }
 
+export interface AccessTokenResponse {
+  access_token: string;
+  expiry_date: number;
+}
+
+export interface GetRefreshHandlerCallback {
+  (): Promise<AccessTokenResponse>;
+}
+
 export interface GetTokenCallback {
   (
     err: GaxiosError | null,
@@ -427,6 +436,8 @@ export class OAuth2Client extends AuthClient {
 
   forceRefreshOnFailure: boolean;
 
+  refreshHandler?: GetRefreshHandlerCallback;
+
   /**
    * Handles OAuth2 flow for Google APIs.
    *
@@ -749,7 +760,18 @@ export class OAuth2Client extends AuthClient {
       !this.credentials.access_token || this.isTokenExpiring();
     if (shouldRefresh) {
       if (!this.credentials.refresh_token) {
-        throw new Error('No refresh token is set.');
+        if (this.refreshHandler) {
+          const refreshedAccessToken =
+            await this.processAndValidateRefreshHandler();
+          if (refreshedAccessToken?.access_token) {
+            this.setCredentials(refreshedAccessToken);
+            return {token: this.credentials.access_token};
+          }
+        } else {
+          throw new Error(
+            'No refresh token or refresh handler callback is set.'
+          );
+        }
       }
 
       const r = await this.refreshAccessTokenAsync();
@@ -781,8 +803,15 @@ export class OAuth2Client extends AuthClient {
     url?: string | null
   ): Promise<RequestMetadataResponse> {
     const thisCreds = this.credentials;
-    if (!thisCreds.access_token && !thisCreds.refresh_token && !this.apiKey) {
-      throw new Error('No access, refresh token or API key is set.');
+    if (
+      !thisCreds.access_token &&
+      !thisCreds.refresh_token &&
+      !this.apiKey &&
+      !this.refreshHandler
+    ) {
+      throw new Error(
+        'No access, refresh token or API key or refresh handler callback is set.'
+      );
     }
 
     if (thisCreds.access_token && !this.isTokenExpiring()) {
@@ -793,6 +822,19 @@ export class OAuth2Client extends AuthClient {
       return {headers: this.addSharedMetadataHeaders(headers)};
     }
 
+    // If refreshHandler exists, call processAndValidateRefreshHandler().
+    if (this.refreshHandler) {
+      const refreshedAccessToken =
+        await this.processAndValidateRefreshHandler();
+      if (refreshedAccessToken?.access_token) {
+        this.setCredentials(refreshedAccessToken);
+        const headers = {
+          Authorization: 'Bearer ' + this.credentials.access_token,
+        };
+        return {headers: this.addSharedMetadataHeaders(headers)};
+      }
+    }
+
     if (this.apiKey) {
       return {headers: {'X-Goog-Api-Key': this.apiKey}};
     }
@@ -945,16 +987,44 @@ export class OAuth2Client extends AuthClient {
         //   fails on the first try because it's expired. Some developers may
         //   choose to enable forceRefreshOnFailure to mitigate time-related
         //   errors.
+        // Or the following criteria are true:
+        // - We haven't already retried.  It only makes sense to retry once.
+        // - The response was a 401 or a 403
+        // - The request didn't send a readableStream
+        // - No refresh_token was available
+        // - An access_token and a refreshHandler callback were available, but
+        //   either no expiry_date was available or the forceRefreshOnFailure
+        //   flag is set. The access_token fails on the first try because it's
+        //   expired. Some developers may choose to enable forceRefreshOnFailure
+        //   to mitigate time-related errors.
         const mayRequireRefresh =
           this.credentials &&
           this.credentials.access_token &&
           this.credentials.refresh_token &&
           (!this.credentials.expiry_date || this.forceRefreshOnFailure);
+        const mayRequireRefreshWithNoRefreshToken =
+          this.credentials &&
+          this.credentials.access_token &&
+          !this.credentials.refresh_token &&
+          (!this.credentials.expiry_date || this.forceRefreshOnFailure) &&
+          this.refreshHandler;
         const isReadableStream = res.config.data instanceof stream.Readable;
         const isAuthErr = statusCode === 401 || statusCode === 403;
         if (!retry && isAuthErr && !isReadableStream && mayRequireRefresh) {
           await this.refreshAccessTokenAsync();
           return this.requestAsync<T>(opts, true);
+        } else if (
+          !retry &&
+          isAuthErr &&
+          !isReadableStream &&
+          mayRequireRefreshWithNoRefreshToken
+        ) {
+          const refreshedAccessToken =
+            await this.processAndValidateRefreshHandler();
+          if (refreshedAccessToken?.access_token) {
+            this.setCredentials(refreshedAccessToken);
+          }
+          return this.requestAsync<T>(opts, true);
         }
       }
       throw e;
@@ -1316,6 +1386,26 @@ export class OAuth2Client extends AuthClient {
     return new LoginTicket(envelope, payload);
   }
 
+  /**
+   * Returns a promise that resolves with AccessTokenResponse type if
+   * refreshHandler is defined.
+   * If not, nothing is returned.
+   */
+  private async processAndValidateRefreshHandler(): Promise<
+    AccessTokenResponse | undefined
+  > {
+    if (this.refreshHandler) {
+      const accessTokenResponse = await this.refreshHandler();
+      if (!accessTokenResponse.access_token) {
+        throw new Error(
+          'No access token is returned by the refreshHandler callback.'
+        );
+      }
+      return accessTokenResponse;
+    }
+    return;
+  }
+
   /**
    * Returns true if a token is expired or will expire within
    * eagerRefreshThresholdMillismilliseconds.

test/test.oauth2.ts

@@ -893,7 +893,7 @@ describe('oauth2', () => {
       client.request({}, (err, result) => {
         assert.strictEqual(
           err!.message,
-          'No access, refresh token or API key is set.'
+          'No access, refresh token or API key or refresh handler callback is set.'
         );
         assert.strictEqual(result, undefined);
         done();
@@ -1149,6 +1149,46 @@ describe('oauth2', () => {
           done();
         });
       });
+
+      it('should call refreshHandler in request() on token expiration and no refresh token available', async () => {
+        const scope = nock('http://example.com')
+          .get('/access')
+          .reply(code, {
+            error: {code, message: 'Invalid Credentials'},
+          });
+        const expectedRefreshedAccessToken = {
+          access_token: 'access_token',
+          expiry_date: new Date().getTime() + 3600 * 1000,
+        };
+        const expectedMetadata = {
+          Authorization: 'Bearer access_token',
+        };
+        client.refreshHandler = async () => {
+          return expectedRefreshedAccessToken;
+        };
+        client.setCredentials({
+          access_token: 'initial-access-token',
+          expiry_date: new Date().getTime() - 1000,
+        });
+
+        client.request({url: 'http://example.com/access'}, () => {
+          scope.done();
+          assert.strictEqual(
+            client.credentials.access_token,
+            expectedRefreshedAccessToken.access_token
+          );
+          assert.strictEqual(
+            client.credentials.expiry_date,
+            expectedRefreshedAccessToken.expiry_date
+          );
+        });
+
+        // Check authorization header has been added.
+        const requestMetaData = await client.getRequestHeaders(
+          'http://example.com/access'
+        );
+        assert.deepStrictEqual(requestMetaData, expectedMetadata);
+      });
     });
 
     it('should not retry requests with streaming data', done => {
@@ -1338,15 +1378,95 @@ describe('oauth2', () => {
       assert.deepStrictEqual(info.scopes, tokenInfo.scope.split(' '));
     });
 
-    it('should throw if tries to refresh but no refresh token is available', async () => {
+    it('should call refreshHandler in getRequestHeaders() on token expiration and refreshHandler is available', async () => {
+      const expectedRefreshedAccessToken = {
+        access_token: 'access_token',
+        expiry_date: new Date().getTime() + 3600 * 1000,
+      };
+      client.refreshHandler = async () => {
+        return expectedRefreshedAccessToken;
+      };
       client.setCredentials({
         access_token: 'initial-access-token',
         expiry_date: new Date().getTime() - 1000,
       });
+      const expectedMetadata = {
+        Authorization: 'Bearer access_token',
+      };
+
+      const requestMetaData = await client.getRequestHeaders(
+        'http://example.com'
+      );
+
+      assert.deepStrictEqual(requestMetaData, expectedMetadata);
+    });
+
+    it('should return authorization header on current access token if it has not expired', async () => {
+      client.setCredentials({
+        access_token: 'initial-access-token',
+        expiry_date: new Date().getTime() + 3600 * 1000,
+      });
+      const expectedMetadata = {
+        Authorization: 'Bearer initial-access-token',
+      };
+
+      const requestMetaData = await client.getRequestHeaders(
+        'http://example.com'
+      );
+
+      assert.deepStrictEqual(requestMetaData, expectedMetadata);
+    });
+
+    it('should throw if tries to refresh but no refreshHandler callback or refresh token is available', async () => {
+      client.setCredentials({
+        access_token: 'initial-access-token',
+        expiry_date: new Date().getTime() - 1000,
+      });
+
       await assert.rejects(
         client.getRequestHeaders('http://example.com'),
         /No refresh token is set./
       );
     });
+
+    it('should call refreshHandler in getAccessToken() on expiration and no refresh token available', async () => {
+      const expectedRefreshedAccessToken = {
+        access_token: 'access_token',
+        expiry_date: new Date().getTime() + 3600 * 1000,
+      };
+      client.refreshHandler = async () => {
+        return expectedRefreshedAccessToken;
+      };
+      client.setCredentials({
+        access_token: 'initial-access-token',
+        expiry_date: new Date().getTime() - 1000,
+      });
+
+      const refreshedAccessToken = await client.getAccessToken();
+
+      assert.strictEqual(
+        refreshedAccessToken.token,
+        expectedRefreshedAccessToken.access_token
+      );
+    });
+
+    it('should throw error if refreshHandler callback response is missing an access token', async () => {
+      const expectedRefreshedAccessToken = {
+        access_token: '',
+        expiry_date: new Date().getTime() + 3600 * 1000,
+      };
+      client.refreshHandler = async () => {
+        return expectedRefreshedAccessToken;
+      };
+      client.setCredentials({
+        access_token: 'initial-access-token',
+        expiry_date: new Date().getTime() - 1000,
+      });
+
+      await assert.rejects(
+        client.getAccessToken(),
+        /No access token is returned by the refreshHandler callback./
+      );
+    });
   });
 });