From e0d088046b6ab3e4f1232c677cc2003f72164b02 Mon Sep 17 00:00:00 2001
From: Harshdeep Singh <6162866+harsh62@users.noreply.github.com>
Date: Wed, 11 Dec 2024 11:22:47 -0500
Subject: [PATCH] fix(auth): fix credential decoding

---
 .../AWSCognitoAuthCredentialStore.swift       |  8 ++--
 .../Models/AuthFlowType.swift                 | 42 ++++++++++++++++---
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/CredentialStorage/AWSCognitoAuthCredentialStore.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/CredentialStorage/AWSCognitoAuthCredentialStore.swift
index 1e8b1f77f5..dd55162b57 100644
--- a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/CredentialStorage/AWSCognitoAuthCredentialStore.swift
+++ b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/CredentialStorage/AWSCognitoAuthCredentialStore.swift
@@ -133,8 +133,8 @@ extension AWSCognitoAuthCredentialStore: AmplifyAuthCredentialStoreBehavior {
     func retrieveCredential() throws -> AmplifyCredentials {
         let authCredentialStoreKey = generateSessionKey(for: authConfiguration)
         let authCredentialData = try keychain._getData(authCredentialStoreKey)
-        let awsCredential: AmplifyCredentials = try decode(data: authCredentialData)
-        return awsCredential
+        let amplifyCredential: AmplifyCredentials = try decode(data: authCredentialData)
+        return amplifyCredential
     }
 
     func deleteCredential() throws {
@@ -191,7 +191,7 @@ private extension AWSCognitoAuthCredentialStore {
         do {
             return try JSONEncoder().encode(object)
         } catch {
-            throw KeychainStoreError.codingError("Error occurred while encoding AWSCredentials", error)
+            throw KeychainStoreError.codingError("Error occurred while encoding credentials", error)
         }
     }
 
@@ -199,7 +199,7 @@ private extension AWSCognitoAuthCredentialStore {
         do {
             return try JSONDecoder().decode(T.self, from: data)
         } catch {
-            throw KeychainStoreError.codingError("Error occurred while decoding AWSCredentials", error)
+            throw KeychainStoreError.codingError("Error occurred while decoding credentials", error)
         }
     }
 
diff --git a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AuthFlowType.swift b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AuthFlowType.swift
index 46b701bcd3..9f611fdd92 100644
--- a/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AuthFlowType.swift
+++ b/AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AuthFlowType.swift
@@ -38,6 +38,8 @@ public enum AuthFlowType {
         switch rawValue {
         case "CUSTOM_AUTH":
             self = .customWithSRP
+        case "CUSTOM_AUTH_WITHOUT_SRP":
+            self = .customWithoutSRP
         case "USER_SRP_AUTH":
             self = .userSRP
         case "USER_PASSWORD_AUTH":
@@ -51,8 +53,10 @@ public enum AuthFlowType {
 
     var rawValue: String {
         switch self {
-        case .custom, .customWithSRP, .customWithoutSRP:
+        case .custom, .customWithSRP:
             return "CUSTOM_AUTH"
+        case .customWithoutSRP:
+            return "CUSTOM_AUTH_WITHOUT_SRP"
         case .userSRP:
             return "USER_SRP_AUTH"
         case .userPassword:
@@ -62,6 +66,23 @@ public enum AuthFlowType {
         }
     }
 
+    internal static func legacyInit(rawValue: String) -> Self? {
+        switch rawValue {
+        case "userSRP":
+            return .userSRP
+        case "userPassword":
+            return .userPassword
+        case "custom":
+            return .custom
+        case "customWithSRP":
+            return .customWithSRP
+        case "customWithoutSRP":
+            return .customWithoutSRP
+        default:
+            return nil
+        }
+    }
+
     public static var userAuth: AuthFlowType {
         return .userAuth(preferredFirstFactor: nil)
     }
@@ -110,9 +131,21 @@ extension AuthFlowType: Codable {
 
     // Decoding the enum
     public init(from decoder: Decoder) throws {
-        let container = try decoder.container(keyedBy: CodingKeys.self)
+        let container: KeyedDecodingContainer<CodingKeys>
+        do {
+            container = try decoder.container(keyedBy: CodingKeys.self)
+        } catch DecodingError.typeMismatch {
+            let legacyContainer = try decoder.singleValueContainer()
+            let type = try legacyContainer.decode(String.self)
+            guard let authFlowType = AuthFlowType.legacyInit(rawValue: type) else {
+                throw DecodingError.dataCorruptedError(in: legacyContainer, debugDescription: "Invalid AuthFlowType value")
+            }
+            self = authFlowType
+            return
+        } catch {
+            throw error
+        }
 
-        // Decode the type (raw value)
         let type = try container.decode(String.self, forKey: .type)
 
         // Initialize based on the type
@@ -130,7 +163,7 @@ extension AuthFlowType: Codable {
             if let preferredFirstFactor = AuthFactorType(rawValue: preferredFirstFactorString) {
                 self = .userAuth(preferredFirstFactor: preferredFirstFactor)
             } else {
-                throw DecodingError.dataCorruptedError(forKey: .type, in: container, debugDescription: "Unable to decode preferredFirstFactor value")
+                throw DecodingError.dataCorruptedError(forKey: .preferredFirstFactor, in: container, debugDescription: "Unable to decode preferredFirstFactor value")
             }
         default:
             throw DecodingError.dataCorruptedError(forKey: .type, in: container, debugDescription: "Invalid AuthFlowType value")
@@ -152,5 +185,4 @@ extension AuthFlowType {
             return .userAuth
         }
     }
-
 }