add modifyCredentials + load stages to LoadClientRegistration flow

Author: samgst-amazon

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

@@ -104,8 +104,17 @@ class DiskCache(
     override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey): ClientRegistration? {
         LOG.debug { "loadClientRegistration for $cacheKey" }
         val inputStream = clientRegistrationCache(cacheKey).tryInputStreamIfExists()
-            ?: return null
-
+        if (inputStream == null) {
+            val stage = LoadCredentialStage.ACCESS_FILE
+            AwsTelemetry.modifyCredentials(
+                credentialModification = CredentialModification.Unknown,
+                result = Result.Failed,
+                reason = "Failed to load ClientRegistration",
+                reasonDesc = "Load Step:$stage failed: previous error writing to file or deleted",
+                source = "loadClientRegistration"
+            )
+            return null
+        }
         return loadClientRegistration(inputStream)
     }
 
@@ -127,15 +136,10 @@ class DiskCache(
                 result = Result.Failed,
                 reason = "Failed to invalidate ClientRegistration",
                 reasonDesc = e.message,
-                source = "DiskCache.invalidateClientRegistration"
+                source = "invalidateClientRegistration"
             )
             throw e
         }
-        AwsTelemetry.modifyCredentials(
-            credentialModification = CredentialModification.Delete,
-            result = Result.Succeeded,
-            source = "DiskCache.invalidateClientRegistration"
-        )
     }
 
     override fun invalidateAccessToken(ssoUrl: String) {
@@ -152,11 +156,6 @@ class DiskCache(
             )
             throw e
         }
-        AwsTelemetry.modifyCredentials(
-            credentialModification = CredentialModification.Delete,
-            result = Result.Succeeded,
-            source = "DiskCache.invalidateAccessToken"
-        )
     }
 
     override fun loadAccessToken(cacheKey: AccessTokenCacheKey): AccessToken? {
@@ -191,11 +190,6 @@ class DiskCache(
             )
             throw e
         }
-        AwsTelemetry.modifyCredentials(
-            credentialModification = CredentialModification.Delete,
-            result = Result.Succeeded,
-            source = "DiskCache.invalidateAccessToken"
-        )
     }
 
     private fun clientRegistrationCache(ssoRegion: String): Path = cacheDir.resolve("aws-toolkit-jetbrains-client-id-$ssoRegion.json")
@@ -222,15 +216,34 @@ class DiskCache(
         return cacheDir.resolve(fileName)
     }
 
-    private fun loadClientRegistration(inputStream: InputStream) =
-        tryOrNull {
+    private fun loadClientRegistration(inputStream: InputStream): ClientRegistration? {
+        var stage = LoadCredentialStage.VALIDATE_CREDENTIALS
+        try{
             val clientRegistration = objectMapper.readValue<ClientRegistration>(inputStream)
+            stage = LoadCredentialStage.CHECK_EXPIRATION
             if (clientRegistration.expiresAt.isNotExpired()) {
-                clientRegistration
+                return clientRegistration
             } else {
-                null
+                AwsTelemetry.modifyCredentials(
+                    credentialModification = CredentialModification.Unknown,
+                    result = Result.Failed,
+                    reason = "Failed to load ClientRegistration",
+                    reasonDesc = "Load Step:$stage failed: ClientRegistration is expired",
+                    source = "loadClientRegistration"
+                )
+                return null
             }
+        } catch (e: Exception) {
+            AwsTelemetry.modifyCredentials(
+                credentialModification = CredentialModification.Unknown,
+                result = Result.Failed,
+                reason = "Failed to load ClientRegistration",
+                reasonDesc = "Load Step:$stage failed: ClientRegistration file is invalid",
+                source = "loadClientRegistration"
+            )
+            return null
         }
+    }
 
     private fun loadAccessToken(inputStream: InputStream) = tryOrNull {
         val accessToken = objectMapper.readValue<AccessToken>(inputStream)
@@ -294,6 +307,12 @@ class DiskCache(
         }
     }
 
+    private enum class LoadCredentialStage {
+        ACCESS_FILE,
+        VALIDATE_CREDENTIALS,
+        CHECK_EXPIRATION,
+    }
+
     companion object {
         val EXPIRATION_THRESHOLD = Duration.ofMinutes(15)
         private val LOG = getLogger<DiskCache>()

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

@@ -432,16 +432,8 @@ class SsoAccessTokenProvider(
 
         stageName = RefreshCredentialStage.VALIDATE_REGISTRATION
         if (registration == null) {
-            val (reason, message) = when {
-                currentToken.expiresAt.isBefore(Instant.now(clock)) -> Pair(
-                    "Reauth Required: $stageName",
-                    "Expired client registration"
-                )
-                else -> Pair(
-                    "Unable to load client registration from cache: $stageName",
-                    "Null client registration"
-                )
-            }
+            val reason = "Unable to load client registration from cache: $stageName"
+            val message ="Null client registration: invalid or expired"
             sendRefreshCredentialsMetric(
                 currentToken,
                 reason = reason,