From 965f0e3eb1814e2cece84cbd6fdc8fa72b3b2d30 Mon Sep 17 00:00:00 2001
From: Nick Cooke <nickcooke@google.com>
Date: Tue, 10 Dec 2024 17:49:51 -0500
Subject: [PATCH] [Auth] Refactor MultiFactor.swift

---
 .../Swift/MultiFactor/MultiFactor.swift       | 74 ++++++++++++++++++-
 1 file changed, 70 insertions(+), 4 deletions(-)

diff --git a/FirebaseAuth/Sources/Swift/MultiFactor/MultiFactor.swift b/FirebaseAuth/Sources/Swift/MultiFactor/MultiFactor.swift
index 5b1675fdfa0..a2d7e43289d 100644
--- a/FirebaseAuth/Sources/Swift/MultiFactor/MultiFactor.swift
+++ b/FirebaseAuth/Sources/Swift/MultiFactor/MultiFactor.swift
@@ -67,6 +67,19 @@ 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.TOTPMultiFactorID
+      else {
+        return
+      }
+
+      guard let user = user, let auth = user.auth else {
+        fatalError("Internal Auth error: failed to get user enrolling in MultiFactor")
+      }
+
+      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,13 +87,11 @@ 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)
+<<<<<<< Updated upstream
           let request = FinalizeMFAEnrollmentRequest(idToken: self.user?.rawAccessToken(),
                                                      displayName: displayName,
                                                      totpVerificationInfo: finalizeMFATOTPRequestInfo,
@@ -106,13 +117,63 @@ import Foundation
                     completion(error)
                   }
                 }
+=======
+          request = FinalizeMFAEnrollmentRequest(idToken: self.user?.rawAccessToken(),
+                                                 displayName: displayName,
+                                                 totpVerificationInfo: finalizeMFATOTPRequestInfo,
+                                                 requestConfiguration: user
+                                                   .requestConfiguration)
+        }
+      } else if assertion.factorID == PhoneMultiFactorInfo.PhoneMultiFactorID {
+        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
+            )
+          request = FinalizeMFAEnrollmentRequest(
+            idToken: self.user?.rawAccessToken(),
+            displayName: displayName,
+            phoneVerificationInfo: finalizeMFAPhoneRequestInfo,
+            requestConfiguration: user.requestConfiguration
+          )
+        }
+      }
+
+      guard let request else {
+        // Assertion is not a phone assertion or TOTP assertion.
+        return
+      }
+
+      Task {
+        do {
+          let response = try await AuthBackend.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)
+>>>>>>> Stashed changes
               }
-            } catch {
+            }
+          } catch {
+            DispatchQueue.main.async {
               if let completion {
                 completion(error)
               }
             }
           }
+<<<<<<< Updated upstream
         }
         return
       } else if assertion.factorID != PhoneMultiFactorInfo.PhoneMultiFactorID {
@@ -162,6 +223,11 @@ import Foundation
             if let completion {
               completion(error)
             }
+=======
+        } catch {
+          if let completion {
+            completion(error)
+>>>>>>> Stashed changes
           }
         }
       }