[RFC] Use akka's built-in authenticate/authorize directives in AuthProvider (#136)

* Use akka's built-in authenticate/authorize directives in AuthProvider

* Move AuthProvider companion object to AuthProvider file, move realm to parameter of AuthProvider

* Add secondary constructor to maintain ABI compat

5 files changed, 88 insertions, 66 deletions

diff --git a/src/main/scala/xyz/driver/core/rest/auth/AuthProvider.scala b/src/main/scala/xyz/driver/core/rest/auth/AuthProvider.scala
index 82edcc7..1fddd45 100644
--- a/src/main/scala/xyz/driver/core/rest/auth/AuthProvider.scala
+++ b/src/main/scala/xyz/driver/core/rest/auth/AuthProvider.scala
@@ -1,23 +1,24 @@
 package xyz.driver.core.rest.auth
 
-import akka.http.scaladsl.model.headers.HttpChallenges
-import akka.http.scaladsl.server.AuthenticationFailedRejection.CredentialsRejected
+import akka.http.scaladsl.server.directives.Credentials
 import com.typesafe.scalalogging.Logger
-import xyz.driver.core._
-import xyz.driver.core.auth.{Permission, User}
-import xyz.driver.core.rest.{AuthorizedServiceRequestContext, ServiceRequestContext, serviceContext}
+import scalaz.OptionT
+import xyz.driver.core.auth.{AuthToken, Permission, User}
+import xyz.driver.core.rest.{AuthorizedServiceRequestContext, ContextHeaders, ServiceRequestContext, serviceContext}
 
 import scala.concurrent.{ExecutionContext, Future}
-import scala.util.{Failure, Success}
-
-import scalaz.Scalaz.futureInstance
-import scalaz.OptionT
 
