Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with
or
.
Download ZIP
Fetching contributors…

Cannot retrieve contributors at this time

executable file 266 lines (233 sloc) 8.935 kb
#!/usr/bin/env ruby
$: << File.join(File.dirname(__FILE__), '../lib')
require 'goliath'
require 'em-mongo'
require 'em-http'
require 'em-synchrony/em-http'
require 'em-synchrony/em-mongo'
require 'yajl/json_gem' if RUBY_PLATFORM != 'java'
require File.join(File.dirname(__FILE__), 'http_log') # Use the HttpLog as our actual endpoint, but include this in the middleware
#
# Usage:
#
# First launch the test rig:
# bundle exec ./examples/test_rig.rb -sv -p 8080 -e prod &
#
# Then launch this script
# bundle exec ./examples/auth_and_rate_limit.rb -sv -p 9000 --config $PWD/examples/config/auth_and_rate_limit.rb
#
# The auth info is returned in the headers:
#
# curl -vv 'http://127.0.0.1:9000/?_apikey=i_am_busy&drop=false' ; echo
# ...snip...
# < X-RateLimit-MaxRequests: 1000
# < X-RateLimit-Requests: 999
# < X-RateLimit-Reset: 1312059600
#
# This user will hit the rate limit after 10 requests:
#
# for foo in 1 2 3 4 5 6 7 8 9 10 11 12 ; do echo -ne $foo "\t" ; curl 'http://127.0.0.1:9000/?_apikey=i_am_limited' ; echo ; done
# 1 {"Special":"Header","Params":"_apikey: i_am_awesome|drop: false","Path":"/","Headers":"User-Agent: ...
# ...
# 11 [:error, "Your request rate (11) is over your limit (10)"]
#
# You can test the barrier (both delays are in fractional seconds):
# * drop=true will drop the request at the remote host
# * auth_db_delay will fake a slow response from the mongo
# * delay will cause a slow response from the remote host
#
# time curl -vv 'http://127.0.0.1:9000/?_apikey=i_am_awesome&drop=false&delay=0.4&auth_db_delay=0.3'
# ...
# X-Tracer: ... received_usage_info: 0.06, received_sleepy: 299.52, received_downstream_resp: 101.67, ..., total: 406.09
# ...
# real 0m0.416s user 0m0.002s sys 0m0.003s pct 1.24
#
# This shows the mongodb response returning quickly, the fake DB delay returning
# after 300ms, and the downstream response returning after an additional 101 ms.
# The total request took 416ms of wall-clock time
#
# This will hold up even in the face of many concurrent connections. Relaunch in
# production (you may have to edit the config/auth_and_rate_limit scripts):
#
# bundle exec ./examples/auth_and_rate_limit.rb -sv -p 9000 -e prod --config $PWD/examples/config/auth_and_rate_limit.rb
#
# On my laptop, with 20 concurrent requests (each firing two db gets, a 400 ms
# http get, and two db writes), the median/90%ile times were 431ms / 457ms:
#
# time ab -c20 -n20 'http://127.0.0.1:9000/?_apikey=i_am_awesome&drop=false&delay=0.4&auth_db_delay=0.3'
# ...
# Percentage of the requests served within a certain time (ms)
# 50% 431
# 90% 457
# real 0m0.460s user 0m0.001s sys 0m0.003s pct 0.85
#
# With 100 concurrent requests, the request latency starts to drop but the
# throughput and variance stand up:
#
# time ab -c100 -n100 'http://127.0.0.1:9000/?_apikey=i_am_awesome&drop=false&delay=0.4&auth_db_delay=0.3'
# ...
# Percentage of the requests served within a certain time (ms)
# 50% 640
# 90% 673
# real 0m0.679s user 0m0.002s sys 0m0.007s pct 1.33
#
# Tracks and enforces account and rate limit policies.
#
# This is like a bouncer who lets townies order a drink while he checks their
# ID, but who's a jerk to college kids.
#
# On GET or HEAD requests, it proxies the request and gets account/usage info
# concurrently; authorizing the account doesn't delay the response.
#
# On a POST or other non-idempotent request, it checks the account/usage info
# *before* allowing the request to fire. This takes longer, but is necessary and
# tolerable.
#
# The magic of BarrierAroundware:
#
# 1) In pre_process (before the request):
# * validate an apikey was given; if not, raise (returning directly)
# * launch requests for the account and rate limit usage
#
# 2) On a POST or other non-GET non-HEAD, we issue `perform`, which barriers
# (allowing other requests to proceed) until the two pending requests
# complete. It then checks the account exists and is valid, and that the rate
# limit is OK
#
# 3) If the auth check fails, we raise an error (later caught by a safely{}
# block and turned into the right 4xx HTTP response.
#
# 4) If the auth check succeeds, or the request is a GET or HEAD, we return
# Goliath::Connection::AsyncResponse, and BarrierAroundwareFactory passes the
# request down the middleware chain
#
# 5) post_process resumes only when both proxied request & auth info are complete
# (it already has of course in the non-lazy scenario)
#
# 6) If we were lazy, the post_process method now checks authorization
#
class AuthBarrier
include Goliath::Rack::BarrierAroundware
include Goliath::Validation
attr_reader :db
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 initialize(env, db_name)
@db = env.config[db_name]
super(env)
end
def pre_process
env.trace('pre_process_beg')
validate_apikey!
# the results of the afirst deferrable will be set right into account_info (and the request into successes)
enqueue_mongo_request(:account_info, { :_id => apikey })
enqueue_mongo_request(:usage_info, { :_id => usage_id })
maybe_fake_delay!
# On non-GET non-HEAD requests, we have to check auth now.
unless lazy_authorization?
perform # yield execution until user_info has arrived
charge_usage
check_authorization!
end
env.trace('pre_process_end')
return Goliath::Connection::AsyncResponse
end
def post_process
env.trace('post_process_beg')
# [:account_info, :usage_info, :status, :headers, :body].each{|attr| env.logger.info(("%23s\t%s" % [attr, self.send(attr).inspect[0..200]])) }
inject_headers
# We have to check auth now, we skipped it before
if lazy_authorization?
charge_usage
check_authorization!
end
env.trace('post_process_end')
[status, headers, body]
end
def lazy_authorization?
(env['REQUEST_METHOD'] == 'GET') || (env['REQUEST_METHOD'] == 'HEAD')
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
# Fake out a delay in the database response if auth_db_delay is given
def maybe_fake_delay!
if (auth_db_delay = env.params['auth_db_delay'].to_f) > 0
enqueue_acceptor(:sleepy){|acc| EM.add_timer(auth_db_delay){ acc.succeed } }
end
end
def accept_response(handle, *args)
env.trace("received_#{handle}")
super(handle, *args)
end
# ===========================================================================
def check_authorization!
check_apikey!
check_rate_limit!
end
def validate_apikey!
if apikey.to_s.empty?
raise MissingApikeyError
end
end
def check_apikey!
unless account_info && (account_info['valid'] == true)
raise InvalidApikeyError
end
end
def check_rate_limit!
self.usage_info ||= {}
rate = usage_info['calls'].to_i + 1
limit = account_info['max_call_rate'].to_i
return true if rate <= limit
raise RateLimitExceededError, "Your request rate (#{rate}) is over your limit (#{limit})"
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!({
'X-RateLimit-MaxRequests' => account_info['max_call_rate'].to_s,
'X-RateLimit-Requests' => usage_info['calls'].to_i.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
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
Jump to Line
Something went wrong with that request. Please try again.