From a96b9c5e7f2122b59315d46f47d1bd726d511f7a Mon Sep 17 00:00:00 2001
From: Luna Mendes <lnmds@protonmail.com>
Date: Sun, 4 Nov 2018 02:23:26 -0300
Subject: [PATCH] ratelimits.handler: five better retry_after and global flag

 - run: add X-RateLimit-Global and Retry-After headers
---
 litecord/ratelimits/handler.py | 14 ++++++++++++--
 litecord/ratelimits/main.py    |  2 +-
 run.py                         | 11 +++++++++++
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/litecord/ratelimits/handler.py b/litecord/ratelimits/handler.py
index db896bf..715a283 100644
--- a/litecord/ratelimits/handler.py
+++ b/litecord/ratelimits/handler.py
@@ -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)
diff --git a/litecord/ratelimits/main.py b/litecord/ratelimits/main.py
index ffcc54d..f3e6bfb 100644
--- a/litecord/ratelimits/main.py
+++ b/litecord/ratelimits/main.py
@@ -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):
diff --git a/run.py b/run.py
index ff3e623..a8059c6 100644
--- a/run.py
+++ b/run.py
@@ -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