Merge pull request #1039 from modelix/fix/MODELIX-1016

fix(model-server): make specifying a key ID optional

Author: slisson

authorization/src/main/kotlin/org/modelix/authorization/AuthorizationConfig.kt

@@ -19,6 +19,7 @@ package org.modelix.authorization
 import com.auth0.jwk.JwkProvider
 import com.auth0.jwk.JwkProviderBuilder
 import com.auth0.jwt.JWT
+import com.auth0.jwt.JWTVerifier
 import com.auth0.jwt.algorithms.Algorithm
 import com.auth0.jwt.interfaces.DecodedJWT
 import io.ktor.server.application.Application
@@ -134,22 +135,23 @@ class ModelixAuthorizationConfig : IModelixAuthorizationConfig {
         hmac384KeyFromEnv?.let { return@lazy Algorithm.HMAC384(it) }
         hmac256KeyFromEnv?.let { return@lazy Algorithm.HMAC256(it) }
 
-        val jwk = cachedJwkProvider?.get(jwkKeyId)
-        if (jwk != null) {
-            val publicKey = jwk.publicKey as? RSAPublicKey ?: error("Invalid key type: ${jwk.publicKey}")
-            return@lazy when (jwk.algorithm) {
-                "RS256" -> Algorithm.RSA256(publicKey, null)
-                "RSA384" -> Algorithm.RSA384(publicKey, null)
-                "RS512" -> Algorithm.RSA512(publicKey, null)
-                else -> error("Unsupported algorithm: ${jwk.algorithm}")
-            }
+        val localJwkProvider = cachedJwkProvider
+        val localJwkKeyId = jwkKeyId
+        if (localJwkProvider == null || localJwkKeyId == null) {
+            return@lazy null
         }
-
-        null
+        return@lazy getAlgorithmFromJwkProviderAndKeyId(localJwkProvider, localJwkKeyId)
     }
 
-    fun getJwtSignatureAlgorithm(): Algorithm {
-        return checkNotNull(algorithm) { "No signature algorithm configured" }
+    private fun getAlgorithmFromJwkProviderAndKeyId(jwkProvider: JwkProvider, jwkKeyId: String): Algorithm {
+        val jwk = jwkProvider.get(jwkKeyId)
+        val publicKey = jwk.publicKey as? RSAPublicKey ?: error("Invalid key type: ${jwk.publicKey}")
+        return when (jwk.algorithm) {
+            "RS256" -> Algorithm.RSA256(publicKey, null)
+            "RSA384" -> Algorithm.RSA384(publicKey, null)
+            "RS512" -> Algorithm.RSA512(publicKey, null)
+            else -> error("Unsupported algorithm: ${jwk.algorithm}")
+        }
     }
 
     fun getJwtSignatureAlgorithmOrNull(): Algorithm? {
@@ -161,10 +163,17 @@ class ModelixAuthorizationConfig : IModelixAuthorizationConfig {
     }
 
     fun verifyTokenSignature(token: DecodedJWT) {
-        val algorithm = getJwtSignatureAlgorithm()
-        val verifier = JWT.require(algorithm)
-            .acceptLeeway(0L)
-            .build()
+        val algorithm = getJwtSignatureAlgorithmOrNull()
+        val jwkProvider = getJwkProvider()
+
+        val verifier = if (algorithm != null) {
+            getVerifierForSpecificAlgorithm(algorithm)
+        } else if (jwkProvider != null) {
+            val algorithmForKeyFromToken = getAlgorithmFromJwkProviderAndKeyId(jwkProvider, token.keyId)
+            getVerifierForSpecificAlgorithm(algorithmForKeyFromToken)
+        } else {
+            error("Either an JWT algorithm or a JWK URI must be configured.")
+        }
         verifier.verify(token)
     }
 
@@ -178,8 +187,19 @@ class ModelixAuthorizationConfig : IModelixAuthorizationConfig {
         }
     }
 
-    fun shouldGenerateFakeTokens() = generateFakeTokens ?: (algorithm == null)
-    fun permissionCheckingEnabled() = permissionChecksEnabled ?: (algorithm != null)
+    // TODO MODELIX-1019 Instead of creating a fake token, we should refactor our code to work without a username
+    // when no authentication and authorization is configured.
+    /**
+     * Whether a fake token should be generated based on the configuration values provided.
+     *
+     * The fake token is generated so that we always have a username that can be used in the server logic.
+     */
+    fun shouldGenerateFakeTokens() = generateFakeTokens ?: (algorithm == null && cachedJwkProvider == null)
+
+    /**
+     * Whether permission checking should be enabled based on the configuration values provided.
+     */
+    fun permissionCheckingEnabled() = permissionChecksEnabled ?: (algorithm != null || cachedJwkProvider != null)
 
     override fun configureForUnitTests() {
         generateFakeTokens = true
@@ -198,3 +218,7 @@ private fun getBooleanFromEnv(name: String): Boolean? {
         throw IllegalArgumentException("Failed to read boolean value $name", ex)
     }
 }
+
+internal fun getVerifierForSpecificAlgorithm(algorithm: Algorithm): JWTVerifier =
+    JWT.require(algorithm)
+        .build()

authorization/src/main/kotlin/org/modelix/authorization/AuthorizationPlugin.kt

@@ -18,6 +18,9 @@ package org.modelix.authorization
 
 import com.auth0.jwt.JWT
 import com.auth0.jwt.algorithms.Algorithm
+import com.auth0.jwt.exceptions.JWTVerificationException
+import com.auth0.jwt.interfaces.DecodedJWT
+import com.auth0.jwt.interfaces.JWTVerifier
 import com.google.common.cache.CacheBuilder
 import io.ktor.http.HttpStatusCode
 import io.ktor.server.application.Application
@@ -80,29 +83,18 @@ object ModelixAuthorization : BaseRouteScopedPlugin<IModelixAuthorizationConfig,
             } else {
                 // "Authorization: Bearer ..." header is provided in the header by OAuth proxy
                 jwt(MODELIX_JWT_AUTH) {
-                    val jwkProvider = config.getJwkProvider()
-                    if (jwkProvider != null) {
-                        verifier(jwkProvider) {}
-                    } else {
-                        verifier(
-                            JWT.require(config.getJwtSignatureAlgorithm())
-                                .build(),
-                        )
-                    }
+                    verifier(config.getVerifier())
                     challenge { _, _ ->
                         call.respond(status = HttpStatusCode.Unauthorized, "No or invalid JWT token provided")
                         // login and token generation is done by OAuth proxy. Only validation is required here.
                     }
                     validate {
                         try {
-                            val token = jwtFromHeaders()
-                            if (token != null) {
-                                return@validate config.nullIfInvalid(token)?.let { AccessTokenPrincipal(it) }
-                            }
+                            jwtFromHeaders()?.let(::AccessTokenPrincipal)
                         } catch (e: Exception) {
                             LOG.warn(e) { "Failed to read JWT token" }
+                            null
                         }
-                        null
                     }
                 }
             }
@@ -201,3 +193,21 @@ class ModelixAuthorizationPluginInstance(val config: ModelixAuthorizationConfig)
         }
     }
 }
