From a2dfdf0ed075995712d760127b989fe2e155301c Mon Sep 17 00:00:00 2001
From: Harsh <6162866+harsh62@users.noreply.github.com>
Date: Wed, 4 Dec 2024 20:27:13 -0500
Subject: [PATCH] fix(auth): use auth flow type correctly from
 amplifyconfiguraiton.json (#3928)

---
 .../Models/AuthFlowType.swift                     | 15 +++++++++++++++
 .../Support/Helpers/ConfigurationHelper.swift     |  9 ++++++---
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AuthFlowType.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AuthFlowType.swift
index 0229bd5285..46b701bcd3 100644
--- a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AuthFlowType.swift
+++ b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AuthFlowType.swift
@@ -34,6 +34,21 @@ public enum AuthFlowType {
     /// - `preferredFirstFactor`: the auth factor type the user should begin signing with if available. If the preferred first factor is not available, the flow would fallback to provide available first factors.
     case userAuth(preferredFirstFactor: AuthFactorType?)
 
+    internal init?(rawValue: String) {
+        switch rawValue {
+        case "CUSTOM_AUTH":
+            self = .customWithSRP
+        case "USER_SRP_AUTH":
+            self = .userSRP
+        case "USER_PASSWORD_AUTH":
+            self = .userPassword
+        case "USER_AUTH":
+            self = .userAuth
+        default:
+            return nil
+        }
+    }
+
     var rawValue: String {
         switch self {
         case .custom, .customWithSRP, .customWithoutSRP:
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Support/Helpers/ConfigurationHelper.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Support/Helpers/ConfigurationHelper.swift
index ae6d367070..6baadb7269 100644
--- a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Support/Helpers/ConfigurationHelper.swift
+++ b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Support/Helpers/ConfigurationHelper.swift
@@ -47,14 +47,17 @@ struct ConfigurationHelper {
 
         // parse `authFlowType`
         var authFlowType: AuthFlowType
+
+        // If Migration path is enabled, auth flow type should always be set to USER_PASSWORD_AUTH
         if case .boolean(let isMigrationEnabled) = cognitoUserPoolJSON.value(at: "MigrationEnabled"),
            isMigrationEnabled == true {
             authFlowType = .userPassword
         } else if let authJson = config.value(at: "Auth.Default"),
-                  case .string(let authFlowTypeJSON) = authJson.value(at: "authenticationFlowType"),
-                  authFlowTypeJSON == "CUSTOM_AUTH" {
-            authFlowType = .customWithSRP
+                  case .string(let authFlowTypeConfigValue) = authJson.value(at: "authenticationFlowType"),
+                  let authFlowTypeFromConfig = AuthFlowType(rawValue: authFlowTypeConfigValue) {
+            authFlowType = authFlowTypeFromConfig
         } else {
+            // if the auth flow type is not found from config, default to SRP
             authFlowType = .userSRP
         }