feat: support objects and lists in request mapping claims (#699)

* fixes #683 and #674 

Co-authored-by: ybelmekk <[email protected]>

build.gradle.kts

@@ -66,6 +66,7 @@ dependencies {
     implementation("com.fasterxml.jackson.module:jackson-module-kotlin:$jacksonVersion")
     implementation("org.freemarker:freemarker:$freemarkerVersion")
     implementation("org.bouncycastle:bcpkix-jdk18on:$bouncyCastleVersion")
+    implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:1.6.3")
     testImplementation("org.assertj:assertj-core:$assertjVersion")
     testImplementation("org.junit.jupiter:junit-jupiter-api:$junitJupiterVersion")
     testImplementation("org.junit.jupiter:junit-jupiter-params:$junitJupiterVersion")

src/main/kotlin/no/nav/security/mock/oauth2/extensions/Template.kt

@@ -0,0 +1,31 @@
+package no.nav.security.mock.oauth2.extensions
+
+/**
+ * Replaces all template values denoted with ${key} in a map with the corresponding values from the templates map.
+ *
+ * @param templates a map of template values
+ * @return a new map with all template values replaced
+ */
+fun Map<String, Any>.replaceValues(templates: Map<String, Any>): Map<String, Any> {
+    fun replaceTemplateString(
+        value: String,
+        templates: Map<String, Any>,
+    ): String {
+        val regex = Regex("""\$\{(\w+)\}""")
+        return regex.replace(value) { matchResult ->
+            val key = matchResult.groupValues[1]
+            templates[key]?.toString() ?: matchResult.value
+        }
+    }
+
+    fun replaceValue(value: Any): Any {
+        return when (value) {
+            is String -> replaceTemplateString(value, templates)
+            is List<*> -> value.map { it?.let { replaceValue(it) } }
+            is Map<*, *> -> value.mapValues { v -> v.value?.let { replaceValue(it) } }
+            else -> value
+        }
+    }
+
+    return this.mapValues { replaceValue(it.value) }
+}

src/main/kotlin/no/nav/security/mock/oauth2/token/OAuth2TokenCallback.kt

@@ -5,6 +5,7 @@ import com.nimbusds.oauth2.sdk.GrantType
 import com.nimbusds.oauth2.sdk.TokenRequest
 import no.nav.security.mock.oauth2.extensions.clientIdAsString
 import no.nav.security.mock.oauth2.extensions.grantType
+import no.nav.security.mock.oauth2.extensions.replaceValues
 import no.nav.security.mock.oauth2.extensions.scopesWithoutOidcScopes
 import no.nav.security.mock.oauth2.extensions.tokenExchangeGrantOrNull
 import java.time.Duration
@@ -89,27 +90,14 @@ data class RequestMappingTokenCallback(
 
     private fun List<RequestMapping>.getClaims(tokenRequest: TokenRequest): Map<String, Any> {
         val claims = firstOrNull { it.isMatch(tokenRequest) }?.claims ?: emptyMap()
-        val customParameters = tokenRequest.customParameters.mapValues { (_, value) -> value.first() }
-        val variables =
-            if (tokenRequest.grantType() == GrantType.CLIENT_CREDENTIALS) {
-                customParameters + ("clientId" to tokenRequest.clientIdAsString())
-            } else {
-                customParameters
-            }
-        return claims.mapValues { (_, value) ->
-            when (value) {
-                is String -> replaceVariables(value, variables)
-                is List<*> ->
-                    value.map { v ->
-                        if (v is String) {
-                            replaceVariables(v, variables)
-                        } else {
-                            v
-                        }
-                    }
-                else -> value
-            }
-        }
+        val templateParams = tokenRequest.toHTTPRequest().bodyAsFormParameters.mapValues { it.value.joinToString(separator = " ") }
+
+        // in case client_id is not set as form param but as basic auth, we add it to the template params in two different formats for backwards compatibility
+        return claims.replaceValues(
+            templateParams +
+                mapOf("clientId" to tokenRequest.clientIdAsString()) +
+                mapOf("client_id" to tokenRequest.clientIdAsString()),
+        )
     }
 
     private inline fun <reified T> List<RequestMapping>.getClaimOrNull(
@@ -118,18 +106,6 @@ data class RequestMappingTokenCallback(
     ): T? = getClaims(tokenRequest)[key] as? T
 
     private fun List<RequestMapping>.getTypeHeader(tokenRequest: TokenRequest) = firstOrNull { it.isMatch(tokenRequest) }?.typeHeader ?: JOSEObjectType.JWT.type
-
-    private fun replaceVariables(
-        input: String,
-        replacements: Map<String, String>,
-    ): String {
-        val pattern = Regex("""\$\{(\w+)}""")
-        return pattern.replace(input) { result ->
-            val variableName = result.groupValues[1]
-            val replacement = replacements[variableName]
-            replacement ?: result.value
-        }
-    }
 }
 
 data class RequestMapping(

src/test/kotlin/no/nav/security/mock/oauth2/extensions/TemplateTest.kt

@@ -0,0 +1,28 @@
+package no.nav.security.mock.oauth2.extensions
+
+import io.kotest.assertions.asClue
+import io.kotest.matchers.shouldBe
+import org.junit.jupiter.api.Test
+
+class TemplateTest {
+    @Test
+    fun `template values in map should be replaced`() {
+        val templates =
+            mapOf(
+                "templateVal1" to "val1",
+                "templateVal2" to "val2",
+                "templateListVal" to "listVal1",
+            )
+
+        mapOf(
+            "object1" to mapOf("key1" to "\${templateVal1}"),
+            "object2" to "\${templateVal2}",
+            "nestedObject" to mapOf("nestedKey" to mapOf("nestedKeyAgain" to "\${templateVal2}")),
+            "list1" to listOf("\${templateListVal}"),
+        ).replaceValues(templates).asClue {
+            it["object1"] shouldBe mapOf("key1" to "val1")
+            it["list1"] shouldBe listOf("listVal1")
+            println(it)
+        }
+    }
+}

src/test/kotlin/no/nav/security/mock/oauth2/token/OAuth2TokenCallbackTest.kt

@@ -124,6 +124,73 @@ internal class OAuth2TokenCallbackTest {
                 issuer1.typeHeader(grantTypeShouldMatch) shouldBe "JWT"
             }
         }
+
+        @Test
+        fun `token request with custom parameters in token request should include claims with placeholder names`() {
+            val request =
+                clientCredentialsRequest(
+                    "scope" to "testscope:something another:scope",
+                    "mock_token_type" to "custom",
+                )
+            RequestMappingTokenCallback(
+                issuerId = "issuer1",
+                requestMappings =
+                    listOf(
+                        RequestMapping(
+                            requestParam = "scope",
+                            match = "testscope:.*",
+                            claims =
+                                mapOf(
+                                    "sub" to "\${clientId}",
+                                    "scope" to "\${scope}",
+                                    "mock_token_type" to "\${mock_token_type}",
+                                ),
+                        ),
+                    ),
+            ).addClaims(request).asClue {
+                it shouldContainAll mapOf("sub" to clientId, "scope" to "testscope:something another:scope", "mock_token_type" to "custom")
+            }
+        }
+    }
+
+    @Test
+    fun `token request with custom parameters in token request should include claims with placeholder names`() {
+        val request =
+            clientCredentialsRequest(
+                "mock_token_type" to "custom",
+                "participantId" to "participantId",
+                "actAs" to "actAs",
+                "readAs" to "readAs",
+            )
+        RequestMappingTokenCallback(
+            issuerId = "issuer1",
+            requestMappings =
+                listOf(
+                    RequestMapping(
+                        requestParam = "mock_token_type",
+                        match = "custom",
+                        claims =
+                            mapOf(
+                                "https://daml.com/ledger-api" to
+                                    mapOf(
+                                        "participantId" to "\${participantId}",
+                                        "actAs" to listOf("\${actAs}"),
+                                        "readAs" to listOf("\${readAs}"),
+                                    ),
+                            ),
+                    ),
+                ),
+        ).addClaims(request).asClue {
+            it shouldContainAll
+                mapOf(
+                    "https://daml.com/ledger-api" to
+                        mapOf(
+                            "participantId" to "participantId",
+                            "actAs" to listOf("actAs"),
+                            "readAs" to listOf("readAs"),
+                        ),
+                )
+        }
     }
 
     @Nested