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

src/auth/oauth2client.ts

@@ -28,6 +28,7 @@ import {BodyResponseCallback} from '../transporters';
 import {AuthClient} from './authclient';
 import {CredentialRequest, Credentials} from './credentials';
 import {LoginTicket, TokenPayload} from './loginticket';
+import {getErrorFromOAuthErrorResponse} from './oauth2common';
 /**
  * The results from the `generateCodeVerifierAsync` method.  To learn more,
  * See the sample:
@@ -298,6 +299,15 @@ export interface GenerateAuthUrlOpts {
   code_challenge?: string;
 }
 
+export interface AccessTokenResponse {
+  access_token: string;
+  expiry_date: number;
+}
+
+export interface GetRefreshHandlerCallack {
+  (): Promise<AccessTokenResponse>;
+}
+
 export interface GetTokenCallback {
   (
     err: GaxiosError | null,
@@ -427,6 +437,8 @@ export class OAuth2Client extends AuthClient {
 
   forceRefreshOnFailure: boolean;
 
+  refreshHandlerCallback?: GetRefreshHandlerCallack;
+
   /**
    * Handles OAuth2 flow for Google APIs.
    *
@@ -749,7 +761,15 @@ 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.refreshHandlerCallback && this.credentials.access_token) {
+          const refreshedAccessToken = await this.refreshHandler();
+          if (refreshedAccessToken && refreshedAccessToken.access_token) {
+            this.setCredentials(refreshedAccessToken);
+            return {token: this.credentials.access_token};
+          }
+        } else {
+          throw new Error('No refresh token is set.');
+        }
       }
 
       const r = await this.refreshAccessTokenAsync();
@@ -781,7 +801,12 @@ export class OAuth2Client extends AuthClient {
     url?: string | null
   ): Promise<RequestMetadataResponse> {
     const thisCreds = this.credentials;
-    if (!thisCreds.access_token && !thisCreds.refresh_token && !this.apiKey) {
+    if (
+      !thisCreds.access_token &&
+      !thisCreds.refresh_token &&
+      !this.apiKey &&
+      !this.refreshHandlerCallback
+    ) {
       throw new Error('No access, refresh token or API key is set.');
     }
 
@@ -793,6 +818,18 @@ export class OAuth2Client extends AuthClient {
       return {headers: this.addSharedMetadataHeaders(headers)};
     }
 
+    // If refreshHandlerCallback exists, refresh
+    if (this.refreshHandlerCallback) {
+      const refreshedAccessToken = await this.refreshHandler();
+      if (refreshedAccessToken && 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 +982,43 @@ 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);
         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 &&
+          this.refreshHandler
+        ) {
+          const refreshedAccessToken = await this.refreshHandler();
+          if (refreshedAccessToken && refreshedAccessToken.access_token) {
+            this.setCredentials(refreshedAccessToken);
+          }
+          return this.requestAsync<T>(opts, true);
         }
       }
       throw e;
@@ -1316,6 +1380,22 @@ export class OAuth2Client extends AuthClient {
     return new LoginTicket(envelope, payload);
   }
 
+  /**
+   * Returns a Promise that resolves with AccessTokenResponse type if
+   * refreshHandlerCallBack is defined.
+   * If it is undefined, nothing returned.
+   */
+  async refreshHandler(): Promise<AccessTokenResponse | void> {
+    if (this.refreshHandlerCallback) {
+      const accessTokenResponse = await this.refreshHandlerCallback();
+      if (!accessTokenResponse.access_token) {
+        throw new Error('There is no access token being returned');
+      }
+      return accessTokenResponse;
+    }
+    return;
+  }
+
   /**
    * Returns true if a token is expired or will expire within
    * eagerRefreshThresholdMillismilliseconds.

test/test.oauth2.ts

@@ -25,6 +25,7 @@ import * as sinon from 'sinon';
 
 import {CodeChallengeMethod, OAuth2Client} from '../src';
 import {LoginTicket} from '../src/auth/loginticket';
+import {AccessTokenResponse} from '../src/auth/oauth2client';
 
 nock.disableNetConnect();
 
@@ -1149,6 +1150,36 @@ describe('oauth2', () => {
           done();
         });
       });
+
+      it('should call refreshHandler and set credential if token expires in request()', async () => {
+        const scope = nock('http://example.com')
+          .get('/access')
+          .reply(code, {
+            error: {code, message: 'Invalid Credentials'},
+          });
+        const expectedRefreshedAccessToken = {
+          access_token: 'access_token',
+          expiry_date: 123456789,
+        };
+        client.refreshHandlerCallback = async () => {
+          return expectedRefreshedAccessToken;
+        };
+        client.setCredentials({
+          access_token: 'initial-access-token',
+          expiry_date: new Date().getTime() - 1000,
+        });
+        client.request({url: 'http://example.com'}, () => {
+          scope.done();
+          assert.strictEqual(
+            client.credentials.access_token,
+            expectedRefreshedAccessToken.access_token
+          );
+          assert.strictEqual(
+            client.credentials.expiry_date,
+            expectedRefreshedAccessToken.expiry_date
+          );
+        });
+      });
     });
 
     it('should not retry requests with streaming data', done => {
@@ -1338,6 +1369,27 @@ describe('oauth2', () => {
       assert.deepStrictEqual(info.scopes, tokenInfo.scope.split(' '));
     });
 
+    it('should refresh request header when refreshHandler is available', async () => {
+      const expectedRefreshedAccessToken = {
+        access_token: 'access_token',
+        expiry_date: 123456789,
+      };
+      client.refreshHandlerCallback = 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 throw if tries to refresh but no refresh token is available', async () => {
       client.setCredentials({
         access_token: 'initial-access-token',
@@ -1348,5 +1400,58 @@ describe('oauth2', () => {
         /No refresh token is set./
       );
     });
+
+    it('should throw error is refreshHandler does not return token', async () => {
+      const expectedRefreshedAccessToken = {
+        access_token: '',
+        expiry_date: 123456789,
+      };
+      client.refreshHandlerCallback = async () => {
+        return expectedRefreshedAccessToken;
+      };
+      await assert.rejects(
+        client.refreshHandler(),
+        'There is no access token being returned.'
+      );
+    });
+
+    it('should call refreshHandler and set credential if providing refresh handler callback', async () => {
+      const expectedRefreshedAccessToken = {
+        access_token: 'access_token',
+        expiry_date: 123456789,
+      };
+      client.refreshHandlerCallback = async () => {
+        return expectedRefreshedAccessToken;
+      };
+      const refreshedAccessToken =
+        (await client.refreshHandler()) as AccessTokenResponse;
+      assert.strictEqual(
+        refreshedAccessToken.access_token,
+        expectedRefreshedAccessToken.access_token
+      );
+      assert.strictEqual(
+        refreshedAccessToken.expiry_date,
+        expectedRefreshedAccessToken.expiry_date
+      );
+    });
+
+    it('should refresh credential when refreshHandler is available in getAccessToken()', async () => {
+      const expectedRefreshedAccessToken = {
+        access_token: 'access_token',
+        expiry_date: 123456789,
+      };
+      client.refreshHandlerCallback = 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
+      );
+    });
   });
 });