Fix: force reauth when client registration is not found in cache

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<DiskCache>()
     }
 }
+
+class ClientRegistrationNotFoundException : RuntimeException("Client registration file not found")

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt

@@ -434,6 +434,10 @@ class SsoAccessTokenProvider(
                 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(

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)
+    }
 }

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