Fix getting credentials possibly deadlocking (#2961)

abrooksv · web-flow · commit f88193de3fdb · 2021-11-17T09:31:23.000-08:00
* Move SSO to jetbrains-core so we get access to ProgressIndicator
* Fetch credentials under a modal progress indicator
* Make SSO token pending respect ProgressIndicator cancel
* Throw ProcessCanceledException when user prompts are canceled
* Remove some coroutines usage that led to losing track of the ProgressIndicator
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/CorrectThreadCredentialsProvider.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/CorrectThreadCredentialsProvider.kt
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/CredentialManager.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/CredentialManager.kt
@@ -19,6 +19,8 @@ import software.aws.toolkits.core.credentials.ToolkitCredentialsProvider
 import software.aws.toolkits.core.region.AwsRegion
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.tryOrNull
+import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
+import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.AwsTelemetry
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.atomic.AtomicInteger
@@ -68,7 +70,9 @@ abstract class CredentialManager : SimpleModificationTracker() {
      * as loading from disk when new values.
      */
     private inner class AwsCredentialProviderProxy(private val providerId: String, private val region: AwsRegion) : AwsCredentialsProvider {
-        override fun resolveCredentials(): AwsCredentials = getOrCreateAwsCredentialsProvider(providerId, region).resolveCredentials()
+        override fun resolveCredentials(): AwsCredentials = runUnderProgressIfNeeded(null, message("credentials.retrieving"), cancelable = true) {
+            getOrCreateAwsCredentialsProvider(providerId, region).resolveCredentials()
+        }
 
         private fun getOrCreateAwsCredentialsProvider(providerId: String, region: AwsRegion): AwsCredentialsProvider {
             // Validate that the provider ID is still valid and get the latest copy
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/MfaSupport.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/MfaSupport.kt
@@ -4,10 +4,9 @@
 package software.aws.toolkits.jetbrains.core.credentials
 
 import com.intellij.openapi.actionSystem.AnAction
+import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.openapi.ui.Messages
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.withContext
-import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineUiContext
+import software.aws.toolkits.jetbrains.utils.computeOnEdt
 import software.aws.toolkits.resources.message
 
 interface MfaRequiredInteractiveCredentials : InteractiveCredential {
@@ -19,12 +18,10 @@ interface MfaRequiredInteractiveCredentials : InteractiveCredential {
     override fun userActionRequired(): Boolean = true
 }
 
-fun promptForMfaToken(name: String, mfaSerial: String): String = runBlocking {
-    withContext(getCoroutineUiContext()) {
-        Messages.showInputDialog(
-            message("credentials.mfa.message", mfaSerial),
-            message("credentials.mfa.title", name),
-            null
-        ) ?: throw IllegalStateException("MFA challenge is required")
-    }
+fun promptForMfaToken(name: String, mfaSerial: String): String = computeOnEdt {
+    Messages.showInputDialog(
+        message("credentials.mfa.message", mfaSerial),
+        message("credentials.mfa.title", name),
+        null
+    ) ?: throw ProcessCanceledException(IllegalStateException("MFA challenge is required"))
 }
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/SsoSupport.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/SsoSupport.kt
@@ -5,13 +5,13 @@ package software.aws.toolkits.jetbrains.core.credentials
 
 import com.intellij.ide.BrowserUtil
 import com.intellij.openapi.actionSystem.AnAction
+import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.openapi.ui.Messages
-import kotlinx.coroutines.withContext
-import software.aws.toolkits.core.credentials.sso.Authorization
-import software.aws.toolkits.core.credentials.sso.DiskCache
-import software.aws.toolkits.core.credentials.sso.SsoCache
-import software.aws.toolkits.core.credentials.sso.SsoLoginCallback
-import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineUiContext
+import software.aws.toolkits.jetbrains.core.credentials.sso.Authorization
+import software.aws.toolkits.jetbrains.core.credentials.sso.DiskCache
+import software.aws.toolkits.jetbrains.core.credentials.sso.SsoCache
+import software.aws.toolkits.jetbrains.core.credentials.sso.SsoLoginCallback
+import software.aws.toolkits.jetbrains.utils.computeOnEdt
 import software.aws.toolkits.jetbrains.utils.notifyError
 import software.aws.toolkits.resources.message
 
@@ -21,8 +21,8 @@ import software.aws.toolkits.resources.message
 val diskCache by lazy { DiskCache() }
 
 object SsoPrompt : SsoLoginCallback {
-    override suspend fun tokenPending(authorization: Authorization) {
-        withContext(getCoroutineUiContext()) {
+    override fun tokenPending(authorization: Authorization) {
+        computeOnEdt {
             val result = Messages.showOkCancelDialog(
                 message("credentials.sso.login.message", authorization.verificationUri, authorization.userCode),
                 message("credentials.sso.login.title"),
@@ -34,7 +34,7 @@ object SsoPrompt : SsoLoginCallback {
             if (result == Messages.OK) {
                 BrowserUtil.browse(authorization.verificationUriComplete)
             } else {
-                throw IllegalStateException(message("credentials.sso.login.cancelled"))
+                throw ProcessCanceledException(IllegalStateException(message("credentials.sso.login.cancelled")))
             }
         }
     }
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/profiles/ProfileCredentialProviderFactory.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/profiles/ProfileCredentialProviderFactory.kt
@@ -19,12 +19,12 @@ import software.aws.toolkits.core.credentials.CredentialSourceId
 import software.aws.toolkits.core.credentials.CredentialType
 import software.aws.toolkits.core.credentials.CredentialsChangeEvent
 import software.aws.toolkits.core.credentials.CredentialsChangeListener
-import software.aws.toolkits.core.credentials.sso.SsoCache
 import software.aws.toolkits.core.region.AwsRegion
 import software.aws.toolkits.jetbrains.core.credentials.MfaRequiredInteractiveCredentials
 import software.aws.toolkits.jetbrains.core.credentials.SsoRequiredInteractiveCredentials
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitCredentialProcessProvider
 import software.aws.toolkits.jetbrains.core.credentials.diskCache
+import software.aws.toolkits.jetbrains.core.credentials.sso.SsoCache
 import software.aws.toolkits.jetbrains.settings.AwsSettings
 import software.aws.toolkits.jetbrains.settings.ProfilesNotification
 import software.aws.toolkits.jetbrains.utils.createNotificationExpiringAction
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/profiles/ProfileSsoProvider.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/profiles/ProfileSsoProvider.kt
@@ -12,11 +12,11 @@ import software.amazon.awssdk.regions.Region
 import software.amazon.awssdk.services.sso.SsoClient
 import software.amazon.awssdk.services.ssooidc.SsoOidcClient
 import software.amazon.awssdk.utils.SdkAutoCloseable
-import software.aws.toolkits.core.credentials.sso.SsoAccessTokenProvider
-import software.aws.toolkits.core.credentials.sso.SsoCredentialProvider
 import software.aws.toolkits.jetbrains.core.AwsClientManager
 import software.aws.toolkits.jetbrains.core.credentials.SsoPrompt
 import software.aws.toolkits.jetbrains.core.credentials.diskCache
+import software.aws.toolkits.jetbrains.core.credentials.sso.SsoAccessTokenProvider
+import software.aws.toolkits.jetbrains.core.credentials.sso.SsoCredentialProvider
 
 class ProfileSsoProvider(profile: Profile) : AwsCredentialsProvider, SdkAutoCloseable {
     private val ssoClient: SsoClient
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/AccessToken.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/AccessToken.kt
@@ -1,7 +1,7 @@
-// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-package software.aws.toolkits.core.credentials.sso
+package software.aws.toolkits.jetbrains.core.credentials.sso
 
 import software.amazon.awssdk.services.sso.SsoClient
 import software.amazon.awssdk.services.ssooidc.SsoOidcClient
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/Authorization.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/Authorization.kt
@@ -1,7 +1,7 @@
 // Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-package software.aws.toolkits.core.credentials.sso
+package software.aws.toolkits.jetbrains.core.credentials.sso
 
 import software.amazon.awssdk.services.ssooidc.SsoOidcClient
 import java.time.Instant
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/ClientRegistration.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/ClientRegistration.kt
@@ -1,7 +1,7 @@
 // Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-package software.aws.toolkits.core.credentials.sso
+package software.aws.toolkits.jetbrains.core.credentials.sso
 
 import software.amazon.awssdk.services.ssooidc.SsoOidcClient
 import java.time.Instant
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt
@@ -1,7 +1,7 @@
-// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-package software.aws.toolkits.core.credentials.sso
+package software.aws.toolkits.jetbrains.core.credentials.sso
 
 import com.fasterxml.jackson.core.JsonParser
 import com.fasterxml.jackson.databind.DeserializationContext
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt
@@ -1,13 +1,15 @@
-// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-package software.aws.toolkits.core.credentials.sso
+package software.aws.toolkits.jetbrains.core.credentials.sso
 
-import kotlinx.coroutines.delay
+import com.intellij.openapi.progress.ProgressManager
 import software.amazon.awssdk.services.ssooidc.SsoOidcClient
 import software.amazon.awssdk.services.ssooidc.model.AuthorizationPendingException
 import software.amazon.awssdk.services.ssooidc.model.InvalidClientException
 import software.amazon.awssdk.services.ssooidc.model.SlowDownException
+import software.aws.toolkits.jetbrains.utils.assertIsNonDispatchThread
+import software.aws.toolkits.jetbrains.utils.sleepWithCancellation
 import java.time.Clock
 import java.time.Duration
 import java.time.Instant
@@ -23,7 +25,9 @@ class SsoAccessTokenProvider(
     private val client: SsoOidcClient,
     private val clock: Clock = Clock.systemUTC()
 ) {
-    suspend fun accessToken(): AccessToken {
+    fun accessToken(): AccessToken {
+        assertIsNonDispatchThread()
+
         cache.loadAccessToken(ssoUrl)?.let {
             return it
         }
@@ -81,7 +85,8 @@ class SsoAccessTokenProvider(
         )
     }
 
-    private suspend fun pollForToken(): AccessToken {
+    private fun pollForToken(): AccessToken {
+        val progressIndicator = ProgressManager.getInstance().progressIndicator
         val registration = registerClient()
         val authorization = authorizeClient(registration)
 
@@ -117,7 +122,7 @@ class SsoAccessTokenProvider(
                 throw e
             }
 
-            delay(backOffTime.toMillis())
+            sleepWithCancellation(backOffTime, progressIndicator)
         }
     }
 
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCache.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCache.kt
@@ -1,7 +1,7 @@
-// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-package software.aws.toolkits.core.credentials.sso
+package software.aws.toolkits.jetbrains.core.credentials.sso
 
 interface SsoCache {
     fun loadClientRegistration(ssoRegion: String): ClientRegistration?
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCredentialProvider.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCredentialProvider.kt
@@ -1,12 +1,10 @@
-// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
 @file:Suppress("BannedImports")
 
-package software.aws.toolkits.core.credentials.sso
+package software.aws.toolkits.jetbrains.core.credentials.sso
 
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.runBlocking
 import software.amazon.awssdk.auth.credentials.AwsCredentials
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
 import software.amazon.awssdk.auth.credentials.AwsSessionCredentials
@@ -15,6 +13,7 @@ import software.amazon.awssdk.services.sso.model.UnauthorizedException
 import software.amazon.awssdk.utils.SdkAutoCloseable
 import software.amazon.awssdk.utils.cache.CachedSupplier
 import software.amazon.awssdk.utils.cache.RefreshResult
+import software.aws.toolkits.jetbrains.utils.assertIsNonDispatchThread
 import java.time.Duration
 import java.time.Instant
 
@@ -34,10 +33,10 @@ class SsoCredentialProvider(
     override fun resolveCredentials(): AwsCredentials = sessionCache.get().credentials
 
     private fun refreshCredentials(): RefreshResult<SsoCredentialsHolder> {
+        assertIsNonDispatchThread()
+
         val roleCredentials = try {
-            val accessToken = runBlocking(Dispatchers.IO) {
-                ssoAccessTokenProvider.accessToken()
-            }
+            val accessToken = ssoAccessTokenProvider.accessToken()
 
             ssoClient.getRoleCredentials {
                 it.accessToken(accessToken.accessToken)
@@ -58,8 +57,7 @@ class SsoCredentialProvider(
 
         val expirationTime = Instant.ofEpochMilli(roleCredentials.roleCredentials().expiration())
 
-        val ssoCredentials =
-            SsoCredentialsHolder(awsCredentials, expirationTime)
+        val ssoCredentials = SsoCredentialsHolder(awsCredentials, expirationTime)
 
         return RefreshResult.builder(ssoCredentials)
             .staleTime(expirationTime.minus(Duration.ofMinutes(1)))
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoLoginCallback.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoLoginCallback.kt
@@ -1,7 +1,7 @@
-// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-package software.aws.toolkits.core.credentials.sso
+package software.aws.toolkits.jetbrains.core.credentials.sso
 
 /**
  * Callback interface to allow for UI elements to react to the different stages of the SSO login flow
@@ -10,7 +10,7 @@ interface SsoLoginCallback {
     /**
      * Called when a new authorization is pending within SSO service. User should be notified so they can perform the login flow.
      */
-    suspend fun tokenPending(authorization: Authorization)
+    fun tokenPending(authorization: Authorization)
 
     /**
      * Called when the user successfully logs into the SSO service.
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/utils/ThreadingUtils.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/utils/ThreadingUtils.kt
@@ -4,9 +4,18 @@
 package software.aws.toolkits.jetbrains.utils
 
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.ModalityState
+import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.progress.util.ProgressIndicatorUtils
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Ref
 import com.intellij.openapi.util.ThrowableComputable
+import com.intellij.util.ExceptionUtil
+import com.intellij.util.concurrency.AppExecutorUtil
+import com.intellij.util.concurrency.Semaphore
+import java.time.Duration
+import java.util.concurrent.TimeUnit
 
 // There is a new/experimental API in IJ SDK, but replicate a simpler one here till we can use it
 fun assertIsNonDispatchThread() {
@@ -20,3 +29,46 @@ fun <T> runUnderProgressIfNeeded(project: Project?, title: String, cancelable: B
     } else {
         task.invoke()
     }
+
+fun <T> computeOnEdt(supplier: () -> T): T {
+    val application = ApplicationManager.getApplication()
+    if (application.isDispatchThread) {
+        return supplier.invoke()
+    }
+    val indicator = ProgressManager.getInstance().progressIndicator
+    val semaphore = Semaphore(1)
+    val result = Ref.create<T>()
+    val error = Ref.create<Throwable>()
+    val runnable = Runnable {
+        try {
+            if (indicator == null || !indicator.isCanceled) {
+                result.set(supplier.invoke())
+            }
+        } catch (ex: Throwable) {
+            error.set(ex)
+        } finally {
+            semaphore.up()
+        }
+    }
+
+    ApplicationManager.getApplication().invokeLater(runnable, ModalityState.any())
+
+    ProgressIndicatorUtils.awaitWithCheckCanceled(semaphore, indicator)
+    ExceptionUtil.rethrowAllAsUnchecked(error.get())
+
+    return result.get()
+}
+
+fun sleepWithCancellation(sleepAmount: Duration, indicator: ProgressIndicator?) {
+    val semaphore = Semaphore(1)
+    val future = AppExecutorUtil.getAppScheduledExecutorService().schedule(
+        { semaphore.up() },
+        sleepAmount.toMillis(),
+        TimeUnit.MILLISECONDS
+    )
+    try {
+        ProgressIndicatorUtils.awaitWithCheckCanceled(semaphore, indicator)
+    } finally {
+        future.cancel(true)
+    }
+}
diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/CredentialManagerTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/CredentialManagerTest.kt
diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/profiles/ProfileCredentialProviderFactoryTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/profiles/ProfileCredentialProviderFactoryTest.kt
diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCacheTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCacheTest.kt
diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProviderTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProviderTest.kt
diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCredentialProviderTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCredentialProviderTest.kt
diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/utils/ThreadingUtilsKtTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/utils/ThreadingUtilsKtTest.kt
