基于Akka-Streams的HTTP代理的实现

来源:互联网 发布:arcgis java 开发手册 编辑:程序博客网 时间:2024/05/21 11:27

Akka-Streams是一个让人激动的Reactive Streams的框架,Akka-Http也是构建在其之上,除了内置背压模式的支持,使用其DSL构建一个Graph也是一个让人惊艳的过程。对于Akka-Streams的介绍会在后续的文章中逐一展开,本文只奉上一段代码,实现的是HTTP代理功能,同时可以对请求和响应做一些修饰。

import akka.NotUsedimport akka.actor.ActorSystemimport akka.http.scaladsl.Httpimport akka.http.scaladsl.coding.{Deflate, Gzip, NoCoding}import akka.http.scaladsl.model._import akka.http.scaladsl.model.headers.HttpEncodingsimport akka.stream._import akka.stream.scaladsl.{Flow, GraphDSL, Sink, Source}import akka.util.ByteStringimport io.circe._import io.circe.parser._import scala.concurrent.duration.Durationimport scala.concurrent.{Await, Future}object HttpProxy extends App {  implicit val actorSystem = ActorSystem(name = "system")  implicit val streamMaterializer = ActorMaterializer()  implicit val executionContext = actorSystem.dispatcher  // The handle flow contains 2 sub-flows: forward flow & decorate flow.  // The forward flow will forward http request to destination server.  // The decorate flow will decorate http response then return to clients.  val handleFlow: Flow[HttpRequest, HttpResponse, _] = Flow.fromGraph(GraphDSL.create() {    implicit builder: GraphDSL.Builder[NotUsed] =>      import GraphDSL.Implicits._      val forwardFlow = Http(actorSystem).outgoingConnection(host = "destination-server", port = 8080)      val forwardRequestFlow: Flow[HttpRequest, (HttpRequest, Future[HttpResponse]), _] = {        Flow[HttpRequest]          .map {            orgRequest => {              val request = orgRequest.copy(                // The uri string value looks like: "http://domain/path",                // this is invalid for some http servers, the acceptance format is: "/path"                // So, we need remove domain, now, The value 21 is the length of                // string: "http://destination-server:8080", it's hard code for short-term,                // later, it will be configurable.                uri = orgRequest.uri.toString().substring(30),                headers =                  for {                    header <- orgRequest.headers                    if header.isNot("timeout-access")                  } yield {                    if (header.is("host")) {                      akka.http.scaladsl.model.headers.Host(Uri.Host("destination-server"), 8080)                    }                    else {                      header                    }                  }              )              val response = Source.single(request).via(forwardFlow).runWith(Sink.head)              (request, response)            }          }      }      val decorateResponseFlow: Flow[(HttpRequest, Future[HttpResponse]), HttpResponse, _] = {        Flow[(HttpRequest, Future[HttpResponse])]          .mapAsync(1) {            (tuple) => {              val request = tuple._1              val responseFuture = tuple._2              val (path, body) = extractPathAndBody(request)              path match {                case "/path/to/decorate" if body.contains("some special key words") =>                  responseFuture.flatMap {                    response =>                      val coder = response.encoding match{                        case HttpEncodings.gzip => Gzip                        case HttpEncodings.deflate => Deflate                        case HttpEncodings.identity => NoCoding                        case _ => throw new RuntimeException("akka not supported encoding "+response.encoding.toString())                      }                      response.entity                        .dataBytes.runFold(ByteString(""))(_ ++ _)                        .map {                          byteString =>                            val gzipDecodeFuture = coder.decode(byteString)                            val gzipDecodeByteString = Await.result(gzipDecodeFuture, Duration.Inf)                            val decodedByteString = gzipDecodeByteString.decodeString("utf-8")                            val json: Json = parse(decodedByteString).getOrElse(Json.Null)                            // do some transformation jobs based on origin json.                            val transformedEntity = ???                            val entity = HttpEntity(ContentTypes.`application/json`,                              coder.encode(ByteString.fromString(transformedEntity)))                            response.copy(entity = entity)                        }                  }                case _ => responseFuture              }            }          }      }      val forwardRequest: FlowShape[HttpRequest, (HttpRequest, Future[HttpResponse])] = builder.add(forwardRequestFlow)      val decorateResponse: FlowShape[(HttpRequest, Future[HttpResponse]), HttpResponse] = builder.add(decorateResponseFlow)      forwardRequest ~> decorateResponse      FlowShape(forwardRequest.in, decorateResponse.out)  })  def extractPathAndBody(request:HttpRequest):(String,String) ={    val path = request.uri.path.toString()    val body = request.entity.dataBytes.map(_.decodeString("utf-8")).runWith(Sink.head)    val bodyStr = Await.result(body, Duration.Inf)    (path,bodyStr)  }  val serverSource = Http(actorSystem).bind(interface = "0.0.0.0", port = 7000)  serverSource.to(Sink.foreach { connection =>    connection.handleWith(handleFlow)  }).run()}
2 0
原创粉丝点击