Limit HTTP request rate for spray/akka-http

I want to limit the request rate in a spray/akka-http route. I'm not an expert with both scala and spray/akka but learning. So I came up with the following custom directive:

trait RequestRateLimiter[T] {
    def check(ctx: RequestContext, realIp: Option[RemoteAddress], proxiedIp: Option[RemoteAddress]) : RequestRateResult[T]
  }

  sealed trait RequestRateResult[T]
  case class Pass[T](result: T) extends RequestRateResult[T]
  case class Block[T](rejection: Rejection) extends RequestRateResult[T]

  def limit[T](requestRateLimiter: RequestRateLimiter[T]): Directive1[T] = {
    def remoteAddress: Directive1[Option[RemoteAddress]] = optionalHeaderValuePF { case `Remote-Address`(address) ⇒ address }
    def xForwardedFor: Directive1[Option[RemoteAddress]] = optionalHeaderValuePF { case `X-Forwarded-For`(Seq(address, _*)) ⇒ address }
    extract(identity) flatMap { ctx =>
      remoteAddress.flatMap { realIp =>
        xForwardedFor.flatMap { proxiedIp =>
          requestRateLimiter.check(ctx, realIp, proxiedIp) match {
            case Pass(result)     => provide(result)
            case Block(rejection) => reject(rejection)
          }
        }
      }
    }
  }

The custom directive takes an instance of an RequestRateLimiter which keeps track of the request rate based on the RequestContext and RemoteAddress. The RequestRateLimiter then decides if the request should be passed or blocked.
I'm not quiet happy with the RequestRateLimiter because it's too generic. It can be anything that either provides something or rejects.

I also posted a Gist (feel free to share, use and improve).

What do you think of the directive? Anything that could be improved?

I think you could rewrite your limit function like this :

def limit[T](requestRateLimiter: RequestRateLimiter[T]) = {
  def remoteAddress = optionalHeaderValuePF { case `Remote-Address`(address) ⇒ address }
  def xForwardedFor = optionalHeaderValuePF { case `X-Forwarded-For`(Seq(address, _*)) ⇒ address }

  for {
    ctx  provide(result)
    case Block(rejection) => reject(rejection)
  }
}

Here are the two main changes that I've done :

• Changed your flatMaps for a "for comprehension" (which imho is more readable)

• Removed the types on the declaration of the two variables.

Overall, I find it more readable.

StackExchange Code Review Q#120733, answer score: 4