From c092c02dd2840782132e61dfc77558729bf63450 Mon Sep 17 00:00:00 2001
From: Nick Cooke <36927374+ncooke3@users.noreply.github.com>
Date: Mon, 6 Jan 2025 12:20:55 -0500
Subject: [PATCH] [Auth] Refactor MultiFactor.swift (#14241)

---
 .../Swift/MultiFactor/MultiFactor.swift       | 160 +++++++++---------
 1 file changed, 77 insertions(+), 83 deletions(-)

diff --git a/FirebaseAuth/Sources/Swift/MultiFactor/MultiFactor.swift b/FirebaseAuth/Sources/Swift/MultiFactor/MultiFactor.swift
index 5b1675fdfa0..17ec6b18731 100644
--- a/FirebaseAuth/Sources/Swift/MultiFactor/MultiFactor.swift
+++ b/FirebaseAuth/Sources/Swift/MultiFactor/MultiFactor.swift
@@ -67,6 +67,53 @@ import Foundation
                      displayName: String?,
                      completion: ((Error?) -> Void)?) {
       // TODO: Refactor classes so this duplicated code isn't necessary for phone and totp.
+
+      guard
+        assertion.factorID == PhoneMultiFactorInfo.TOTPMultiFactorID ||
+        assertion.factorID == PhoneMultiFactorInfo.PhoneMultiFactorID
+      else {
+        return
+      }
+
+      guard let user, let auth = user.auth else {
+        fatalError("Internal Auth error: failed to get user enrolling in MultiFactor")
+      }
+
+      let request = Self.enrollmentFinalizationRequest(
+        with: assertion,
+        displayName: displayName,
+        user: user,
+        auth: auth
+      )
+
+      Task {
+        do {
+          let response = try await auth.backend.call(with: request)
+          let user = try await auth.completeSignIn(withAccessToken: response.idToken,
+                                                   accessTokenExpirationDate: nil,
+                                                   refreshToken: response.refreshToken,
+                                                   anonymous: false)
+          try auth.updateCurrentUser(user, byForce: false, savingToDisk: true)
+          if let completion {
+            DispatchQueue.main.async {
+              completion(nil)
+            }
+          }
+        } catch {
+          if let completion {
+            DispatchQueue.main.async {
+              completion(error)
+            }
+          }
+        }
+      }
+    }
+
+    private static func enrollmentFinalizationRequest(with assertion: MultiFactorAssertion,
+                                                      displayName: String?,
+                                                      user: User,
+                                                      auth: Auth) -> FinalizeMFAEnrollmentRequest {
+      var request: FinalizeMFAEnrollmentRequest? = nil
       if assertion.factorID == PhoneMultiFactorInfo.TOTPMultiFactorID {
         guard let totpAssertion = assertion as? TOTPMultiFactorAssertion else {
           fatalError("Auth Internal Error: Failed to find TOTPMultiFactorAssertion")
@@ -74,97 +121,44 @@ import Foundation
         switch totpAssertion.secretOrID {
         case .enrollmentID: fatalError("Missing secret in totpAssertion")
         case let .secret(secret):
-          guard let user = user, let auth = user.auth else {
-            fatalError("Internal Auth error: failed to get user enrolling in MultiFactor")
-          }
           let finalizeMFATOTPRequestInfo =
             AuthProtoFinalizeMFATOTPEnrollmentRequestInfo(sessionInfo: secret.sessionInfo,
                                                           verificationCode: totpAssertion
                                                             .oneTimePassword)
-          let request = FinalizeMFAEnrollmentRequest(idToken: self.user?.rawAccessToken(),
-                                                     displayName: displayName,
-                                                     totpVerificationInfo: finalizeMFATOTPRequestInfo,
-                                                     requestConfiguration: user
-                                                       .requestConfiguration)
-          Task {
-            do {
-              let response = try await auth.backend.call(with: request)
-              do {
-                let user = try await auth.completeSignIn(withAccessToken: response.idToken,
-                                                         accessTokenExpirationDate: nil,
-                                                         refreshToken: response.refreshToken,
-                                                         anonymous: false)
-                try auth.updateCurrentUser(user, byForce: false, savingToDisk: true)
-                if let completion {
-                  DispatchQueue.main.async {
-                    completion(nil)
-                  }
-                }
-              } catch {
-                DispatchQueue.main.async {
-                  if let completion {
-                    completion(error)
-                  }
-                }
-              }
-            } catch {
-              if let completion {
-                completion(error)
-              }
-            }
-          }
+          request = FinalizeMFAEnrollmentRequest(idToken: user.rawAccessToken(),
+                                                 displayName: displayName,
+                                                 totpVerificationInfo: finalizeMFATOTPRequestInfo,
+                                                 requestConfiguration: user
+                                                   .requestConfiguration)
         }
-        return
-      } else if assertion.factorID != PhoneMultiFactorInfo.PhoneMultiFactorID {
-        return
-      }
-      let phoneAssertion = assertion as? PhoneMultiFactorAssertion
-      guard let credential = phoneAssertion?.authCredential else {
-        fatalError("Internal Error: Missing credential")
-      }
-      switch credential.credentialKind {
-      case .phoneNumber: fatalError("Internal Error: Missing verificationCode")
-      case let .verification(verificationID, code):
-        let finalizeMFAPhoneRequestInfo =
-          AuthProtoFinalizeMFAPhoneRequestInfo(sessionInfo: verificationID, verificationCode: code)
-        guard let user = user, let auth = user.auth else {
-          fatalError("Internal Auth error: failed to get user enrolling in MultiFactor")
+      } else if assertion.factorID == PhoneMultiFactorInfo.PhoneMultiFactorID {
+        let phoneAssertion = assertion as? PhoneMultiFactorAssertion
+        guard let credential = phoneAssertion?.authCredential else {
+          fatalError("Internal Error: Missing credential")
         }
-        let request = FinalizeMFAEnrollmentRequest(
-          idToken: self.user?.rawAccessToken(),
-          displayName: displayName,
-          phoneVerificationInfo: finalizeMFAPhoneRequestInfo,
-          requestConfiguration: user.requestConfiguration
-        )
-
-        Task {
-          do {
-            let response = try await auth.backend.call(with: request)
-            do {
-              let user = try await auth.completeSignIn(withAccessToken: response.idToken,
-                                                       accessTokenExpirationDate: nil,
-                                                       refreshToken: response.refreshToken,
-                                                       anonymous: false)
-              try auth.updateCurrentUser(user, byForce: false, savingToDisk: true)
-              if let completion {
-                DispatchQueue.main.async {
-                  completion(nil)
-                }
-              }
-            } catch {
-              DispatchQueue.main.async {
-                if let completion {
-                  completion(error)
-                }
-              }
-            }
-          } catch {
-            if let completion {
-              completion(error)
-            }
-          }
+        switch credential.credentialKind {
+        case .phoneNumber: fatalError("Internal Error: Missing verificationCode")
+        case let .verification(verificationID, code):
+          let finalizeMFAPhoneRequestInfo =
+            AuthProtoFinalizeMFAPhoneRequestInfo(
+              sessionInfo: verificationID,
+              verificationCode: code
+            )
+          request = FinalizeMFAEnrollmentRequest(
+            idToken: user.rawAccessToken(),
+            displayName: displayName,
+            phoneVerificationInfo: finalizeMFAPhoneRequestInfo,
+            requestConfiguration: user.requestConfiguration
+          )
         }
       }
+
+      guard let request else {
+        // Assertion is not a phone assertion or TOTP assertion.
+        fatalError("Internal Error: Unsupported assertion with factor ID: \(assertion.factorID).")
+      }
+
+      return request
     }
 
     /// Enrolls a second factor as identified by the `MultiFactorAssertion` parameter for the