ratelimits.handler: five better retry_after and global flag

- run: add X-RateLimit-Global and Retry-After headers
This commit is contained in:
Luna Mendes 2018-11-04 02:23:26 -03:00
parent 33f893c0ff
commit a96b9c5e7f
3 changed files with 24 additions and 3 deletions

View File

@ -10,8 +10,11 @@ async def _check_bucket(bucket):
request.bucket = bucket
if retry_after:
raise Ratelimited('You are being ratelimited.', {
'retry_after': retry_after
request.retry_after = retry_after
raise Ratelimited('You are being rate limited.', {
'retry_after': int(retry_after * 1000),
'global': request.bucket_global,
})
@ -22,6 +25,7 @@ async def _handle_global(ratelimit):
except Unauthorized:
user_id = request.remote_addr
request.bucket_global = True
bucket = ratelimit.get_bucket(user_id)
await _check_bucket(bucket)
@ -59,6 +63,12 @@ async def ratelimit_handler():
# methods have different ratelimits
rule_path = rule.endpoint
# some request ratelimit context.
# TODO: maybe put those in a namedtuple or contextvar of sorts?
request.bucket = None
request.retry_after = None
request.bucket_global = False
try:
ratelimit = app.ratelimiter.get_ratelimit(rule_path)
await _handle_specific(ratelimit)

View File

@ -41,7 +41,7 @@ class RatelimitManager:
"""Manager for the bucket managers"""
def __init__(self):
self._ratelimiters = {}
self.global_bucket = Ratelimit(50, 1)
self.global_bucket = Ratelimit(1, 1)
self._fill_rtl()
def _fill_rtl(self):

11
run.py
View File

@ -122,9 +122,20 @@ async def app_set_ratelimit_headers(resp):
"""Set the specific ratelimit headers."""
try:
bucket = request.bucket
if bucket is None:
raise AttributeError()
resp.headers['X-RateLimit-Limit'] = str(bucket.requests)
resp.headers['X-RateLimit-Remaining'] = str(bucket._tokens)
resp.headers['X-RateLimit-Reset'] = str(bucket._window + bucket.second)
resp.headers['X-RateLimit-Global'] = str(request.bucket_global).lower()
# only add Retry-After if we actually hit a ratelimit
retry_after = request.retry_after
if request.retry_after:
resp.headers['Retry-After'] = str(retry_after)
except AttributeError:
pass