From 6cb8f1f4053f57fa28060ae18fe3809e544f8929 Mon Sep 17 00:00:00 2001 From: Sam Stewart Date: Mon, 18 Nov 2024 13:05:35 -0800 Subject: [PATCH 1/4] initial commit --- .../core/credentials/sso/DiskCache.kt | 5 +++ .../credentials/sso/SsoAccessTokenProvider.kt | 32 +++++++++++-------- 2 files changed, 24 insertions(+), 13 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 e39a9f27c8b..fd9290427cb 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 @@ -114,6 +114,9 @@ class DiskCache( reason = "Failed to load Client Registration", reasonDesc = "Load Step:$stage failed. Cache file does not exist" ) + if (source == SsoAccessTokenProvider.SourceOfLoadRegistration.REFRESH_TOKEN.toString()) { + throw ClientRegistrationNotFoundException() + } return null } return loadClientRegistration(inputStream) @@ -320,3 +323,5 @@ class DiskCache( private val LOG = getLogger() } } + +class ClientRegistrationNotFoundException : Exception("Client registration file not found") 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 7fd718a8220..f6e9aafc76c 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 @@ -429,20 +429,26 @@ class SsoAccessTokenProvider( } stageName = RefreshCredentialStage.LOAD_REGISTRATION - val registration = try { - when (currentToken) { - is DeviceAuthorizationGrantToken -> loadDagClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) - is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) + val registration = run { + try { + when (currentToken) { + is DeviceAuthorizationGrantToken -> loadDagClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) + is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) + } + } catch (e: ClientRegistrationNotFoundException) { + //invalidate tokens to force a reauth + invalidate() + null + } catch (e: Exception) { + val message = e.message ?: "$stageName: ${e::class.java.name}" + sendRefreshCredentialsMetric( + currentToken, + reason = "Refresh access token request failed: $stageName", + reasonDesc = message, + result = Result.Failed + ) + throw InvalidClientException.builder().message(message).cause(e).build() } - } catch (e: Exception) { - val message = e.message ?: "$stageName: ${e::class.java.name}" - sendRefreshCredentialsMetric( - currentToken, - reason = "Refresh access token request failed: $stageName", - reasonDesc = message, - result = Result.Failed - ) - throw InvalidClientException.builder().message(message).cause(e).build() } stageName = RefreshCredentialStage.VALIDATE_REGISTRATION From f1d29ef8132acf5fa7c2a671ed489d85b7fecfba Mon Sep 17 00:00:00 2001 From: Sam Stewart Date: Mon, 18 Nov 2024 13:08:19 -0800 Subject: [PATCH 2/4] 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 f6e9aafc76c..f9091d1909f 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 @@ -436,7 +436,7 @@ class SsoAccessTokenProvider( is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) } } catch (e: ClientRegistrationNotFoundException) { - //invalidate tokens to force a reauth + // invalidate tokens to force a reauth invalidate() null } catch (e: Exception) { From dafd506fc86bd8396507e8668f3f73cfb615d250 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Wed, 11 Dec 2024 16:44:59 -0800 Subject: [PATCH 3/4] diskCacheTests --- .../core/credentials/sso/DiskCacheTest.kt | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) 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 e5ac374d365..05492fbf46d 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 @@ -7,6 +7,7 @@ import com.intellij.openapi.util.SystemInfo import com.intellij.openapi.util.io.NioFiles import com.intellij.testFramework.ApplicationExtension import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test import org.junit.jupiter.api.condition.DisabledOnOs @@ -714,4 +715,31 @@ class DiskCacheTest { .usingRecursiveComparison() .isEqualTo(sut.loadAccessToken(key2)) } + + @Test + fun `loadClientRegistration returns null when file not found during registration`() { + val key = DeviceAuthorizationClientRegistrationCacheKey( + startUrl = ssoUrl, + scopes = scopes, + region = ssoRegion + ) + + assertThat(sut.loadClientRegistration(key, SsoAccessTokenProvider.SourceOfLoadRegistration.REGISTER_CLIENT.toString())).isNull() + } + + @Test + fun `loadClientRegistration throws exception when file not found during refresh`() { + val key = DeviceAuthorizationClientRegistrationCacheKey( + startUrl = ssoUrl, + scopes = scopes, + region = ssoRegion + ) + + assertThatThrownBy { + sut.loadClientRegistration( + key, + SsoAccessTokenProvider.SourceOfLoadRegistration.REFRESH_TOKEN.toString() + ) + }.isInstanceOf(ClientRegistrationNotFoundException::class.java) + } } From 81098faeb5654b3a5d1bff77a2f39788b454c884 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Thu, 12 Dec 2024 12:42:19 -0800 Subject: [PATCH 4/4] SsoAccessTokenProvider tests --- .../core/credentials/sso/DiskCache.kt | 2 +- .../credentials/sso/SsoAccessTokenProvider.kt | 36 +++++++++---------- .../sso/SsoAccessTokenProviderTest.kt | 33 +++++++++++++++++ 3 files changed, 51 insertions(+), 20 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 fd9290427cb..bd532eb980d 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 @@ -324,4 +324,4 @@ class DiskCache( } } -class ClientRegistrationNotFoundException : Exception("Client registration file not found") +class ClientRegistrationNotFoundException : RuntimeException("Client registration file not found") 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 f9091d1909f..46df6d2290c 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 @@ -429,26 +429,24 @@ class SsoAccessTokenProvider( } stageName = RefreshCredentialStage.LOAD_REGISTRATION - val registration = run { - try { - when (currentToken) { - is DeviceAuthorizationGrantToken -> loadDagClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) - is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) - } - } catch (e: ClientRegistrationNotFoundException) { - // invalidate tokens to force a reauth - invalidate() - null - } catch (e: Exception) { - val message = e.message ?: "$stageName: ${e::class.java.name}" - sendRefreshCredentialsMetric( - currentToken, - reason = "Refresh access token request failed: $stageName", - reasonDesc = message, - result = Result.Failed - ) - throw InvalidClientException.builder().message(message).cause(e).build() + val registration = try { + when (currentToken) { + is DeviceAuthorizationGrantToken -> loadDagClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) + is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString()) } + } catch (e: ClientRegistrationNotFoundException) { + // invalidate tokens to force a reauth + invalidate() + null + } catch (e: Exception) { + val message = e.message ?: "$stageName: ${e::class.java.name}" + sendRefreshCredentialsMetric( + currentToken, + reason = "Refresh access token request failed: $stageName", + reasonDesc = message, + result = Result.Failed + ) + throw InvalidClientException.builder().message(message).cause(e).build() } stageName = RefreshCredentialStage.VALIDATE_REGISTRATION 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 5f6f2676870..5a255636777 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 @@ -479,6 +479,39 @@ class SsoAccessTokenProviderTest { verify(ssoCache).invalidateAccessToken(ssoUrl) } + @Test + fun `refreshToken invalidates tokens when client registration not found during refresh`() { + setPkceTrue() + + val accessToken = PKCEAuthorizationGrantToken( + ssoUrl, + ssoRegion, + "dummyToken", + "refreshToken", + clock.instant(), + clock.instant() + ) + + ssoCache.stub { + on(ssoCache.loadAccessToken(any())) + .thenReturn(accessToken) + on( + ssoCache.loadClientRegistration( + any(), + eq(SsoAccessTokenProvider.SourceOfLoadRegistration.REFRESH_TOKEN.toString()) + ) + ).thenThrow(ClientRegistrationNotFoundException()) + } + + assertThatThrownBy { + runBlocking { + sut.refreshToken(sut.accessToken()) + } + }.isInstanceOf(InvalidClientException::class.java) + + verify(ssoCache, times(2)).invalidateAccessToken(any()) + } + private fun setupCacheStub(expirationClientRegistration: Instant) { setupCacheStub(DeviceAuthorizationClientRegistration(clientId, clientSecret, expirationClientRegistration)) }