CustomStage.scala

class CustomStage extends GraphStage[FlowShape[ByteString, ByteString]] {
    val in = Inlet[ByteString]("CustomStage.in")
    val out = Outlet[ByteString]("CustomStage.out")
    override val shape = FlowShape.of(in, out)

    override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) with InHandler with OutHandler {

        setHandlers(in, out, this)

        override def onPush() = emit(out, grab(in))

        override def onPull() = pull(in)
    }
}