API key authentication + rate limiting in Goliath using MongoDB (incomplete sketch)

auth_and_rate_limit.rb

#!/usr/bin/env ruby
$: << File.join(File.dirname(__FILE__), '../lib')
require 'goliath'
require 'em-mongo'
require 'em-http'
require 'em-synchrony/em-http'
require 'yajl/json_gem'

require 'goliath/synchrony/mongo_receiver'            # has the aroundware logic for talking to mongodb
require File.join(File.dirname(__FILE__), 'auth_receiver')

# Usage:
#
# First launch a dummy responder, like hello_world.rb or test_rig.rb:
# ruby ./examples/hello_world.rb -sv -p 8080 -e prod &
#
# Then launch this script
# ruby ./examples/auth_and_rate_limit.rb -sv -p 9000 --config $PWD/auth_and_rate_limit_config.rb
#
# Rate limit!
#
#     $ curl  'http://127.0.0.1:9000/?_apikey=i_am_busy' ; echo
#     hello world
#     $ curl  'http://127.0.0.1:9000/?_apikey=i_am_busy' ; echo
#     hello world
#     $ curl  'http://127.0.0.1:9000/?_apikey=i_am_busy' ; echo
#     [:error, "Forbidden"]

# taken from examples/http_log.rb

class AuthAndRateLimit < Goliath::API
  use Goliath::Rack::Tracer, 'X-Tracer'
  use Goliath::Rack::Params             # parse & merge query and body parameters
  use Goliath::Rack::AsyncAroundware, AuthReceiver, 'api_auth_db'


  # Capture the headers when they roll in, to replay for the remote target
  def on_headers(env, headers)
    env['client-headers'] = headers
  end

  # Pass the request on to host given in config[:forwarder]
  def response(env)
    env.trace :response_beg

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

    url = "#{forwarder}#{env[Goliath::Request::REQUEST_PATH]}"
    env.logger.info ['proxy', url].join("\t")
    req = EM::HttpRequest.new(url)
    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

    env.trace :response_end
    [resp.response_header.status, response_header_hash(resp), resp.response]
  end

  # Need to convert from the CONTENT_TYPE we'll get back from the server
  # to the normal Content-Type header
  def response_header_hash(resp)
    hsh = {}
    resp.response_header.each_pair do |k, v|
      hsh[to_http_header(k)] = v
    end
    hsh
  end

  def to_http_header(k)
    k.downcase.split('_').map{|e| e.capitalize }.join('-')
  end
end

auth_and_rate_limit_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('buzzkill_test')
  end

  # for demo purposes, some dummy accounts
  timebin = ((Time.now.to_i / 3600).floor * 3600)

  # This user's calls should all go through
  config['api_auth_db'].collection('AccountInfo').save({
      :_id => 'i_am_awesome', 'valid' => true,  'max_call_rate' => 1_000_000 })

  # this user's account is disabled
  config['api_auth_db'].collection('AccountInfo').save({
      :_id => 'i_am_lame',    'valid' => false, 'max_call_rate' => 1_000 })

  # this user has not been seen, but will very quickly hit their limit
  config['api_auth_db'].collection('AccountInfo').save({
      :_id => 'i_am_limited', 'valid' => true, 'max_call_rate' =>     10 })

  # fakes a user with a bunch of calls already made this hour -- two more = no yuo
  config['api_auth_db'].collection('AccountInfo').save({
      :_id => 'i_am_busy',    'valid' => true, 'max_call_rate' =>  1_000 })
  config['api_auth_db'].collection('UsageInfo').save({
      :_id => "i_am_busy-#{timebin}", 'calls' =>  999 })
end

config['forwarder'] = 'http://localhost:8080'

auth_receiver.rb

#
# Tracks and enforces account and rate limit policies.
#
# Before the request:
#
# * validates the apikey exists
# * launches requests for the account and current usage (hourly rate limit, etc)
#
# It then passes the request down the middleware chain; execution resumes only
# when both the remote request and the auth info have returned.
#
# After remote request and auth info return:
#
# * Check the account exists and is valid
# * Check the rate limit is OK
#
# If it passes all those checks, the request goes through; otherwise we raise an
# error that Goliath::Rack::Validator turns into a 4xx response
#
# WARNING: Since this passes ALL requests through to the responder, it's only
# suitable for idempotent requests (GET, typically).  You may need to handle
# POST/PUT/DELETE requests differently.
#
#
class AuthReceiver < Goliath::Synchrony::MongoReceiver
  include Goliath::Validation
  include Goliath::Rack::Validator
  attr_accessor :account_info, :usage_info

  # time period to aggregate stats over, in seconds
  TIMEBIN_SIZE = 60 * 60

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

  def pre_process
    validate_apikey!
    first('AccountInfo', { :_id => apikey   }){|res| self.account_info = res }
    first('UsageInfo',   { :_id => usage_id }){|res| self.usage_info   = res }
    env.trace('pre_process_end')
  end

  def post_process
    env.trace('post_process_beg')
    env.logger.info [account_info, usage_info].inspect
    self.account_info ||= {}
    self.usage_info   ||= {}

    inject_headers

    EM.next_tick do
      safely(env){ charge_usage }
    end

    safely(env, headers) do
      check_apikey!
      check_rate_limit!

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

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

  def validate_apikey!
    if apikey.to_s.empty?
      raise MissingApikeyError
    end
  end

  def check_apikey!
    unless account_info['valid'] == true
      raise InvalidApikeyError
    end
  end

  def check_rate_limit!
    return true if usage_info['calls'].to_f <= account_info['max_call_rate'].to_f
    raise RateLimitExceededError
  end

  def charge_usage
    update('UsageInfo', { :_id => usage_id },
      { '$inc' => { :calls   => 1 } }, :upsert => true)
  end

  def inject_headers
    headers.merge!({
        'X-RateLimit-MaxRequests' => account_info['max_call_rate'].to_s,
        'X-RateLimit-Requests'    => usage_info['calls'].to_s,
        'X-RateLimit-Reset'       => timebin_end.to_s,
      })
  end

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

  def apikey
    env.params['_apikey']
  end

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

  def timebin
    @timebin ||= timebin_beg
  end

  def timebin_beg
    ((Time.now.to_i / TIMEBIN_SIZE).floor * TIMEBIN_SIZE)
  end

  def timebin_end
    timebin_beg + TIMEBIN_SIZE
  end
end