Akka HTTP (Scala) Keycloak token verifier #scala #akka-http #keycloak

KeycloakAuthorization.scala

import java.math.BigInteger
import java.security.spec.RSAPublicKeySpec
import java.security.{KeyFactory, PublicKey}
import java.util.Base64

import akka.actor.ActorSystem
import akka.event.LoggingAdapter
import akka.http.scaladsl.Http
import akka.http.scaladsl.model.HttpRequest
import akka.http.scaladsl.model.headers.{Authorization, OAuth2BearerToken}
import akka.http.scaladsl.server.Directives.{onComplete, optionalCookie, optionalHeaderValueByType, provide, reject}
import akka.http.scaladsl.server.{AuthorizationFailedRejection, Directive1}
import akka.http.scaladsl.unmarshalling.{PredefinedFromEntityUnmarshallers, Unmarshal}
import akka.stream.ActorMaterializer
import de.heikoseeberger.akkahttpcirce.ErrorAccumulatingCirceSupport
import io.circe._
import org.keycloak.TokenVerifier
import org.keycloak.adapters.KeycloakDeployment
import org.keycloak.common.VerificationException
import org.keycloak.jose.jws.{AlgorithmType, JWSHeader}
import org.keycloak.representations.AccessToken

import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.Future

class KeycloakAuthorization(logger: LoggingAdapter, keycloakDeployment: KeycloakDeployment,
                            realm: String, implicit val system: ActorSystem) extends ErrorAccumulatingCirceSupport {
  private implicit val mat: ActorMaterializer = ActorMaterializer()

  def authorized: Directive1[AccessToken] = {
    bearerToken.flatMap {
      case Some(token) =>
        onComplete(verifyToken(token)).flatMap {
          _.map(accessToken => provide(accessToken))
            .recover {
              case e =>
                logger.error(e, "Couldn't log in using provided authorization token")
                reject(AuthorizationFailedRejection)
                  .toDirective[Tuple1[AccessToken]]
            }
            .get
        }
      case None =>
        reject(AuthorizationFailedRejection)
    }
  }

  /**
    * Obtain Bearer Token from Authentication Header.
    * Fallback to X-Authorization-Token Cookie on failure.
    *
    * @return
    *         The Bearer Token.
    */
  private def bearerToken: Directive1[Option[String]] = {
    for {
      authBearerHeader <- optionalHeaderValueByType(classOf[Authorization])
        .map(authHeader => authHeader.collect {
          case Authorization(OAuth2BearerToken(token)) => token
        })
      xAuthCookie <- optionalCookie("X-Authorization-Token")
        .map(_.map(cookie => cookie.value))
    } yield authBearerHeader.orElse(xAuthCookie)
  }

  private def verifyToken(token: String): Future[AccessToken] = {
    val tokenVerifier = TokenVerifier.create(token, classOf[AccessToken])
      .withDefaultChecks()
    getPublicKey(tokenVerifier.getHeader).map {
      case Left(error) =>
        throw new VerificationException("Error deserializing public key from Keycloak.", error)
      case Right(option) => option match {
        case Some(publicKey) =>
          tokenVerifier.publicKey(publicKey).verify().getToken
        case None =>
          throw new VerificationException("No matching public key found.")
      }
    }
  }

  /**
    * Because Keycloak supports key rotation, we cannot hardcode the public key.
    * This method performs a HTTP GET request to Keycloak's OpenID Connect's
    * certification endpoint to obtain the public keys, then uses the kid header
    * in the JWS Header to find the public key that corresponds with the
    * private key used to generate the token.
    *
    * This example uses akka-http-json and circe to unmarshal the JSON response
    * from the GET request.
    *
    * @see
    *      https://gist.github.com/thomasdarimont/52152ed68486c65b50a04fcf7bd9bbde
    *
    * @param jwsHeader
    *                   The JSON Web Signature header.
    * @return
    *         The current public key.
    */
  private def getPublicKey(jwsHeader: JWSHeader): Future[Either[DecodingFailure, Option[PublicKey]]] = {
    Http().singleRequest(HttpRequest(uri = keycloakDeployment.getJwksUrl))
      .flatMap(response => {PredefinedFromEntityUnmarshallers
        Unmarshal(response).to[Json]
          .map(json => json.hcursor
            .get[List[Map[String, String]]]("keys")
            .map(keys => keys
              .find(key => key("kid") == jwsHeader.getKeyId)
              .map(key => {
                val keyFactory = KeyFactory.getInstance(AlgorithmType.RSA.toString)
                val modulusBase64 = key("n")
                val exponentBase64 = key("e")

                val urlDecoder = Base64.getUrlDecoder
                val modulus = new BigInteger(1, urlDecoder.decode(modulusBase64))
                val publicExponent = new BigInteger(1, urlDecoder.decode(exponentBase64))

                keyFactory.generatePublic(new RSAPublicKeySpec(modulus, publicExponent))
              })
            )
          )
      })
  }
}