Centralize bearer token auth logic and add more logs (#3538)

rli · web-flow · commit eef50ee409d7 · 2023-03-23T16:01:24.000-07:00
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ConnectionSettingsMenuBuilder.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ConnectionSettingsMenuBuilder.kt
@@ -241,7 +241,7 @@ class ConnectionSettingsMenuBuilder private constructor() {
             addAll(
                 object : DumbAwareAction(message("credentials.individual_identity.reconnect")) {
                     override fun actionPerformed(e: AnActionEvent) {
-                        reauthProviderIfNeeded(value)
+                        reauthProviderIfNeeded(e.project, value)
                     }
                 },
 
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt
@@ -15,6 +15,7 @@ import software.aws.toolkits.core.TokenConnectionSettings
 import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
+import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.core.credentials.pinning.FeatureWithPinnedConnection
 import software.aws.toolkits.jetbrains.core.credentials.sono.ALL_SONO_SCOPES
 import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
@@ -165,7 +166,7 @@ fun loginSso(project: Project?, startUrl: String, scopes: List<String> = ALL_SON
 }
 
 private fun reauthConnection(project: Project?, connection: ToolkitConnection): BearerTokenProvider {
-    val provider = reauthProviderIfNeeded(connection)
+    val provider = reauthProviderIfNeeded(project, connection)
 
     ToolkitConnectionManager.getInstance(project).switchConnection(connection)
 
@@ -204,23 +205,42 @@ fun AwsBearerTokenConnection.lazyIsUnauthedBearerConnection(): Boolean {
     return false
 }
 
-fun reauthProviderIfNeeded(connection: ToolkitConnection): BearerTokenProvider {
+fun reauthProviderIfNeeded(project: Project?, connection: ToolkitConnection): BearerTokenProvider {
     val tokenProvider = (connection.getConnectionSettings() as TokenConnectionSettings).tokenProvider.delegate as BearerTokenProvider
+
+    return reauthProviderIfNeeded(project, tokenProvider)
+}
+
+fun reauthProviderIfNeeded(project: Project?, tokenProvider: BearerTokenProvider): BearerTokenProvider {
+    maybeReauthProviderIfNeeded(project, tokenProvider) {
+        runUnderProgressIfNeeded(project, message("credentials.sono.login.pending"), true) {
+            tokenProvider.reauthenticate()
+        }
+    }
+
+    return tokenProvider
+}
+
+fun maybeReauthProviderIfNeeded(project: Project?, tokenProvider: BearerTokenProvider, onReauthRequired: (SsoOidcException?) -> Any) {
     val state = tokenProvider.state()
-    runUnderProgressIfNeeded(null, message("settings.states.validating.short"), true) {
-        if (state == BearerTokenAuthState.NEEDS_REFRESH) {
+    when (state) {
+        BearerTokenAuthState.NOT_AUTHENTICATED -> {
+            getLogger<ToolkitAuthManager>().info { "Token provider NOT_AUTHENTICATED, requesting login" }
+            onReauthRequired(null)
+        }
+
+        BearerTokenAuthState.NEEDS_REFRESH -> {
             try {
-                tokenProvider.resolveToken()
-                BearerTokenProviderListener.notifyCredUpdate(tokenProvider.id)
+                runUnderProgressIfNeeded(project, message("credentials.sono.login.refreshing"), true) {
+                    tokenProvider.resolveToken()
+                    BearerTokenProviderListener.notifyCredUpdate(tokenProvider.id)
+                }
             } catch (e: SsoOidcException) {
-                tokenProvider.reauthenticate()
+                getLogger<ToolkitAuthManager>().warn(e) { "Redriving AWS Builder ID login flow since token could not be refreshed" }
+                onReauthRequired(e)
             }
-        } else if (state == BearerTokenAuthState.NOT_AUTHENTICATED) {
-            tokenProvider.reauthenticate()
         }
 
-        Unit
+        BearerTokenAuthState.AUTHORIZED -> {}
     }
-
-    return tokenProvider
 }
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sono/SonoCredentialManager.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sono/SonoCredentialManager.kt
@@ -5,22 +5,18 @@ package software.aws.toolkits.jetbrains.core.credentials.sono
 
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
-import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
 import software.aws.toolkits.core.TokenConnectionSettings
 import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
 import software.aws.toolkits.core.telemetry.DefaultMetricEvent
-import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.tryOrNull
-import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.core.AwsResourceCache
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
-import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.loginSso
 import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeCatalystConnection
+import software.aws.toolkits.jetbrains.core.credentials.reauthProviderIfNeeded
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
-import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
 import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
 import software.aws.toolkits.jetbrains.services.caws.CawsResources
 import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
@@ -72,35 +68,11 @@ class SonoCredentialManager {
     fun getProviderAndPromptAuth(): BearerTokenProvider {
         val provider = provider()
         return when (provider?.state()) {
-            null -> runUnderProgressIfNeeded(null, message("credentials.sono.login.pending"), true) {
+            null -> runUnderProgressIfNeeded(project, message("credentials.sono.login.pending"), true) {
                 loginSso(project, SONO_URL, ALL_SONO_SCOPES)
             }
 
-            BearerTokenAuthState.NOT_AUTHENTICATED -> {
-                runUnderProgressIfNeeded(null, message("credentials.sono.login.pending"), true) {
-                    provider.reauthenticate()
-                }
-
-                return provider
-            }
-
-            BearerTokenAuthState.NEEDS_REFRESH -> {
-                try {
-                    runUnderProgressIfNeeded(null, message("credentials.sono.login.refreshing"), true) {
-                        provider.resolveToken()
-                        BearerTokenProviderListener.notifyCredUpdate(provider.id)
-                    }
-                } catch (e: SsoOidcException) {
-                    LOG.warn(e) { "Redriving AWS Builder ID login flow since token could not be refreshed" }
-                    runUnderProgressIfNeeded(null, message("credentials.sono.login.pending"), true) {
-                        provider.reauthenticate()
-                    }
-                }
-
-                return provider
-            }
-
-            BearerTokenAuthState.AUTHORIZED -> provider
+            else -> reauthProviderIfNeeded(project, provider)
         }
     }
 
@@ -112,13 +84,9 @@ class SonoCredentialManager {
         fun getInstance(project: Project? = null) = project?.let { it.service<SonoCredentialManager>() } ?: service()
 
         fun loginSono(project: Project?) {
-            val provider = SonoCredentialManager.getInstance(project).getProviderAndPromptAuth()
-            val connection = ToolkitAuthManager.getInstance().getConnection(provider.id) ?: error("Mismatch between provider id and connection id")
-            ToolkitConnectionManager.getInstance(project).switchConnection(connection)
+            SonoCredentialManager.getInstance(project).getProviderAndPromptAuth()
         }
 
-        private val LOG = getLogger<SonoCredentialManager>()
-
         private val SSO_REGION by lazy {
             with(AwsRegionProvider.getInstance()) {
                 get(SONO_REGION) ?: throw RuntimeException("AwsRegionProvider was unable to provide SSO_REGION for AWS Builder ID")
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
@@ -141,10 +141,10 @@ class SsoAccessTokenProvider(
 
     fun refreshToken(currentToken: AccessToken): AccessToken {
         if (currentToken.refreshToken == null) {
-            throw InvalidRequestException.builder().build()
+            throw InvalidRequestException.builder().message("Requested token refresh, but refresh token was null").build()
         }
 
-        val registration = loadClientRegistration() ?: throw InvalidClientException.builder().build()
+        val registration = loadClientRegistration() ?: throw InvalidClientException.builder().message("Unable to load client registration").build()
 
         val newToken = client.createToken {
             it.clientId(registration.clientId)
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt
@@ -9,14 +9,14 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.project.Project
 import software.amazon.awssdk.services.codewhisperer.model.Recommendation
-import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.BearerSsoConnection
 import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.loginSso
 import software.aws.toolkits.jetbrains.core.credentials.logoutFromSsoConnection
+import software.aws.toolkits.jetbrains.core.credentials.maybeReauthProviderIfNeeded
 import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
@@ -29,7 +29,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.actions.DoNotShowA
 import software.aws.toolkits.jetbrains.utils.notifyError
 import software.aws.toolkits.jetbrains.utils.notifyInfo
 import software.aws.toolkits.jetbrains.utils.notifyWarn
-import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.CodewhispererCompletionType
 import java.net.URI
@@ -104,19 +103,7 @@ object CodeWhispererUtil {
         val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
         if (connection !is BearerSsoConnection) return
         val tokenProvider = tokenProvider(project) ?: return
-        val state = tokenProvider.state()
-        if (state == BearerTokenAuthState.NEEDS_REFRESH) {
-            try {
-                runUnderProgressIfNeeded(null, message("settings.states.validating.short"), false) {
-                    tokenProvider.resolveToken()
-                }
-            } catch (e: SsoOidcException) {
-                runInEdt {
-                    notifyConnectionExpired(project, connection)
-                    callback()
-                }
-            }
-        } else if (state == BearerTokenAuthState.NOT_AUTHENTICATED) {
+        maybeReauthProviderIfNeeded(project, tokenProvider) {
             runInEdt {
                 notifyConnectionExpired(project, connection)
                 callback()
diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/DefaultToolkitAuthManagerTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/DefaultToolkitAuthManagerTest.kt
@@ -291,6 +291,7 @@ class DefaultToolkitAuthManagerTest {
 
         mockConstruction(InteractiveBearerTokenProvider::class.java) { context, _ ->
             doNothing().whenever(context).reauthenticate()
+            whenever(context.state()).thenReturn(BearerTokenAuthState.NOT_AUTHENTICATED)
         }.use {
             // before
             assertThat(sut.listConnections()).hasSize(0)

Original file line number	Diff line number	Diff line change
`@@ -241,7 +241,7 @@ class ConnectionSettingsMenuBuilder private constructor() {`
`241`	`241`	`addAll(`
`242`	`242`	`object : DumbAwareAction(message("credentials.individual_identity.reconnect")) {`
`243`	`243`	`override fun actionPerformed(e: AnActionEvent) {`
`244`		`- reauthProviderIfNeeded(value)`
	`244`	`+ reauthProviderIfNeeded(e.project, value)`
`245`	`245`	`}`
`246`	`246`	`},`
`247`	`247`