From 7aaaf6b31cafe0dd4a50262bcda825d6020ab6ee Mon Sep 17 00:00:00 2001
From: Kevin Wallimann <kevin.wallimann@absa.africa>
Date: Tue, 22 Aug 2023 11:24:26 +0200
Subject: [PATCH] JWKS demo

---
 project/Dependencies.scala                     |  6 ++++++
 .../co/absa/loginsvc/rest/SecurityConfig.scala |  1 +
 .../rest/controller/TokenController.scala      | 18 ++++++++++++++++++
 .../loginsvc/rest/service/JWTService.scala     | 13 +++++++++++++
 4 files changed, 38 insertions(+)

diff --git a/project/Dependencies.scala b/project/Dependencies.scala
index 41b7fde..ca2c968 100644
--- a/project/Dependencies.scala
+++ b/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,
 
diff --git a/service/src/main/scala/za/co/absa/loginsvc/rest/SecurityConfig.scala b/service/src/main/scala/za/co/absa/loginsvc/rest/SecurityConfig.scala
index c000670..38680d0 100644
--- a/service/src/main/scala/za/co/absa/loginsvc/rest/SecurityConfig.scala
+++ b/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()
diff --git a/service/src/main/scala/za/co/absa/loginsvc/rest/controller/TokenController.scala b/service/src/main/scala/za/co/absa/loginsvc/rest/controller/TokenController.scala
index 1afbbf0..5c761fd 100644
--- a/service/src/main/scala/za/co/absa/loginsvc/rest/controller/TokenController.scala
+++ b/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)
+  }
 }
diff --git a/service/src/main/scala/za/co/absa/loginsvc/rest/service/JWTService.scala b/service/src/main/scala/za/co/absa/loginsvc/rest/service/JWTService.scala
index c50558c..3ff808e 100644
--- a/service/src/main/scala/za/co/absa/loginsvc/rest/service/JWTService.scala
+++ b/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 {