Scala
Akka
JWT
Akka-HTTP

Scala、AkkaHttp、JsonWebToken(Jwt)を使用して認証を実装したい!パート3

公開鍵や無効トークンリストを取得してチェックしたい!

パート1では鍵は定数で保持していましたし、無効リストのチェックは実装していませんでした。

なので次は
公開鍵と、すでに無効になっているアクセストークンリストを動的に取得してその都度トークンの有効性をチェックするようにしていきます。

AkkaHttpを使って、外部の認証サーバーにGetリクエストできるように認証のみで使用するリクエストを作成しておきます。

qiita.scala
  private def getRequest[A](uri: String): Future[String] = {
    val request = HttpRequest(GET, uri = Uri(uri))
    val res = Http().singleRequest(request)
    res.flatMap { r =>
      Unmarshal(r.entity).to[String]
    }
  }

取得してきたbiteをStringに治す方法なんですけど、

qiita.scala
・ネットでよく見る例
val body: Future[String] = res.entity.dataBytes.runFold(ByteString.empty)(_ ++ _).map(_.utf8String)

・簡単な例
val body: Future[String] = Unmarshal(res.entity).to[String]

Unmarshalを使った方が全然短くかけるので、使っていきましょう。

受け取ったJsonをStringからScalaの型にparseしよう!

これで認証サーバーからのresponseからFuture[String]形式で取得することができるようになりました。

ただJsonのStringのままだと

qiita.scala
val rawJson = 
'{
  "user_account_id": "3961",
  "user_name": "MLKK0002",
  "scope": [
    "asaas" 
  ],
  "session_id": "f05d6ab96204e7b250dbde6ce7824251",
  "exp": 1526406373,
  "authorities": [
    "OFFICE_ADMIN" 
  ],
  "jti": "6ed59891-c569-4558-92ad-9713db577819",
  "client_id": "asaas_client" 
}'

のように非常に使い勝手が悪いので、Scalaの型に直していきたいと思います。

parseの仕方はこちらを参考にしました。
http://takkkun.hatenablog.com/entry/2017/04/04/Scala%E3%81%AEJSON%E3%83%A9%E3%82%A4%E3%83%96%E3%83%A9%E3%83%AAcirce%E3%81%AE%E4%BD%BF%E3%81%84%E6%96%B9

上の記事を参考にしてgetRequestを拡張します。
parseの独自の型にコンバートする関数を渡せるようにしました。

qiita.scala
  private def getRequest[A](uri: String)(f: Json => Decoder.Result[A]): Future[Either[Error,A]] = {
    val request = HttpRequest(GET, uri = Uri(uri))
    val res = Http().singleRequest(request)
    res.flatMap { r =>
      Unmarshal(r.entity).to[String].map{ b => parse(b).right.flatMap(f)}
    }
  }
qiita.scala
parse(b).right.flatMap(_.as[PublicKey])

 この一行でPublicKeyという型に格納してます。
PublicKey以外にも無効リストを取得したい時もあるので、parseする型は関数で渡せるようにしておきます。

qiita.scala
case class PublicKey(
  alg: String,
  value: String
)

使うときは

qiita.scala
val publicKey: Future[Either[Error, A]] = getRequest(uri = PublicKeyPath)(_.as[PublicKey])

こんな感じで呼び出します。
これで公開鍵や無効トークンリストの取得ができるようになりました。

注意:この時にas[PublicKey]のようにDecoder.Result[A]を渡すのですが、そこにStringやLongのような型ではなく、独自の型を設定するとimplicit Decoderを自分で設定しないとコンパイルエラーになります。

型が色々ごちゃごちゃしてきたからEitherTを使いたい!

Future[Either[Error,A]]と色々ごちゃごちゃしているので、EitherTを使おうと思います。

qiita.scala
  type ProcessResult[T] = EitherT[Future, Error, T]

このようにEitherTをエイリアス貼って、使いやすいようにしておきます。

さっきのgetRequestの返りの型をProcessResultにしてみます。

qiita.scala
  private def getRequest[A](uri: String)(f: Json => Decoder.Result[A]): ProcessResult[A] = {
    val request = HttpRequest(GET, uri = Uri(uri))
    val res = Http().singleRequest(request)
    EitherT{
      res.flatMap { r =>
        Unmarshal(r.entity).to[String].map{ b => parse(b).right.flatMap(f)}
      }
    }
  }

これでEitherTのコンテキストでflatMapが使えるようになったので、Leftになった瞬間に処理を中断できる形になりました。
パッと実装していきます。

