rack-throttle to limit API request

For a API project I am using rack-throttle gem, which is a Rack middleware that provides logic for rate-limiting incoming HTTP requests to Rack applications.

Now, the usage was pretty straight forward as it is explained in the gem page but the requirements I had were different. So I had to extend the gem and customize the code to support my requirements. rack-throttle gem check limitation per host, I wanted to it per authentication token. Each API call must have a authentication token passed to get API result back, and I wanted to set different limit for different authentication-token, or you can say different user.

I used memcache to keep the counts. And I wanted the limit per day, so I extended Rack::Throttle::Daily.

Here is my code snippet in /lib/api_defender.rb

At First, in the initialization I want to check the Database for all the authentication tokens and their limits and want to populate memcache with those values:

require 'rack/throttle'
class ApiDefender < Rack::Throttle::Daily 
  def initialize(app)
     options = {
       :code => 403,
       :message => "Rate Limit Exceeded"
     }
     @app, @options = app, options
     api_auth_limits = User.select('authentication_token, api_req_limit')
     api_auth_limits.each do |a_l|
        Rails.cache.write(a_l.authentication_token+'limit', a_l.api_req_limit, expires_in: TOKEN_LIMIT_EXPIRATION_MIN)
     end
     rescue
       nil
   end

After that, the allowed method check if the request link eligible for this limit check? In my case, if the request is coming from a browser I do the limit per host, and if it not browser and pass authentication token I check the authentication token limit for this request. I also keep increasing the limit in memcache ot keep track and once it reaches the final limit per day I return forbidden.

  def allowed?(request)
    if request_is_browser?(request).blank?
      token_val = get_token_val(request)
      max_per_window = Rails.cache.read(token_val+'limit')
      count_increase = cache_incr(request, max_per_window)
      if max_per_window.blank?
        api_auth_limits = User.select('authentication_token, req_limit').where(:authentication_token => token_val)
        Rails.cache.write(api_auth_limits[0].authentication_token+'limit', api_auth_limits[0].req_limit, expires_in: TOKEN_LIMIT_EXPIRATION_MIN)
        max_per_window =  api_auth_limits[0].req_limit
        max_per_window = 0 if max_per_window.blank?
      end
      need_defense?(request) ? count_increase <= max_per_window : true
    else
      if need_defense?(request)
        token_expire_seconds = TOKEN_COUNTER_EXPIRATION_MIN*60
        token_expire_seconds_for_time = TOKEN_COUNTER_EXPIRATION_MIN*60 + 120
        host_ip = request.ip.to_s
        host_ip_limit = Rails.cache.read(host_ip+'limit')
        host_ip_count = Rails.cache.read(host_ip)
        if host_ip_limit.blank?
          Rails.cache.write(host_ip+'limit', BROWSER_API_REQ_LIMIT, expires_in: TOKEN_LIMIT_EXP_MIN)
          host_ip_limit = BROWSER_API_REQ_LIMIT
        end
        if host_ip_count.blank?
          current_time = Time.now
          Rails.cache.write(host_ip+'time', current_time, expires_in: token_expire_seconds_for_time)
          Rails.cache.write(host_ip, 1, expires_in: token_expire_seconds)
          host_ip_count = 1
        elsif host_ip_count <= host_ip_limit
          host_ip_count = host_ip_count + 1
          current_time = Time.now
          last_update_time = Rails.cache.read(host_ip+'time')
          seconds_left = ((token_expire_seconds  - (current_time - last_update_time)))
          Rails.cache.write(host_ip, host_ip_count, expires_in: seconds_left)
        end
        host_ip_count  max_per_window)

    token_expire_seconds = TOKEN_COUNTER_EXPIRATION_MIN*60
    token_expire_seconds_for_time = TOKEN_COUNTER_EXPIRATION_MIN*60 + 120
    if token_count.blank?
      current_time = Time.now
      Rails.cache.write(token_val, 1, expires_in: token_expire_seconds)
      Rails.cache.write(token_val+'time', current_time, expires_in: token_expire_seconds_for_time)
      count = 0
    else
      current_time = Time.now
      last_update_time = Rails.cache.read(token_val+'time')
      seconds_left = ((token_expire_seconds  - (current_time - last_update_time)))
      count = token_count
      Rails.cache.write(token_val, count + 1, expires_in: seconds_left)
    end
    count + 1
    rescue
      true
  end

And the need_defense check if this path should be validated for request.

  def need_defense?(request)
    request.fullpath.include? "api/v1/"
  end

The call() method is the one which return forbidden if limit exceeds or pass handle to the API application if the validation of limit pass.

  def call(env)
    request = Rack::Request.new(env)
    token = get_token_val(request)
    existing_time = Rails.cache.read(token+'time')
    token = request.ip.to_s if existing_time.blank?
    token = 'Test' if token.blank?
    allowed?(request) ? app.call(env) : rate_limit_exceeded(token)
  end

The rate_limit_exceeded method is customized to return results the way I want it to.

  def rate_limit_exceeded(token)
    existing_time = Rails.cache.read(token+'time')
    time_left = ((Time.now - existing_time)/60).round(2) if !existing_time.blank?
    time_attr = "min"
    time_attr = "time"  if time_left.blank?
    time_left = "later" if time_left.blank?
    http_error("{'error':[{'message':'Sorry, Rate Limit Exceeded, please try again after #{time_left} #{time_attr}','code':403}]}")
  rescue
    time_left = "later time"
    http_error("{'error':[{'message':'Sorry, Rate Limit Exceeded, please try again after #{time_left}','code':403}]}")
 end

Also, I had to add this custom class in the environment files:

  config.middleware.insert_after Rack::Lock, ApiDefender