From 3c6b272033341c38b2cb1852a7b48ea7befb7002 Mon Sep 17 00:00:00 2001
From: vlad <vlad@drivergrp.com>
Date: Wed, 9 Nov 2016 11:14:04 -0800
Subject: DIR-135 Return tracing headers in all responses

---
 src/main/scala/xyz/driver/core/app.scala      | 23 +++++++++--------------
 src/main/scala/xyz/driver/core/database.scala |  4 ++--
 src/main/scala/xyz/driver/core/rest.scala     | 10 ++++++++--
 3 files changed, 19 insertions(+), 18 deletions(-)

(limited to 'src/main/scala/xyz/driver')

diff --git a/src/main/scala/xyz/driver/core/app.scala b/src/main/scala/xyz/driver/core/app.scala
index dd4f661..227be57 100644
--- a/src/main/scala/xyz/driver/core/app.scala
+++ b/src/main/scala/xyz/driver/core/app.scala
@@ -67,12 +67,14 @@ object app {
       val _ = Future {
         http.bindAndHandle(route2HandlerFlow(handleExceptions(ExceptionHandler(exceptionHandler)) { ctx =>
           val trackingId = rest.extractTrackingId(ctx)
+          log.audit(s"Received request ${ctx.request} with tracking id $trackingId")
+
           val contextWithTrackingId =
             ctx.withRequest(ctx.request.addHeader(RawHeader(ContextHeaders.TrackingIdHeader, trackingId)))
 
-          log.audit(s"Received request ${ctx.request} with tracking id $trackingId")
-
-          modules.map(_.route).foldLeft(versionRt ~ healthRoute ~ swaggerRoutes)(_ ~ _)(contextWithTrackingId)
+          respondWithHeaders(List(RawHeader(ContextHeaders.TrackingIdHeader, trackingId))) {
+            modules.map(_.route).foldLeft(versionRt ~ healthRoute ~ swaggerRoutes)(_ ~ _)
+          }(contextWithTrackingId)
         }), interface, port)(materializer)
       }
     }
@@ -83,27 +85,20 @@ object app {
         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}" }"""))(
-            ctx)
+          complete(HttpResponse(BadRequest, entity = is.getMessage))(ctx)
 
       case cm: ConcurrentModificationException =>
         ctx =>
-          val trackingId                    = rest.extractTrackingId(ctx)
-          val concurrentModificationMessage = "Resource was changed concurrently, try requesting a newer version"
+          val trackingId = rest.extractTrackingId(ctx)
           log.audit(s"Concurrent modification of the resource ${ctx.request.uri} ($trackingId)", cm)
           complete(
-            HttpResponse(
-              Conflict,
-              entity = s"""{ "trackingId": "$trackingId", "message": "$concurrentModificationMessage" }"""))(ctx)
+            HttpResponse(Conflict, entity = "Resource was changed concurrently, try requesting a newer version"))(ctx)
 
       case t: Throwable =>
         ctx =>
           val trackingId = rest.extractTrackingId(ctx)
           log.error(s"Request to ${ctx.request.uri} could not be handled normally ($trackingId)", t)
-          complete(
-            HttpResponse(InternalServerError,
-                         entity = s"""{ "trackingId": "$trackingId", "message": "${t.getMessage}" }"""))(ctx)
+          complete(HttpResponse(InternalServerError, entity = t.getMessage))(ctx)
     }
 
     protected def versionRoute(version: String, gitHash: String, startupTime: Time): Route = {
diff --git a/src/main/scala/xyz/driver/core/database.scala b/src/main/scala/xyz/driver/core/database.scala
index 1039be4..45d1b6d 100644
--- a/src/main/scala/xyz/driver/core/database.scala
+++ b/src/main/scala/xyz/driver/core/database.scala
@@ -1,10 +1,10 @@
 package xyz.driver.core
 
+import slick.backend.DatabaseConfig
+import slick.driver.JdbcProfile
 import xyz.driver.core.time.Time
 
 import scala.concurrent.Future
-import slick.backend.DatabaseConfig
-import slick.driver.JdbcProfile
 
 object database {
 
diff --git a/src/main/scala/xyz/driver/core/rest.scala b/src/main/scala/xyz/driver/core/rest.scala
index 1988cf4..f05a800 100644
--- a/src/main/scala/xyz/driver/core/rest.scala
+++ b/src/main/scala/xyz/driver/core/rest.scala
@@ -22,7 +22,12 @@ object rest {
 
   object ContextHeaders {
     val AuthenticationTokenHeader = "WWW-Authenticate"
-    val TrackingIdHeader          = "l5d-ctx-trace" // https://linkerd.io/doc/0.7.4/linkerd/protocol-http/
+    val TrackingIdHeader          = "X-Trace"
+
+    object LinkerD {
+      // https://linkerd.io/doc/0.7.4/linkerd/protocol-http/
+      def isLinkerD(headerName: String) = headerName.startsWith("l5d-")
+    }
   }
 
   final case class ServiceRequestContext(
@@ -45,7 +50,8 @@ object rest {
 
   def extractContextHeaders(ctx: RequestContext): Map[String, String] = {
     ctx.request.headers.filter { h =>
-      h.lowercaseName.startsWith("l5d-") || h.name === ContextHeaders.AuthenticationTokenHeader
+      h.name === ContextHeaders.AuthenticationTokenHeader || h.name === ContextHeaders.TrackingIdHeader
+      // || ContextHeaders.LinkerD.isLinkerD(h.lowercaseName)
     } map { header =>
       header.name -> header.value
     } toMap
-- 
cgit v1.2.3