From 410deae471ebfcc07f31af0574852501b9526308 Mon Sep 17 00:00:00 2001 From: Tyler Roach Date: Fri, 28 Jul 2023 14:08:08 -0400 Subject: [PATCH 1/3] prevent continuation from returning multiple times. --- .../auth/cognito/CredentialStoreClient.kt | 82 +++++++++++-------- .../auth/cognito/CredentialStoreClientTest.kt | 54 ++++++++++++ 2 files changed, 104 insertions(+), 32 deletions(-) create mode 100644 aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/CredentialStoreClientTest.kt diff --git a/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/CredentialStoreClient.kt b/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/CredentialStoreClient.kt index 1354dead92..b7d43d03e7 100644 --- a/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/CredentialStoreClient.kt +++ b/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) -> Unit, onError: (Exception) -> Unit ) { - var capturedSuccess: Result? = 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,40 @@ internal class CredentialStoreClient(configuration: AuthConfiguration, context: ) } } + + internal class OneShotCredentialStoreStateListener( + val onSuccess: (Result) -> Unit, + val onError: (Exception) -> Unit, + val logger: Logger + ) { + private var capturedSuccess: Result? = 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 + } + } + } } diff --git a/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/CredentialStoreClientTest.kt b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/CredentialStoreClientTest.kt new file mode 100644 index 0000000000..c56139544d --- /dev/null +++ b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/CredentialStoreClientTest.kt @@ -0,0 +1,54 @@ +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) + } +} From 173a7e9a5f08860d457a57b6f06bfd4d1b34615b Mon Sep 17 00:00:00 2001 From: Tyler Roach Date: Wed, 9 Aug 2023 08:20:44 -0400 Subject: [PATCH 2/3] add license --- .../auth/cognito/CredentialStoreClientTest.kt | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/CredentialStoreClientTest.kt b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/CredentialStoreClientTest.kt index c56139544d..094c97a644 100644 --- a/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/CredentialStoreClientTest.kt +++ b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/CredentialStoreClientTest.kt @@ -1,3 +1,18 @@ +/* + * 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 From be50107e24acda656fb5cd2c830c64964566a31f Mon Sep 17 00:00:00 2001 From: Tyler Roach Date: Wed, 9 Aug 2023 09:46:09 -0400 Subject: [PATCH 3/3] add comment --- .../amplifyframework/auth/cognito/CredentialStoreClient.kt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/CredentialStoreClient.kt b/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/CredentialStoreClient.kt index b7d43d03e7..f7182f966f 100644 --- a/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/CredentialStoreClient.kt +++ b/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/CredentialStoreClient.kt @@ -104,6 +104,10 @@ 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) -> Unit, val onError: (Exception) -> Unit,