From 91b01c6bdc905db9db7874fb0a66574463a0d484 Mon Sep 17 00:00:00 2001
From: Ivan Topolnjak <ivantopo@gmail.com>
Date: Fri, 2 Feb 2018 03:35:02 +0100
Subject: allow arbitrary header names for automatic broadcast string keys

---
 .../src/main/scala/kamon/context/Codecs.scala      | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

(limited to 'kamon-core/src/main/scala')

diff --git a/kamon-core/src/main/scala/kamon/context/Codecs.scala b/kamon-core/src/main/scala/kamon/context/Codecs.scala
index 8ba70157..c5d237a9 100644
--- a/kamon-core/src/main/scala/kamon/context/Codecs.scala
+++ b/kamon-core/src/main/scala/kamon/context/Codecs.scala
@@ -43,7 +43,7 @@ class Codecs(initialConfig: Config) {
     import scala.collection.JavaConverters._
     try {
       val codecsConfig = config.getConfig("kamon.context.codecs")
-      val stringKeys = codecsConfig.getStringList("string-keys").asScala
+      val stringKeys = readStringKeysConfig(codecsConfig.getConfig("string-keys"))
       val knownHttpHeaderCodecs = readEntryCodecs[TextMap]("http-headers-keys", codecsConfig) ++ stringHeaderCodecs(stringKeys)
       val knownBinaryCodecs = readEntryCodecs[ByteBuffer]("binary-keys", codecsConfig) ++ stringBinaryCodecs(stringKeys)
 
@@ -72,13 +72,14 @@ class Codecs(initialConfig: Config) {
     entries.result()
   }
 
-  private def stringHeaderCodecs(keys: Seq[String]): Map[String, Codecs.ForEntry[TextMap]] = {
-    keys.map(key => (key, new Codecs.StringHeadersCodec(key))).toMap
-  }
+  private def readStringKeysConfig(config: Config): Map[String, String] =
+    config.topLevelKeys.map(key => (key, config.getString(key))).toMap
 
-  private def stringBinaryCodecs(keys: Seq[String]): Map[String, Codecs.ForEntry[ByteBuffer]] = {
-    keys.map(key => (key, new Codecs.StringBinaryCodec(key))).toMap
-  }
+  private def stringHeaderCodecs(keys: Map[String, String]): Map[String, Codecs.ForEntry[TextMap]] =
+    keys.map { case (key, header) => (key, new Codecs.StringHeadersCodec(key, header)) }
+
+  private def stringBinaryCodecs(keys: Map[String, String]): Map[String, Codecs.ForEntry[ByteBuffer]] =
+    keys.map { case (key, _) => (key, new Codecs.StringBinaryCodec(key)) }
 }
 
 object Codecs {
@@ -224,21 +225,20 @@ object Codecs {
     }
   }
 
-  private class StringHeadersCodec(key: String) extends Codecs.ForEntry[TextMap] {
-    private val dataKey = "X-KamonContext-" + key
+  private class StringHeadersCodec(key: String, headerName: String) extends Codecs.ForEntry[TextMap] {
     private val contextKey = Key.broadcast[Option[String]](key, None)
 
     override def encode(context: Context): TextMap = {
       val textMap = TextMap.Default()
       context.get(contextKey).foreach { value =>
-        textMap.put(dataKey, value)
+        textMap.put(headerName, value)
       }
 
       textMap
     }
 
     override def decode(carrier: TextMap, context: Context): Context = {
-      carrier.get(dataKey) match {
+      carrier.get(headerName) match {
         case value @ Some(_) => context.withKey(contextKey, value)
         case None            => context
       }
-- 
cgit v1.2.3