auth_proxy_example

auth_and_rate_limit.rb

require 'rubygems'
require 'bundler'
Bundler.require

require 'goliath'
require 'em-mongo'
require 'em-http'
require 'em-synchrony/em-http'
require 'em-synchrony/em-mongo'

require File.join(File.dirname(__FILE__), 'http_log')

class AuthBarrier
  include Goliath::Rack::BarrierAroundware
  include Goliath::Validation
  attr_reader   :db
  attr_accessor :usage_info

  class MissingApikeyError     < BadRequestError   ; end
  class RateLimitExceededError < ForbiddenError    ; end
  class InvalidApikeyError     < UnauthorizedError ; end

  def initialize(env, db_name)
    @db = env.config[db_name]
    super(env)
  end

  def pre_process
    env.trace('pre_process_beg')

    enqueue_mongo_request(:usage_info,   { :_id => usage_id })

    perform     # yield execution until user_info has arrived
    charge_usage
    check_authorization!

    env.trace('pre_process_end')
    return Goliath::Connection::AsyncResponse
  end

  def post_process
    env.trace('post_process_beg')

    inject_headers

    env.trace('post_process_end')
    [status, headers, body]
  end

  if defined?(EM::Mongo::Cursor)
    # em-mongo > 0.3.6 gives us a deferrable back. nice and clean.
    def enqueue_mongo_request(handle, query)
      enqueue handle, db.collection(handle).afirst(query)
    end
  else
    # em-mongo <= 0.3.6 makes us fake a deferrable response.
    def enqueue_mongo_request(handle, query)
      enqueue_acceptor(handle) do |acc|
        db.collection(handle).afind(query){|resp| acc.succeed(resp.first) }
      end
    end
  end

  def accept_response(handle, *args)
    env.trace("received_#{handle}")
    super(handle, *args)
  end

  # ===========================================================================

  def check_authorization!
    check_rate_limit!
  end

  def check_rate_limit!
    self.usage_info ||= {}
    rate  = usage_info['calls'].to_i + 1
    return true if rate <= env.config['rate_limit_per_timebin']
    raise RateLimitExceededError, "Your request rate (#{rate}) is over your limit (#{env.config['rate_limit_per_timebin']})"
  end

  def current_request_count
    if usage_info.nil?
      1
    else
      usage_info['calls'].to_i + 1
    end
  end

  def charge_usage
    EM.next_tick do
      safely(env){ db.collection(:usage_info).update({ :_id => usage_id },
          { '$inc' => { :calls   => 1 } }, :upsert => true) }
    end
  end

  def inject_headers
    headers.merge!({
      'Access-Control-Allow-Origin' => '*',
      'Content-Type' => 'application/json; charset=utf-8',
      'X-RateLimit-Limit' => "#{env.config['rate_limit_per_timebin']}",
      'X-RateLimit-Remaining' => "#{remaining_credits}",
      'X-RateLimit-Reset' => "#{timebin_end}",
    })
  end

  def ip
    env['HTTP_X_REAL_IP'] || env['REMOTE_ADDR']
  end
  def remaining_credits
    count = (env.config['rate_limit_per_timebin'] - current_request_count)
    count < 0 ? 0 : count
  end

  def usage_id
    "#{ip}-#{timebin}"
  end

  def timebin
    @timebin ||= timebin_beg
  end

  def timebin_beg
    ((Time.now.to_i / env.config['rate_limit_timebin']).floor * env.config['rate_limit_timebin'])
  end

  def timebin_end
    timebin_beg + env.config['rate_limit_timebin']
  end
end

class AuthAndRateLimit < HttpLog
  use Goliath::Rack::Tracer, 'X-Tracer'
  use Goliath::Rack::Params             # parse & merge query and body parameters
  use Goliath::Rack::BarrierAroundwareFactory, AuthBarrier, 'api_auth_db'
end

config.rb

environment(:development) do
  config['api_auth_db'] = EventMachine::Synchrony::ConnectionPool.new(:size => 20) do
    conn = EM::Mongo::Connection.new('localhost', 27017, 1, {:reconnect_in => 1})
    conn.db('rate_limits')
  end
  config['backend'] = "http://localhost:4568"
  config['rate_limit_timebin'] = 10
  config['rate_limit_per_timebin'] = 4000
end

Gemfile

source 'https://rubygems.org'

gem "bson_ext", "1.9.2"
gem "mongo", "1.9.2"

gem "em-http-request"
gem "em-mongo"
gem "em-synchrony"
gem "goliath", "1.0.3"
gem "oj"

http_log.rb

require 'rubygems'
require 'bundler'
Bundler.require

require 'goliath'
require 'em-mongo'
require 'em-http'
require 'em-synchrony/em-http'
require 'pp'

class HttpLog < Goliath::API
  use Goliath::Rack::Params

  def on_headers(env, headers)
    env.logger.info 'proxying new request: ' + headers.inspect
    env['client-headers'] = headers
  end

  def response(env)
    start_time = Time.now.to_f

    params = {:head => env['client-headers'], :query => env.params}

    req = EM::HttpRequest.new("#{env.config['backend']}#{env[Goliath::Request::REQUEST_PATH]}")
    resp = case(env[Goliath::Request::REQUEST_METHOD])
      when 'GET'  then req.get(params)
      when 'POST' then req.post(params.merge(:body => env[Goliath::Request::RACK_INPUT].read))
      when 'HEAD' then req.head(params)
      else p "UNKNOWN METHOD #{env[Goliath::Request::REQUEST_METHOD]}"
    end

    process_time = Time.now.to_f - start_time

    response_headers = {}
    resp.response_header.each_pair do |k, v|
      response_headers[to_http_header(k)] = v
    end

    [resp.response_header.status, response_headers, resp.response]
  end

  # Need to convert from the CONTENT_TYPE we'll get back from the server
  # to the normal Content-Type header
  def to_http_header(k)
    k.downcase.split('_').collect { |e| e.capitalize }.join('-')
  end
end