Merge pull request #35 from drivergroup/zsmith/PDW-600-query-params

PDW-600 Cached Permissions JWT

5 files changed, 208 insertions, 75 deletions

diff --git a/build.sbt b/build.sbt
index c4de456..0963f6c 100644
--- a/build.sbt
+++ b/build.sbt
@@ -10,6 +10,7 @@ lazy val core = (project in file("."))
     "com.typesafe.akka"            %% "akka-http-core"       % akkaHttpV,
     "com.typesafe.akka"            %% "akka-http-spray-json" % akkaHttpV,
     "com.typesafe.akka"            %% "akka-http-testkit"    % akkaHttpV,
+    "com.pauldijou"                %% "jwt-core"             % "0.9.2",
     "org.scalatest"                % "scalatest_2.11"        % "2.2.6" % "test",
     "org.scalacheck"               %% "scalacheck"           % "1.12.5" % "test",
     "org.mockito"                  % "mockito-core"          % "1.9.5" % "test",
diff --git a/src/main/scala/xyz/driver/core/auth.scala b/src/main/scala/xyz/driver/core/auth.scala
index f9a1a57..5dea2db 100644
--- a/src/main/scala/xyz/driver/core/auth.scala
+++ b/src/main/scala/xyz/driver/core/auth.scala
@@ -23,6 +23,7 @@ object auth {
 
   final case class AuthToken(value: String)
   final case class RefreshToken(value: String)
+  final case class PermissionsToken(value: String)
 
   final case class PasswordHash(value: String)
 
diff --git a/src/main/scala/xyz/driver/core/file/S3Storage.scala b/src/main/scala/xyz/driver/core/file/S3Storage.scala
index 50bfe85..933b01a 100644
--- a/src/main/scala/xyz/driver/core/file/S3Storage.scala
+++ b/src/main/scala/xyz/driver/core/file/S3Storage.scala
@@ -53,11 +53,13 @@ class S3Storage(s3: AmazonS3, bucket: Name[Bucket], executionContext: ExecutionC
         result.isTruncated
       } flatMap { result =>
         result.getObjectSummaries.asScala.toList.map { summary =>
-          FileLink(Name[File](summary.getKey),
-                   Paths.get(path.toString + "/" + summary.getKey),
-                   Revision[File](summary.getETag),
-                   Time(summary.getLastModified.getTime),
-                   summary.getSize)
+          FileLink(
+            Name[File](summary.getKey),
+            Paths.get(path.toString + "/" + summary.getKey),
+            Revision[File](summary.getETag),
+            Time(summary.getLastModified.getTime),
+            summary.getSize
+          )
         } filterNot isInSubFolder(path)
       } toList
     })
diff --git a/src/main/scala/xyz/driver/core/rest.scala b/src/main/scala/xyz/driver/core/rest.scala
index f1eab45..1db9d09 100644
--- a/src/main/scala/xyz/driver/core/rest.scala
+++ b/src/main/scala/xyz/driver/core/rest.scala
@@ -1,21 +1,24 @@
 package xyz.driver.core
 
+import java.nio.file.{Files, Path}
+import java.security.spec.X509EncodedKeySpec
+import java.security.{KeyFactory, PublicKey}
+
 import akka.actor.ActorSystem
 import akka.http.scaladsl.Http
 import akka.http.scaladsl.model._
 import akka.http.scaladsl.model.headers.{HttpChallenges, RawHeader}
 import akka.http.scaladsl.server.AuthenticationFailedRejection.CredentialsRejected
-import akka.http.scaladsl.server.Directive0
-import com.typesafe.scalalogging.Logger
-import akka.http.scaladsl.unmarshalling.Unmarshal
-import akka.http.scaladsl.unmarshalling.Unmarshaller
+import akka.http.scaladsl.unmarshalling.{Unmarshal, Unmarshaller}
 import akka.stream.ActorMaterializer
 import akka.stream.scaladsl.Flow
 import akka.util.ByteString
 import com.github.swagger.akka.model._
 import com.github.swagger.akka.{HasActorSystem, SwaggerHttpService}
 import com.typesafe.config.Config
+import com.typesafe.scalalogging.Logger
 import io.swagger.models.Scheme
+import pdi.jwt.{Jwt, JwtAlgorithm}
 import xyz.driver.core.auth._
 import xyz.driver.core.time.provider.TimeProvider
 
