finagle-core: make broadcast context keys case insensitive

Problem:
When marshalling context keys over http keys, the id is marshalled into the
http header which is canonicalized. For keys with non-lowercase letters, this
can break when unmarshalling the headers.

i.e `com.myKey -> Finagle-Ctx-Com.mykey`

Solution / Result:
Lets make the keys case insensitive during lookups. This change will sit behind
a toggle that will be at 100% since this will be backwards compatible (unit tests).
The toggle will act as a "kill switch" that we can trigger at runtime in case anything
goes wrong

JIRA Issues: CSL-10942

Differential Revision: https://phabricator.twitter.biz/D665209

CHANGELOG.rst

@@ -7,6 +7,14 @@ Note that ``PHAB_ID=#`` and ``RB_ID=#`` correspond to associated messages in com
 Unreleased
 ----------
 
+Runtime Behavior Changes
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+* finagle-core: Broadcast context keys lookups are now case insensitive. This change is backwards
+  compatible as the marshalled key id is unchanged. Although enabled by default, this change will
+  be temporarily sitting behind a toggle, `com.twitter.finagle.context.MarshalledContextLookupId`
+  that can be used to turn off this change. ``PHAB_ID=D665209``.
+
 21.5.0
 ------
 

finagle-benchmark/src/main/scala/com/twitter/finagle/context/ContextBenchmark.scala

@@ -22,7 +22,7 @@ class ContextBenchmark extends StdBenchAnnotations {
   @Param(Array("5"))
   var depth: Int = 0
 
-  var env: Map[Buf, broadcast.Cell] = _
+  var env: Map[String, broadcast.Cell] = _
 
   // Build up a context with some keys that won't be used
   @Setup(Level.Iteration)

finagle-core/src/main/resources/com/twitter/toggles/configs/com.twitter.finagle.core.json

@@ -4,6 +4,11 @@
       "id": "com.twitter.finagle.OffloadEarly",
       "description": "Hand off work to the Offload pool right after we enter the Finagle stack",
       "fraction": 0.0
+    },
+    {
+      "id": "com.twitter.finagle.context.MarshalledContextLookupId",
+      "description": "Case insenstive broadcast context key lookups",
+      "fraction": 1.0
     }
   ]
 }

finagle-core/src/main/scala/com/twitter/finagle/context/MarshalledContext.scala

@@ -1,9 +1,18 @@
 package com.twitter.finagle.context
 
+import com.twitter.finagle.CoreToggles
+import com.twitter.finagle.server.ServerInfo
 import com.twitter.io.Buf
 import com.twitter.logging.Logger
 import com.twitter.util.{Local, Return, Throw, Try}
 import java.security.{MessageDigest, NoSuchAlgorithmException}
