JWKS demo

project/Dependencies.scala

@@ -28,6 +28,8 @@ object Dependencies {
 
     val jjwt = "0.11.5"
 
+    val nimbusJoseJwt = "9.31"
+
     val scalatest = "3.2.15"
 
     val pureConfig = "0.17.2"
@@ -47,6 +49,8 @@ object Dependencies {
   lazy val jjwtImpl = "io.jsonwebtoken" % "jjwt-impl" % Versions.jjwt % Runtime
   lazy val jjwtJackson = "io.jsonwebtoken" % "jjwt-jackson" % Versions.jjwt % Runtime
 
+  lazy val nimbusJoseJwt = "com.nimbusds" % "nimbus-jose-jwt" % Versions.nimbusJoseJwt
+
   lazy val pureConfig = "com.github.pureconfig" %% "pureconfig" % Versions.pureConfig
   lazy val pureConfigYaml = "com.github.pureconfig" %% "pureconfig-yaml" % Versions.pureConfig
 
@@ -78,6 +82,8 @@ object Dependencies {
     jjwtImpl,
     jjwtJackson,
 
+    nimbusJoseJwt,
+
     pureConfig,
     pureConfigYaml,
 

service/src/main/scala/za/co/absa/loginsvc/rest/SecurityConfig.scala

@@ -39,6 +39,7 @@ class SecurityConfig {
         "/swagger-ui/**", "/swagger-ui.html", // "/swagger-ui.html" redirects to "/swagger-ui/index.html
         "/swagger-resources/**", "/v3/api-docs/**", // swagger needs these
         "/actuator/**",
+        "/token/public-key-jwks",
         "/token/public-key").permitAll()
         .anyRequest().authenticated()
         .and()

service/src/main/scala/za/co/absa/loginsvc/rest/controller/TokenController.scala

@@ -107,4 +107,22 @@ class TokenController @Autowired()(jwtService: JWTService) {
     Future.successful(PublicKeyWrapper(publicKeyBase64))
   }
 
+  @Tags(Array(new Tag(name = "token")))
+  @Operation(
+    summary = "Gives payload with the RSA256 public key in JWKS format",
+    description = "Returns the same information as /token/public-key, but as a JSON Web Key Set",
+    responses = Array(
+    new ApiResponse(responseCode = "200", description = "Success", content = Array(new Content(examples = Array(new ExampleObject(value = """{"keys":[{"kty": "EC","crv": "P-256","x": "MKBCTNIcKUSDii11ySs3526iDZ8AiTo7Tu6KPAqv7D4","y": "4Etl6SRW2YiLUrN5vfvVHuhp7x8PxltmWWlbbM4IFyM","use": "enc","kid": "1"},{"kty": "RSA","n": "0vx7agoebGcQSuuPiLJXZptN9nndrQmbXEps2aiAFbWhM78LhWx4cbbfAAtVT86zwu1RK7aPFFxuhDR1L6tSoc_BJECPebWKRXjBZCiFV4n3oknjhMstn64tZ_2W-5JsGY4Hc5n9yBXArwl93lqt7_RN5w6Cf0h4QyQ5v-65YGjQR0_FDW2QvzqY368QQMicAtaSqzs8KJZgnYb9c7d0zgdAZHzu6qMQvRL5hajrn1n91CbOpbISD08qNLyrdkt-bFTWhAI4vMQFh6WeZu0fM4lFd2NcRwr3XPksINHaQ-G_xBniIqbw0Ls1jF44-csFCur-kEgU8awapJzKnqDKgw","e": "AQAB","alg": "RS256","kid": "2011-04-29"}]}"""))))),
+  ))
+  @GetMapping(
+    path = Array("/public-key-jwks"),
+    produces = Array(MediaType.APPLICATION_JSON_VALUE)
+  )
+  @ResponseStatus(HttpStatus.OK)
+  def getPublicKeyJwks(): CompletableFuture[Map[String, AnyRef]] = {
+    val jwks = jwtService.jwks
+
+    import scala.collection.JavaConverters._
+    Future.successful(jwks.toJSONObject(true).asScala.toMap)
+  }
 }

service/src/main/scala/za/co/absa/loginsvc/rest/service/JWTService.scala

@@ -16,6 +16,8 @@
 
 package za.co.absa.loginsvc.rest.service
 
+import com.nimbusds.jose.JWSAlgorithm
+import com.nimbusds.jose.jwk.{JWKSet, KeyUse, RSAKey}
 import io.jsonwebtoken.security.Keys
 import io.jsonwebtoken.{JwtBuilder, Jwts, SignatureAlgorithm}
 import org.springframework.beans.factory.annotation.Autowired
@@ -25,6 +27,7 @@ import za.co.absa.loginsvc.rest.config.provider.JwtConfigProvider
 import za.co.absa.loginsvc.rest.service.JWTService.JwtBuilderExt
 import za.co.absa.loginsvc.utils.OptionExt
 
+import java.security.interfaces.RSAPublicKey
 import java.security.{KeyPair, PublicKey}
 import java.time.Instant
 import java.time.temporal.ChronoUnit
@@ -61,6 +64,16 @@ class JWTService @Autowired()(jwtConfigProvider: JwtConfigProvider) {
 
   def publicKey: PublicKey = rsaKeyPair.getPublic
 
+  def jwks: JWKSet = {
+    val jwk = publicKey match {
+      case rsaKey: RSAPublicKey => new RSAKey.Builder(rsaKey)
+        .keyUse(KeyUse.SIGNATURE)
+        .algorithm(JWSAlgorithm.parse(jwtConfig.algName))
+        .build()
+      case _ => throw new IllegalArgumentException("Unsupported public key type")
+    }
+    new JWKSet(jwk).toPublicJWKSet
+  }
 }
 
 object JWTService {

service/src/test/scala/za/co/absa/loginsvc/rest/controller/TokenControllerTest.scala

@@ -16,6 +16,7 @@
 
 package za.co.absa.loginsvc.rest.controller
 
+import com.nimbusds.jose.jwk.{JWKSet, RSAKey}
 import io.jsonwebtoken.SignatureAlgorithm
 import io.jsonwebtoken.security.Keys
 import org.mockito.ArgumentMatchers.any
@@ -30,6 +31,7 @@ import za.co.absa.loginsvc.model.User
 import za.co.absa.loginsvc.rest.{FakeAuthentication, SecurityConfig}
 import za.co.absa.loginsvc.rest.service.JWTService
 
+import java.security.interfaces.RSAPublicKey
 import java.util.Base64
 
 @Import(Array(classOf[SecurityConfig]))
@@ -95,4 +97,58 @@ class TokenControllerTest extends AnyFlatSpec with ControllerIntegrationTestBase
     )(FakeAuthentication.fakeAnonymousAuthentication)
   }
 
+  behavior of "getPublicKeyJwks"
+
+  it should "return a JWKS from JWTService when user is authenticated" in {
+    val publicKey = Keys.keyPairFor(SignatureAlgorithm.RS256).getPublic
+    val jwk = new RSAKey.Builder(publicKey.asInstanceOf[RSAPublicKey]).build()
+    val jwks = new JWKSet(jwk)
+
+    when(jwtService.jwks).thenReturn(jwks)
+
+    val expectedResponse = s"""
+         |{
+         |  "keys": [
+         |    {
+         |      "kty":"${jwk.getKeyType}",
+         |      "e":"${jwk.getPublicExponent}",
+         |      "n":"${jwk.getModulus}"
+         |    }
+         |  ]
+         |}
+         |""".stripMargin
+
+    assertOkAndResultBodyJsonEquals(
+      "/token/public-key-jwks",
+      Get(),
+      expectedResponse
+    )(FakeAuthentication.fakeUserAuthentication)
+  }
+
+  it should "return a JWKS from JWTService when user is not authenticated" in {
+    val publicKey = Keys.keyPairFor(SignatureAlgorithm.RS256).getPublic
+    val jwk = new RSAKey.Builder(publicKey.asInstanceOf[RSAPublicKey]).build()
+    val jwks = new JWKSet(jwk)
+
+    when(jwtService.jwks).thenReturn(jwks)
+
+    val expectedResponse =
+      s"""
+         |{
+         |  "keys": [
+         |    {
+         |      "kty":"${jwk.getKeyType}",
+         |      "e":"${jwk.getPublicExponent}",
+         |      "n":"${jwk.getModulus}"
+         |    }
+         |  ]
+         |}
+         |""".stripMargin
+
+    assertOkAndResultBodyJsonEquals(
+      "/token/public-key-jwks",
+      Get(),
+      expectedResponse
+    )(FakeAuthentication.fakeAnonymousAuthentication)
+  }
 }

