Save refreshed PKCE token as PKCE token instead of DAG token (#4470)

Author: rli

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

@@ -124,7 +124,7 @@ class SsoAccessTokenProvider(
         val registerResponse = client.registerClient {
             it.clientType(PUBLIC_CLIENT_REGISTRATION_TYPE)
             it.scopes(scopes)
-            it.clientName("AWS IDE Plugins for JetBrains")
+            it.clientName(PKCE_CLIENT_NAME)
         }
 
         val registeredClient = DeviceAuthorizationClientRegistration(
@@ -310,7 +310,11 @@ class SsoAccessTokenProvider(
             it.refreshToken(currentToken.refreshToken)
         }
 
-        val token = newToken.toDAGAccessToken(currentToken.createdAt)
+        val token = when (currentToken) {
+            is DeviceAuthorizationGrantToken -> newToken.toDAGAccessToken(currentToken.createdAt)
+            is PKCEAuthorizationGrantToken -> newToken.toPKCEAccessToken(currentToken.createdAt)
+        }
+
         saveAccessToken(token)
 
         return token
@@ -399,7 +403,7 @@ class SsoAccessTokenProvider(
         }
     }
 
-    private fun CreateTokenResponse.toDAGAccessToken(creationTime: Instant): AccessToken {
+    private fun CreateTokenResponse.toDAGAccessToken(creationTime: Instant): DeviceAuthorizationGrantToken {
         val expirationTime = Instant.now(clock).plusSeconds(expiresIn().toLong())
 
         return DeviceAuthorizationGrantToken(
@@ -412,6 +416,19 @@ class SsoAccessTokenProvider(
         )
     }
 
+    private fun CreateTokenResponse.toPKCEAccessToken(creationTime: Instant): PKCEAuthorizationGrantToken {
+        val expirationTime = Instant.now(clock).plusSeconds(expiresIn().toLong())
+
+        return PKCEAuthorizationGrantToken(
+            issuerUrl = ssoUrl,
+            region = ssoRegion,
+            accessToken = accessToken(),
+            refreshToken = refreshToken(),
+            expiresAt = expirationTime,
+            createdAt = creationTime
+        )
+    }
+
     private companion object {
         const val PUBLIC_CLIENT_REGISTRATION_TYPE = "public"
         const val DEVICE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:device_code"

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

@@ -293,6 +293,54 @@ class SsoAccessTokenProviderTest {
         verify(ssoCache).saveAccessToken(ssoUrl, refreshedToken)
     }
 
+    @Test
+    fun `PKCE refresh access token saves PKCE token`() {
+        val sut = SsoAccessTokenProvider(ssoUrl, "us-east-1", ssoCache, ssoOidcClient, scopes = listOf("dummy:scope"), clock = clock)
+
+        val expirationClientRegistration = clock.instant().plusSeconds(120)
+        setupCacheStub(expirationClientRegistration)
+
+        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>())
+            ).thenReturn(
+                PKCEClientRegistration(
+                    clientType = "public",
+                    redirectUris = listOf("uri"),
+                    grantTypes = listOf("grant"),
+                    issuerUrl = ssoUrl,
+                    region = ssoRegion,
+                    scopes = listOf("dummy:scope"),
+                    clientId = clientId,
+                    clientSecret = clientSecret,
+                    expiresAt = clock.instant()
+                )
+            )
+        }
+
+        ssoOidcClient.stub {
+            on(
+                ssoOidcClient.createToken(refreshTokenRequest())
+            ).thenReturn(
+                refreshTokenResponse()
+            )
+        }
+
+        val refreshedToken = runBlocking { sut.refreshToken(sut.accessToken()) }
+
+        verify(ssoCache).loadAccessToken(any<PKCEAccessTokenCacheKey>())
+        verify(ssoCache).loadClientRegistration(any<PKCEClientRegistrationCacheKey>())
+        verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
+        verify(ssoCache).saveAccessToken(any<PKCEAccessTokenCacheKey>(), eq(refreshedToken))
+    }
+
     @Test
     fun exceptionStopsPolling() {
         val expirationClientRegistration = clock.instant().plusSeconds(120)