AWS Cognito Swift: Current Token Not Updating After First Login

[thanhnd10kn]

Describe the bug

I am using AWS Cognito for authentication in my iOS app with Swift. Below is the loginCognito function I use to authenticate users:

private func loginCognito(data: AWSConfigData) async throws {
    let loginsKey = "cognito-idp.\(data.region).amazonaws.com/\(data.userPoolId)"
    let logins = [loginsKey: data.idToken]

    let identityProviderManager = IdentityProviderManager(loginMaps: logins)

    let credentialsProvider = AWSCognitoCredentialsProvider(
        regionType: data.regionType,
        identityPoolId: data.identityPoolId,
        identityProviderManager: identityProviderManager
    )

    return try await withCheckedThrowingContinuation { continuation in
        credentialsProvider.identityProvider.logins().continueWith { task in
            if let result = task.result, let token = result[loginsKey] as? String {
                print("Login with Token:", token)

                self.awsConfigData = data
                let configuration = AWSServiceConfiguration(
                    region: data.regionType,
                    credentialsProvider: credentialsProvider
                )
                AWSServiceManager.default()?.defaultServiceConfiguration = configuration

                Task {
                    let currentToken = try await self.getLoginsToken()
                    print("Current Token:", currentToken)
                }
                return continuation.resume(returning: ())
            } else if let error = task.error {
                return continuation.resume(throwing: error)
            } else {
                let error = NSError(domain: "S3Error", code: -1, userInfo: [NSLocalizedDescriptionKey: "Get login error"])
                return continuation.resume(throwing: error)
            }
        }
    }
}

/// IdentityProviderManager for AWS Cognito
private class IdentityProviderManager: NSObject, AWSIdentityProviderManager {
    private let loginMaps: [String: String]

    init(loginMaps: [String: String]) {
        self.loginMaps = loginMaps
    }

    func logins() -> AWSTask<NSDictionary> {
        return AWSTask(result: loginMaps as NSDictionary)
    }
}

This is my getLoginsToken function to retrieve the current token:

private func getLoginsToken() async throws -> String {
    guard let configuration = AWSServiceManager.default().defaultServiceConfiguration,
          let credentialsProvider = configuration.credentialsProvider as? AWSCognitoCredentialsProvider,
          let awsConfigData else {
        let error = NSError(domain: "S3Error", code: -1, userInfo: [NSLocalizedDescriptionKey: "Cannot get credentialsProvider"])
        print("Token Error:", error)
        throw error
    }

    return try await withCheckedThrowingContinuation { continuation in
        credentialsProvider.identityProvider.logins().continueWith { task in
            let loginsKey = "cognito-idp.\(awsConfigData.region).amazonaws.com/\(awsConfigData.userPoolId)"

            if let result = task.result, let token = result[loginsKey] as? String {
                print("Token:", token)
                return continuation.resume(returning: token)
            } else if let error = task.error {
                print("Token Error:", error)
                return continuation.resume(throwing: error)
            } else {
                let error = NSError(domain: "S3Error", code: -1, userInfo: [NSLocalizedDescriptionKey: "Get login error"])
                return continuation.resume(throwing: error)
            }
        }
    }
}

---

Issue:

• On the first login, Login with Token and Current Token are the same.
• However, from the second login onwards, Login with Token is updated, but Current Token remains the same as the first login.
• I have tried using credentialsProvider.clearCredentials(), but it did not solve the issue.

How can I ensure that getLoginsToken() always retrieves the latest token from Cognito?

Any help would be greatly appreciated! Thanks in advance! 🙏

Steps To Reproduce

1. **Perform the first login**  
   - Call `loginCognito(data: AWSConfigData)`.  
   - Observe that `Login with Token` and `Current Token` are the same.  

2. **Perform the second login**  
   - Call `loginCognito(data: AWSConfigData)` again with a new `idToken`.  
   - Observe that `Login with Token` updates correctly, but `Current Token` still returns the token from the first login.  

3. **Retrieve the current token**  
   - Call `getLoginsToken()`.  
   - Observe that it returns the token from the first login instead of the latest one.

Expected behavior

• On the first login, Login with Token and Current Token should be the same.
• On subsequent logins, Login with Token should update with a new token, and Current Token should also reflect the latest token.
• Calling getLoginsToken() should always return the most recent token issued by AWS Cognito after each login.
• The AWS credentials should refresh correctly, ensuring that the new token is used for authentication in subsequent API calls.

Amplify Framework Version

2.40.1

Amplify Categories

Storage

Dependency manager

Swift PM

Swift version

5.9

CLI version

I don't use CLI

Xcode version

16.1

Relevant log output

Is this a regression?

Yes

Regression additional context

No response

Platforms

iOS

OS Version

iOS 17.6.1

Device

iPad Gen 9th

Specific to simulators

No response

Additional context

No response

[tylerjroach]

It appears that you are using the AWS iOS SDK. I'll go ahead and transfer your issue to the proper library so our team can take a look.

[thanhnd10kn]

@tylerjroach Thanks for your help! I made some small updates to the code. It seems to be working now, but I wanted to ask if there's a better way to do it. Thanks again!

Login func

    private func loginCognito(data: AWSConfigData) {
        let loginsKey = "cognito-idp.\(data.region).amazonaws.com/\(data.userPoolId)"
        let logins = [loginsKey: data.idToken]

        self.awsConfigData = data

        printAWS(addition: "Login with Token:", message: "\(data.idToken)")

        // If defaultServiceConfiguration has configured -> update loginMaps to IdentityProviderManager
        if let existingCredentialsProvider = AWSServiceManager.default()?.defaultServiceConfiguration?.credentialsProvider,
            let existingCredentialsProvider = existingCredentialsProvider as? AWSCognitoCredentialsProvider,
            let existingIdentityProviderManager = existingCredentialsProvider.identityProvider.identityProviderManager,
            let existingIdentityProviderManager = existingIdentityProviderManager as? IdentityProviderManager {
            existingIdentityProviderManager.setLogins(loginMaps: logins)
            return
        }

        let identityProviderManager = IdentityProviderManager(loginMaps: logins)
        let credentialsProvider = AWSCognitoCredentialsProvider(
            regionType: data.regionType,
            identityPoolId: data.identityPoolId,
            identityProviderManager: identityProviderManager
        )
        // Set AWS default configuration
        let configuration = AWSServiceConfiguration(
            region: data.regionType,
            credentialsProvider: credentialsProvider
        )

        AWSServiceManager.default()?.defaultServiceConfiguration = configuration
    }

IdentityProviderManager

private class IdentityProviderManager: NSObject, AWSIdentityProviderManager {
    private var loginMaps: [String: String]

    init(loginMaps: [String: String]) {
        self.loginMaps = loginMaps
    }

    func setLogins(loginMaps: [String: String]) {
        self.loginMaps = loginMaps
    }

    func logins() -> AWSTask<NSDictionary> {
        return AWSTask(result: loginMaps as NSDictionary)
    }
}