implemented headers for moving-window + redis

TODO: improve test coverage.

doc/source/index.rst

@@ -164,6 +164,7 @@ The following flask configuration values are honored by
                                while memcached relies on the `pymemcache`_ package.
 ``RATELIMIT_STRATEGY``         The rate limiting strategy to use.  :ref:`ratelimit-strategy`
                                for details.
+``RATELIMIT_HEADERS_ENABLED``  Enables returning :ref:`ratelimit-headers`
 ``RATELIMIT_ENABLED``          Overall killswitch for ratelimits. Defaults to ``True``
 ============================== ================================================
 
@@ -242,6 +243,39 @@ rate limit as the window for each limit is not fixed at the start and end of eac
 however a higher memory cost associated with this strategy as it requires ``N`` items to
 be maintained in memory per resource and rate limit.
 
+.. _ratelimit-headers:
+
+Rate-limiting Headers
+=====================
+
+If the configuration is enabled, information about the rate limit with respect to the
+route being requested will be written as part of the response. Since multiple rate limits
+can be active for a given route - the rate limit with the lowest time granularity will be
+used in the scenario when the request does not breach any rate limits.
+
+.. tabularcolumns:: |p{6.5cm}|p{8.5cm}|
+
+============================== ================================================
+``X-RateLimit-Limit``          The total number of requests allowed for the
+                               active window
+``X-RateLimit-Remaining``      The number of requests remaining in the active
+                               window.
+``X-RateLimit-Reset``          UTC seconds since epoch when the window will be
+                               reset.
+============================== ================================================
+
+Depending on the :ref:`ratelimit-strategy` chosen, the meaning of the headers
+may differ. For example, with a moving window strategy there is no actual
+reset for the window, and therefore the value of ``X-RateLimit-Reset`` is always
+``now() + 1``.
+
+.. warning:: Enabling the headers has an additional with certain storage / strategy combinations.
+
+    * Memcached + Fixed Window: an extra key per rate limit is stored to calculate
+      ``X-RateLimit-Reset``
+    * Redis + Moving Window: an extra call to redis is involved during every request
+      to calculate ``X-RateLimit-Remaining``
+
 .. _keyfunc-customization:
 
 Customization

flask_limiter/extension.py

@@ -20,6 +20,7 @@ class Limiter(object):
      limits to apply to all routes. :ref:`ratelimit-string` for  more details.
     :param function key_func: a callable that returns the domain to rate limit by.
      Defaults to the remote address of the request.
+    :param bool headers_enabled: whether ``X-RateLimit`` response headers are written.
     """
 
     def __init__(self, app=None

flask_limiter/storage.py

@@ -178,6 +178,20 @@ def __init__(self, redis_url):
         self.storage = get_dependency("redis").from_url(redis_url)
         if not self.storage.ping():
             raise ConfigurationError("unable to connect to redis at %s" % redis_url) # pragma: no cover
+        script = """
+        local items = redis.call('lrange', KEYS[1], 0, tonumber(ARGV[2]))
+        local expiry = tonumber(ARGV[1])
+        local a = 0
+        for idx=1,#items do
+            if tonumber(items[idx]) >= expiry then
+                a = a + 1
+            else
+                break
+            end
+        end
+        return a
+        """
+        self.script_hash = self.storage.script_load(script)
         super(RedisStorage, self).__init__()
 
     def incr(self, key, expiry, elastic_expiry=False):
@@ -229,7 +243,10 @@ def get_acquirable(self, key, limit, expiry):
         :param int limit: amount of entries allowed
         :param int expiry: expiry of the entry
         """
-        raise NotImplementedError
+        timestamp = time.time()
+        return limit - self.storage.evalsha(
+            self.script_hash, 1, key, int(timestamp - expiry), limit
+        )
 
     def get_expiry(self, key):
         """

tests/test_flask_ext.py

@@ -9,10 +9,13 @@
 import time
 from datetime import datetime
 from flask.ext.limiter.extension import Limiter
-
+import redis
 
 
 class FlaskExtTests(unittest.TestCase):
+    def setUp(self):
+        redis.Redis().flushall()
+
     def test_combined_rate_limits(self):
         app = Flask(__name__)
         app.config.setdefault("RATELIMIT_GLOBAL", "1 per hour;10 per day")
@@ -378,31 +381,31 @@ def t():
                     str(int(time.time() + 49))
                 )
 
-    def test_headers_moving_window(self):
+    def test_headers_moving_window_redis(self):
+
         app = Flask(__name__)
         app.config["RATELIMIT_STRATEGY"] = "moving-window"
+        app.config["RATELIMIT_STORAGE_URL"] = "redis://localhost:6379"
         limiter = Limiter(app, global_limits=["10/minute"], headers_enabled=True)
 
         @app.route("/t1")
-        @limiter.limit("2/second; 10 per minute; 20/hour")
+        @limiter.limit("10/second; 20per minute")
         def t():
             return "test"
 
-        with hiro.Timeline().freeze(datetime(2012,12,12,0,0,0)) as timeline:
-            with app.test_client() as cli:
-                for i in range(11):
-                    resp = cli.get("/t1")
-                    timeline.forward(1)
-                self.assertEqual(
-                    resp.headers.get('X-RateLimit-Limit'),
-                    '10'
-                )
-                self.assertEqual(
-                    resp.headers.get('X-RateLimit-Remaining'),
-                    '0'
-                )
-                self.assertEqual(
-                    resp.headers.get('X-RateLimit-Reset'),
-                    str(int(time.time()))
-                )
-
+        with app.test_client() as cli:
+            for i in range(21):
+                resp = cli.get("/t1")
+                time.sleep(0.1)
+            self.assertEqual(
+                resp.headers.get('X-RateLimit-Limit'),
+                '20'
+            )
+            self.assertEqual(
+                resp.headers.get('X-RateLimit-Remaining'),
+                '0'
+            )
+            self.assertEqual(
+                resp.headers.get('X-RateLimit-Reset'),
+                str(int(time.time() + 1))
+            )