+import java.util.Locale
+
+private object MarshalledContext {
+  val toggle = CoreToggles("com.twitter.finagle.context.MarshalledContextLookupId")
+
+  def caseInsensitiveLookupId(id: String): String = id.toLowerCase(Locale.US)
+}
 
 /**
  * A marshalled context contains bindings that may be
@@ -14,10 +23,15 @@ import java.security.{MessageDigest, NoSuchAlgorithmException}
  * tree.
  */
 final class MarshalledContext private[context] extends Context {
+  import MarshalledContext._
+
+  private[this] val useCaseInsensitiveLookup: Boolean =
+    if (toggle.isDefined) toggle(ServerInfo().id.hashCode)
+    else true
 
   private[this] val log = Logger.get()
 
-  private[this] val local = new Local[Map[Buf, Cell]]
+  private[this] val local = new Local[Map[String, Cell]]
 
   // Cell and its sub types are package visible for testing purposes only.
   private[context] sealed trait Cell
@@ -53,10 +67,12 @@ final class MarshalledContext private[context] extends Context {
 
   /**
    * Keys in MarshalledContext must provide a marshaller
-   * and unmarshaller. The key `id` is used for marshalling and
+   * and unmarshaller. The `key` is used for marshalling and
    * unmarshalling and thus must be unique. The behavior of the
    * MarshalledContext when using two keys with the same key `id`
    * is undefined.
+   *
+   * @Note the key's `id` is treated without case sensitivity.
    */
   abstract class Key[A](val id: String) {
 
@@ -67,6 +83,14 @@ final class MarshalledContext private[context] extends Context {
      */
     final val marshalId: Buf = Buf.ByteBuffer.coerce(Buf.Utf8(id))
 
+    /**
+     * The identifier used to lookup the key in the stored context.
+     */
+    private[context] final val lookupId: String = {
+      if (useCaseInsensitiveLookup) MarshalledContext.caseInsensitiveLookupId(id)
+      else id
+    }
+
     /**
      * Marshal an A-typed value into a Buf.
      */
@@ -78,35 +102,35 @@ final class MarshalledContext private[context] extends Context {
     def tryUnmarshal(buf: Buf): Try[A]
   }
 
-  def get[A](key: Key[A]): Option[A] = env.get(key.marshalId) match {
+  def get[A](key: Key[A]): Option[A] = env.get(key.lookupId) match {
     case Some(Real(_, someValue)) => someValue.asInstanceOf[Some[A]]
     case Some(t: Translucent) => t.unmarshal(key)
     case None => None
   }
 
   def let[A, R](key: Key[A], value: A)(fn: => R): R =
-    letLocal(env.updated(key.marshalId, Real(key, Some(value))))(fn)
+    letLocal(env.updated(key.lookupId, Real(key, Some(value))))(fn)
 
   def let[A, B, R](key1: Key[A], value1: A, key2: Key[B], value2: B)(fn: => R): R = {
     val next = env
-      .updated(key1.marshalId, Real(key1, Some(value1)))
-      .updated(key2.marshalId, Real(key2, Some(value2)))
+      .updated(key1.lookupId, Real(key1, Some(value1)))
+      .updated(key2.lookupId, Real(key2, Some(value2)))
     letLocal(next)(fn)
   }
 
   def let[R](pairs: Iterable[KeyValuePair[_]])(fn: => R): R = {
     val next = pairs.foldLeft(env) {
       case (e, KeyValuePair(k, v)) =>
-        e.updated(k.marshalId, Real(k, Some(v)))
+        e.updated(k.lookupId, Real(k, Some(v)))
     }
     letLocal(next)(fn)
   }
 
   def letClear[R](key: Key[_])(fn: => R): R =
-    letLocal(env - key.marshalId)(fn)
+    letLocal(env - key.lookupId)(fn)
 
   def letClear[R](keys: Iterable[Key[_]])(fn: => R): R = {
-    val next = keys.foldLeft(env) { case (e, k) => e - k.marshalId }
+    val next = keys.foldLeft(env) { case (e, k) => e - k.lookupId }
     letLocal(next)(fn)
   }
 
@@ -127,27 +151,31 @@ final class MarshalledContext private[context] extends Context {
   def marshal(): Iterable[(Buf, Buf)] = marshal(env)
 
   // Exposed for testing
-  private[context] def marshal(env: Map[Buf, Cell]): Iterable[(Buf, Buf)] = {
-    env.mapValues {
-      case Real(k, v) => k.marshal(v.get)
-      case Translucent(_, vBuf) => vBuf
+  private[context] def marshal(env: Map[String, Cell]): Iterable[(Buf, Buf)] = {
+    // Since we internally store our keys in a case insensitive format, we want to
+    // make sure we marshal our keys using the original `marshalId` which preserves
+    // the casing of the context key.
+    env.map {
+      case (_, Real(k, v)) => (k.marshalId, k.marshal(v.get))
+      case (_, Translucent(k, vBuf)) => (k, vBuf)
     }
   }
 
   // Exposed for testing
-  private[context] def env: Map[Buf, Cell] = local() match {
+  private[context] def env: Map[String, Cell] = local() match {
     case Some(env) => env
     case None => Map.empty
   }
 
   // Exposed for testing
   private[context] def doUnmarshal(
-    env: Map[Buf, Cell],
+    env: Map[String, Cell],
     contexts: Iterable[(Buf, Buf)]
-  ): Map[Buf, Cell] = {
+  ): Map[String, Cell] = {
     contexts.foldLeft(env) {
-      case (env, (k, v)) =>
-        env.updated(k, Translucent(k, v))
+      case (env, (k @ Buf.Utf8(id), v)) =>
+        if (useCaseInsensitiveLookup) env.updated(caseInsensitiveLookupId(id), Translucent(k, v))
+        else env.updated(id, Translucent(k, v))
     }
   }
 
@@ -172,5 +200,5 @@ final class MarshalledContext private[context] extends Context {
   }
 
   // Exposed for testing
-  private[context] def letLocal[T](env: Map[Buf, Cell])(fn: => T): T = local.let(env)(fn)
+  private[context] def letLocal[T](env: Map[String, Cell])(fn: => T): T = local.let(env)(fn)
 }

finagle-core/src/test/scala/com/twitter/finagle/context/MarshalledContextTest.scala

@@ -21,6 +21,134 @@ class MarshalledContextTest extends AbstractContextTest {
     }
   }
 
+  def withCaseInsensitiveOverride[R](enabled: Boolean)(f: => R): R = {
+    val fraction = if (enabled) 1.0 else 0.0
+    com.twitter.finagle.toggle.flag.overrides
+      .let("com.twitter.finagle.context.MarshalledContextLookupId", fraction) {
+        f
+      }
+  }
+
+  test("Key#marshalId is used when marshalling") {
+    val key = new ctx.Key[String]("C.kEy") {
+      def marshal(value: String) = Buf.Utf8(value)
+      def tryUnmarshal(buf: Buf) = buf match {
+        case Buf.Utf8(value) => Return(value)
+      }
+    }
+
+    ctx.let(key, "bar") {
+      assert(
+        ctx.marshal() == Map(
+          key.marshalId -> Buf.Utf8("bar")
+        )
+      )
+    }
+  }
+
+  test("When disabled, key lookups are case sensitive") {
+    withCaseInsensitiveOverride(false) {
+      val ctx = new MarshalledContext
+      val lowerKey = new ctx.Key[String]("lowerkey") {
+        def marshal(value: String) = Buf.Utf8(value)
+        def tryUnmarshal(buf: Buf) = buf match {
+          case Buf.Utf8(value) => Return(value)
+        }
+      }
+
+      val upperKey = new ctx.Key[String]("LOWERKEY") {
+        def marshal(value: String) = Buf.Utf8(value)
+        def tryUnmarshal(buf: Buf) = buf match {
+          case Buf.Utf8(value) => Return(value)
+        }
+      }
+
+      ctx.let(lowerKey, "bar") {
+        val lowerValue = ctx.get(lowerKey)
+        assert(lowerValue.isDefined)
+
+        val upperValue = ctx.get(upperKey)
+        assert(upperValue == None)
+      }
+    }
+  }
+
+  test("When enabled, keys unmarshal a into a case insensitive format") {
+    withCaseInsensitiveOverride(true) {
+      val ctx = new MarshalledContext
+      val lowerKey = new ctx.Key[String]("foo") {
+        def marshal(value: String) = Buf.Utf8(value)
+        def tryUnmarshal(buf: Buf) = buf match {
+          case Buf.Utf8(value) => Return(value)
+        }
+      }
+      val upperKey = new ctx.Key[String]("FOO") {
+        def marshal(value: String) = Buf.Utf8(value)
+        def tryUnmarshal(buf: Buf) = buf match {
+          case Buf.Utf8(value) => Return(value)
+        }
+      }
+
+      val data = ctx.let(upperKey, "bar") { ctx.marshal() }
+      ctx.letUnmarshal(data) {
+        val value = ctx.get(lowerKey)
+        assert(value.isDefined)
+        assert(value.get == "bar")
+      }
+    }
+  }
+
+  test("when enabled, key lookups are case insenstive") {
+    withCaseInsensitiveOverride(true) {
+      val ctx = new MarshalledContext
+      val lowerKey = new ctx.Key[String]("foo") {
+        def marshal(value: String) = Buf.Utf8(value)
+        def tryUnmarshal(buf: Buf) = buf match {
+          case Buf.Utf8(value) => Return(value)
+        }
+      }
+      val upperKey = new ctx.Key[String]("FOO") {
+        def marshal(value: String) = Buf.Utf8(value)
+        def tryUnmarshal(buf: Buf) = buf match {
+          case Buf.Utf8(value) => Return(value)
+        }
+      }
+
+      ctx.let(lowerKey, "hello") {
+        val value = ctx.get(upperKey)
+        assert(value.isDefined)
+        assert(value.get == "hello")
+
+        val originalKeyValue = ctx.get(lowerKey)
+        assert(originalKeyValue.isDefined)
+        assert(originalKeyValue.get == "hello")
+      }
+    }
+  }
+
+  test("Flipping toggle does not change existing context") {
+    val ctx = withCaseInsensitiveOverride(false) { new MarshalledContext }
+    val lowerKey = new ctx.Key[String]("lowerkey") {
+      def marshal(value: String) = Buf.Utf8(value)
+      def tryUnmarshal(buf: Buf) = buf match {
+        case Buf.Utf8(value) => Return(value)
+      }
+    }
+
+    val upperKey = new ctx.Key[String]("LOWERKEY") {
+      def marshal(value: String) = Buf.Utf8(value)
+      def tryUnmarshal(buf: Buf) = buf match {
+        case Buf.Utf8(value) => Return(value)
+      }
+    }
+
+    withCaseInsensitiveOverride(true) {
+      ctx.let(lowerKey, "foo") {
+        assert(ctx.get(upperKey) == None)
+      }
+    }
+  }
+
   test("Translucency: pass through, replace") {
     ctx.let(b, 333) {
       ctx.letUnmarshal(Seq(Buf.Utf8("bleep") -> Buf.Utf8("bloop"))) {
@@ -88,7 +216,7 @@ class MarshalledContextTest extends AbstractContextTest {
       val roundTrip = ctx.doUnmarshal(Map.empty, ctx.marshal())
 
       def checkKey(key: ctx.Key[_]): Unit = {
-        roundTrip(key.marshalId) match {
+        roundTrip(key.lookupId) match {
           case t: ctx.Translucent => assert(t.unmarshal(key) == ctx.get(key))
           case other => fail(s"Unexpected structure: $other")
         }