@@ -33,7 +36,7 @@ package rest {
     def serviceContext: Directive1[ServiceRequestContext] = extract(ctx => extractServiceContext(ctx.request))
 
     def extractServiceContext(request: HttpRequest): ServiceRequestContext =
-      ServiceRequestContext(extractTrackingId(request), extractContextHeaders(request))
+      new ServiceRequestContext(extractTrackingId(request), extractContextHeaders(request))
 
     def extractTrackingId(request: HttpRequest): String = {
       request.headers
@@ -43,7 +46,8 @@ package rest {
 
     def extractContextHeaders(request: HttpRequest): Map[String, String] = {
       request.headers.filter { h =>
-        h.name === ContextHeaders.AuthenticationTokenHeader || h.name === ContextHeaders.TrackingIdHeader
+        h.name === ContextHeaders.AuthenticationTokenHeader || h.name === ContextHeaders.TrackingIdHeader ||
+        h.name === ContextHeaders.PermissionsTokenHeader
       } map { header =>
         if (header.name === ContextHeaders.AuthenticationTokenHeader) {
           header.name -> header.value.stripPrefix(ContextHeaders.AuthenticationHeaderPrefix).trim
@@ -91,14 +95,59 @@ package rest {
     }
   }
 
-  final case class ServiceRequestContext(trackingId: String = generators.nextUuid().toString,
-                                         contextHeaders: Map[String, String] = Map.empty[String, String]) {
-
+  class ServiceRequestContext(val trackingId: String = generators.nextUuid().toString,
+                              val contextHeaders: Map[String, String] = Map.empty[String, String]) {
     def authToken: Option[AuthToken] =
       contextHeaders.get(AuthProvider.AuthenticationTokenHeader).map(AuthToken.apply)
 
+    def permissionsToken: Option[PermissionsToken] =
+      contextHeaders.get(AuthProvider.PermissionsTokenHeader).map(PermissionsToken.apply)
+
     def withAuthToken(authToken: AuthToken): ServiceRequestContext =
-      copy(contextHeaders = contextHeaders.updated(AuthProvider.AuthenticationTokenHeader, authToken.value))
+      new ServiceRequestContext(
+        trackingId,
+        contextHeaders.updated(AuthProvider.AuthenticationTokenHeader, authToken.value)
+      )
+
+    def withAuthenticatedUser[U <: User](authToken: AuthToken, user: U): AuthorizedServiceRequestContext[U] =
+      new AuthorizedServiceRequestContext(
+        trackingId,
+        contextHeaders.updated(AuthProvider.AuthenticationTokenHeader, authToken.value),
+        user
+      )
+
+    override def hashCode(): Int =
+      Seq[Any](trackingId, contextHeaders).foldLeft(31)((result, obj) => 31 * result + obj.hashCode())
+
+    override def equals(obj: Any): Boolean = obj match {
+      case ctx: ServiceRequestContext => trackingId === ctx.trackingId && contextHeaders === ctx.contextHeaders
+      case _                          => false
+    }
+
+    override def toString: String = s"ServiceRequestContext($trackingId, $contextHeaders)"
+  }
+
+  class AuthorizedServiceRequestContext[U <: User](override val trackingId: String = generators.nextUuid().toString,
+                                                   override val contextHeaders: Map[String, String] =
+                                                     Map.empty[String, String],
+                                                   val authenticatedUser: U)
+      extends ServiceRequestContext {
+
+    def withPermissionsToken(permissionsToken: PermissionsToken): AuthorizedServiceRequestContext[U] =
+      new AuthorizedServiceRequestContext[U](
+        trackingId,
+        contextHeaders.updated(AuthProvider.PermissionsTokenHeader, permissionsToken.value),
+        authenticatedUser)
+
+    override def hashCode(): Int = 31 * super.hashCode() + authenticatedUser.hashCode()
+
+    override def equals(obj: Any): Boolean = obj match {
+      case ctx: AuthorizedServiceRequestContext[U] => super.equals(ctx) && ctx.authenticatedUser == authenticatedUser
+      case _                                       => false
+    }
+
+    override def toString: String =
+      s"AuthorizedServiceRequestContext($trackingId, $contextHeaders, $authenticatedUser)"
   }
 
   object ContextHeaders {
@@ -115,18 +164,78 @@ package rest {
     val SetPermissionsTokenHeader    = "set-permissions"
   }
 
-  trait Authorization {
-    def userHasPermission(user: User, permission: Permission)(implicit ctx: ServiceRequestContext): Future[Boolean]
+  final case class AuthorizationResult(authorized: Boolean, token: Option[PermissionsToken])
+  object AuthorizationResult {
+    val unauthorized: AuthorizationResult = AuthorizationResult(authorized = false, None)
+  }
+
+  trait Authorization[U <: User] {
+    def userHasPermissions(user: U, permissions: Seq[Permission])(
+            implicit ctx: ServiceRequestContext): Future[AuthorizationResult]
+  }
+
+  class AlwaysAllowAuthorization[U <: User](implicit execution: ExecutionContext) extends Authorization[U] {
+    override def userHasPermissions(user: U, permissions: Seq[Permission])(
+            implicit ctx: ServiceRequestContext): Future[AuthorizationResult] =
+      Future.successful(AuthorizationResult(authorized = true, ctx.permissionsToken))
   }
 
-  class AlwaysAllowAuthorization extends Authorization {
-    override def userHasPermission(user: User, permission: Permission)(
-            implicit ctx: ServiceRequestContext): Future[Boolean] = {
-      Future.successful(true)
+  class CachedTokenAuthorization[U <: User](publicKey: => PublicKey, issuer: String) extends Authorization[U] {
+    override def userHasPermissions(user: U, permissions: Seq[Permission])(
+            implicit ctx: ServiceRequestContext): Future[AuthorizationResult] = {
+      import spray.json._
+
+      def extractPermissionsFromTokenJSON(tokenObject: JsObject): Option[Map[String, Boolean]] =
+        tokenObject.fields.get("permissions").collect {
+          case JsObject(fields) =>
+            fields.collect {
+              case (key, JsBoolean(value)) => key -> value
+            }
+        }
+
+      val result = for {
+        token <- ctx.permissionsToken
+        jwt   <- Jwt.decode(token.value, publicKey, Seq(JwtAlgorithm.RS256)).toOption
+        jwtJson = jwt.parseJson.asJsObject
+
+        // Ensure jwt is for the currently authenticated user and the correct issuer, otherwise return None
+        _ <- jwtJson.fields.get("sub").contains(JsString(user.id.value)).option(())
+        _ <- jwtJson.fields.get("iss").contains(JsString(issuer)).option(())
+
+        permissionsMap <- extractPermissionsFromTokenJSON(jwtJson)
+
+        authorized = permissions.forall(p => permissionsMap.get(p.toString).contains(true))
+      } yield AuthorizationResult(authorized, Some(token))
+
+      Future.successful(result.getOrElse(AuthorizationResult.unauthorized))
+    }
+  }
+
+  object CachedTokenAuthorization {
+    def apply[U <: User](publicKeyFile: Path, issuer: String): CachedTokenAuthorization[U] = {
+      lazy val publicKey: PublicKey = {
+        val publicKeyBytes = Files.readAllBytes(publicKeyFile)
+        val spec           = new X509EncodedKeySpec(publicKeyBytes)
+        KeyFactory.getInstance("RSA").generatePublic(spec)
+      }
+      new CachedTokenAuthorization[U](publicKey, issuer)
+    }
+  }
+
+  class ChainedAuthorization[U <: User](authorizations: Authorization[U]*)(implicit execution: ExecutionContext)
+      extends Authorization[U] {
+
+    override def userHasPermissions(user: U, permissions: Seq[Permission])(
+            implicit ctx: ServiceRequestContext): Future[AuthorizationResult] = {
+      authorizations.toList.foldLeftM[Future, AuthorizationResult](AuthorizationResult.unauthorized) {
+        (authResult, authorization) =>
+          if (authResult.authorized) Future.successful(authResult)
+          else authorization.userHasPermissions(user, permissions)
+      }
     }
   }
 
-  abstract class AuthProvider[U <: User](val authorization: Authorization, log: Logger)(
+  abstract class AuthProvider[U <: User](val authorization: Authorization[U], log: Logger)(
           implicit execution: ExecutionContext) {
 
     import akka.http.scaladsl.server._
@@ -142,43 +251,30 @@ package rest {
     def authenticatedUser(implicit ctx: ServiceRequestContext): OptionT[Future, U]
 
     /**
-      * Specific implementation can verify session expiration and single sign out
-      * to verify if session is still valid
-      */
-    def isSessionValid(user: U)(implicit ctx: ServiceRequestContext): Future[Boolean]
-
-    /**
       * Verifies if request is authenticated and authorized to have `permissions`
       */
-    def authorize(permissions: Permission*): Directive1[U] = {
+    def authorize(permissions: Permission*): Directive1[AuthorizedServiceRequestContext[U]] = {
       serviceContext flatMap { ctx =>
-        onComplete(authenticatedUser(ctx).run flatMap { userOption =>
-          userOption.traverseM[Future, (U, Boolean)] { user =>
-            isSessionValid(user)(ctx).flatMap { sessionValid =>
-              if (sessionValid) {
-                permissions.toList
-                  .traverse[Future, Boolean](authorization.userHasPermission(user, _)(ctx))
-                  .map(results => Option(user -> results.forall(identity)))
-              } else {
-                Future.successful(Option.empty[(U, Boolean)])
-              }
-            }
-          }
-        }).flatMap {
-          case Success(Some((user, authorizationResult))) =>
-            if (authorizationResult) provide(user)
-            else {
-              val challenge =
-                HttpChallenges.basic(s"User does not have the required permissions: ${permissions.mkString(", ")}")
-              log.warn(s"User $user does not have the required permissions: ${permissions.mkString(", ")}")
-              reject(AuthenticationFailedRejection(CredentialsRejected, challenge))
-            }
-
+        onComplete {
+          (for {
+            authToken <- OptionT.optionT(Future.successful(ctx.authToken))
+            user      <- authenticatedUser(ctx)
+            authCtx = ctx.withAuthenticatedUser(authToken, user)
+            authorizationResult <- authorization.userHasPermissions(user, permissions)(authCtx).toOptionT
+            cachedPermissionsAuthCtx = authorizationResult.token.fold(authCtx)(authCtx.withPermissionsToken)
+          } yield (cachedPermissionsAuthCtx, authorizationResult.authorized)).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) =>
             log.warn(
               s"Wasn't able to find authenticated user for the token provided to verify ${permissions.mkString(", ")}")
             reject(ValidationRejection(s"Wasn't able to find authenticated user for the token provided"))
-
           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)))
@@ -193,7 +289,6 @@ package rest {
 
     import akka.http.scaladsl.marshallers.sprayjson.SprayJsonSupport._
     import spray.json._
-    import DefaultJsonProtocol._
 
     protected implicit val exec: ExecutionContext
     protected implicit val materializer: ActorMaterializer
@@ -217,10 +312,7 @@ package rest {
     protected def jsonEntity(json: JsValue): RequestEntity =
       HttpEntity(ContentTypes.`application/json`, json.compactPrint)
 
-    protected def get(baseUri: Uri, path: String) =
-      HttpRequest(HttpMethods.GET, endpointUri(baseUri, path))
-
-    protected def get(baseUri: Uri, path: String, query: Map[String, String]) =
+    protected def get(baseUri: Uri, path: String, query: Seq[(String, String)] = Seq.empty) =
       HttpRequest(HttpMethods.GET, endpointUri(baseUri, path, query))
 
     protected def post(baseUri: Uri, path: String, httpEntity: RequestEntity) =
@@ -235,8 +327,8 @@ package rest {
     protected def endpointUri(baseUri: Uri, path: String) =
       baseUri.withPath(Uri.Path(path))
 
-    protected def endpointUri(baseUri: Uri, path: String, query: Map[String, String]) =
-      baseUri.withPath(Uri.Path(path)).withQuery(Uri.Query(query))
+    protected def endpointUri(baseUri: Uri, path: String, query: Seq[(String, String)]) =
+      baseUri.withPath(Uri.Path(path)).withQuery(Uri.Query(query: _*))
   }
 
   trait ServiceTransport {
diff --git a/src/test/scala/xyz/driver/core/AuthTest.scala b/src/test/scala/xyz/driver/core/AuthTest.scala
index ad8cec8..bf776df 100644
--- a/src/test/scala/xyz/driver/core/AuthTest.scala
+++ b/src/test/scala/xyz/driver/core/AuthTest.scala
@@ -7,34 +7,47 @@ import akka.http.scaladsl.server._
 import akka.http.scaladsl.testkit.ScalatestRouteTest
 import org.scalatest.mock.MockitoSugar
 import org.scalatest.{FlatSpec, Matchers}
+import pdi.jwt.{Jwt, JwtAlgorithm}
 import xyz.driver.core.auth._
 import xyz.driver.core.logging._
-import xyz.driver.core.rest.{AuthProvider, Authorization, ServiceRequestContext}
+import xyz.driver.core.rest._
 
 import scala.concurrent.Future
 import scalaz.OptionT
 
 class AuthTest extends FlatSpec with Matchers with MockitoSugar with ScalatestRouteTest {
 
-  case object TestRoleAllowedPermission    extends Permission
-  case object TestRoleNotAllowedPermission extends Permission
+  case object TestRoleAllowedPermission        extends Permission
+  case object TestRoleAllowedByTokenPermission extends Permission
+  case object TestRoleNotAllowedPermission     extends Permission
 
   val TestRole = Role(Id("1"), Name("testRole"))
 
-  implicit val exec = scala.concurrent.ExecutionContext.global
+  val (publicKey, privateKey) = {
+    import java.security.KeyPairGenerator
 
-  val authorization: Authorization = new Authorization {
-    override def userHasPermission(user: User, permission: Permission)(
-            implicit ctx: ServiceRequestContext): Future[Boolean] = {
-      Future.successful(permission === TestRoleAllowedPermission)
+    val keygen = KeyPairGenerator.getInstance("RSA")
+    keygen.initialize(2048)
+
+    val keyPair = keygen.generateKeyPair()
+    (keyPair.getPublic, keyPair.getPrivate)
+  }
+
+  val basicAuthorization: Authorization[User] = new Authorization[User] {
+
+    override def userHasPermissions(user: User, permissions: Seq[Permission])(
+            implicit ctx: ServiceRequestContext): Future[AuthorizationResult] = {
+      val authorized = permissions.forall(_ === TestRoleAllowedPermission)
+      Future.successful(AuthorizationResult(authorized, ctx.permissionsToken))
     }
   }
 
-  val authStatusService = new AuthProvider[User](authorization, NoLogger) {
+  val tokenIssuer        = "users"
+  val tokenAuthorization = new CachedTokenAuthorization[User](publicKey, tokenIssuer)
 
-    override def isSessionValid(user: User)(implicit ctx: ServiceRequestContext): Future[Boolean] =
-      Future.successful(true)
+  val authorization = new ChainedAuthorization[User](tokenAuthorization, basicAuthorization)
 
+  val authStatusService = new AuthProvider[User](authorization, NoLogger) {
     override def authenticatedUser(implicit ctx: ServiceRequestContext): OptionT[Future, User] =
       OptionT.optionT[Future] {
         if (ctx.contextHeaders.keySet.contains(AuthProvider.AuthenticationTokenHeader)) {
@@ -47,7 +60,7 @@ class AuthTest extends FlatSpec with Matchers with MockitoSugar with ScalatestRo
 
   import authStatusService._
 
-  "'authorize' directive" should "throw error is auth token is not in the request" in {
+  "'authorize' directive" should "throw error if auth token is not in the request" in {
 
     Get("/naive/attempt") ~>
       authorize(TestRoleAllowedPermission) { user =>
@@ -59,7 +72,7 @@ class AuthTest extends FlatSpec with Matchers with MockitoSugar with ScalatestRo
       }
   }
 
-  it should "throw error is authorized user is not having the requested permission" in {
+  it should "throw error if authorized user does not have the requested permission" in {
 
     val referenceAuthToken = AuthToken("I am a test role's token")
 
@@ -85,12 +98,36 @@ class AuthTest extends FlatSpec with Matchers with MockitoSugar with ScalatestRo
     Get("/valid/attempt/?a=2&b=5").addHeader(
       RawHeader(AuthProvider.AuthenticationTokenHeader, referenceAuthToken.value)
     ) ~>
-      authorize(TestRoleAllowedPermission) { user =>
-        complete("Alright, user \"" + user.id + "\" is authorized")
+      authorize(TestRoleAllowedPermission) { ctx =>
+        complete(s"Alright, user ${ctx.authenticatedUser.id} is authorized")
+      } ~>
+      check {
+        handled shouldBe true
+        responseAs[String] shouldBe "Alright, user 1 is authorized"
+      }
+  }
+
+  it should "authorize permission found in permissions token" in {
+    import spray.json._
+
+    val claim = JsObject(
+      Map(
+        "iss"         -> JsString(tokenIssuer),
+        "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")
+
+    Get("/alic/attempt/?a=2&b=5")
+      .addHeader(RawHeader(AuthProvider.AuthenticationTokenHeader, referenceAuthToken.value))
+      .addHeader(RawHeader(AuthProvider.PermissionsTokenHeader, permissionsToken.value)) ~>
+      authorize(TestRoleAllowedByTokenPermission) { ctx =>
+        complete(s"Alright, user ${ctx.authenticatedUser.id} is authorized by permissions token")
       } ~>
       check {
         handled shouldBe true
-        responseAs[String] shouldBe "Alright, user \"1\" is authorized"
+        responseAs[String] shouldBe "Alright, user 1 is authorized by permissions token"
       }
   }
 }