fix(auth): Prevent credential continuation from returning multiple times.

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/CredentialStoreClient.kt

@@ -25,6 +25,7 @@ import com.amplifyframework.statemachine.codegen.data.AuthConfiguration
 import com.amplifyframework.statemachine.codegen.data.CredentialType
 import com.amplifyframework.statemachine.codegen.events.CredentialStoreEvent
 import com.amplifyframework.statemachine.codegen.states.CredentialStoreState
+import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.coroutines.resumeWithException
 import kotlin.coroutines.suspendCoroutine
 
@@ -54,40 +55,21 @@ internal class CredentialStoreClient(configuration: AuthConfiguration, context:
         onSuccess: (Result<AmplifyCredential>) -> Unit,
         onError: (Exception) -> Unit
     ) {
-        var capturedSuccess: Result<AmplifyCredential>? = null
-        var capturedError: Exception? = null
-        var token: StateChangeListenerToken? = StateChangeListenerToken()
-        credentialStoreStateMachine.listen(
-            token as StateChangeListenerToken,
-            { storeState ->
-                logger.verbose("Credential Store State Change: $storeState")
-                when (storeState) {
-                    is CredentialStoreState.Success -> {
-                        capturedSuccess = Result.success(storeState.storedCredentials)
-                    }
-                    is CredentialStoreState.Error -> {
-                        capturedError = storeState.error
-                    }
-                    is CredentialStoreState.Idle -> {
-                        val success = capturedSuccess
-                        val error = capturedError
-                        if (success != null && token != null) {
-                            credentialStoreStateMachine.cancel(token!!)
-                            token = null
-                            onSuccess(success)
-                        } else if (error != null && token != null) {
-                            credentialStoreStateMachine.cancel(token!!)
-                            token = null
-                            onError(error)
-                        }
-                    }
-                    else -> Unit
-                }
-            },
+        val token = StateChangeListenerToken()
+        val credentialStoreStateListener = OneShotCredentialStoreStateListener(
             {
-                credentialStoreStateMachine.send(event)
-            }
+                credentialStoreStateMachine.cancel(token)
+                onSuccess(it)
+            }, {
+            credentialStoreStateMachine.cancel(token)
+            onError(it)
+        },
+            logger
         )
+        credentialStoreStateMachine.listen(
+            token,
+            credentialStoreStateListener::listen
+        ) { credentialStoreStateMachine.send(event) }
     }
 
     override suspend fun loadCredentials(credentialType: CredentialType): AmplifyCredential {
@@ -121,4 +103,44 @@ internal class CredentialStoreClient(configuration: AuthConfiguration, context:
             )
         }
     }
+
+    /*
+    This class is a necessary workaround due to undesirable threading issues within the Auth State Machine. If state
+    machine threading is improved, this class should be considered for removal.
+     */
+    internal class OneShotCredentialStoreStateListener(
+        val onSuccess: (Result<AmplifyCredential>) -> Unit,
+        val onError: (Exception) -> Unit,
+        val logger: Logger
+    ) {
+        private var capturedSuccess: Result<AmplifyCredential>? = null
+        private var capturedError: Exception? = null
+        private val isActive = AtomicBoolean(true)
+        fun listen(storeState: CredentialStoreState) {
+            logger.verbose("Credential Store State Change: $storeState")
+            when (storeState) {
+                is CredentialStoreState.Success -> {
+                    capturedSuccess = Result.success(storeState.storedCredentials)
+                }
+
+                is CredentialStoreState.Error -> {
+                    capturedError = storeState.error
+                }
+
+                is CredentialStoreState.Idle -> {
+                    val success = capturedSuccess
+                    val error = capturedError
+
+                    if ((success != null || error != null) && isActive.getAndSet(false)) {
+                        if (success != null) {
+                            onSuccess(success)
+                        } else if (error != null) {
+                            onError(error)
+                        }
+                    }
+                }
+                else -> Unit
+            }
+        }
+    }
 }

aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/CredentialStoreClientTest.kt

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.amplifyframework.auth.cognito
+
+import com.amplifyframework.statemachine.codegen.states.CredentialStoreState
+import io.mockk.mockk
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.random.Random
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
+
+internal class CredentialStoreClientTest {
+
+    /**
+     * This test has been verified to regularly fail if the OneShotCredentialStoreStateListener isActive
+     * field is non-atomic.
+     */
+    @Test
+    fun one_shot_listener_fires_once() {
+        val attempts = 10_000
+        val timesFired = AtomicInteger(0)
+        var timesFailed = 0
+        val listener = CredentialStoreClient.OneShotCredentialStoreStateListener(
+            {
+                if (timesFired.incrementAndGet() != 1) {
+                    timesFailed += 1
+                }
+            }, {
+            if (timesFired.incrementAndGet() != 1) {
+                timesFailed += 1
+            }
+        },
+            mockk(relaxed = true)
+        )
+
+        for (i in 0..attempts) {
+            for (x in 0..5) {
+                if (Random.nextBoolean()) {
+                    listener.listen(CredentialStoreState.Success(mockk()))
+                } else {
+                    listener.listen(CredentialStoreState.Error(mockk()))
+                }
+                CoroutineScope(Dispatchers.IO).launch {
+                    for (y in 0..5) {
+                        listener.listen(CredentialStoreState.Idle())
+                    }
+                }
+            }
+        }
+
+        assertEquals(0, timesFailed)
+    }
+}