load function fixes

Author: samgst-amazon

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

@@ -109,24 +109,27 @@ class DiskCache(
 
     override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey): ClientRegistration? {
         LOG.debug { "loadClientRegistration for $cacheKey" }
-        val inputStream = clientRegistrationCache(cacheKey).tryInputStreamIfExists()
-            ?: // try to load from in memory cache
-            return InMemoryCache.get(clientRegistrationCache(cacheKey).toString())?.let { data ->
-                ByteArrayInputStream(data).use { memoryStream ->
-                    loadClientRegistration(memoryStream)
-                }
-            } ?: run {
-                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"
-                )
-                null
+        val cacheFile = clientRegistrationCache(cacheKey)
+        // try InMemoryCacheFirst in case of stale registration on full disk
+        InMemoryCache.get(cacheFile.toString())?.let { data ->
+            ByteArrayInputStream(data).use { memoryStream ->
+                return loadClientRegistration(memoryStream)
             }
+        }
+
+        val inputStream = cacheFile.tryInputStreamIfExists()
+        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)
     }
 
@@ -172,21 +175,22 @@ class DiskCache(
         }
     }
 
+
     override fun loadAccessToken(cacheKey: AccessTokenCacheKey): AccessToken? {
         LOG.debug { "loadAccessToken for $cacheKey" }
         val cacheFile = accessTokenCache(cacheKey)
-        // If file exists, returns InputStream, if not returns null
-        return cacheFile.tryInputStreamIfExists()
-            // try to load and parse access token, returns AccessToken or null if expired
-            ?.let { loadAccessToken(it) }
-            // If file doesn't exist or loadAccessToken failed, try in-memory cache
-            ?: InMemoryCache.get(cacheFile.toString())?.let { data ->
-                // If in-memory cache has data, create stream and try to load token
-                ByteArrayInputStream(data).use { memoryStream ->
-                    loadAccessToken(memoryStream)
-                }
-                // If both file system and in-memory cache attempts fail, returns null
+        // try InMemoryCacheFirst in case of stale token on full disk
+        InMemoryCache.get(cacheFile.toString())?.let { data ->
+            ByteArrayInputStream(data).use { memoryStream ->
+                return loadAccessToken(memoryStream)
             }
+        }
+
+        val inputStream = cacheFile.tryInputStreamIfExists() ?: return null
+
+        val token = loadAccessToken(inputStream)
+
+        return token
     }
 
     override fun saveAccessToken(cacheKey: AccessTokenCacheKey, accessToken: AccessToken) {

plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCacheTest.kt

@@ -829,4 +829,106 @@ class DiskCacheTest {
         val loadedRegistration = sut.loadClientRegistration(key)
         assertNull(loadedRegistration)
     }
+
+    @Test
+    fun `test client registration update with disk error falls back to memory cache`() {
+        // Create a cache key and initial registration
+        val key = PKCEClientRegistrationCacheKey(
+            issuerUrl = ssoUrl,
+            scopes = scopes,
+            region = ssoRegion,
+            clientType = "public",
+            grantTypes = listOf("authorization_code", "refresh_token"),
+            redirectUris = listOf("http://127.0.0.1/oauth/callback")
+        )
+        val initialRegistration = PKCEClientRegistration(
+            "stale_id",
+            "stale_client_secret",
+            Instant.now().plusSeconds(3600),
+            scopes,
+            ssoUrl,
+            ssoRegion,
+            "public",
+            listOf("authorization_code", "refresh_token"),
+            listOf("http://127.0.0.1/oauth/callback")
+        )
+
+        // Save initial registration (should save to disk)
+        sut.saveClientRegistration(key, initialRegistration)
+
+        // Verify initial save
+        val loadedInitial = sut.loadClientRegistration(key)
+        assertNotNull(loadedInitial)
+        assertEquals(initialRegistration, loadedInitial)
+
+        // Setup mock to throw IOException for future disk writes
+        setupMockOutputStreamThrowingIOException()
+
+        // Create updated registration
+        val updatedRegistration = PKCEClientRegistration(
+    "fresh_ID",
+ "fresh_client_secret",
+            Instant.now().plusSeconds(3600),
+            scopes,
+            ssoUrl,
+            ssoRegion,
+            "public",
+            listOf("authorization_code", "refresh_token"),
+            listOf("http://127.0.0.1/oauth/callback")
+        )
+
+        // Try to save updated registration (should fall back to in-memory cache)
+        sut.saveClientRegistration(key, updatedRegistration)
+
+        // Load registration again
+        val loadedUpdated = sut.loadClientRegistration(key)
+
+        // Verify that we get the updated registration, not the initial one
+        assertNotNull(loadedUpdated)
+        assertEquals(updatedRegistration, loadedUpdated)
+    }
+    
+    @Test
+    fun `test access token update with disk error falls back to memory cache`() {
+        val key = PKCEAccessTokenCacheKey(ssoUrl, ssoRegion, scopes)
+        val initialToken = PKCEAuthorizationGrantToken(
+            ssoUrl,
+            ssoRegion,
+            "stale_access_token",
+            "stale_refresh_token",
+            Instant.now().plusSeconds(3600),
+            Instant.now()
+        )
+
+        // Save initial token (should save to disk)
+        sut.saveAccessToken(key, initialToken)
+
+        // Verify initial save
+        val loadedInitial = sut.loadAccessToken(key)
+        assertNotNull(loadedInitial)
+        assertEquals(initialToken, loadedInitial)
+
+        // Setup mock to throw IOException for future disk writes
+        setupMockOutputStreamThrowingIOException()
+
+        // Create updated token
+        val updatedToken = PKCEAuthorizationGrantToken(
+            ssoUrl,
+            ssoRegion,
+            "fresh_access_token",
+            "fresh_refresh_token",
+            Instant.now().plusSeconds(3600),
+            Instant.now()
+        )
+
+        // Try to save updated token (should fall back to in-memory cache)
+        sut.saveAccessToken(key, updatedToken)
+
+        // Load token again
+        val loadedUpdated = sut.loadAccessToken(key)
+
+        // Verify that we get the updated token, not the initial one
+        assertNotNull(loadedUpdated)
+        assertEquals(updatedToken, loadedUpdated)
+    }
 }