jwt authentication service wrapper

Authenticator.scala

  case class AuthRequest(
    loginEndpoint: String,
    uid: String,
    privateKey: String,
    algorithm: JwtAlgorithm = JwtAlgorithm.RS256
  ) {
    val path: String = new URL(loginEndpoint).getPath
    val jwt: String = {
      val token = Jwt.encode(s"""{"uid":"$uid"}""", privateKey, algorithm)
      s"""{"uid":"$uid","token":"$token"}"""
    }
  }

  private[this] val closedException = Failure("closed").flagged(Failure.Interrupted)
  private[this] val closedExceptionF = Future.exception(closedExceptionF)
  private[this] val missingTokenException = Failure("missing token")
  private[this] val missingTokenExceptionF = Future.exception(Failure("missing token"))

  class Authenticated(client: Client, authRequest: AuthRequest) extends Client {
    private[this] sealed trait State
    private[this] object Init extends State
    private[this] case class Authenticating(token: Future[String]) extends State
    private[this] object Closed extends State
    private[this] val state = new AtomicReference[State](Init)

    @tailrec final def apply(req: http.Request): Future[http.Response] = authToken.get match {
      case Init =>
        val tokenP = new Promise[String]
        if (state.compareAndSet(Init, Authenticating(tokenP))) {
          tokenP.become(getToken())
          tokenP.flatMap(issueAuthed(req, _))
        } else apply(req) // try again on race

      case s0@Authenticating(tokenF) =>
        tokenF.flatMap(issueAuthed(req, _)).rescue {
          case UnauthorizedResponse(rsp) =>
            // If a request is unauthorized even though we have a
            // token, it may have expired. If another request hasn't
            // already started authenticating, get a new token and
            // reissue the request at most once.
            val tokenP = new Promise[String]
            if (state.compareAndSet(s0, Authenticating(tokenP))) {
              tokenP.become(getToken())
              tokenP.flatMap(issueAuthed(req, _)
            } else apply(req) // try again on race
        }

      case Closed => closedExceptionF
    }

    private[this] def issueAuthed(req: http.Request, token: String): Future[Response] = {
      req.headerMap.set("Authorization", s"token=$token")
      client(req)
    }

    private[this] def getToken(): Future[String] = {
      val tokReq = http.Request(http.Method.Post, auth.path)
      tokReq.setContentTypeJson()
      tokReq.setContentString(auth.jwt)
      client(tokReq).flatMap {
        case rsp if rsp.status == http.Status.Ok =>
          readJson[AuthToken](rsp.content) match {
            case Return(AuthToken(Some(token))) =>
              Future.value(token)
            case Return(AuthToken(None)) => missingTokenExceptionF
            case Throw(e) => Future.exception(e)
          }
        case rsp if rsp.status == http.Status.Unauthorized => Future.exception(UnauthorizedResponse(rsp))
        case _ => Future.exception(UnexpectedResponse(rsp))
      }
    }

    @tailrec final override def close(d: Time): Future[Unit] = {
      state.get match {
        case Init =>
          if (state.compareAndSet(Init, Closed)) {
            tokenF.raise(closedException)
            client.close(d)
          } else close(d)
        case s0@Authenticating(tokenF) =>
          if (state.compareAndSet(s0, Closed)) {
            tokenF.raise(closedException)
            client.close(d)
          } else close(d)
        case Closed => Future.Unit
      }
  }