upgrade.scala

import com.fasterxml.jackson.annotation.JsonIgnore
import com.twitter.finagle.Stack.Params
import com.twitter.finagle.mux.{Request, Response}
import com.twitter.finagle.{Filter, Path, Service}
import com.twitter.io.Buf
import com.twitter.util.Future
import io.buoyant.linkerd.protocol.ThriftServerConfig
import io.buoyant.linkerd.{ClientConfig, ProtocolInitializer, RouterConfig}
import io.buoyant.linkerd._
import io.buoyant.router.Thrift

class Thrift2MuxUpgradeInitializer extends ProtocolInitializer {
  val name = "thrift2mux"

  type Bytes = Array[Byte]
  protected type RouterReq = Request
  protected type RouterRsp = Response
  protected type ServerReq = Bytes
  protected type ServerRsp = Bytes

  protected val defaultRouter = io.buoyant.router.Mux.router

  protected val defaultServer = io.buoyant.router.Thrift.server

  val defaultServerPort: Int = 4114

  val configClass = classOf[Thrift2MuxConfig]

  protected val adapter = new Filter[Bytes, Bytes, Request, Response] {
    // ??? ClientId should be automatically propagated ???
    override def apply(request: Bytes, service: Service[Request, Response]): Future[Bytes] = {
      val req = Request(Path.empty, Buf.ByteArray.Owned(request))
      service(req).map(rep => Buf.ByteArray.Owned.extract(rep.body))
    }
  }
}

object Thrift2MuxUpgradeInitializer extends Thrift2MuxUpgradeInitializer

case class Thrift2MuxConfig(
  thriftMethodInDst: Option[Boolean]
) extends RouterConfig {

  var servers: Seq[ThriftServerConfig] = Nil
  var client: Option[ClientConfig] = None

  @JsonIgnore
  override def protocol = Thrift2MuxUpgradeInitializer

  override def routerParams: Params = super.routerParams
    .maybeWith(thriftMethodInDst.map(Thrift.param.MethodInDst(_)))
}