Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.project.DumbAwareAction
import com.intellij.openapi.wm.ToolWindowManager
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
Expand Down Expand Up @@ -39,7 +40,7 @@ abstract class SignInToQActionBase(actionName: String) : DumbAwareAction(actionN
UiTelemetry.click(project, "auth_start_Q")
val connectionManager = ToolkitConnectionManager.getInstance(project)
connectionManager.activeConnectionForFeature(QConnection.getInstance())?.let {
reauthConnectionIfNeeded(project, it, isReAuth = true)
reauthConnectionIfNeeded(project, it, isReAuth = true, reauthSource = ReauthSource.CODEWHISPERER_STATUSBAR)
} ?: run {
runInEdt {
if (requestCredentialsForQ(project, isReauth = false)) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.maybeReauthProviderIfNeeded
Expand Down Expand Up @@ -171,7 +172,7 @@ object CodeWhispererUtil {
if (!isQExpired(project)) return false
val tokenProvider = tokenProvider(project) ?: return false
return try {
maybeReauthProviderIfNeeded(project, tokenProvider) {
maybeReauthProviderIfNeeded(project, ReauthSource.CODEWHISPERER, tokenProvider) {
runInEdt {
if (!CodeWhispererService.hasReAuthPromptBeenShown()) {
notifyConnectionExpiredRequestReauth(project)
Expand Down Expand Up @@ -256,7 +257,7 @@ object CodeWhispererUtil {
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
if (connection !is ManagedBearerSsoConnection) return
pluginAwareExecuteOnPooledThread {
reauthConnectionIfNeeded(project, connection, isReAuth = true)
reauthConnectionIfNeeded(project, connection, isReAuth = true, reauthSource = ReauthSource.CODEWHISPERER)
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import org.junit.jupiter.api.Test
import org.junit.jupiter.api.extension.RegisterExtension
import software.aws.toolkits.core.utils.test.aString
import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
Expand Down Expand Up @@ -51,7 +52,7 @@ class CodeWhispererUtilTest {
CodeWhispererUtil.reconnectCodeWhisperer(projectExtension.project)

verify {
reauthConnectionIfNeeded(projectExtension.project, mockConnection, isReAuth = true)
reauthConnectionIfNeeded(projectExtension.project, mockConnection, isReAuth = true, reauthSource = ReauthSource.CODEWHISPERER)
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -273,6 +273,11 @@
"type": "string",
"description": "Comma-delimited list of enabled auths."
},
{
"name": "authRefreshSource",
"type": "string",
"description": "Source triggering token refresh"
},
{
"name": "component",
"allowedValues": [
Expand Down Expand Up @@ -813,6 +818,17 @@
"required": true
}
]
},
{
"name": "auth_sourceOfRefresh",
"description": "Source of user triggered refresh",
"metadata": [
{
"type": "authRefreshSource",
"required": false
}
]
}

]
}
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,7 @@ fun authAndUpdateConfig(

val connection = try {
ToolkitAuthManager.getInstance().tryCreateTransientSsoConnection(updatedProfile) { connection ->
reauthConnectionIfNeeded(project, connection, onPendingToken)
reauthConnectionIfNeeded(project, connection, onPendingToken, reauthSource = ReauthSource.FRESH_AUTH)
}
} catch (e: Exception) {
onError(e)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenPr
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.InteractiveBearerTokenProvider
import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
import software.aws.toolkits.resources.AwsCoreBundle
import software.aws.toolkits.telemetry.AuthTelemetry
import software.aws.toolkits.telemetry.CredentialSourceId
import software.aws.toolkits.telemetry.CredentialType
import software.aws.toolkits.telemetry.Result
Expand Down Expand Up @@ -135,6 +136,7 @@ fun loginSso(
onPendingToken = onPendingToken,
isReAuth = false,
source = metadata?.sourceId,
reauthSource = ReauthSource.FRESH_AUTH
)
}
} catch (e: Exception) {
Expand Down Expand Up @@ -182,6 +184,7 @@ fun loginSso(
onPendingToken = onPendingToken,
isReAuth = true,
source = metadata?.sourceId,
reauthSource = ReauthSource.COMMON_LOGIN
)
return@let connection
}
Expand Down Expand Up @@ -240,6 +243,7 @@ fun reauthConnectionIfNeeded(
onPendingToken: (InteractiveBearerTokenProvider) -> Unit = {},
isReAuth: Boolean = false,
source: String? = null,
reauthSource: ReauthSource? = ReauthSource.TOOLKIT,
): BearerTokenProvider {
val tokenProvider = (connection.getConnectionSettings() as TokenConnectionSettings).tokenProvider.delegate as BearerTokenProvider
if (tokenProvider is InteractiveBearerTokenProvider) {
Expand All @@ -248,7 +252,7 @@ fun reauthConnectionIfNeeded(

val startUrl = (connection as AwsBearerTokenConnection).startUrl
var didReauth = false
maybeReauthProviderIfNeeded(project, tokenProvider) {
maybeReauthProviderIfNeeded(project, reauthSource, tokenProvider) {
didReauth = true
runUnderProgressIfNeeded(project, AwsCoreBundle.message("credentials.pending.title"), true) {
try {
Expand Down Expand Up @@ -301,6 +305,7 @@ fun reauthConnectionIfNeeded(
// Return true if need to re-auth, false otherwise
fun maybeReauthProviderIfNeeded(
project: Project?,
reauthSource: ReauthSource? = ReauthSource.TOOLKIT,
tokenProvider: BearerTokenProvider,
onReauthRequired: (SsoOidcException?) -> Any,
): Boolean {
Expand All @@ -314,12 +319,14 @@ fun maybeReauthProviderIfNeeded(

BearerTokenAuthState.NEEDS_REFRESH -> {
try {
getLogger<ToolkitAuthManager>().warn { "Starting token refresh" }
return runUnderProgressIfNeeded(project, AwsCoreBundle.message("credentials.refreshing"), true) {
tokenProvider.resolveToken()
BearerTokenProviderListener.notifyCredUpdate(tokenProvider.id)
return@runUnderProgressIfNeeded false
}
} catch (e: SsoOidcException) {
AuthTelemetry.sourceOfRefresh(authRefreshSource = reauthSource.toString())
getLogger<ToolkitAuthManager>().warn(e) { "Redriving bearer token login flow since token could not be refreshed" }
onReauthRequired(e)
return true
Expand All @@ -332,6 +339,17 @@ fun maybeReauthProviderIfNeeded(
}
}

enum class ReauthSource {
CODEWHISPERER,
TOOLKIT,
Q_CHAT,
LOGIN_BROWSER,
CODEWHISPERER_STATUSBAR,
CODECATALYST,
COMMON_LOGIN,
FRESH_AUTH,
}

fun deleteSsoConnection(connection: ProfileSsoManagedBearerSsoConnection) =
deleteSsoConnection(connection.configSessionName)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ import software.aws.toolkits.jetbrains.core.credentials.CredentialManager
import software.aws.toolkits.jetbrains.core.credentials.InteractiveCredential
import software.aws.toolkits.jetbrains.core.credentials.MfaRequiredInteractiveCredentials
import software.aws.toolkits.jetbrains.core.credentials.PostValidateInteractiveCredential
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
import software.aws.toolkits.jetbrains.core.credentials.RefreshConnectionAction
import software.aws.toolkits.jetbrains.core.credentials.SsoRequiredInteractiveCredentials
import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
Expand Down Expand Up @@ -133,7 +134,7 @@ class ProfileCredentialsIdentifierSso @TestOnly constructor(
scopes = session.scopes.toList()
)
)
reauthConnectionIfNeeded(e.project, connection)
reauthConnectionIfNeeded(e.project, connection, reauthSource = ReauthSource.TOOLKIT)
RefreshConnectionAction().actionPerformed(e)
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,6 @@ class DiskCache(
LOG.debug { "loadClientRegistration for $cacheKey" }
val inputStream = clientRegistrationCache(cacheKey).tryInputStreamIfExists()
?: return null

return loadClientRegistration(inputStream)
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -432,6 +432,7 @@ class SsoAccessTokenProvider(
requestId = requestId,
result = Result.Failed
)
getLogger<SsoAccessTokenProvider>().warn("RefreshAccessTokenFailed: ${e.message}")
throw e
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import software.aws.toolkits.core.utils.tryOrNull
import software.aws.toolkits.jetbrains.core.credentials.LegacyManagedBearerSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.ProfileSsoManagedBearerSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.loginSso
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
Expand Down Expand Up @@ -230,7 +231,7 @@ fun reauthenticateWithQ(project: Project) {
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
if (connection !is ManagedBearerSsoConnection) return
pluginAwareExecuteOnPooledThread {
reauthConnectionIfNeeded(project, connection, isReAuth = true)
reauthConnectionIfNeeded(project, connection, isReAuth = true, reauthSource = ReauthSource.Q_CHAT)
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import software.aws.toolkits.core.utils.tryOrNull
import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.Login
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeCatalystConnection
Expand Down Expand Up @@ -380,7 +381,7 @@ abstract class LoginBrowser(
protected fun reauth(connection: ToolkitConnection?) {
if (connection is AwsBearerTokenConnection) {
loginWithBackgroundContext {
reauthConnectionIfNeeded(project, connection, onPendingToken, isReAuth = true)
reauthConnectionIfNeeded(project, connection, onPendingToken, isReAuth = true, reauthSource = ReauthSource.LOGIN_BROWSER)
}
stopAndClearBrowserOpenTimer()
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,7 @@ class ConnectionSettingsMenuBuilder private constructor() {
addAll(
object : DumbAwareAction(message("credentials.individual_identity.reconnect")) {
override fun actionPerformed(e: AnActionEvent) {
reauthConnectionIfNeeded(e.project, value, isReAuth = true)
reauthConnectionIfNeeded(e.project, value, isReAuth = true, reauthSource = ReauthSource.TOOLKIT)

e.project?.let { ToolkitConnectionManager.getInstance(it).switchConnection(value) }
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import software.aws.toolkits.core.telemetry.DefaultMetricEvent
import software.aws.toolkits.core.utils.tryOrNull
import software.aws.toolkits.jetbrains.core.AwsResourceCache
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.logoutFromSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.maybeReauthProviderIfNeeded
Expand Down Expand Up @@ -78,7 +79,7 @@ class CodeCatalystCredentialManager {
fun promptAuth(): BearerTokenProvider? {
connection()?.let {
val tokenProvider = provider(it)
val reauthRequired = maybeReauthProviderIfNeeded(project, tokenProvider) {}
val reauthRequired = maybeReauthProviderIfNeeded(project, ReauthSource.CODECATALYST, tokenProvider) {}
if (reauthRequired) {
val useCurrentCredentials =
computeOnEdt {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import com.intellij.icons.AllIcons
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.project.DumbAwareAction
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeCatalystConnection
import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
Expand All @@ -23,7 +24,7 @@ class SonoLogin : DumbAwareAction(AllIcons.Actions.Execute) {
pluginAwareExecuteOnPooledThread {
val connectionManager = ToolkitConnectionManager.getInstance(project)
connectionManager.activeConnectionForFeature(CodeCatalystConnection.getInstance())?.let {
reauthConnectionIfNeeded(project, it, isReAuth = true)
reauthConnectionIfNeeded(project, it, isReAuth = true, reauthSource = ReauthSource.CODECATALYST)
project.refreshDevToolTree()
} ?: run {
runInEdt {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import com.intellij.openapi.ui.Messages
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.ConfigFilesFacade
import software.aws.toolkits.jetbrains.core.credentials.DefaultConfigFilesFacade
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
import software.aws.toolkits.jetbrains.core.credentials.UserConfigSsoSessionProfile
import software.aws.toolkits.jetbrains.core.credentials.authAndUpdateConfig
import software.aws.toolkits.jetbrains.core.credentials.profiles.SsoSessionConstants
Expand Down Expand Up @@ -42,7 +43,7 @@ fun rolePopupFromConnection(
Messages.showErrorDialog(project, e.message, message("gettingstarted.explorer.iam.add"))
} ?: return@runInEdt
} else {
reauthConnectionIfNeeded(project, connection)
reauthConnectionIfNeeded(project, connection, reauthSource = ReauthSource.TOOLKIT)
connection
}.getConnectionSettings().tokenProvider

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ 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.ReauthSource
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
Expand Down Expand Up @@ -117,7 +118,7 @@ class IamAuth : DatabaseAuthProviderCompatabilityAdapter {
)
)

reauthConnectionIfNeeded(project, ssoConnection)
reauthConnectionIfNeeded(project, ssoConnection, reauthSource = ReauthSource.TOOLKIT)
return connection
}

Expand Down
Loading