-abstract class AuthProvider[U <: User](val authorization: Authorization[U], log: Logger)(
-    implicit execution: ExecutionContext) {
+abstract class AuthProvider[U <: User](
+    val authorization: Authorization[U],
+    log: Logger,
+    val realm: String
+)(implicit execution: ExecutionContext) {
 
   import akka.http.scaladsl.server._
-  import Directives._
+  import Directives.{authorize => akkaAuthorize, _}
+
+  def this(authorization: Authorization[U], log: Logger)(implicit executionContext: ExecutionContext) =
+    this(authorization, log, "driver.xyz")
 
   /**
     * Specific implementation on how to extract user from request context,
@@ -28,37 +29,30 @@ abstract class AuthProvider[U <: User](val authorization: Authorization[U], log:
     */
   def authenticatedUser(implicit ctx: ServiceRequestContext): OptionT[Future, U]
 
+  protected def authenticator(context: ServiceRequestContext): AsyncAuthenticator[U] = {
+    case Credentials.Missing =>
+      log.info(s"Request (${context.trackingId}) missing authentication credentials")
+      Future.successful(None)
+    case Credentials.Provided(authToken) =>
+      authenticatedUser(context.withAuthToken(AuthToken(authToken))).run
+  }
+
   /**
-    * Verifies if a service context is authenticated and authorized to have `permissions`
+    * Verifies that a user agent is properly authenticated, and (optionally) authorized with the specified permissions
     */
   def authorize(
       context: ServiceRequestContext,
       permissions: Permission*): Directive1[AuthorizedServiceRequestContext[U]] = {
-    onComplete {
-      (for {
-        authToken <- OptionT.optionT(Future.successful(context.authToken))
-        user      <- authenticatedUser(context)
-        authCtx = context.withAuthenticatedUser(authToken, user)
-        authorizationResult <- authorization.userHasPermissions(user, permissions)(authCtx).toOptionT
-
-        cachedPermissionsAuthCtx = authorizationResult.token.fold(authCtx)(authCtx.withPermissionsToken)
-        allAuthorized            = permissions.forall(authorizationResult.authorized.getOrElse(_, false))
-      } yield (cachedPermissionsAuthCtx, allAuthorized)).run
-    } flatMap {
-      case Success(Some((authCtx, true))) => provide(authCtx)
-      case Success(Some((authCtx, false))) =>
-        val challenge =
-          HttpChallenges.basic(s"User does not have the required permissions: ${permissions.mkString(", ")}")
-        log.warn(
-          s"User ${authCtx.authenticatedUser} does not have the required permissions: ${permissions.mkString(", ")}")
-        reject(AuthenticationFailedRejection(CredentialsRejected, challenge))
-      case Success(None) =>
-        val challenge = HttpChallenges.basic("Failed to authenticate user")
-        log.warn(s"Failed to authenticate user to verify ${permissions.mkString(", ")}")
-        reject(AuthenticationFailedRejection(CredentialsRejected, challenge))
-      case Failure(t) =>
-        log.warn(s"Wasn't able to verify token for authenticated user to verify ${permissions.mkString(", ")}", t)
-        reject(ValidationRejection(s"Wasn't able to verify token for authenticated user", Some(t)))
+    authenticateOAuth2Async[U](realm, authenticator(context)) flatMap { authenticatedUser =>
+      val authCtx = context.withAuthenticatedUser(context.authToken.get, authenticatedUser)
+      onSuccess(authorization.userHasPermissions(authenticatedUser, permissions)(authCtx)) flatMap {
+        case AuthorizationResult(authorized, token) =>
+          val allAuthorized = permissions.forall(authorized.getOrElse(_, false))
+          akkaAuthorize(allAuthorized) tflatMap { _ =>
+            val cachedPermissionsCtx = token.fold(authCtx)(authCtx.withPermissionsToken)
+            provide(cachedPermissionsCtx)
+          }
+      }
     }
   }
 
@@ -66,8 +60,13 @@ abstract class AuthProvider[U <: User](val authorization: Authorization[U], log:
     * Verifies if request is authenticated and authorized to have `permissions`
     */
   def authorize(permissions: Permission*): Directive1[AuthorizedServiceRequestContext[U]] = {
-    serviceContext flatMap { ctx =>
-      authorize(ctx, permissions: _*)
-    }
+    serviceContext flatMap (authorize(_, permissions: _*))
   }
 }
+
+object AuthProvider {
+  val AuthenticationTokenHeader: String    = ContextHeaders.AuthenticationTokenHeader
+  val PermissionsTokenHeader: String       = ContextHeaders.PermissionsTokenHeader
+  val SetAuthenticationTokenHeader: String = "set-authorization"
+  val SetPermissionsTokenHeader: String    = "set-permissions"
+}
diff --git a/src/main/scala/xyz/driver/core/rest/package.scala b/src/main/scala/xyz/driver/core/rest/package.scala
index d4d01df..7d67138 100644
--- a/src/main/scala/xyz/driver/core/rest/package.scala
+++ b/src/main/scala/xyz/driver/core/rest/package.scala
@@ -14,6 +14,7 @@ import akka.util.ByteString
 import scalaz.Scalaz.{intInstance, stringInstance}
 import scalaz.syntax.equal._
 import scalaz.{Functor, OptionT}
+import xyz.driver.core.rest.auth.AuthProvider
 import xyz.driver.tracing.TracingDirectives
 
 import scala.concurrent.Future
@@ -90,13 +91,6 @@ object `package` {
     val SpanHeaderName: String             = TracingDirectives.SpanHeaderName
   }
 
-  object AuthProvider {
-    val AuthenticationTokenHeader: String    = ContextHeaders.AuthenticationTokenHeader
-    val PermissionsTokenHeader: String       = ContextHeaders.PermissionsTokenHeader
-    val SetAuthenticationTokenHeader: String = "set-authorization"
-    val SetPermissionsTokenHeader: String    = "set-permissions"
-  }
-
   val AllowedHeaders: Seq[String] =
     Seq(
       "Origin",
@@ -131,8 +125,18 @@ object `package` {
       originHeader.fold[HttpOriginRange](HttpOriginRange.*)(h => HttpOriginRange(h.origins: _*)))
 
   def serviceContext: Directive1[ServiceRequestContext] = {
+    def fixAuthorizationHeader(headers: Seq[HttpHeader]): collection.immutable.Seq[HttpHeader] = {
+      headers.map({ header =>
+        if (header.name === ContextHeaders.AuthenticationTokenHeader && !header.value.startsWith(
+              ContextHeaders.AuthenticationHeaderPrefix)) {
+          Authorization(OAuth2BearerToken(header.value))
+        } else header
+      })(collection.breakOut)
+    }
     extractClientIP flatMap { remoteAddress =>
-      extract(ctx => extractServiceContext(ctx.request, remoteAddress))
+      mapRequest(req => req.withHeaders(fixAuthorizationHeader(req.headers))) tflatMap { _ =>
+        extract(ctx => extractServiceContext(ctx.request, remoteAddress))
+      }
     }
   }
 
diff --git a/src/main/scala/xyz/driver/core/rest/serviceRequestContext.scala b/src/main/scala/xyz/driver/core/rest/serviceRequestContext.scala
index 775106e..76f5a0d 100644
--- a/src/main/scala/xyz/driver/core/rest/serviceRequestContext.scala
+++ b/src/main/scala/xyz/driver/core/rest/serviceRequestContext.scala
@@ -4,9 +4,9 @@ import java.net.InetAddress
 
 import xyz.driver.core.auth.{AuthToken, PermissionsToken, User}
 import xyz.driver.core.generators
-
 import scalaz.Scalaz.{mapEqual, stringInstance}
 import scalaz.syntax.equal._
+import xyz.driver.core.rest.auth.AuthProvider
 
 class ServiceRequestContext(
     val trackingId: String = generators.nextUuid().toString,
diff --git a/src/test/scala/xyz/driver/core/AuthTest.scala b/src/test/scala/xyz/driver/core/AuthTest.scala
index a7707aa..2e772fb 100644
--- a/src/test/scala/xyz/driver/core/AuthTest.scala
+++ b/src/test/scala/xyz/driver/core/AuthTest.scala
@@ -1,7 +1,11 @@
 package xyz.driver.core
 
-import akka.http.scaladsl.model.headers.{HttpChallenges, RawHeader}
-import akka.http.scaladsl.server.AuthenticationFailedRejection.CredentialsRejected
+import akka.http.scaladsl.model.headers.{
+  HttpChallenges,
+  OAuth2BearerToken,
+  RawHeader,
+  Authorization => AkkaAuthorization
+}
 import akka.http.scaladsl.server.Directives._
 import akka.http.scaladsl.server._
 import akka.http.scaladsl.testkit.ScalatestRouteTest
@@ -79,36 +83,51 @@ class AuthTest extends FlatSpec with Matchers with ScalatestRouteTest {
       } ~>
       check {
         // handled shouldBe false
-        val challenge = HttpChallenges.basic("Failed to authenticate user")
-        rejections should contain(AuthenticationFailedRejection(CredentialsRejected, challenge))
+        rejections should contain(
+          AuthenticationFailedRejection(
+            AuthenticationFailedRejection.CredentialsMissing,
+            HttpChallenges.oAuth2(authStatusService.realm)))
       }
   }
 
   it should "throw error if authorized user does not have the requested permission" in {
 
-    val referenceAuthToken = AuthToken("I am a test role's token")
+    val referenceAuthToken  = AuthToken("I am a test role's token")
+    val referenceAuthHeader = AkkaAuthorization(OAuth2BearerToken(referenceAuthToken.value))
 
     Post("/administration/attempt").addHeader(
-      RawHeader(AuthProvider.AuthenticationTokenHeader, referenceAuthToken.value)
+      referenceAuthHeader
     ) ~>
       authorize(TestRoleNotAllowedPermission) { user =>
         complete("Never going to get here")
       } ~>
       check {
         handled shouldBe false
-        rejections should contain(
-          AuthenticationFailedRejection(
-            CredentialsRejected,
-            HttpChallenges.basic("User does not have the required permissions: TestRoleNotAllowedPermission")))
+        rejections should contain(AuthorizationFailedRejection)
       }
   }
 
   it should "pass and retrieve the token to client code, if token is in request and user has permission" in {
+    val referenceAuthToken  = AuthToken("I am token")
+    val referenceAuthHeader = AkkaAuthorization(OAuth2BearerToken(referenceAuthToken.value))
+
+    Get("/valid/attempt/?a=2&b=5").addHeader(
+      referenceAuthHeader
+    ) ~>
+      authorize(TestRoleAllowedPermission) { ctx =>
+        complete(s"Alright, user ${ctx.authenticatedUser.id} is authorized")
+      } ~>
+      check {
+        handled shouldBe true
+        responseAs[String] shouldBe "Alright, user 1 is authorized"
+      }
+  }
 
-    val referenceAuthToken = AuthToken("I am token")
+  it should "authenticate correctly even without the 'Bearer' prefix on the Authorization header" in {
+    val referenceAuthToken = AuthToken("unprefixed_token")
 
     Get("/valid/attempt/?a=2&b=5").addHeader(
-      RawHeader(AuthProvider.AuthenticationTokenHeader, referenceAuthToken.value)
+      RawHeader(ContextHeaders.AuthenticationTokenHeader, referenceAuthToken.value)
     ) ~>
       authorize(TestRoleAllowedPermission) { ctx =>
         complete(s"Alright, user ${ctx.authenticatedUser.id} is authorized")
@@ -128,11 +147,12 @@ class AuthTest extends FlatSpec with Matchers with ScalatestRouteTest {
         "sub"         -> JsString("1"),
         "permissions" -> JsObject(Map(TestRoleAllowedByTokenPermission.toString -> JsBoolean(true)))
       )).prettyPrint
-    val permissionsToken   = PermissionsToken(Jwt.encode(claim, privateKey, JwtAlgorithm.RS256))
-    val referenceAuthToken = AuthToken("I am token")
+    val permissionsToken    = PermissionsToken(Jwt.encode(claim, privateKey, JwtAlgorithm.RS256))
+    val referenceAuthToken  = AuthToken("I am token")
+    val referenceAuthHeader = AkkaAuthorization(OAuth2BearerToken(referenceAuthToken.value))
 
     Get("/alic/attempt/?a=2&b=5")
-      .addHeader(RawHeader(AuthProvider.AuthenticationTokenHeader, referenceAuthToken.value))
+      .addHeader(referenceAuthHeader)
       .addHeader(RawHeader(AuthProvider.PermissionsTokenHeader, permissionsToken.value)) ~>
       authorize(TestRoleAllowedByTokenPermission) { ctx =>
         complete(s"Alright, user ${ctx.authenticatedUser.id} is authorized by permissions token")
diff --git a/src/test/scala/xyz/driver/core/rest/RestTest.scala b/src/test/scala/xyz/driver/core/rest/RestTest.scala
index e742462..19e4ed1 100644
--- a/src/test/scala/xyz/driver/core/rest/RestTest.scala
+++ b/src/test/scala/xyz/driver/core/rest/RestTest.scala
@@ -1,6 +1,5 @@
 package xyz.driver.core.rest
 
-import akka.http.javadsl.server.MalformedRequestContentRejection
 import akka.http.scaladsl.model.StatusCodes
 import akka.http.scaladsl.server.{Directives, Route, ValidationRejection}
 import akka.http.scaladsl.testkit.ScalatestRouteTest