JWKS demo (#63)

* JWKS demo * Add tests * Update service/src/test/scala/za/co/absa/loginsvc/rest/service/JWTServiceTest.scala Co-authored-by: Daniel K <[email protected]> * Feature/jwks kid (#64) * jwks: kid added * jwks: kid test added * jwks: kid changed to be a public rsakey thumprint * remove static val kid (#66) --------- Co-authored-by: Daniel K <[email protected]>
AbsaOSS · Aug 23, 2023 · 07fd00e · 07fd00e
1 parent 828ee98
commit 07fd00e
Show file tree

Hide file tree

Showing 6 changed files with 137 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -52,6 +55,7 @@ class JWTService @Autowired()(jwtConfigProvider: JwtConfigProvider) {
       .setSubject(user.name)
       .setExpiration(expiration)
       .setIssuedAt(issuedAt)
+      .claim("kid", publicKeyThumbprint)
       .claim("groups", groupsClaim)
       .applyIfDefined(user.email, (builder, value: String) => builder.claim("email", value))
       .applyIfDefined(user.displayName, (builder, value: String) => builder.claim("displayname", value))
@@ -61,6 +65,23 @@ class JWTService @Autowired()(jwtConfigProvider: JwtConfigProvider) {
 
   def publicKey: PublicKey = rsaKeyPair.getPublic
 
+  private def rsaPublicKey: RSAKey = {
+    publicKey match {
+      case rsaKey: RSAPublicKey => new RSAKey.Builder(rsaKey)
+        .keyUse(KeyUse.SIGNATURE)
+        .algorithm(JWSAlgorithm.parse(jwtConfig.algName))
+        .keyIDFromThumbprint()
+        .build()
+      case _ => throw new IllegalArgumentException("Unsupported public key type")
+    }
+  }
+
+  def publicKeyThumbprint: String = rsaPublicKey.getKeyID
+
+  def jwks: JWKSet = {
+    val jwk = rsaPublicKey
+    new JWKSet(jwk).toPublicJWKSet
+  }
 }
 
 object JWTService {

diff --git a/service/src/test/scala/za/co/absa/loginsvc/rest/controller/TokenControllerTest.scala b/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)
+  }
 }
diff --git a/service/src/test/scala/za/co/absa/loginsvc/rest/service/JWTServiceTest.scala b/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
@@ -89,6 +91,17 @@ class JWTServiceTest extends AnyFlatSpec {
     }
   }
 
+  it should "return a JWT kid" in {
+    val jwt = jwtService.generateToken(userWithoutEmailAndGroups)
+    val parsedJWT = parseJWT(jwt)
+
+    assert(parsedJWT.isSuccess)
+    parsedJWT.foreach { jwt =>
+      val kid = jwt.getBody.get("kid")
+      assert(kid === jwtService.publicKeyThumbprint)
+    }
+  }
+
   it should "turn groups into empty `groups` claim for user without groups" in {
     import scala.collection.JavaConverters._
 
@@ -115,4 +128,26 @@ 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 keys = jwtService.jwks.getKeys.asScala
+    assert(keys.length == 1, "One JWK is expected to be generated now")
+
+    val jwk = keys.head
+    assert(jwk.getAlgorithm == JWSAlgorithm.RS256)
+    assert(jwk.getKeyUse == KeyUse.SIGNATURE)
+  }
 }