Improve Client Registration Error Collection (#4993)

* Add try-catch blocks for refresh steps.

* add metric that indicates when token was successfully written to the json & stored in cache

* remove comment

* add more metrics to catch blocks.

* Added reauth for expired client registration

* fix try-catch Exception blocks. remove bad reauth logic

* add metrics to saveClientRegistration

* make metrics same in saveClientRegistration

* Add modifyCredentials to invalidateClientRegistration

* Move save/modify credentials to diskCache

* linter

* saveCredentials moved to writeKey

* saveCredentials moved to writeKey

* capture saveCredentials per saveAccessToken/ClientRegistration calls

* add modifyCredentials + load stages to LoadClientRegistration flow

* remove success metrics

* add logs

* add auth_ModifyConnection to Override

* Fix Tests

* nit fixes

* consistent metrics

* null e message, clean up unused enums

* add more modifyConnection

* nits

* add enums to diskCache stages

* try block success metric

* detekt

---------

Co-authored-by: Manodnya Jaydeep Bhoite <[email protected]>

Author: samgst-amazon

plugins/core/jetbrains-community/resources/telemetryOverride.json

@@ -64,6 +64,11 @@
             ],
             "description": "Explict user intent associated with a chat message"
         },
+        {
+            "name": "connectionState",
+            "type": "string",
+            "description": "A detailed state of a specific auth connection. Use `authStatus` for a higher level view of an extension's general connection."
+        },
         {
             "name": "cwsprChatHasCodeSnippet",
             "type": "boolean",
@@ -703,6 +708,53 @@
                 }
             ]
         },
+        {
+            "name": "auth_modifyConnection",
+            "description": "An auth connection was modified in some way, e.g. deleted, updated",
+            "metadata": [
+                {
+                    "type": "action",
+                    "required": true
+                },
+                {
+                    "type": "authScopes",
+                    "required": false
+                },
+                {
+                    "type": "authStatus",
+                    "required": false
+                },
+                {
+                    "type": "connectionState",
+                    "required": false
+                },
+                {
+                    "type": "credentialStartUrl",
+                    "required": false
+                },
+                {
+                    "type": "sessionDuration",
+                    "required": false
+                },
+                {
+                    "type": "source",
+                    "required": true
+                },
+                {
+                    "type": "ssoRegistrationClientId",
+                    "required": false
+                },
+                {
+                    "type": "ssoRegistrationExpiresAt",
+                    "required": false
+                },
+                {
+                    "type": "id",
+                    "required": false
+                }
+            ],
+            "passive": true
+        },
         {
             "name": "auth_userState",
             "description": "The state of the user's authentication.",

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

@@ -32,6 +32,8 @@ import software.aws.toolkits.core.utils.touch
 import software.aws.toolkits.core.utils.tryDirOp
 import software.aws.toolkits.core.utils.tryFileOp
 import software.aws.toolkits.core.utils.tryOrNull
+import software.aws.toolkits.telemetry.AuthTelemetry
+import software.aws.toolkits.telemetry.Result
 import java.io.InputStream
 import java.io.OutputStream
 import java.nio.file.Path
@@ -101,7 +103,18 @@ 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
+            LOG.warn("Failed to load Client Registration: cache file does not exist")
+            AuthTelemetry.modifyConnection(
+                action = "Load cache file",
+                source = "loadClientRegistration",
+                result = Result.Failed,
+                reason = "Failed to load Client Registration",
+                reasonDesc = "Load Step:$stage failed. Unable to load file"
+            )
+            return null
+        }
         return loadClientRegistration(inputStream)
     }
 
