DIR-135 Consistent request context extractionv0.9.8

5 files changed, 80 insertions, 122 deletions

diff --git a/src/main/scala/xyz/driver/core/app.scala b/src/main/scala/xyz/driver/core/app.scala
index 8f892e8..f972158 100644
--- a/src/main/scala/xyz/driver/core/app.scala
+++ b/src/main/scala/xyz/driver/core/app.scala
@@ -4,18 +4,18 @@ import akka.actor.ActorSystem
 import akka.http.scaladsl.Http
 import akka.http.scaladsl.marshallers.sprayjson.SprayJsonSupport
 import akka.http.scaladsl.model.StatusCodes._
+import akka.http.scaladsl.model.headers.RawHeader
 import akka.http.scaladsl.model.{HttpResponse, StatusCodes}
 import akka.http.scaladsl.server.Directives._
 import akka.http.scaladsl.server.RouteResult._
-import akka.http.scaladsl.server.{ExceptionHandler, RequestContext, Route, RouteConcatenation}
+import akka.http.scaladsl.server.{ExceptionHandler, Route, RouteConcatenation}
 import akka.stream.ActorMaterializer
 import com.typesafe.config.Config
 import org.slf4j.LoggerFactory
 import spray.json.DefaultJsonProtocol
 import xyz.driver.core
-import xyz.driver.core.auth.AuthService
 import xyz.driver.core.logging.{Logger, TypesafeScalaLogger}
-import xyz.driver.core.rest.Swagger
+import xyz.driver.core.rest.{ContextHeaders, Swagger}
 import xyz.driver.core.stats.SystemStats
 import xyz.driver.core.time.Time
 import xyz.driver.core.time.provider.{SystemTimeProvider, TimeProvider}
