Skip to content

Commit e75c8d6

Browse files
manodnyabmanodnyab
authored
Add flag to indicate reauth in metrics (aws#4819)
* Add flag to indicate reauth * feedback * unused lines deleted * fixed tests
1 parent 35a316a commit e75c8d6

File tree

10 files changed

+84
-53
lines changed

10 files changed

+84
-53
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/explorerActions/SignInToQAction.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ abstract class SignInToQActionBase(actionName: String) : DumbAwareAction(actionN
3939
UiTelemetry.click(project, "auth_start_Q")
4040
val connectionManager = ToolkitConnectionManager.getInstance(project)
4141
connectionManager.activeConnectionForFeature(QConnection.getInstance())?.let {
42-
reauthConnectionIfNeeded(project, it)
42+
reauthConnectionIfNeeded(project, it, isReAuth = true)
4343
} ?: run {
4444
runInEdt {
4545
if (requestCredentialsForQ(project)) {

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -273,7 +273,7 @@ object CodeWhispererUtil {
273273
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
274274
if (connection !is ManagedBearerSsoConnection) return
275275
pluginAwareExecuteOnPooledThread {
276-
reauthConnectionIfNeeded(project, connection)
276+
reauthConnectionIfNeeded(project, connection, isReAuth = true)
277277
}
278278
}
279279

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtilTest.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ class CodeWhispererUtilTest {
5151
CodeWhispererUtil.reconnectCodeWhisperer(projectExtension.project)
5252

5353
verify {
54-
reauthConnectionIfNeeded(projectExtension.project, mockConnection)
54+
reauthConnectionIfNeeded(projectExtension.project, mockConnection, isReAuth = true)
5555
}
5656
}
5757
}

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/LoginUtils.kt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import software.aws.toolkits.jetbrains.core.AwsClientManager
2222
import software.aws.toolkits.jetbrains.core.credentials.profiles.SsoSessionConstants
2323
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_REGION
2424
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL
25+
import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
2526
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.InteractiveBearerTokenProvider
2627
import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
2728
import software.aws.toolkits.resources.AwsCoreBundle
@@ -245,3 +246,6 @@ fun messageFromConfigFacadeError(e: Exception): Pair<String, String> {
245246

246247
return errorMessage to errorType
247248
}
249+
250+
fun getCredentialIdForTelemetry(connection: ToolkitConnection): CredentialSourceId =
251+
if (connection.isSono()) CredentialSourceId.AwsId else CredentialSourceId.IamIdentityCenter

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt‎

Lines changed: 34 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,7 @@ fun loginSso(
120120
onSuccess: () -> Unit = {},
121121
metadata: ConnectionMetadata? = null
122122
): AwsBearerTokenConnection? {
123+
val source = metadata
123124
fun createAndAuthNewConnection(profile: AuthProfile): AwsBearerTokenConnection? {
124125
val authManager = ToolkitAuthManager.getInstance()
125126
val connection = try {
@@ -129,19 +130,6 @@ fun loginSso(
129130
connection = transientConnection,
130131
onPendingToken = onPendingToken,
131132
)
132-
recordLoginWithBrowser(
133-
credentialStartUrl = transientConnection.startUrl,
134-
credentialSourceId = CredentialSourceId.AwsId,
135-
isReAuth = true,
136-
result = Result.Succeeded,
137-
source = metadata?.sourceId
138-
)
139-
recordAddConnection(
140-
credentialSourceId = CredentialSourceId.AwsId,
141-
isReAuth = true,
142-
result = Result.Failed,
143-
source = metadata?.sourceId
144-
)
145133
}
146134
} catch (e: Exception) {
147135
onError(e)
@@ -191,6 +179,7 @@ fun loginSso(
191179
reauthConnectionIfNeeded(
192180
project = project,
193181
connection = connection,
182+
isReAuth = true
194183
)
195184
return connection
196185
} ?: run {
@@ -243,15 +232,46 @@ fun reauthConnectionIfNeeded(
243232
project: Project?,
244233
connection: ToolkitConnection,
245234
onPendingToken: (InteractiveBearerTokenProvider) -> Unit = {},
235+
isReAuth: Boolean = false
246236
): BearerTokenProvider {
247237
val tokenProvider = (connection.getConnectionSettings() as TokenConnectionSettings).tokenProvider.delegate as BearerTokenProvider
248238
if (tokenProvider is InteractiveBearerTokenProvider) {
249239
onPendingToken(tokenProvider)
250240
}
251241

242+
val startUrl = (connection as AwsBearerTokenConnection).startUrl
252243
maybeReauthProviderIfNeeded(project, tokenProvider) {
253244
runUnderProgressIfNeeded(project, AwsCoreBundle.message("credentials.pending.title"), true) {
254-
tokenProvider.reauthenticate()
245+
try {
246+
tokenProvider.reauthenticate()
247+
if (isReAuth) {
248+
recordLoginWithBrowser(
249+
credentialStartUrl = startUrl,
250+
credentialSourceId = getCredentialIdForTelemetry(connection),
251+
isReAuth = true,
252+
result = Result.Succeeded
253+
)
254+
recordAddConnection(
255+
credentialSourceId = getCredentialIdForTelemetry(connection),
256+
isReAuth = true,
257+
result = Result.Succeeded
258+
)
259+
}
260+
} catch (e: Exception) {
261+
if (isReAuth) {
262+
recordLoginWithBrowser(
263+
credentialStartUrl = startUrl,
264+
credentialSourceId = getCredentialIdForTelemetry(connection),
265+
isReAuth = true,
266+
result = Result.Failed
267+
)
268+
recordAddConnection(
269+
credentialSourceId = getCredentialIdForTelemetry(connection),
270+
isReAuth = true,
271+
result = Result.Failed
272+
)
273+
}
274+
}
255275
}
256276
}
257277
return tokenProvider

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/gettingstarted/GettingStartedAuthUtils.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ fun reauthenticateWithQ(project: Project) {
225225
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
226226
if (connection !is ManagedBearerSsoConnection) return
227227
pluginAwareExecuteOnPooledThread {
228-
reauthConnectionIfNeeded(project, connection)
228+
reauthConnectionIfNeeded(project, connection, isReAuth = true)
229229
}
230230
}
231231

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/LoginBrowser.kt‎

Lines changed: 39 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeCatalystConn
3333
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
3434
import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
3535
import software.aws.toolkits.jetbrains.core.credentials.sono.CODECATALYST_SCOPES
36+
import software.aws.toolkits.jetbrains.core.credentials.sono.IDENTITY_CENTER_ROLE_ACCESS_SCOPE
3637
import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES
3738
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL
3839
import software.aws.toolkits.jetbrains.core.credentials.sso.PendingAuthorization
@@ -154,21 +155,28 @@ abstract class LoginBrowser(
154155
else -> ""
155156
}
156157

157-
private fun isReAuth(scopes: List<String>): Boolean = ToolkitConnectionManager.getInstance(project)
158+
private fun isReAuth(scopes: List<String>, requestedStartUrl: String): Boolean = ToolkitConnectionManager.getInstance(project)
158159
.let {
160+
val qConnection = it.activeConnectionForFeature(QConnection.getInstance()) as AwsBearerTokenConnection?
161+
val codeCatalystConnection = it.activeConnectionForFeature(CodeCatalystConnection.getInstance()) as AwsBearerTokenConnection?
159162
if (scopes.toSet().intersect(Q_SCOPES.toSet()).isNotEmpty()) {
160-
it.activeConnectionForFeature(QConnection.getInstance()) != null
163+
qConnection != null && qConnection.startUrl == requestedStartUrl
161164
} else if (scopes.toSet().intersect(CODECATALYST_SCOPES.toSet()).isNotEmpty()) {
162-
it.activeConnectionForFeature(CodeCatalystConnection.getInstance()) != null
165+
codeCatalystConnection != null && codeCatalystConnection.startUrl == requestedStartUrl
163166
} else {
164-
val activeCon = it.activeConnection()
165-
val ccCon = it.activeConnectionForFeature(CodeCatalystConnection.getInstance())
166-
val qCon = it.activeConnectionForFeature(QConnection.getInstance())
167-
activeCon != null && activeCon != ccCon && activeCon != qCon
167+
if (it.activeConnection() is AwsBearerTokenConnection) {
168+
val activeCon = it.activeConnection() as? AwsBearerTokenConnection
169+
return activeCon?.scopes?.toSet()?.intersect(IDENTITY_CENTER_ROLE_ACCESS_SCOPE.toSet())?.isNotEmpty() == true &&
170+
activeCon != qConnection &&
171+
activeCon != codeCatalystConnection
172+
} else {
173+
return false
174+
}
168175
}
169176
}
170177

171178
open fun loginBuilderId(scopes: List<String>) {
179+
val isReauth = isReAuth(scopes, SONO_URL)
172180
val onError: (Exception) -> Unit = { e ->
173181
stopAndClearBrowserOpenTimer()
174182
isUserCancellation(e)
@@ -177,12 +185,14 @@ abstract class LoginBrowser(
177185
credentialStartUrl = SONO_URL,
178186
result = Result.Failed,
179187
reason = e.message,
180-
credentialSourceId = CredentialSourceId.AwsId
188+
credentialSourceId = CredentialSourceId.AwsId,
189+
isReAuth = isReauth
181190
)
182191
AuthTelemetry.addConnection(
183192
result = Result.Failed,
184193
credentialSourceId = CredentialSourceId.AwsId,
185-
reason = e.message
194+
reason = e.message,
195+
isReAuth = isReauth
186196
)
187197
}
188198
val onSuccess: () -> Unit = {
@@ -191,11 +201,13 @@ abstract class LoginBrowser(
191201
project = null,
192202
credentialStartUrl = SONO_URL,
193203
result = Result.Succeeded,
194-
credentialSourceId = CredentialSourceId.AwsId
204+
credentialSourceId = CredentialSourceId.AwsId,
205+
isReAuth = isReauth
195206
)
196207
AuthTelemetry.addConnection(
197208
result = Result.Succeeded,
198-
credentialSourceId = CredentialSourceId.AwsId
209+
credentialSourceId = CredentialSourceId.AwsId,
210+
isReAuth = isReauth
199211
)
200212
}
201213

@@ -208,7 +220,20 @@ abstract class LoginBrowser(
208220

209221
open fun loginIdC(url: String, region: AwsRegion, scopes: List<String>) {
210222
// assumes scopes contains either Q or non-Q permissions but not both
211-
val isReAuth = isReAuth(scopes)
223+
val (onError: (Exception) -> Unit, onSuccess: () -> Unit) = getSuccessAndErrorActionsForIdcLogin(scopes, url, region)
224+
loginWithBackgroundContext {
225+
Login
226+
.IdC(url, region, scopes, onPendingToken, onSuccess, onError)
227+
.login(project)
228+
}
229+
}
230+
231+
fun getSuccessAndErrorActionsForIdcLogin(
232+
scopes: List<String>,
233+
url: String,
234+
region: AwsRegion
235+
): Pair<(Exception) -> Unit, () -> Unit> {
236+
val isReAuth = isReAuth(scopes, url)
212237

213238
val onError: (Exception) -> Unit = { e ->
214239
stopAndClearBrowserOpenTimer()
@@ -257,11 +282,7 @@ abstract class LoginBrowser(
257282
credentialSourceId = CredentialSourceId.IamIdentityCenter
258283
)
259284
}
260-
loginWithBackgroundContext {
261-
Login
262-
.IdC(url, region, scopes, onPendingToken, onSuccess, onError)
263-
.login(project)
264-
}
285+
return Pair(onError, onSuccess)
265286
}
266287

267288
open fun loginIAM(profileName: String, accessKey: String, secretKey: String) {
@@ -330,7 +351,7 @@ abstract class LoginBrowser(
330351
protected fun reauth(connection: ToolkitConnection?) {
331352
if (connection is AwsBearerTokenConnection) {
332353
loginWithBackgroundContext {
333-
reauthConnectionIfNeeded(project, connection, onPendingToken)
354+
reauthConnectionIfNeeded(project, connection, onPendingToken, isReAuth = true)
334355
}
335356
stopAndClearBrowserOpenTimer()
336357
}

‎plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ConnectionSettingsMenuBuilder.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,7 @@ class ConnectionSettingsMenuBuilder private constructor() {
244244
addAll(
245245
object : DumbAwareAction(message("credentials.individual_identity.reconnect")) {
246246
override fun actionPerformed(e: AnActionEvent) {
247-
reauthConnectionIfNeeded(e.project, value)
247+
reauthConnectionIfNeeded(e.project, value, isReAuth = true)
248248

249249
ToolkitConnectionManager.getInstance(e.project).switchConnection(value)
250250
}

‎plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/core/explorer/devToolsTab/nodes/actions/SonoLogin.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ class SonoLogin : DumbAwareAction(AllIcons.Actions.Execute) {
2323
pluginAwareExecuteOnPooledThread {
2424
val connectionManager = ToolkitConnectionManager.getInstance(project)
2525
connectionManager.activeConnectionForFeature(CodeCatalystConnection.getInstance())?.let {
26-
reauthConnectionIfNeeded(project, it)
26+
reauthConnectionIfNeeded(project, it, isReAuth = true)
2727
project.refreshDevToolTree()
2828
} ?: run {
2929
runInEdt {

‎plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/core/explorer/webview/ToolkitLoginWebview.kt‎

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ import com.intellij.openapi.application.runInEdt
1212
import com.intellij.openapi.components.Service
1313
import com.intellij.openapi.components.service
1414
import com.intellij.openapi.project.Project
15-
import com.intellij.openapi.ui.Messages
1615
import com.intellij.openapi.util.Disposer
1716
import com.intellij.ui.JBColor
1817
import com.intellij.ui.components.JBTextArea
@@ -47,7 +46,6 @@ import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeCatalystConn
4746
import software.aws.toolkits.jetbrains.core.credentials.sono.CODECATALYST_SCOPES
4847
import software.aws.toolkits.jetbrains.core.credentials.sono.IDENTITY_CENTER_ROLE_ACCESS_SCOPE
4948
import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
50-
import software.aws.toolkits.jetbrains.core.credentials.ssoErrorMessageFromException
5149
import software.aws.toolkits.jetbrains.core.explorer.showExplorerTree
5250
import software.aws.toolkits.jetbrains.core.gettingstarted.IdcRolePopup
5351
import software.aws.toolkits.jetbrains.core.gettingstarted.IdcRolePopupState
@@ -307,19 +305,7 @@ class ToolkitWebviewBrowser(val project: Project, private val parentDisposable:
307305
}
308306

309307
override fun loginIdC(url: String, region: AwsRegion, scopes: List<String>) {
310-
val onIdCError: (Exception) -> Unit = { e ->
311-
stopAndClearBrowserOpenTimer()
312-
if (!isUserCancellation(e)) {
313-
runInEdt {
314-
Messages.showErrorDialog(jcefBrowser.component, ssoErrorMessageFromException(e), "Failed to Authenticate")
315-
}
316-
}
317-
// TODO: telemetry
318-
}
319-
val onIdCSuccess: () -> Unit = {
320-
stopAndClearBrowserOpenTimer()
321-
// TODO: telemetry
322-
}
308+
val (onIdCError: (Exception) -> Unit, onIdCSuccess: () -> Unit) = getSuccessAndErrorActionsForIdcLogin(scopes, url, region)
323309

324310
val login = Login.IdC(url, region, scopes, onPendingToken, onIdCSuccess, onIdCError)
325311

0 commit comments

Comments
 (0)