Substitute headers

src/main/scala/com/twitter/diffy/ApiController.scala

@@ -33,7 +33,7 @@ class ApiController @Inject()(
     JoinedEndpoint: JoinedEndpoint,
     includeWeights: Boolean,
     excludeNoise: Boolean
-  ) =
+  ) : Map[String, Map[String, Any]] =
     Map(
       "endpoint" -> Renderer.endpoint(JoinedEndpoint.endpoint),
       "fields" ->

src/main/scala/com/twitter/diffy/proxy/ClientService.scala

@@ -1,7 +1,7 @@
 package com.twitter.diffy.proxy
 
-import com.twitter.finagle.{Addr, Name, Resolver, Service}
 import com.twitter.finagle.thrift.ThriftClientRequest
+import com.twitter.finagle.{Addr, Name, Service}
 import com.twitter.util.{Time, Var}
 import org.jboss.netty.handler.codec.http.{HttpRequest, HttpResponse}
 
@@ -42,7 +42,7 @@ case class ThriftService(
     }
   }
 
-  private[this] def sizeChange(size: Int) {
+  private[this] def sizeChange(size: Int) : Unit = {
     changeCount += 1
     if (changeCount > 1) {
       changedAt = Some(Time.now)

src/main/scala/com/twitter/diffy/proxy/CloneHttpRequestFilter.scala

@@ -0,0 +1,27 @@
+package com.twitter.diffy.proxy
+
+import com.twitter.finagle.Filter
+import com.twitter.util.Future
+import org.jboss.netty.handler.codec.http.{HttpHeaders, DefaultHttpRequest, HttpResponse, HttpRequest}
+import scala.collection.JavaConverters._
+
+object CloneHttpRequestFilter {
+
+  type Req = HttpRequest
+  type Res = HttpResponse
+
+  private def mkCloneRequest(reqIn:Req, filter: (Req => Future[Res])):Future[Res] = {
+    def copyHeader(headers:HttpHeaders)(k: String, v:String) : Unit = { headers.add(k, v)}
+
+    val reqOut: DefaultHttpRequest = new DefaultHttpRequest(reqIn.getProtocolVersion, reqIn.getMethod, reqIn.getUri)
+    reqOut.setChunked(reqIn.isChunked)
+    reqOut.setContent(reqIn.getContent)
+
+    val headers = for {entry <- reqIn.headers().asScala} yield (entry.getKey, entry.getValue)
+    headers.foreach((copyHeader(reqOut.headers)_).tupled)
+
+    filter(reqOut)
+  }
+
+  def apply(): Filter[Req, Res, Req, Res] = Filter.mk(mkCloneRequest)
+}

src/main/scala/com/twitter/diffy/proxy/HttpDifferenceProxy.scala

@@ -2,13 +2,13 @@ package com.twitter.diffy.proxy
 
 import java.net.SocketAddress
 
-import com.twitter.diffy.analysis.{DifferenceAnalyzer, JoinedDifferences, InMemoryDifferenceCollector}
+import com.twitter.diffy.analysis.{DifferenceAnalyzer, InMemoryDifferenceCollector, JoinedDifferences}
 import com.twitter.diffy.lifter.{HttpLifter, Message}
 import com.twitter.diffy.proxy.DifferenceProxy.NoResponseException
-import com.twitter.finagle.{Service, Http, Filter}
-import com.twitter.finagle.http.{Status, Response, Method, Request}
-import com.twitter.util.{Try, Future}
-import org.jboss.netty.handler.codec.http.{HttpResponse, HttpRequest}
+import com.twitter.finagle.http.{Method, Request, Response, Status}
+import com.twitter.finagle.{Filter, Http, Service}
+import com.twitter.util.{Future, Try}
+import org.jboss.netty.handler.codec.http.{HttpRequest, HttpResponse}
 
 object HttpDifferenceProxy {
   val okResponse = Future.value(Response(Status.Ok))
@@ -34,8 +34,12 @@ trait HttpDifferenceProxy extends DifferenceProxy {
   override type Rep = HttpResponse
   override type Srv = HttpService
 
+
   override def serviceFactory(serverset: String, label: String) =
-    HttpService(Http.newClient(serverset, label).toService)
+    HttpService(
+      CloneHttpRequestFilter.apply.
+        andThen(RefineHttpHeadersByLabel(label)).
+        andThen(Http.newClient(serverset, label).toService))
 
   override lazy val server =
     Http.serve(

src/main/scala/com/twitter/diffy/proxy/RefineHttpHeadersByLabel.scala

@@ -0,0 +1,54 @@
+package com.twitter.diffy.proxy
+
+import com.twitter.finagle.Filter
+import com.twitter.finagle.http.Request
+import com.twitter.util.Future
+import org.jboss.netty.handler.codec.http.{HttpResponse, HttpHeaders, HttpRequest}
+import scala.collection.JavaConverters._
+
+object RefineHttpHeadersByLabel {
+  val knownLabels: Seq[String] = Seq("primary", "secondary", "candidate")
+
+  def substitute(inclusions: Seq[String], exclusion: Seq[String])(req: HttpRequest) : HttpRequest = {
+
+    def orFn[T](f: T => Boolean, g: T => Boolean): T => Boolean = { t =>  f(t) || g(t) }
+    def constFalse[T](t: T): Boolean = false
+
+    def renameHeaders(headers: HttpHeaders)(nameOld: String, nameNew: String): Unit = {
+      val temp = headers.get(nameOld)
+      headers.remove(nameOld)
+      headers.set(nameNew, temp)
+    }
+
+
+    val removeExclusion  = exclusion.
+      map(_+"_").
+      map( x => (_:String).startsWith(x)).
+      foldLeft(constFalse[String]_)(orFn)
+
+    val reqOut = Request.apply(req)
+    val headers: HttpHeaders = reqOut.headers
+    //k => Option[Function[k => k]]
+    val incomingHeaders = headers.names().asScala.toSet
+
+    val rewrites = for {
+      name <- incomingHeaders
+      inclusion <- inclusions if name.startsWith(inclusion + "_")
+    } yield name->name.substring(inclusion.length+1)
+
+    incomingHeaders.filter( removeExclusion).foreach( headers.remove )
+    rewrites.foreach((renameHeaders(headers) _ ).tupled)
+
+    reqOut
+  }
+
+  def apply(label:String): Filter[HttpRequest, HttpResponse, HttpRequest, HttpResponse] = {
+    def mkFilter(inclusions: Seq[String], exclusion: Seq[String])
+                (req:HttpRequest, filter: (HttpRequest => Future[HttpResponse]) ) =
+      filter(substitute(inclusions, exclusion)(req))
+
+    val (inclusions, exclusion) = knownLabels.partition(_.equalsIgnoreCase(label))
+    Filter.mk( mkFilter(inclusions, exclusion))
+  }
+
+}