Fix AWS SSO connection with RDS using IAM (#4278)

* Fix AWS SSO connection with RDS using IAM

* Removed unUseed message in messageBundles file

* addressing comments

* Addressed test case

* Exception Updated

* detekt failed

* fix failed test

Author: vchikoti1998

.changes/next-release/bugfix-7d06545e-eb51-4a84-9333-96a7add7e78c.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Fix AWS SSO connection when authenticating with RDS using an IAM Identity Center profile (#4145)"
+}

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/profiles/ProfileCredentialProviderFactory.kt

@@ -51,6 +51,7 @@ import software.aws.toolkits.jetbrains.core.credentials.profiles.SsoSessionConst
 import software.aws.toolkits.jetbrains.core.credentials.profiles.SsoSessionConstants.SSO_SESSION_SECTION_NAME
 import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
 import software.aws.toolkits.jetbrains.core.credentials.sso.SsoCache
+import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.NoTokenInitializedException
 import software.aws.toolkits.jetbrains.settings.AwsSettings
 import software.aws.toolkits.jetbrains.settings.ProfilesNotification
 import software.aws.toolkits.jetbrains.utils.createNotificationExpiringAction
@@ -93,7 +94,7 @@ class ProfileCredentialsIdentifierSso @TestOnly constructor(
 
     override fun handleValidationException(e: Exception): ConnectionState.RequiresUserAction? {
         // in the new SSO flow, we must attempt validation before knowing if user action is truly required
-        if (findUpException<SsoOidcException>(e) || findUpException<IllegalStateException>(e)) {
+        if (findUpException<SsoOidcException>(e) || findUpException<IllegalStateException>(e) || findUpException<NoTokenInitializedException>(e)) {
             return ConnectionState.RequiresUserAction(
                 object : InteractiveCredential, CredentialIdentifier by this {
                     override val userActionDisplayMessage = message("credentials.sso.display", displayName)

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/bearer/BearerTokenProvider.kt

@@ -137,7 +137,7 @@ class InteractiveBearerTokenProvider(
     }
 
     private fun refreshToken(): RefreshResult<out SdkToken> {
-        val lastToken = lastToken.get() ?: error("Token refresh started before session initialized")
+        val lastToken = lastToken.get() ?: throw NoTokenInitializedException("Token refresh started before session initialized")
         val token = if (Duration.between(Instant.now(), lastToken.expiresAt) > Duration.ofMinutes(30)) {
             lastToken
         } else {
@@ -167,7 +167,7 @@ class InteractiveBearerTokenProvider(
      * Only use if you know what you're doing.
      */
     override fun refresh(): AccessToken {
-        val lastToken = lastToken.get() ?: error("Token refresh started before session initialized")
+        val lastToken = lastToken.get() ?: throw NoTokenInitializedException("Token refresh started before session initialized")
         return accessTokenProvider.refreshToken(lastToken).also {
             this.lastToken.set(it)
         }
@@ -189,6 +189,8 @@ class InteractiveBearerTokenProvider(
     }
 }
 
+class NoTokenInitializedException(message: String) : Exception(message)
+
 public enum class BearerTokenAuthState {
     AUTHORIZED,
     NEEDS_REFRESH,

plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/profiles/ProfileCredentialsIdentifierSsoTest.kt

@@ -15,6 +15,7 @@ import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
 import software.aws.toolkits.jetbrains.core.MockClientManagerExtension
 import software.aws.toolkits.jetbrains.core.credentials.sso.DiskCache
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.InteractiveBearerTokenProvider
+import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.NoTokenInitializedException
 
 @ExtendWith(ApplicationExtension::class)
 class ProfileCredentialsIdentifierSsoTest {
@@ -46,7 +47,7 @@ class ProfileCredentialsIdentifierSsoTest {
         mockClientManager.create<SsoOidcClient>()
 
         // IllegalStateException instead of more general base Exception so we know if the type changes
-        val exception = assertThrows<IllegalStateException> {
+        val exception = assertThrows<NoTokenInitializedException> {
             InteractiveBearerTokenProvider("", "us-east-1", emptyList(), cache = cache, id = "test").resolveToken()
         }
         assertThat(sut.handleValidationException(exception)).isNotNull()

plugins/core/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/core/credentials/MockCredentialsManager.kt

@@ -16,6 +16,7 @@ import software.aws.toolkits.core.credentials.CredentialIdentifierBase
 import software.aws.toolkits.core.credentials.CredentialProviderFactory
 import software.aws.toolkits.core.credentials.CredentialSourceId
 import software.aws.toolkits.core.credentials.CredentialsChangeListener
+import software.aws.toolkits.core.credentials.SsoSessionIdentifier
 import software.aws.toolkits.core.credentials.ToolkitCredentialsProvider
 import software.aws.toolkits.core.region.AwsRegion
 import software.aws.toolkits.core.utils.test.aString
@@ -49,6 +50,20 @@ class MockCredentialsManager : CredentialManager() {
         addProvider(it)
     }
 
+    fun addCredentials(
+        credentialIdentifier: CredentialIdentifier
+    ): CredentialIdentifier {
+        addProvider(credentialIdentifier)
+        return credentialIdentifier
+    }
+
+    fun addSsoProvider(
+        ssoSessionIdentifier: SsoSessionIdentifier
+    ): SsoSessionIdentifier {
+        super.addSsoSession(ssoSessionIdentifier)
+        return ssoSessionIdentifier
+    }
+
     fun createCredentialProvider(
         id: String = aString(),
         credentials: AwsCredentials,
@@ -113,6 +128,14 @@ open class MockCredentialManagerRule : ApplicationRule() {
         regionId: String? = null
     ): MockCredentialsManager.MockCredentialIdentifier = credentialManager.addCredentials(id, credentials, regionId)
 
+    fun addCredentials(
+        credentialIdentifier: CredentialIdentifier
+    ): CredentialIdentifier = credentialManager.addCredentials(credentialIdentifier)
+
+    fun addSsoProvider(
+        ssoSessionIdentifier: SsoSessionIdentifier
+    ): SsoSessionIdentifier = credentialManager.addSsoProvider(ssoSessionIdentifier)
+
     fun createCredentialProvider(
         id: String = aString(),
         credentials: AwsCredentials = AwsBasicCredentials.create("Access", "Secret"),

plugins/toolkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/services/rds/auth/IamAuth.kt

@@ -5,10 +5,14 @@ package software.aws.toolkits.jetbrains.services.rds.auth
 
 import com.intellij.credentialStore.Credentials
 import com.intellij.database.access.DatabaseCredentials
+import com.intellij.database.connection.throwable.KnownDatabaseException
+import com.intellij.database.connection.throwable.info.ErrorInfo
+import com.intellij.database.connection.throwable.info.SimpleErrorInfo
 import com.intellij.database.dataSource.DatabaseAuthProvider.AuthWidget
 import com.intellij.database.dataSource.DatabaseConnectionInterceptor.ProtoConnection
 import com.intellij.database.dataSource.DatabaseCredentialsAuthProvider
 import com.intellij.database.dataSource.LocalDataSource
+import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.project.Project
 import kotlinx.coroutines.future.future
 import software.amazon.awssdk.regions.Region
@@ -17,6 +21,12 @@ import software.aws.toolkits.core.ConnectionSettings
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
+import software.aws.toolkits.jetbrains.core.credentials.CredentialManager
+import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
+import software.aws.toolkits.jetbrains.core.credentials.UserConfigSsoSessionProfile
+import software.aws.toolkits.jetbrains.core.credentials.profiles.ProfileCredentialsIdentifierSso
+import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
+import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.NoTokenInitializedException
 import software.aws.toolkits.jetbrains.datagrip.auth.compatability.DatabaseAuthProviderCompatabilityAdapter
 import software.aws.toolkits.jetbrains.datagrip.auth.compatability.project
 import software.aws.toolkits.jetbrains.datagrip.getAwsConnectionSettings
@@ -49,6 +59,17 @@ class IamAuth : DatabaseAuthProviderCompatabilityAdapter {
 
     override fun createWidget(project: Project?, credentials: DatabaseCredentials, dataSource: LocalDataSource): AuthWidget? = IamAuthWidget()
 
+    inner class SsoNoTokenFix(val project: Project, val connection: ProtoConnection) : ErrorInfo.Fix {
+
+        override fun getName(): String = message("credentials.sso.login.session", getAuthInformation(connection).connectionSettings.credentials.id)
+
+        override fun isSilent(): Boolean = false
+
+        override fun apply(dataContext: DataContext) {
+            handleSsoAuthentication(project, connection)
+        }
+    }
+
     override fun intercept(
         connection: ProtoConnection,
         silent: Boolean
@@ -63,13 +84,43 @@ class IamAuth : DatabaseAuthProviderCompatabilityAdapter {
                 DatabaseCredentialsAuthProvider.applyCredentials(connection, credentials, true)
             } catch (e: Throwable) {
                 result = Result.Failed
-                throw e
+                if (e is NoTokenInitializedException) {
+                    val simpleErrorInfo = SimpleErrorInfo(
+                        message("rds.validation.iam_sso_connection.error_info"),
+                        e,
+                        listOf(SsoNoTokenFix(project, connection))
+                    )
+                    throw KnownDatabaseException(simpleErrorInfo)
+                } else {
+                    throw e
+                }
             } finally {
                 RdsTelemetry.getCredentials(project = project, result = result, databaseCredentials = IAM, databaseEngine = connection.getDatabaseEngine())
             }
         }
     }
 
+    fun handleSsoAuthentication(project: Project, connection: ProtoConnection): ProtoConnection {
+        val authInformation = getAuthInformation(connection)
+        val profileCredentials =
+            CredentialManager.getInstance().getCredentialIdentifierById(authInformation.connectionSettings.credentials.id) as ProfileCredentialsIdentifierSso
+
+        val session = CredentialManager.getInstance()
+            .getSsoSessionIdentifiers()
+            .first { it.id == profileCredentials.sessionIdentifier }
+        val ssoConnection = ToolkitAuthManager.getInstance().getOrCreateSsoConnection(
+            UserConfigSsoSessionProfile(
+                configSessionName = profileCredentials.ssoSessionName,
+                ssoRegion = session.ssoRegion,
+                startUrl = session.startUrl,
+                scopes = session.scopes.toList()
+            )
+        )
+
+        reauthConnectionIfNeeded(project, ssoConnection)
+        return connection
+    }
+
     internal fun getAuthInformation(connection: ProtoConnection): RdsAuth {
         validateIamConfiguration(connection)
         val signingUrl = connection.connectionPoint.additionalProperties[RDS_SIGNING_HOST_PROPERTY]

plugins/toolkit/jetbrains-ultimate/tst/software/aws/toolkits/jetbrains/services/rds/auth/IamAuthTest.kt

@@ -19,16 +19,25 @@ import org.mockito.kotlin.doAnswer
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.mock
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials
+import software.amazon.awssdk.services.ssooidc.SsoOidcClient
+import software.aws.toolkits.core.credentials.CredentialType
 import software.aws.toolkits.core.region.AwsRegion
 import software.aws.toolkits.core.utils.RuleUtils
 import software.aws.toolkits.core.utils.unwrap
+import software.aws.toolkits.jetbrains.core.MockClientManagerRule
 import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule
+import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule
+import software.aws.toolkits.jetbrains.core.credentials.diskCache
+import software.aws.toolkits.jetbrains.core.credentials.profiles.ProfileCredentialsIdentifierSso
+import software.aws.toolkits.jetbrains.core.credentials.profiles.ProfileSsoSessionIdentifier
+import software.aws.toolkits.jetbrains.core.credentials.sso.DeviceAuthorizationGrantToken
 import software.aws.toolkits.jetbrains.core.region.MockRegionProviderRule
 import software.aws.toolkits.jetbrains.datagrip.CREDENTIAL_ID_PROPERTY
 import software.aws.toolkits.jetbrains.datagrip.REGION_ID_PROPERTY
 import software.aws.toolkits.jetbrains.datagrip.RequireSsl
 import software.aws.toolkits.jetbrains.datagrip.auth.compatability.project
 import software.aws.toolkits.resources.message
+import java.time.Instant
 
 class IamAuthTest {
     @Rule
@@ -53,10 +62,34 @@ class IamAuthTest {
     @JvmField
     val credentialManager = MockCredentialManagerRule()
 
+    @Rule
+    @JvmField
+    val authManager = MockToolkitAuthManagerRule()
+
+    @JvmField
+    @Rule
+    val mockClientManager = MockClientManagerRule()
+
+    private lateinit var ssoClient: SsoOidcClient
+
     @Before
     fun setUp() {
         credentialManager.addCredentials(credentialId, mockCreds)
         regionProvider.addRegion(AwsRegion(defaultRegion, RuleUtils.randomName(), RuleUtils.randomName()))
+        ssoClient = mockClientManager.create()
+    }
+
+    @Test
+    fun `Handle Sso authentication no token present`() {
+        val noTokenCredentialId = RuleUtils.randomName()
+        val ssoUrl = RuleUtils.randomName()
+        diskCache.saveAccessToken(ssoUrl, DeviceAuthorizationGrantToken(ssoUrl, "us-east-1", "access1", "refresh1", Instant.MAX))
+        credentialManager.addCredentials(ProfileCredentialsIdentifierSso(noTokenCredentialId, noTokenCredentialId, "us-east-1", CredentialType.SsoProfile))
+        credentialManager.addSsoProvider(ProfileSsoSessionIdentifier(noTokenCredentialId, ssoUrl, "us-east-1", setOf()))
+        val conneciton = buildConnection(hasCredentials = true, credentialId = "profile:" + noTokenCredentialId)
+
+        val connection = iamAuth.handleSsoAuthentication(projectRule.project, conneciton)
+        assertThat(connection).isNotNull
     }
 
     @Test
@@ -168,7 +201,8 @@ class IamAuthTest {
         hasCredentials: Boolean = true,
         hasBadHost: Boolean = false,
         hasSslConfig: Boolean = true,
-        dbmsType: Dbms = Dbms.POSTGRES
+        dbmsType: Dbms = Dbms.POSTGRES,
+        credentialId: String = this.credentialId
     ): ProtoConnection {
         val mockConnection = mock<LocalDataSource> {
             on { url } doReturn "jdbc:postgresql://$dbHost:$connectionPort/dev"

plugins/toolkit/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties

@@ -1455,6 +1455,7 @@ rds.port=RDS Port:
 rds.postgres=PostgreSQL
 rds.url=RDS Host:
 rds.validation.aurora_mysql_ssl_required=Aurora MySQL requires SSL to be enabled to connect
+rds.validation.iam_sso_connection.error_info=No token found, please login and try again
 rds.validation.no_iam_auth=Database {0} does not have IAM authentication enabled
 rds.validation.no_instance_host=No RDS database host specified
 rds.validation.no_instance_port=No RDS database port specified