From 22f1369f12c42065948e6080066b9ef728dcb4eb Mon Sep 17 00:00:00 2001
From: renkelvin <chuanr@google.com>
Date: Wed, 11 Oct 2023 10:53:55 -0700
Subject: [PATCH] Add handleRecaptchaFlow

---
 .../recaptcha_enterprise_verifier.test.ts     | 93 ++++++++++++++++++-
 .../recaptcha_enterprise_verifier.ts          | 40 ++++++++
 2 files changed, 131 insertions(+), 2 deletions(-)

diff --git a/packages/auth/src/platform_browser/recaptcha/recaptcha_enterprise_verifier.test.ts b/packages/auth/src/platform_browser/recaptcha/recaptcha_enterprise_verifier.test.ts
index b1cbd959c4c..7853e7ea342 100644
--- a/packages/auth/src/platform_browser/recaptcha/recaptcha_enterprise_verifier.test.ts
+++ b/packages/auth/src/platform_browser/recaptcha/recaptcha_enterprise_verifier.test.ts
@@ -20,16 +20,23 @@ import chaiAsPromised from 'chai-as-promised';
 import * as sinon from 'sinon';
 import sinonChai from 'sinon-chai';
 
-import { Endpoint, RecaptchaClientType, RecaptchaVersion } from '../../api';
+import {
+  Endpoint,
+  RecaptchaClientType,
+  RecaptchaVersion,
+  RecaptchaActionName
+} from '../../api';
 import { mockEndpointWithParams } from '../../../test/helpers/api/helper';
 import { testAuth, TestAuth } from '../../../test/helpers/mock_auth';
 import * as mockFetch from '../../../test/helpers/mock_fetch';
 import { ServerError } from '../../api/errors';
+import { AuthInternal } from '../../model/auth';
 
 import { MockGreCAPTCHATopLevel } from './recaptcha_mock';
 import {
   RecaptchaEnterpriseVerifier,
-  FAKE_TOKEN
+  FAKE_TOKEN,
+  handleRecaptchaFlow
 } from './recaptcha_enterprise_verifier';
 
 use(chaiAsPromised);
@@ -117,4 +124,86 @@ describe('platform_browser/recaptcha/recaptcha_enterprise_verifier', () => {
       expect(await verifier.verify()).to.eq(FAKE_TOKEN);
     });
   });
+
+  context('handleRecaptchaFlow', () => {
+    let mockAuthInstance: AuthInternal;
+    let mockRequest: any;
+    let mockActionMethod: sinon.SinonStub;
+
+    beforeEach(async () => {
+      mockAuthInstance = await testAuth();
+      mockRequest = {};
+      mockActionMethod = sinon.stub();
+    });
+
+    afterEach(() => {
+      sinon.restore();
+    });
+
+    it('should handle recaptcha when emailPasswordEnabled is true', async () => {
+      if (typeof window === 'undefined') {
+        return;
+      }
+      sinon.stub(mockAuthInstance, '_getRecaptchaConfig').returns({
+        emailPasswordEnabled: true,
+        siteKey: 'mock_site_key'
+      });
+      mockActionMethod.resolves('success');
+
+      const result = await handleRecaptchaFlow(
+        mockAuthInstance,
+        mockRequest,
+        RecaptchaActionName.GET_OOB_CODE,
+        mockActionMethod
+      );
+
+      expect(result).to.equal('success');
+      expect(mockActionMethod).to.have.been.calledOnce;
+    });
+
+    it('should handle action without recaptcha when emailPasswordEnabled is false and no error', async () => {
+      if (typeof window === 'undefined') {
+        return;
+      }
+      sinon.stub(mockAuthInstance, '_getRecaptchaConfig').returns({
+        emailPasswordEnabled: false,
+        siteKey: 'mock_site_key'
+      });
+      mockActionMethod.resolves('success');
+
+      const result = await handleRecaptchaFlow(
+        mockAuthInstance,
+        mockRequest,
+        RecaptchaActionName.GET_OOB_CODE,
+        mockActionMethod
+      );
+
+      expect(result).to.equal('success');
+      expect(mockActionMethod).to.have.been.calledOnce;
+    });
+
+    it('should handle MISSING_RECAPTCHA_TOKEN error when emailPasswordEnabled is false', async () => {
+      if (typeof window === 'undefined') {
+        return;
+      }
+      sinon.stub(mockAuthInstance, '_getRecaptchaConfig').returns({
+        emailPasswordEnabled: false,
+        siteKey: 'mock_site_key'
+      });
+      mockActionMethod.onFirstCall().rejects({
+        code: 'auth/MISSING_RECAPTCHA_TOKEN'
+      });
+      mockActionMethod.onSecondCall().resolves('success-after-recaptcha');
+
+      const result = await handleRecaptchaFlow(
+        mockAuthInstance,
+        mockRequest,
+        RecaptchaActionName.GET_OOB_CODE,
+        mockActionMethod
+      );
+
+      expect(result).to.equal('success-after-recaptcha');
+      expect(mockActionMethod).to.have.been.calledTwice;
+    });
+  });
 });
diff --git a/packages/auth/src/platform_browser/recaptcha/recaptcha_enterprise_verifier.ts b/packages/auth/src/platform_browser/recaptcha/recaptcha_enterprise_verifier.ts
index c2c0303088a..63ec91fceae 100644
--- a/packages/auth/src/platform_browser/recaptcha/recaptcha_enterprise_verifier.ts
+++ b/packages/auth/src/platform_browser/recaptcha/recaptcha_enterprise_verifier.ts
@@ -28,6 +28,7 @@ import { Auth } from '../../model/public_types';
 import { AuthInternal } from '../../model/auth';
 import { _castAuth } from '../../core/auth/auth_impl';
 import * as jsHelpers from '../load_js';
+import { AuthErrorCode } from '../../core/errors';
 
 const RECAPTCHA_ENTERPRISE_URL =
   'https://www.google.com/recaptcha/enterprise.js?render=';
@@ -175,6 +176,45 @@ export async function injectRecaptchaFields<T>(
   return newRequest;
 }
 
+type ActionMethod<TRequest, TResponse> = (
+  auth: Auth,
+  request: TRequest
+) => Promise<TResponse>;
+
+export async function handleRecaptchaFlow<TRequest, TResponse>(
+  authInstance: AuthInternal,
+  request: TRequest,
+  actionName: RecaptchaActionName,
+  actionMethod: ActionMethod<TRequest, TResponse>
+): Promise<TResponse> {
+  if (authInstance._getRecaptchaConfig()?.emailPasswordEnabled) {
+    const requestWithRecaptcha = await injectRecaptchaFields(
+      authInstance,
+      request,
+      actionName,
+      actionName === RecaptchaActionName.GET_OOB_CODE
+    );
+    return actionMethod(authInstance, requestWithRecaptcha);
+  } else {
+    return actionMethod(authInstance, request).catch(async error => {
+      if (error.code === `auth/${AuthErrorCode.MISSING_RECAPTCHA_TOKEN}`) {
+        console.log(
+          `${actionName} is protected by reCAPTCHA Enterprise for this project. Automatically triggering the reCAPTCHA flow and restarting the flow.`
+        );
+        const requestWithRecaptcha = await injectRecaptchaFields(
+          authInstance,
+          request,
+          actionName,
+          actionName === RecaptchaActionName.GET_OOB_CODE
+        );
+        return actionMethod(authInstance, requestWithRecaptcha);
+      } else {
+        return Promise.reject(error);
+      }
+    });
+  }
+}
+
 export async function _initializeRecaptchaConfig(auth: Auth): Promise<void> {
   const authInternal = _castAuth(auth);