+
+/**
+ * Returns an [JWTVerifier] that wraps our common authorization logic,
+ * so that it can be configured in the verification with Ktor's JWT authorization.
+ */
+internal fun ModelixAuthorizationConfig.getVerifier() = object : JWTVerifier {
+    override fun verify(token: String?): DecodedJWT {
+        val jwt = JWT.decode(token)
+        return verify(jwt)
+    }
+
+    override fun verify(jwt: DecodedJWT?): DecodedJWT {
+        if (jwt == null) {
+            throw JWTVerificationException("No JWT provided.")
+        }
+        return this@getVerifier.nullIfInvalid(jwt) ?: throw JWTVerificationException("JWT invalid.")
+    }
+}

docs/global/modules/core/pages/reference/component-model-server.adoc

@@ -64,11 +64,11 @@ To enable it you can specify the following environment variables.
 |Variable |Description
 
 |MODELIX_PERMISSION_CHECKS_ENABLED
-|By default, permission checking is enabled when an algorithm for the JWT signature is configured.
+|By default, permission checking is enabled when an algorithm for the JWT signature or a `MODELIX_JWK_URI` is configured.
  This variable can be set explicitly to `true` or `false` to avoid security issues by a misconfigured algorithm.
 
 |MODELIX_GENERATE_FAKE_JWT
-|By default, if no signature algorithm is configured,
+|By default, if no signature algorithm and no `MODELIX_JWK_URI` is configured,
  a token is generated for all requests with the identity `[email protected]` and no permissions.
  This option can be set to `true` or `false` to enable/disable this behaviour explicitly.
 
@@ -90,7 +90,10 @@ To enable it you can specify the following environment variables.
 |MODELIX_JWK_URI
 |If keys are created and signed by some OpenID connect server the public keys are provided via HTTP.
  Here you can specify the URI of the key set.
- Only RSA (256, 284 and 512) keys are currently supported.
+ Only RSA (256, 384 and 512) keys are currently supported.
+
+|MODELIX_JWK_KEY_ID
+|Optional key ID that can be used together with `MODELIX_JWK_URI`. If specified, it ensures that only tokens that use the specified key are valid. If not specified, a token can use any RSA (256, 384 and 512) key provided by `MODELIX_JWK_URI`.
 
 |KEYCLOAK_BASE_URL
  KEYCLOAK_REALM

gradle/libs.versions.toml

@@ -41,6 +41,7 @@ detekt = "1.23.7"
 xmlunit = "2.10.0"
 kotest = "5.9.1"
 testcontainers = "1.20.3"
+keycloak = "26.0.2"
 
 [libraries]
 
@@ -85,7 +86,8 @@ ktor-client-auth = { group = "io.ktor", name = "ktor-client-auth", version.ref =
 ktor-serialization = { group = "io.ktor", name = "ktor-serialization", version.ref = "ktor" }
 ktor-serialization-json = { group = "io.ktor", name = "ktor-serialization-kotlinx-json", version.ref = "ktor" }
 
-keycloak-authz-client = { group = "org.keycloak", name = "keycloak-authz-client", version = "26.0.2" }
+keycloak-authz-client = { group = "org.keycloak", name = "keycloak-authz-client", version.ref = "keycloak" }
+keycloak-admin-client = { group = "org.keycloak", name = "keycloak-admin-client", version.ref = "keycloak" }
 
 kotest-assertions-core = { group = "io.kotest", name = "kotest-assertions-core", version.ref = "kotest" }
 kotest-assertions-coreJvm = { group = "io.kotest", name = "kotest-assertions-core-jvm", version.ref = "kotest" }

model-server/build.gradle.kts

@@ -85,6 +85,8 @@ dependencies {
     testImplementation(project(":modelql-untyped"))
     testImplementation(libs.testcontainers)
     testImplementation(libs.testcontainers.postgresql)
+    testImplementation(libs.keycloak.authz.client)
+    testImplementation(libs.keycloak.admin.client)
 }
 
 tasks.named<ShadowJar>("shadowJar") {