SsoAccessTokenProvider tests

Author: samgst-amazon

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")

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

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<PKCEAccessTokenCacheKey>()))
+                .thenReturn(accessToken)
+            on(
+                ssoCache.loadClientRegistration(
+                    any<PKCEClientRegistrationCacheKey>(),
+                    eq(SsoAccessTokenProvider.SourceOfLoadRegistration.REFRESH_TOKEN.toString())
+                )
+            ).thenThrow(ClientRegistrationNotFoundException())
+        }
+
+        assertThatThrownBy {
+            runBlocking {
+                sut.refreshToken(sut.accessToken())
+            }
+        }.isInstanceOf(InvalidClientException::class.java)
+
+        verify(ssoCache, times(2)).invalidateAccessToken(any<AccessTokenCacheKey>())
+    }
+
     private fun setupCacheStub(expirationClientRegistration: Instant) {
         setupCacheStub(DeviceAuthorizationClientRegistration(clientId, clientSecret, expirationClientRegistration))
     }