@@ -65,26 +65,23 @@ object app {
       val versionRt      = versionRoute(version, gitHash, time.currentTime())
 
       val _ = Future {
-        http.bindAndHandle(route2HandlerFlow(handleExceptions(exceptionHandler) { ctx =>
-          log.audit(s"Received request ${ctx.request}")
-          modules.map(_.route).foldLeft(versionRt ~ healthRoute ~ swaggerRoutes)(_ ~ _)(ctx)
+        http.bindAndHandle(route2HandlerFlow(handleExceptions(ExceptionHandler(exceptionHandler)) { ctx =>
+          val trackingId = rest.extractTrackingId(ctx)
+          val contextWithTrackingId =
+            ctx.withRequest(ctx.request.withHeaders(RawHeader(ContextHeaders.TrackingIdHeader, trackingId)))
+
+          log.audit(s"Received request ${ctx.request} with tracking id $trackingId")
+
+          modules.map(_.route).foldLeft(versionRt ~ healthRoute ~ swaggerRoutes)(_ ~ _)(contextWithTrackingId)
         }), interface, port)(materializer)
       }
     }
 
-    protected def extractTrackingId(ctx: RequestContext) = {
-      ctx.request.headers
-        .find(_.name == AuthService.TrackingIdHeader)
-        .map(_.value())
-        .getOrElse(java.util.UUID.randomUUID.toString)
-      // TODO: In the case when absent, should be taken the same generated id, as in `authorize`
-    }
-
-    protected def exceptionHandler = ExceptionHandler {
+    protected def exceptionHandler = PartialFunction[Throwable, Route] {
 
       case is: IllegalStateException =>
         ctx =>
-          val trackingId = extractTrackingId(ctx)
+          val trackingId = rest.extractTrackingId(ctx)
           log.debug(s"Request is not allowed to ${ctx.request.uri} ($trackingId)", is)
           complete(
             HttpResponse(BadRequest, entity = s"""{ "trackingId": "$trackingId", "message": "${is.getMessage}" }"""))(
@@ -92,7 +89,7 @@ object app {
 
       case cm: ConcurrentModificationException =>
         ctx =>
-          val trackingId                    = extractTrackingId(ctx)
+          val trackingId                    = rest.extractTrackingId(ctx)
           val concurrentModificationMessage = "Resource was changed concurrently, try requesting a newer version"
           log.audit(s"Concurrent modification of the resource ${ctx.request.uri} ($trackingId)", cm)
           complete(
@@ -102,7 +99,7 @@ object app {
 
       case t: Throwable =>
         ctx =>
-          val trackingId = extractTrackingId(ctx)
+          val trackingId = rest.extractTrackingId(ctx)
           log.error(s"Request to ${ctx.request.uri} could not be handled normally ($trackingId)", t)
           complete(
             HttpResponse(InternalServerError,
diff --git a/src/main/scala/xyz/driver/core/auth.scala b/src/main/scala/xyz/driver/core/auth.scala
index 17f89c0..3dd21d9 100644
--- a/src/main/scala/xyz/driver/core/auth.scala
+++ b/src/main/scala/xyz/driver/core/auth.scala
@@ -4,7 +4,7 @@ import akka.http.scaladsl.model.headers.HttpChallenges
 import akka.http.scaladsl.server.AuthenticationFailedRejection.CredentialsRejected
 
 import scala.concurrent.Future
-import scala.util.{Failure, Success, Try}
+import scala.util.{Failure, Success}
 import scalaz.OptionT
 
 object auth {
@@ -68,17 +68,12 @@ object auth {
     def permissions: Set[Permission] = roles.flatMap(_.permissions)
   }
 
-  final case class Macaroon(value: String)
-
-  final case class Base64[T](value: String)
-
-  final case class AuthToken(value: Base64[Macaroon], trackingId: String)
+  final case class AuthToken(value: String)
 
   final case class PasswordHash(value: String)
 
   object AuthService {
     val AuthenticationTokenHeader = "WWW-Authenticate"
-    val TrackingIdHeader          = "l5d-ctx-trace" // https://linkerd.io/doc/0.7.4/linkerd/protocol-http/
   }
 
   trait AuthService[U <: User] {
@@ -88,49 +83,25 @@ object auth {
 
     protected def authStatus(authToken: AuthToken): OptionT[Future, U]
 
-    def authorize(permissions: Permission*): Directive1[(AuthToken, U)] = {
-      parameters('authToken.?).flatMap { parameterTokenValue =>
-        optionalHeaderValueByName(AuthService.AuthenticationTokenHeader).flatMap { headerTokenValue =>
-          optionalHeaderValueByName(AuthService.TrackingIdHeader).flatMap { trackingIdValue =>
-            verifyAuthToken(headerTokenValue.orElse(parameterTokenValue), trackingIdValue, permissions.toSet)
-          }
-        }
-      }
-    }
-
-    private def verifyAuthToken(tokenOption: Option[String],
-                                trackingIdValue: Option[String],
-                                permissions: Set[Permission]): Directive1[(AuthToken, U)] =
-      tokenOption match {
-        case Some(tokenValue) =>
-          val trackingId = trackingIdValue.getOrElse(java.util.UUID.randomUUID.toString)
-          val token      = AuthToken(Base64[Macaroon](tokenValue), trackingId)
+    def authorize(permissions: Permission*): Directive1[U] = {
+      headerValueByName(AuthService.AuthenticationTokenHeader).flatMap { tokenValue =>
+        val token = AuthToken(tokenValue)
 
-          onComplete(authStatus(token).run).flatMap { tokenUserResult =>
-            checkPermissions(tokenUserResult, permissions, token)
-          }
+        onComplete(authStatus(token).run).flatMap {
+          case Success(Some(user)) =>
+            if (permissions.forall(user.permissions.contains)) provide(user)
+            else {
+              val challenge =
+                HttpChallenges.basic(s"User does not have the required permissions: ${permissions.mkString(", ")}")
+              reject(AuthenticationFailedRejection(CredentialsRejected, challenge))
+            }
 
-        case None =>
-          reject(MissingHeaderRejection(AuthService.AuthenticationTokenHeader))
-      }
+          case Success(None) =>
+            reject(ValidationRejection(s"Wasn't able to find authenticated user for the token provided"))
 
-    private def checkPermissions(userResult: Try[Option[U]],
-                                 permissions: Set[Permission],
-                                 token: AuthToken): Directive1[(AuthToken, U)] = {
-      userResult match {
-        case Success(Some(user)) =>
-          if (permissions.forall(user.permissions.contains)) provide(token -> user)
-          else {
-            val challenge =
-              HttpChallenges.basic(s"User does not have the required permissions: ${permissions.mkString(", ")}")
-            reject(AuthenticationFailedRejection(CredentialsRejected, challenge))
-          }
-
-        case Success(None) =>
-          reject(ValidationRejection(s"Wasn't able to find authenticated user for the token provided"))
-
-        case Failure(t) =>
-          reject(ValidationRejection(s"Wasn't able to verify token for authenticated user", Some(t)))
+          case Failure(t) =>
+            reject(ValidationRejection(s"Wasn't able to verify token for authenticated user", Some(t)))
+        }
       }
     }
   }
diff --git a/src/main/scala/xyz/driver/core/crypto.scala b/src/main/scala/xyz/driver/core/crypto.scala
deleted file mode 100644
index d001e0f..0000000
--- a/src/main/scala/xyz/driver/core/crypto.scala
+++ /dev/null
@@ -1,27 +0,0 @@
-package xyz.driver.core
-
-import xyz.driver.core.auth.AuthToken
-
-object crypto {
-
-  final case class EncryptionKey(value: String)
-
-  final case class DecryptionKey(value: String)
-
-  trait Crypto {
-
-    def keyForToken(authToken: AuthToken): EncryptionKey
-
-    def encrypt(encryptionKey: EncryptionKey)(message: Array[Byte]): Array[Byte]
-
-    def decrypt(decryptionKey: EncryptionKey)(message: Array[Byte]): Array[Byte]
-  }
-
-  object NoCrypto extends Crypto {
-
-    override def keyForToken(authToken: AuthToken): EncryptionKey = EncryptionKey(authToken.value.value)
-
-    override def decrypt(decryptionKey: EncryptionKey)(message: Array[Byte]): Array[Byte] = message
-    override def encrypt(encryptionKey: EncryptionKey)(message: Array[Byte]): Array[Byte] = message
-  }
-}
diff --git a/src/main/scala/xyz/driver/core/rest.scala b/src/main/scala/xyz/driver/core/rest.scala
index eaf97db..c52d9e0 100644
--- a/src/main/scala/xyz/driver/core/rest.scala
+++ b/src/main/scala/xyz/driver/core/rest.scala
@@ -4,15 +4,13 @@ import akka.actor.ActorSystem
 import akka.http.scaladsl.Http
 import akka.http.scaladsl.model._
 import akka.http.scaladsl.model.headers.RawHeader
+import akka.http.scaladsl.server.RequestContext
 import akka.http.scaladsl.unmarshalling.Unmarshal
 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 xyz.driver.core.auth.{AuthService, AuthToken}
-import xyz.driver.core.crypto.Crypto
+import xyz.driver.core.auth.AuthService
 import xyz.driver.core.logging.Logger
 import xyz.driver.core.stats.Stats
 import xyz.driver.core.time.TimeRange
@@ -20,15 +18,41 @@ import xyz.driver.core.time.provider.TimeProvider
 
 import scala.concurrent.{ExecutionContext, Future}
 import scala.util.{Failure, Success}
-import scalaz.{Failure => _, Success => _}
+import scalaz.Scalaz.{Id => _, _}
 
 object rest {
 
+  object ContextHeaders {
+    val AuthenticationTokenHeader = AuthService.AuthenticationTokenHeader
+    val TrackingIdHeader          = "l5d-ctx-trace" // https://linkerd.io/doc/0.7.4/linkerd/protocol-http/
+  }
+
+  final case class ServiceRequestContext(trackingId: String, contextHeaders: Map[String, String])
+
+  def serviceContext(ctx: RequestContext): ServiceRequestContext = {
+    ServiceRequestContext(extractTrackingId(ctx), extractContextHeaders(ctx))
+  }
+
+  def extractTrackingId(ctx: RequestContext): String = {
+    ctx.request.headers
+      .find(_.name == ContextHeaders.TrackingIdHeader)
+      .fold(java.util.UUID.randomUUID.toString)(_.value())
+  }
+
+  def extractContextHeaders(ctx: RequestContext): Map[String, String] = {
+    ctx.request.headers.filter { h =>
+      h.lowercaseName.startsWith("l5d-") || h.name === ContextHeaders.AuthenticationTokenHeader
+    } map { header =>
+      header.name -> header.value
+    } toMap
+  }
+
+
   trait Service
 
   trait ServiceTransport {
 
-    def sendRequest(authToken: AuthToken)(requestStub: HttpRequest): Future[Unmarshal[ResponseEntity]]
+    def sendRequest(context: ServiceRequestContext)(requestStub: HttpRequest): Future[Unmarshal[ResponseEntity]]
   }
 
   trait ServiceDiscovery {
@@ -37,25 +61,18 @@ object rest {
   }
 
   class HttpRestServiceTransport(actorSystem: ActorSystem, executionContext: ExecutionContext,
-                                 crypto: Crypto, log: Logger, stats: Stats, time: TimeProvider) extends ServiceTransport {
+                                 log: Logger, stats: Stats, time: TimeProvider) extends ServiceTransport {
 
     protected implicit val materializer = ActorMaterializer()(actorSystem)
     protected implicit val execution = executionContext
 
-    def sendRequest(authToken: AuthToken)(requestStub: HttpRequest): Future[Unmarshal[ResponseEntity]] = {
+    def sendRequest(context: ServiceRequestContext)(requestStub: HttpRequest): Future[Unmarshal[ResponseEntity]] = {
 
       val requestTime = time.currentTime()
-      val encryptionFlow = Flow[ByteString] map { bytes =>
-        ByteString(crypto.encrypt(crypto.keyForToken(authToken))(bytes.toArray))
-      }
-      val decryptionFlow = Flow[ByteString] map { bytes =>
-        ByteString(crypto.decrypt(crypto.keyForToken(authToken))(bytes.toArray))
-      }
 
-      val request = (if(requestStub.entity.isKnownEmpty()) requestStub else {
-        requestStub.withEntity(requestStub.entity.transformDataBytes(encryptionFlow))
-      }).withHeaders(RawHeader(AuthService.AuthenticationTokenHeader, authToken.value.value),
-        RawHeader(AuthService.TrackingIdHeader, authToken.trackingId))
+      val request = requestStub
+        .withHeaders(RawHeader(ContextHeaders.TrackingIdHeader, context.trackingId))
+        .withHeaders(context.contextHeaders.toSeq.map { h => RawHeader(h._1, h._2): HttpHeader }: _*)
 
       log.audit(s"Sending to ${request.uri} request $request")
 
@@ -65,7 +82,7 @@ object rest {
         } else if(response.status.isFailure()) {
           throw new Exception("Http status is failure " + response.status)
         } else {
-          Unmarshal(response.entity.transformDataBytes(decryptionFlow))
+          Unmarshal(response.entity)
         }
       }
 
diff --git a/src/test/scala/xyz/driver/core/AuthTest.scala b/src/test/scala/xyz/driver/core/AuthTest.scala
index 97279de..ca7e019 100644
--- a/src/test/scala/xyz/driver/core/AuthTest.scala
+++ b/src/test/scala/xyz/driver/core/AuthTest.scala
@@ -16,10 +16,10 @@ class AuthTest extends FlatSpec with Matchers with MockitoSugar with ScalatestRo
 
   val authStatusService: AuthService[User] = new AuthService[User] {
     override def authStatus(authToken: AuthToken): OptionT[Future, User] = OptionT.optionT[Future] {
-      Future.successful(Some(new User() {
+      Future.successful(Some(new User {
         override def id: Id[User]     = Id[User](1L)
         override def roles: Set[Role] = Set(PathologistRole)
-      }))
+      }: User))
     }
   }
 
@@ -29,7 +29,7 @@ class AuthTest extends FlatSpec with Matchers with MockitoSugar with ScalatestRo
 
     Get("/naive/attempt") ~>
       authorize(CanSignOutReport) {
-        case (authToken, user) =>
+        case user =>
           complete("Never going to be here")
       } ~>
       check {
@@ -40,13 +40,13 @@ class AuthTest extends FlatSpec with Matchers with MockitoSugar with ScalatestRo
 
   it should "throw error is authorized user is not having the requested permission" in {
 
-    val referenceAuthToken = AuthToken(Base64("I am a pathologist's token"), "BC131CD")
+    val referenceAuthToken = AuthToken("I am a pathologist's token")
 
     Post("/administration/attempt").addHeader(
-      RawHeader(AuthService.AuthenticationTokenHeader, referenceAuthToken.value.value)
+      RawHeader(AuthService.AuthenticationTokenHeader, referenceAuthToken.value)
     ) ~>
       authorize(CanAssignRoles) {
-        case (authToken, user) =>
+        case user =>
           complete("Never going to get here")
       } ~>
       check {
@@ -60,18 +60,18 @@ class AuthTest extends FlatSpec with Matchers with MockitoSugar with ScalatestRo
 
   it should "pass and retrieve the token to client code, if token is in request and user has permission" in {
 
-    val referenceAuthToken = AuthToken(Base64("I am token"), "AAADDDFFF")
+    val referenceAuthToken = AuthToken("I am token")
 
     Get("/valid/attempt/?a=2&b=5").addHeader(
-      RawHeader(AuthService.AuthenticationTokenHeader, referenceAuthToken.value.value)
+      RawHeader(AuthService.AuthenticationTokenHeader, referenceAuthToken.value)
     ) ~>
       authorize(CanSignOutReport) {
-        case (authToken, user) =>
-          complete("Alright, \"" + authToken.value.value + "\" is handled")
+        case user =>
+          complete("Alright, user \"" + user.id + "\" is authorized")
       } ~>
       check {
         handled shouldBe true
-        responseAs[String] shouldBe "Alright, \"I am token\" is handled"
+        responseAs[String] shouldBe "Alright, user \"1\" is authorized"
       }
   }
 }