feat: support refreshtoken (as JWT) from keycloak (#242)

* feat: support keycloak refresh token format
* includes nonce from auth request in a plain JWT
* see #210

Co-authored-by: Youssef Bel Mekki <[email protected]>

src/main/kotlin/no/nav/security/mock/oauth2/grant/AuthorizationCodeGrantHandler.kt

@@ -76,7 +76,7 @@ internal class AuthorizationCodeHandler(
         val loginTokenCallbackOrDefault = getLoginTokenCallbackOrDefault(code, oAuth2TokenCallback)
         val idToken: SignedJWT = tokenProvider.idToken(tokenRequest, issuerUrl, loginTokenCallbackOrDefault, nonce)
         val accessToken: SignedJWT = tokenProvider.accessToken(tokenRequest, issuerUrl, loginTokenCallbackOrDefault, nonce)
-        val refreshToken: RefreshToken = refreshTokenManager.refreshToken(loginTokenCallbackOrDefault)
+        val refreshToken: RefreshToken = refreshTokenManager.refreshToken(loginTokenCallbackOrDefault, nonce)
 
         return OAuth2TokenResponse(
             tokenType = "Bearer",

src/main/kotlin/no/nav/security/mock/oauth2/grant/RefreshTokenManager.kt

@@ -1,7 +1,9 @@
 package no.nav.security.mock.oauth2.grant
 
-import java.util.UUID
+import com.nimbusds.jwt.JWTClaimsSet
+import com.nimbusds.jwt.PlainJWT
 import no.nav.security.mock.oauth2.token.OAuth2TokenCallback
+import java.util.UUID
 
 typealias RefreshToken = String
 
@@ -10,9 +12,21 @@ internal data class RefreshTokenManager(
 ) {
     operator fun get(refreshToken: RefreshToken) = cache[refreshToken]
 
-    fun refreshToken(tokenCallback: OAuth2TokenCallback): RefreshToken {
-        val refreshToken = UUID.randomUUID().toString()
+    fun refreshToken(tokenCallback: OAuth2TokenCallback, nonce: String?): RefreshToken {
+        val jti = UUID.randomUUID().toString()
+        // added for compatibility with keycloak js client which expects a jwt with nonce
+        val refreshToken = nonce?.let { plainJWT(jti, nonce) } ?: jti
         cache[refreshToken] = tokenCallback
         return refreshToken
     }
+
+    private fun plainJWT(jti: String, nonce: String?): String =
+        PlainJWT(
+            JWTClaimsSet.parse(
+                mapOf(
+                    "jti" to jti,
+                    "nonce" to nonce
+                )
+            )
+        ).serialize()
 }

src/test/kotlin/no/nav/security/mock/oauth2/grant/AuthorizationCodeHandlerTest.kt

@@ -1,6 +1,7 @@
 package no.nav.security.mock.oauth2.grant
 
 import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
+import com.nimbusds.jwt.PlainJWT
 import com.nimbusds.jwt.SignedJWT
 import com.nimbusds.oauth2.sdk.ResponseMode
 import com.nimbusds.openid.connect.sdk.AuthenticationRequest
@@ -99,6 +100,16 @@ internal class AuthorizationCodeHandlerTest {
         }
     }
 
+    @Test
+    fun `auth request with nonce should result in a token response with refresh token as a JWT containing the nonce`() {
+        val code: String = handler.retrieveAuthorizationCode(Login("foo"))
+
+        handler.tokenResponse(tokenRequest(code = code), "http://myissuer".toHttpUrl(), DefaultOAuth2TokenCallback()).asClue {
+            val claims = PlainJWT.parse(it.refreshToken).jwtClaimsSet.claims
+            claims["nonce"] shouldBe "5678"
+        }
+    }
+
     private fun AuthorizationCodeHandler.retrieveAuthorizationCode(login: Login): String =
         authorizationCodeResponse(
             authenticationRequest = "http://authorizationendpoint".toHttpUrl().authenticationRequest().asNimbusAuthRequest(),

src/test/kotlin/no/nav/security/mock/oauth2/grant/RefreshTokenManagerTest.kt

@@ -0,0 +1,35 @@
+package no.nav.security.mock.oauth2.grant
+
+import com.nimbusds.jwt.PlainJWT
+import io.kotest.assertions.asClue
+import io.kotest.matchers.shouldBe
+import io.kotest.matchers.shouldNotBe
+import no.nav.security.mock.oauth2.token.DefaultOAuth2TokenCallback
+import org.junit.jupiter.api.Test
+
+internal class RefreshTokenManagerTest {
+
+    @Test
+    fun `refresh token should be a jwt with nonce included if nonce is not null (for keycloak compatibility)`() {
+        val mgr = RefreshTokenManager()
+        val tokenCallback = DefaultOAuth2TokenCallback()
+
+        mgr.refreshToken(tokenCallback, "nonce123").asClue {
+            val claims = PlainJWT.parse(it).jwtClaimsSet.claims
+
+            claims["nonce"] shouldBe "nonce123"
+            claims["jti"] shouldNotBe null
+        }
+    }
+
+    @Test
+    fun `tokencallback should be available in cache for specific refresh token`() {
+        val mgr = RefreshTokenManager()
+        val tokenCallback = DefaultOAuth2TokenCallback()
+
+        val refreshToken = mgr.refreshToken(tokenCallback, null)
+        mgr[refreshToken] shouldBe tokenCallback
+        val refreshToken2 = mgr.refreshToken(tokenCallback, "nonce123")
+        mgr[refreshToken2] shouldBe tokenCallback
+    }
+}