From 0b696f9bef33e593ebb9b80a4a41ca2ee3435d45 Mon Sep 17 00:00:00 2001 From: Sam Stewart Date: Mon, 11 Nov 2024 10:35:51 -0800 Subject: [PATCH 1/5] initial commit --- .../jetbrains/core/credentials/sso/DiskCache.kt | 2 +- .../credentials/sso/SsoAccessTokenProvider.kt | 15 ++++++++++----- .../jetbrains/core/credentials/sso/SsoCache.kt | 2 +- 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt index fd38ef3dd98..8a460966ad1 100644 --- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt +++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt @@ -101,7 +101,7 @@ class DiskCache( clientRegistrationCache(ssoRegion).tryDeleteIfExists() } - override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey): ClientRegistration? { + override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey, source: String): ClientRegistration? { LOG.info { "loadClientRegistration for $cacheKey" } val inputStream = clientRegistrationCache(cacheKey).tryInputStreamIfExists() if (inputStream == null) { diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt index bb694f3042b..87f74fb85ba 100644 --- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt +++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt @@ -194,7 +194,7 @@ class SsoAccessTokenProvider( @Deprecated("Device authorization grant flow is deprecated") private fun registerDAGClient(): ClientRegistration { - loadDagClientRegistration()?.let { + loadDagClientRegistration(SourceOf)?.let { return it } @@ -505,6 +505,11 @@ class SsoAccessTokenProvider( } } + enum class SourceOfLoadRegistration{ + REGISTER_CLIENT, + REFRESH_TOKEN, + } + private enum class RefreshCredentialStage { VALIDATE_REFRESH_TOKEN, LOAD_REGISTRATION, @@ -514,13 +519,13 @@ class SsoAccessTokenProvider( SAVE_TOKEN, } - private fun loadDagClientRegistration(): ClientRegistration? = - cache.loadClientRegistration(dagClientRegistrationCacheKey)?.let { + private fun loadDagClientRegistration(source: String): ClientRegistration? = + cache.loadClientRegistration(dagClientRegistrationCacheKey, source)?.let { return it } - private fun loadPkceClientRegistration(): PKCEClientRegistration? = - cache.loadClientRegistration(pkceClientRegistrationCacheKey)?.let { + private fun loadPkceClientRegistration(source: String): PKCEClientRegistration? = + cache.loadClientRegistration(pkceClientRegistrationCacheKey, source)?.let { return it as PKCEClientRegistration } diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCache.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCache.kt index 4aa845e3883..b3a74969355 100644 --- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCache.kt +++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCache.kt @@ -7,7 +7,7 @@ interface SsoCache { fun invalidateClientRegistration(ssoRegion: String) fun invalidateAccessToken(ssoUrl: String) - fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey): ClientRegistration? + fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey, source: String): ClientRegistration? fun saveClientRegistration(cacheKey: ClientRegistrationCacheKey, registration: ClientRegistration) fun invalidateClientRegistration(cacheKey: ClientRegistrationCacheKey) From 21a49b75117259ecb676b0886bff472fba99a588 Mon Sep 17 00:00:00 2001 From: Sam Stewart Date: Mon, 11 Nov 2024 11:16:28 -0800 Subject: [PATCH 2/5] add source to loading metrics --- .../toolkits/jetbrains/core/credentials/sso/DiskCache.kt | 4 ++-- .../core/credentials/sso/SsoAccessTokenProvider.kt | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt index 8a460966ad1..e39a9f27c8b 100644 --- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt +++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt @@ -102,14 +102,14 @@ class DiskCache( } override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey, source: String): ClientRegistration? { - LOG.info { "loadClientRegistration for $cacheKey" } + LOG.info { "loadClientRegistration:$source for $cacheKey" } val inputStream = clientRegistrationCache(cacheKey).tryInputStreamIfExists() if (inputStream == null) { val stage = LoadCredentialStage.ACCESS_FILE LOG.info { "Failed to load Client Registration: cache file does not exist" } AuthTelemetry.modifyConnection( action = "Load cache file", - source = "loadClientRegistration", + source = "loadClientRegistration:$source", result = Result.Failed, reason = "Failed to load Client Registration", reasonDesc = "Load Step:$stage failed. Cache file does not exist" diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt index 87f74fb85ba..c8193124307 100644 --- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt +++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt @@ -194,7 +194,7 @@ class SsoAccessTokenProvider( @Deprecated("Device authorization grant flow is deprecated") private fun registerDAGClient(): ClientRegistration { - loadDagClientRegistration(SourceOf)?.let { + loadDagClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT.toString())?.let { return it } @@ -235,7 +235,7 @@ class SsoAccessTokenProvider( } private fun registerPkceClient(): PKCEClientRegistration { - loadPkceClientRegistration()?.let { + loadPkceClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT.toString())?.let { return it } @@ -431,8 +431,8 @@ class SsoAccessTokenProvider( stageName = RefreshCredentialStage.LOAD_REGISTRATION val registration = try { when (currentToken) { - is DeviceAuthorizationGrantToken -> loadDagClientRegistration() - is PKCEAuthorizationGrantToken -> loadPkceClientRegistration() + is DeviceAuthorizationGrantToken -> loadDagClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) + is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) } } catch (e: Exception) { val message = e.message ?: "$stageName: ${e::class.java.name}" From 9f5cb703741765d4a157be1b83425ef902aa913b Mon Sep 17 00:00:00 2001 From: Sam Stewart Date: Mon, 11 Nov 2024 11:25:56 -0800 Subject: [PATCH 3/5] detekt --- .../jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt index c8193124307..7fd718a8220 100644 --- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt +++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt @@ -505,7 +505,7 @@ class SsoAccessTokenProvider( } } - enum class SourceOfLoadRegistration{ + enum class SourceOfLoadRegistration { REGISTER_CLIENT, REFRESH_TOKEN, } From dbc0119c0a64120f5c9aeff8f68f4d73703d257c Mon Sep 17 00:00:00 2001 From: Sam Stewart Date: Mon, 11 Nov 2024 12:11:04 -0800 Subject: [PATCH 4/5] fix tests --- .../core/credentials/sso/DiskCacheTest.kt | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCacheTest.kt b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCacheTest.kt index bdb30af2104..e5ac374d365 100644 --- a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCacheTest.kt +++ b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCacheTest.kt @@ -57,7 +57,8 @@ class DiskCacheTest { startUrl = ssoUrl, scopes = scopes, region = ssoRegion - ) + ), + "testSource" ) ).isNull() } @@ -71,7 +72,7 @@ class DiskCacheTest { ) cacheLocation.resolve("223224b6f0b4702c1a984be8284fe2c9d9718759.json").writeText("badData") - assertThat(sut.loadClientRegistration(key)).isNull() + assertThat(sut.loadClientRegistration(key, "testSource")).isNull() } @Test @@ -91,7 +92,7 @@ class DiskCacheTest { """.trimIndent() ) - assertThat(sut.loadClientRegistration(key)).isNull() + assertThat(sut.loadClientRegistration(key, "testSource")).isNull() } @Test @@ -112,7 +113,7 @@ class DiskCacheTest { """.trimIndent() ) - assertThat(sut.loadClientRegistration(key)).isNull() + assertThat(sut.loadClientRegistration(key, "testSource")).isNull() } @Test @@ -134,7 +135,7 @@ class DiskCacheTest { """.trimIndent() ) - assertThat(sut.loadClientRegistration(key)) + assertThat(sut.loadClientRegistration(key, "testSource")) .usingRecursiveComparison() .isEqualTo( DeviceAuthorizationClientRegistration( @@ -217,7 +218,7 @@ class DiskCacheTest { """.trimIndent() ) - assertThat(sut.loadClientRegistration(key)) + assertThat(sut.loadClientRegistration(key, "testSource")) .usingRecursiveComparison() .isEqualTo( PKCEClientRegistration( @@ -323,10 +324,10 @@ class DiskCacheTest { ) ) - assertThat(sut.loadClientRegistration(key1)) + assertThat(sut.loadClientRegistration(key1, "testSource")) .usingRecursiveComparison() .isEqualTo( - sut.loadClientRegistration(key2) + sut.loadClientRegistration(key2, "testSource") ) } @@ -350,11 +351,11 @@ class DiskCacheTest { region = ssoRegion ) - assertThat(sut.loadClientRegistration(key)).isNotNull() + assertThat(sut.loadClientRegistration(key, "testSource")).isNotNull() sut.invalidateClientRegistration(key) - assertThat(sut.loadClientRegistration(key)).isNull() + assertThat(sut.loadClientRegistration(key, "testSource")).isNull() assertThat(cacheFile).doesNotExist() } @@ -619,7 +620,7 @@ class DiskCacheTest { registration.setPosixFilePermissions(emptySet()) assertPosixPermissions(registration, "---------") - assertThat(sut.loadClientRegistration(key)).isNotNull() + assertThat(sut.loadClientRegistration(key, "testSource")).isNotNull() assertPosixPermissions(registration, "rw-------") } From dce27a961fcd52d04b257d06143d8ecaefe8aa39 Mon Sep 17 00:00:00 2001 From: Sam Stewart Date: Mon, 11 Nov 2024 15:39:10 -0800 Subject: [PATCH 5/5] fix more tests --- .../sso/SsoAccessTokenProviderTest.kt | 20 +++++++++---------- .../InteractiveBearerTokenProviderTest.kt | 3 ++- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProviderTest.kt b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProviderTest.kt index 7291d561ab4..5f6f2676870 100644 --- a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProviderTest.kt +++ b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProviderTest.kt @@ -124,7 +124,7 @@ class SsoAccessTokenProviderTest { verify(ssoOidcClient).startDeviceAuthorization(any()) verify(ssoOidcClient).createToken(any()) verify(ssoCache).loadAccessToken(argThat { startUrl == ssoUrl }) - verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }) + verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any()) verify(ssoCache).saveAccessToken(argThat { startUrl == ssoUrl }, eq(accessToken)) } @@ -170,7 +170,7 @@ class SsoAccessTokenProviderTest { verify(ssoOidcClient).startDeviceAuthorization(any()) verify(ssoOidcClient).createToken(any()) verify(ssoCache).loadAccessToken(argThat { startUrl == ssoUrl }) - verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }) + verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any()) verify(ssoCache).saveClientRegistration(argThat { region == ssoRegion }, any()) verify(ssoCache).saveAccessToken(argThat { startUrl == ssoUrl }, eq(accessToken)) } @@ -267,7 +267,7 @@ class SsoAccessTokenProviderTest { verify(ssoOidcClient).startDeviceAuthorization(any()) verify(ssoOidcClient, times(2)).createToken(any()) verify(ssoCache).loadAccessToken(argThat { startUrl == ssoUrl }) - verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }) + verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any()) verify(ssoCache).saveAccessToken(argThat { startUrl == ssoUrl }, eq(accessToken)) } @@ -296,7 +296,7 @@ class SsoAccessTokenProviderTest { val refreshedToken = runBlocking { sut.refreshToken(sut.accessToken()) } verify(ssoCache).loadAccessToken(argThat { startUrl == ssoUrl }) - verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }) + verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any()) verify(ssoOidcClient).createToken(any()) verify(ssoCache).saveAccessToken(argThat { startUrl == ssoUrl }, eq(refreshedToken)) } @@ -342,7 +342,7 @@ class SsoAccessTokenProviderTest { ) on( - ssoCache.loadClientRegistration(any()) + ssoCache.loadClientRegistration(any(), any()) ).thenReturn( PKCEClientRegistration( clientType = "public", @@ -369,7 +369,7 @@ class SsoAccessTokenProviderTest { val refreshedToken = runBlocking { sut.refreshToken(sut.accessToken()) } verify(ssoCache).loadAccessToken(any()) - verify(ssoCache).loadClientRegistration(any()) + verify(ssoCache).loadClientRegistration(any(), any()) verify(ssoOidcClient).createToken(any()) verify(ssoCache).saveAccessToken(any(), eq(refreshedToken)) } @@ -390,7 +390,7 @@ class SsoAccessTokenProviderTest { verify(ssoOidcClient).startDeviceAuthorization(any()) verify(ssoOidcClient).createToken(any()) verify(ssoCache).loadAccessToken(argThat { startUrl == ssoUrl }) - verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }) + verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any()) } @Test @@ -432,7 +432,7 @@ class SsoAccessTokenProviderTest { verify(ssoOidcClient).startDeviceAuthorization(any()) verify(ssoOidcClient, times(2)).createToken(any()) verify(ssoCache).loadAccessToken(argThat { startUrl == ssoUrl }) - verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }) + verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any()) verify(ssoCache).saveAccessToken(argThat { startUrl == ssoUrl }, eq(accessToken)) } @@ -452,7 +452,7 @@ class SsoAccessTokenProviderTest { verify(ssoOidcClient).registerClient(any()) verify(ssoCache).loadAccessToken(any()) - verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }) + verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any()) } @Test @@ -492,7 +492,7 @@ class SsoAccessTokenProviderTest { ) on( - ssoCache.loadClientRegistration(argThat { region == ssoRegion }) + ssoCache.loadClientRegistration(argThat { region == ssoRegion }, any()) ).thenReturn( returnValue ) diff --git a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/bearer/InteractiveBearerTokenProviderTest.kt b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/bearer/InteractiveBearerTokenProviderTest.kt index 4ca91a5fade..ca5661b1f23 100644 --- a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/bearer/InteractiveBearerTokenProviderTest.kt +++ b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/bearer/InteractiveBearerTokenProviderTest.kt @@ -16,6 +16,7 @@ import org.junit.jupiter.api.assertThrows import org.mockito.Mockito import org.mockito.kotlin.any import org.mockito.kotlin.argThat +import org.mockito.kotlin.eq import org.mockito.kotlin.mock import org.mockito.kotlin.spy import org.mockito.kotlin.times @@ -273,7 +274,7 @@ class InteractiveBearerTokenProviderTest { ) private fun stubClientRegistration() { - whenever(diskCache.loadClientRegistration(any())).thenReturn( + whenever(diskCache.loadClientRegistration(any(), eq("testSource"))).thenReturn( DeviceAuthorizationClientRegistration( "", "",