service/src/test/scala/za/co/absa/loginsvc/rest/service/JWTServiceTest.scala

@@ -16,6 +16,8 @@
 
 package za.co.absa.loginsvc.rest.service
 
+import com.nimbusds.jose.JWSAlgorithm
+import com.nimbusds.jose.jwk.KeyUse
 import io.jsonwebtoken.{Claims, Jws, Jwts}
 import org.scalatest.flatspec.AnyFlatSpec
 import za.co.absa.loginsvc.model.User
@@ -115,4 +117,23 @@ class JWTServiceTest extends AnyFlatSpec {
     assert(actualGroups === userWithGroups.groups)
   }
 
+  behavior of "jwks"
+
+  it should "return a JWK that is equivalent to the `publicKey`" in {
+    import scala.collection.JavaConverters._
+
+    val publicKey = jwtService.publicKey
+    val jwks = jwtService.jwks
+    val rsaKey = jwks.getKeys.asScala.head.toRSAKey
+
+    assert(publicKey == rsaKey.toPublicKey)
+  }
+
+  it should "return a JWK with parameters" in {
+    import scala.collection.JavaConverters._
+
+    val jwk = jwtService.jwks.getKeys.asScala.head
+    assert(jwk.getAlgorithm == JWSAlgorithm.RS256)
+    assert(jwk.getKeyUse == KeyUse.SIGNATURE)
+  }
 }