From 747d92f28c592db9e5c1143dfb242679e1fb22fb Mon Sep 17 00:00:00 2001
From: Richard Li <richali@amazon.com>
Date: Thu, 26 Sep 2024 10:45:15 -0700
Subject: [PATCH] Block webviews from allowing multiple parallel
 authorization_grant flows

---
 .../toolkits/jetbrains/core/credentials/LoginUtils.kt | 11 +++++++++++
 .../core/credentials/sso/pkce/ToolkitOAuthService.kt  |  2 ++
 .../aws/toolkits/resources/MessagesBundle.properties  |  1 +
 3 files changed, 14 insertions(+)

diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/LoginUtils.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/LoginUtils.kt
index d1e470195fe..54c8bad9b63 100644
--- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/LoginUtils.kt
+++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/LoginUtils.kt
@@ -18,24 +18,35 @@ import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
 import software.amazon.awssdk.services.sts.StsClient
 import software.aws.toolkits.core.credentials.validatedSsoIdentifierFromUrl
 import software.aws.toolkits.core.region.AwsRegion
+import software.aws.toolkits.core.utils.debug
+import software.aws.toolkits.core.utils.getLogger
+import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.core.AwsClientManager
 import software.aws.toolkits.jetbrains.core.credentials.profiles.SsoSessionConstants
 import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_REGION
 import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL
 import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.InteractiveBearerTokenProvider
+import software.aws.toolkits.jetbrains.core.credentials.sso.pkce.ToolkitOAuthService
 import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
 import software.aws.toolkits.resources.AwsCoreBundle
 import software.aws.toolkits.telemetry.CredentialSourceId
 import java.io.IOException
 
+private val LOG = getLogger<Login<*>>()
+
 sealed class Login<T> {
     abstract val id: CredentialSourceId
     abstract val onError: (Exception) -> Unit
     protected abstract fun doLogin(project: Project): T
 
     fun login(project: Project): T {
+        LOG.debug { "Starting login with request: $this" }
         try {
+            check(!ToolkitOAuthService.getInstance().hasPendingRequest()) {
+                LOG.warn { "$this attempt initiated with pending request: ${ToolkitOAuthService.getInstance().pendingRequest()}" }
+                AwsCoreBundle.message("toolkit.login.singleton")
+            }
             return doLogin(project)
         } catch (e: Exception) {
             onError(e)
diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/pkce/ToolkitOAuthService.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/pkce/ToolkitOAuthService.kt
index 468e103f134..e62e6a43bfa 100644
--- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/pkce/ToolkitOAuthService.kt
+++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/pkce/ToolkitOAuthService.kt
@@ -43,6 +43,8 @@ const val PKCE_CLIENT_NAME = "AWS IDE Plugins for JetBrains"
 class ToolkitOAuthService : OAuthServiceBase<AccessToken>() {
     override val name: String = "aws/toolkit"
 
+    internal fun pendingRequest() = (currentRequest.get()?.request as? ToolkitOAuthRequest)?.registration
+
     fun hasPendingRequest() = currentRequest.get() != null
 
     fun authorize(registration: PKCEClientRegistration): CompletableFuture<AccessToken> {
diff --git a/plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties b/plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties
index 867cf7178b0..e1dfff2d02b 100644
--- a/plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties
+++ b/plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties
@@ -1954,6 +1954,7 @@ toolkit.login.dialog.sso.text_field.region=Region:
 toolkit.login.dialog.sso.text_field.start_url=Start URL:
 toolkit.login.dialog.sso.title=Connect using AWS IAM Identity Center
 toolkit.login.dialog.title=AWS Toolkit: Add Connection
+toolkit.login.singleton=Only one browser authorization flow may be active at once
 toolkit.sso_expire.dialog.cancel_button=Cancel
 toolkit.sso_expire.dialog.no_button=Don't show again
 toolkit.sso_expire.dialog.title=Connection Expired