@@ -115,12 +128,34 @@ class DiskCache(
 
     override fun invalidateClientRegistration(cacheKey: ClientRegistrationCacheKey) {
         LOG.debug { "invalidateClientRegistration for $cacheKey" }
-        clientRegistrationCache(cacheKey).tryDeleteIfExists()
+        try {
+            clientRegistrationCache(cacheKey).tryDeleteIfExists()
+        } catch (e: Exception) {
+            AuthTelemetry.modifyConnection(
+                action = "Delete cache file",
+                source = "invalidateClientRegistration",
+                result = Result.Failed,
+                reason = "Failed to invalidate Client Registration",
+                reasonDesc = e.message ?: e::class.java.name
+            )
+            throw e
+        }
     }
 
     override fun invalidateAccessToken(ssoUrl: String) {
         LOG.debug { "invalidateAccessToken for $ssoUrl" }
-        accessTokenCache(ssoUrl).tryDeleteIfExists()
+        try {
+            accessTokenCache(ssoUrl).tryDeleteIfExists()
+        } catch (e: Exception) {
+            AuthTelemetry.modifyConnection(
+                action = "Delete cache file",
+                source = "invalidateAccessToken",
+                result = Result.Failed,
+                reason = "Failed to invalidate Access Token",
+                reasonDesc = e.message ?: e::class.java.name
+            )
+            throw e
+        }
     }
 
     override fun loadAccessToken(cacheKey: AccessTokenCacheKey): AccessToken? {
@@ -143,7 +178,18 @@ class DiskCache(
 
     override fun invalidateAccessToken(cacheKey: AccessTokenCacheKey) {
         LOG.debug { "invalidateAccessToken for $cacheKey" }
-        accessTokenCache(cacheKey).tryDeleteIfExists()
+        try {
+            accessTokenCache(cacheKey).tryDeleteIfExists()
+        } catch (e: Exception) {
+            AuthTelemetry.modifyConnection(
+                action = "Delete cache file",
+                source = "invalidateAccessToken",
+                result = Result.Failed,
+                reason = "Failed to invalidate Access Token",
+                reasonDesc = e.message ?: e::class.java.name
+            )
+            throw e
+        }
     }
 
     private fun clientRegistrationCache(ssoRegion: String): Path = cacheDir.resolve("aws-toolkit-jetbrains-client-id-$ssoRegion.json")
@@ -170,15 +216,36 @@ 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
+                LOG.warn("Client Registration is expired")
+                AuthTelemetry.modifyConnection(
+                    action = "Validate Credentials",
+                    source = "loadClientRegistration",
+                    result = Result.Failed,
+                    reason = "Failed to load Client Registration",
+                    reasonDesc = "Load Step:$stage failed: Client Registration is expired"
+                )
+                return null
             }
+        } catch (e: Exception) {
+            LOG.warn("Client Registration could not be read")
+            AuthTelemetry.modifyConnection(
+                action = "Validate Credentials",
+                source = "loadClientRegistration",
+                result = Result.Failed,
+                reason = "Failed to load Client Registration",
+                reasonDesc = "Load Step:$stage failed: File could not be read"
+            )
+            return null
         }
+    }
 
     private fun loadAccessToken(inputStream: InputStream) = tryOrNull {
         val accessToken = objectMapper.readValue<AccessToken>(inputStream)
@@ -202,11 +269,22 @@ class DiskCache(
 
     private fun writeKey(path: Path, consumer: (OutputStream) -> Unit) {
         LOG.debug { "writing to $path" }
-        path.tryDirOp(LOG) { createParentDirectories() }
+        try {
+            path.tryDirOp(LOG) { createParentDirectories() }
 
-        path.tryFileOp(LOG) {
-            touch(restrictToOwner = true)
-            outputStream().use(consumer)
+            path.tryFileOp(LOG) {
+                touch(restrictToOwner = true)
+                outputStream().use(consumer)
+            }
+        } catch (e: Exception) {
+            AuthTelemetry.modifyConnection(
+                action = "Write file",
+                source = "writeKey",
+                result = Result.Failed,
+                reason = "Failed to write to cache",
+                reasonDesc = e.message ?: e::class.java.name
+            )
+            throw e
         }
     }
 
@@ -230,6 +308,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>()