qiita.scala
  private var publicKey = Option.empty[(PublicKey,LocalDate)] #サーバー起動時からの初回リクエストを識別したい

  override def authenticate(credentials: Credentials): Future[Option[AuthContext]] = {
    credentials match {
      case Credentials.Provided(jwt) => validateJwt(jwt)
      case _ => Future { None }
    }
  }

  private def validateJwt(jwt: String): Future[Option[AuthContext]] = {
    (for {
      pub <- saveAndGetPublicKey(jwt)
      auth <- generateAuthContext(jwt, pub._1)
    }yield{
      auth
    }).value.map(_.toOption)
  }

  private def saveAndGetPublicKey(jwt: String): ProcessResult[(PublicKey,LocalDate)] = {
    if (checkPublicKey)
      for {
        pubKey <- getRequest(uri = PublicKeyPath)(_.as[PublicKey])
        _ <- ProcessResult.wrapE(decodeJwt(jwt, pubKey))
        _ <- ProcessResult.wrapSuccess(publicKey = (PublicKey(alg = pubKey.alg, value = pubKey.value), LocalDate.now()).some)
        publicKey <- ProcessResult.wrapE(Either.fromOption(publicKey, err(AuthError)))
      } yield {
        publicKey
      }
    else
      ProcessResult.wrapE {
        Either.fromOption(publicKey, err(AuthError))
      }
  }

  private def generateAuthContext(jwt: String, pub: PublicKey): ProcessResult[AuthContext] =
    for {
      decodedJwt <- ProcessResult.wrapE(decodeJwt(jwt, pub))
      payload <- ProcessResult.wrap{parse(decodedJwt).right.flatMap(_.as[PayloadInput])}
      revokedList <- getRequest(uri = RevokedTokenListPath)(_.as[RevokedTokenInput])
      auth <- ProcessResult.wrapE {
        if (!revokedList.revokedTokens.exists(_.id == payload.jti)) {
          AuthContext(
            userAccountId = payload.user_account_id.toLong.some,
            roles = payload.authorities.map(convertRole),
            sessionId = payload.session_id.some
          ).asRight}
        else err(InvalidJsonSessionId).asLeft
      }
    } yield {
      auth
    }

実装できました。

saveAndGetPublicKeyでは、サーバー起動してから初回のリクエスト、または公開鍵の有効期限が切れている時に取得しに行って、その公開鍵を状態保持する機能です。

またProcessResult.wrapとかProcessResult.wrapEとかは独自の拡張で、中身がただのEitherとかの時にEitherT[Future, Error, T]の形にラップしてくれるものです。

これでただのFuture[Either[Error,A]]の時に比べて、いちいち

qiita.rb
p match {
    case Right(r) => ???
    case Left(e) => ???
}

のように処理を考えなくて良くなりました。
これで実装をすると何が嫌って、

qiita.scala
Future[Either[Error,Either[Error,Option[A]]]]

みたいなめんどくさい型を整理していかないといけないんですよね。

その点EitherTのfor,yieldで書いた場合はすごく楽です。
だってLeftの時は即処理が止まるから、全ての型がProcessResultのコンテキストになるようにfor,yieldを書いていけば何も問題がないわけです。

qiita.scala
  private def validateJwt(jwt: String): Future[Option[AuthContext]] = {
    (for {
      pub <- saveAndGetPublicKey(jwt)
      auth <- generateAuthContext(jwt, pub._1)
    }yield{
      auth
    }).value.map(_.toOption)
  }

でも、Eitherでエラーハンドリングしているのに、最後にtoOptionしてたら意味ないじゃんってことで
EitherTを使うなら、authenticatorもEitherを使うものにしよう!ってことで下にしようとしたんですが、、、

qiita.scala
  private def extractCredentialsAndAuthenticateOrRejectWithChallenge[C <: HttpCredentials, T](
    extractCredentials: Directive1[Option[C]],
    authenticator:      Option[C]  Future[AuthenticationResult[T]]): AuthenticationDirective[T] =
    extractCredentials.flatMap { cred 
      onSuccess(authenticator(cred)).flatMap {
        case Right(user)  provide(user)
        case Left(challenge) 
          val cause = if (cred.isEmpty) CredentialsMissing else CredentialsRejected
          reject(AuthenticationFailedRejection(cause, challenge)): Directive1[T]
      }
    }

よくよく考えたらErrorの内容をユーザーが意識する必要ないんですよねぇ。。。
Noneを返せば、一律「認証エラー」という形でレスポンスが帰るので、別にもう良いかなって(ヤグニの精神大事)

改善点

現状の仕様だと、認証サーバーから公開鍵を取得してから1日経過してたら、鍵を取りに行くという微妙な仕様のためにvar publicという形で取得日時を保持しています。

そもそも公開鍵が変更されたら、アプリサーバー側では知りようがないのが微妙なところなので、公開鍵の変更をコールバックで受け付けるようにするか、もしくは公開鍵でのデコードが失敗したら、もう一度取り直しに行くという仕様に変更を考えてます。

あとはアプリサーバーのsingletonのvarでpublicKeyを保持しているのですが、スレッドセーフな仕様になっていないので
Actorを使用してスレッドセーフな構築にするつもりです。