add anyToken function to MockOAuth2Server.kt (#26)

* feat: add anyToken function to MockOAuth2Server.kt
* can issue any token with the servers jwk, accepting claims as input to function

Author: tommytroen

src/main/kotlin/no/nav/security/mock/oauth2/MockOAuth2Server.kt

@@ -1,15 +1,19 @@
 package no.nav.security.mock.oauth2
 
+import com.nimbusds.jwt.JWTClaimsSet
 import com.nimbusds.jwt.SignedJWT
 import com.nimbusds.oauth2.sdk.AuthorizationCode
 import com.nimbusds.oauth2.sdk.AuthorizationCodeGrant
+import com.nimbusds.oauth2.sdk.AuthorizationGrant
+import com.nimbusds.oauth2.sdk.GrantType
 import com.nimbusds.oauth2.sdk.TokenRequest
 import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic
 import com.nimbusds.oauth2.sdk.auth.Secret
 import com.nimbusds.oauth2.sdk.id.ClientID
 import java.io.IOException
 import java.net.InetAddress
 import java.net.URI
+import java.time.Duration
 import java.util.UUID
 import mu.KotlinLogging
 import no.nav.security.mock.oauth2.extensions.toAuthorizationEndpointUrl
@@ -112,8 +116,33 @@ open class MockOAuth2Server(
             expiry
         )
     )
+
+    @JvmOverloads
+    fun anyToken(issuerUrl: HttpUrl, claims: Map<String, Any>, expiry: Duration = Duration.ofHours(1)): SignedJWT {
+        val jwtClaimsSet = claims.toJwtClaimsSet()
+        val mockGrant: AuthorizationGrant = object : AuthorizationGrant(GrantType("MockGrant")) {
+            override fun toParameters(): MutableMap<String, MutableList<String>> = mutableMapOf()
+        }
+        return this.config.tokenProvider.exchangeAccessToken(
+            TokenRequest(URI.create("http://mockgrant"), mockGrant),
+            issuerUrl,
+            jwtClaimsSet,
+            DefaultOAuth2TokenCallback(
+                audience = jwtClaimsSet.audience,
+                expiry = expiry.toMillis()
+            )
+        )
+    }
 }
 
+internal fun Map<String, Any>.toJwtClaimsSet(): JWTClaimsSet =
+    JWTClaimsSet.Builder()
+        .apply {
+            this@toJwtClaimsSet.forEach {
+                this.claim(it.key, it.value)
+            }
+        }.build()
+
 fun <R> withMockOAuth2Server(
     test: MockOAuth2Server.() -> R
 ): R {

src/test/kotlin/no/nav/security/mock/oauth2/MockOAuth2ServerTest.kt

@@ -7,14 +7,17 @@ import com.nimbusds.jwt.JWTClaimsSet
 import com.nimbusds.jwt.SignedJWT
 import com.nimbusds.oauth2.sdk.GrantType
 import com.nimbusds.oauth2.sdk.id.Issuer
+import io.kotest.matchers.maps.shouldContainAll
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.string.shouldStartWith
 import java.net.URLEncoder
+import java.time.Duration
 import no.nav.security.mock.oauth2.extensions.verifySignatureAndIssuer
 import no.nav.security.mock.oauth2.http.OAuth2HttpResponse
 import no.nav.security.mock.oauth2.http.OAuth2TokenResponse
 import no.nav.security.mock.oauth2.http.WellKnown
 import no.nav.security.mock.oauth2.http.route
+import no.nav.security.mock.oauth2.testutils.claims
 import no.nav.security.mock.oauth2.testutils.get
 import no.nav.security.mock.oauth2.testutils.post
 import no.nav.security.mock.oauth2.token.DefaultOAuth2TokenCallback
@@ -388,6 +391,29 @@ class MockOAuth2ServerTest {
         assertThat(jwtClaimsSet.getClaim("someclaim")).isEqualTo("claimvalue")
     }
 
+    @Test
+    fun `anyToken should issue token with claims from input and be verifyable by servers keys`() {
+        val issuerId = "issuer1"
+        val token = server.anyToken(
+            server.issuerUrl(issuerId),
+            mutableMapOf(
+                "sub" to "yolo",
+                "aud" to listOf("myapp"),
+                "customInt" to 123,
+                "customList" to listOf(1, 2, 3)
+            ),
+            Duration.ofSeconds(10)
+        )
+        val jwkSet: JWKSet = retrieveJwks(server.jwksUrl(issuerId).toString())
+        token.verifySignatureAndIssuer(Issuer(server.issuerUrl("issuer1").toString()), jwkSet)
+        token.claims shouldContainAll mutableMapOf(
+            "sub" to "yolo",
+            "aud" to listOf("myapp"),
+            "customInt" to 123,
+            "customList" to listOf(1, 2, 3)
+        )
+    }
+
     private fun retrieveJwks(jwksUri: String): JWKSet {
         return client.newCall(
             Request.Builder()