From 7b05356e64809bfe812ffacadc0df2f18396e9a6 Mon Sep 17 00:00:00 2001 From: jay-parikh Date: Thu, 16 Jul 2020 16:06:35 +0530 Subject: [PATCH 1/4] feat: Ignore Rate Limit for Whitelisted Clients --- redis_rate_limit/__init__.py | 29 +++++++++++++++-------------- tests/rate_limit_test.py | 13 +++++++++++++ 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/redis_rate_limit/__init__.py b/redis_rate_limit/__init__.py index 147db8f..9c7bd7e 100644 --- a/redis_rate_limit/__init__.py +++ b/redis_rate_limit/__init__.py @@ -43,21 +43,23 @@ class RateLimit(object): This class offers an abstraction of a Rate Limit algorithm implemented on top of Redis >= 2.6.0. """ - def __init__(self, resource, client, max_requests, expire=None, redis_pool=REDIS_POOL): + def __init__(self, resource, client, max_requests, whitelisted_clients=[], expire=None, redis_pool=REDIS_POOL): """ Class initialization method checks if the Rate Limit algorithm is actually supported by the installed Redis version and sets some useful properties. - If Rate Limit is not supported, it raises an Exception. - :param resource: resource identifier string (i.e. ‘user_pictures’) :param client: client identifier string (i.e. ‘192.168.0.10’) :param max_requests: integer (i.e. ‘10’) + :param whitelisted_clients: list of ip addresses (i.e. ['127.0.0.1']) :param expire: seconds to wait before resetting counters (i.e. ‘60’) :param redis_pool: instance of redis.ConnectionPool. Default: ConnectionPool(host='127.0.0.1', port=6379, db=0) """ + self.client = client + self.whitelisted_clients = whitelisted_clients + self._redis = Redis(connection_pool=redis_pool) if not self._is_rate_limit_supported(): raise RedisVersionNotSupported() @@ -88,7 +90,6 @@ def get_usage(self): """ Returns actual resource usage by client. Note that it could be greater than the maximum number of requests set. - :return: integer: current usage """ return int(self._redis.get(self._rate_limit_key) or 0) @@ -97,7 +98,6 @@ def get_wait_time(self): """ Returns estimated optimal wait time for subsequent requests. If limit has already been reached, return wait time until it gets reset. - :return: float: wait time in seconds """ expire = self._redis.pttl(self._rate_limit_key) @@ -111,7 +111,6 @@ def get_wait_time(self): def has_been_reached(self): """ Checks if Rate Limit has been reached. - :return: bool: True if limit has been reached or False otherwise """ return self.get_usage() >= self._max_requests @@ -119,14 +118,11 @@ def has_been_reached(self): def increment_usage(self, increment_by=1): """ Calls a LUA script that should increment the resource usage by client. - If the resource limit overflows the maximum number of requests, this method raises an Exception. - :param increment_by: The count to increment the rate limiter by. This is by default 1, but higher values are provided for more flexible rate-limiting schemes. - :return: integer: current usage """ if increment_by > self._max_requests: @@ -140,8 +136,11 @@ def increment_usage(self, increment_by=1): .format(increment_by=increment_by)) try: - current_usage = self._redis.evalsha( - INCREMENT_SCRIPT_HASH, 1, self._rate_limit_key, self._expire, increment_by) + if self.client not in self.whitelisted_clients: + current_usage = self._redis.evalsha( + INCREMENT_SCRIPT_HASH, 1, self._rate_limit_key, self._expire, increment_by) + else: + current_usage = 0 except NoScriptError: current_usage = self._redis.eval( INCREMENT_SCRIPT, 1, self._rate_limit_key, self._expire, increment_by) @@ -155,7 +154,6 @@ def _is_rate_limit_supported(self): """ Checks if Rate Limit is supported which can basically be found by looking at Redis database version that should be 2.6.0 or greater. - :return: bool """ redis_version = self._redis.info()['redis_version'] @@ -172,17 +170,19 @@ def _reset(self): class RateLimiter(object): - def __init__(self, resource, max_requests, expire=None, redis_pool=REDIS_POOL): + def __init__(self, resource, max_requests, whitelisted_clients=[], expire=None, redis_pool=REDIS_POOL): """ Rate limit factory. Checks if RateLimit is supported when limit is called. :param resource: resource identifier string (i.e. ‘user_pictures’) :param max_requests: integer (i.e. ‘10’) + :param whitelisted_clients: list of ip addresses (i.e. ['127.0.0.1']) :param expire: seconds to wait before resetting counters (i.e. ‘60’) :param redis_pool: instance of redis.ConnectionPool. Default: ConnectionPool(host='127.0.0.1', port=6379, db=0) """ self.resource = resource self.max_requests = max_requests + self.whitelisted_clients = whitelisted_clients self.expire = expire self.redis_pool = redis_pool @@ -194,6 +194,7 @@ def limit(self, client): resource=self.resource, client=client, max_requests=self.max_requests, + whitelisted_clients=self.whitelisted_clients, expire=self.expire, redis_pool=self.redis_pool, - ) + ) \ No newline at end of file diff --git a/tests/rate_limit_test.py b/tests/rate_limit_test.py index 58c5fe6..2c871f4 100644 --- a/tests/rate_limit_test.py +++ b/tests/rate_limit_test.py @@ -41,6 +41,19 @@ def test_limit_10_max_request(self): self.assertEqual(self.rate_limit.get_usage(), 11) self.assertEqual(self.rate_limit.has_been_reached(), True) + def test_whitelisted_clients(self): + """ + Should not increment counter if client is part of whitelisted_clients list. + """ + self.rate_limit = RateLimit(resource='test', client='localhost', whitelisted_clients=['localhost'], + max_requests=10, expire=2) + self.assertEqual(self.rate_limit.get_usage(), 0) + self.assertEqual(self.rate_limit.has_been_reached(), False) + + self._make_10_requests() + self.assertEqual(self.rate_limit.get_usage(), 0) + self.assertEqual(self.rate_limit.has_been_reached(), False) + def test_expire(self): """ Should not raise TooManyRequests Exception when trying to increment for From 8dd5605608fcc35a4bc3832751e52ddab95d104f Mon Sep 17 00:00:00 2001 From: jay-parikh Date: Fri, 17 Jul 2020 11:24:22 +0530 Subject: [PATCH 2/4] fix: renamed whitelisted_clients to ignored_clients and other suggested changes --- redis_rate_limit/__init__.py | 32 ++++++++++++++++++++------------ tests/rate_limit_test.py | 6 +++--- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/redis_rate_limit/__init__.py b/redis_rate_limit/__init__.py index 9c7bd7e..3b79321 100644 --- a/redis_rate_limit/__init__.py +++ b/redis_rate_limit/__init__.py @@ -43,22 +43,24 @@ class RateLimit(object): This class offers an abstraction of a Rate Limit algorithm implemented on top of Redis >= 2.6.0. """ - def __init__(self, resource, client, max_requests, whitelisted_clients=[], expire=None, redis_pool=REDIS_POOL): + def __init__(self, resource, client, max_requests, ignored_clients=None, expire=None, redis_pool=REDIS_POOL): """ Class initialization method checks if the Rate Limit algorithm is actually supported by the installed Redis version and sets some useful properties. + If Rate Limit is not supported, it raises an Exception. + :param resource: resource identifier string (i.e. ‘user_pictures’) :param client: client identifier string (i.e. ‘192.168.0.10’) :param max_requests: integer (i.e. ‘10’) - :param whitelisted_clients: list of ip addresses (i.e. ['127.0.0.1']) + :param ignored_clients: list of ip addresses (i.e. ['127.0.0.1']) :param expire: seconds to wait before resetting counters (i.e. ‘60’) :param redis_pool: instance of redis.ConnectionPool. Default: ConnectionPool(host='127.0.0.1', port=6379, db=0) """ self.client = client - self.whitelisted_clients = whitelisted_clients + self.ignored_clients = ignored_clients self._redis = Redis(connection_pool=redis_pool) if not self._is_rate_limit_supported(): @@ -98,6 +100,7 @@ def get_wait_time(self): """ Returns estimated optimal wait time for subsequent requests. If limit has already been reached, return wait time until it gets reset. + :return: float: wait time in seconds """ expire = self._redis.pttl(self._rate_limit_key) @@ -111,6 +114,7 @@ def get_wait_time(self): def has_been_reached(self): """ Checks if Rate Limit has been reached. + :return: bool: True if limit has been reached or False otherwise """ return self.get_usage() >= self._max_requests @@ -118,13 +122,19 @@ def has_been_reached(self): def increment_usage(self, increment_by=1): """ Calls a LUA script that should increment the resource usage by client. + If the resource limit overflows the maximum number of requests, this method raises an Exception. + :param increment_by: The count to increment the rate limiter by. This is by default 1, but higher values are provided for more flexible rate-limiting schemes. + :return: integer: current usage """ + if self.ignored_clients and self.client in self.ignored_clients: + return 0 + if increment_by > self._max_requests: raise ValueError('increment_by {increment_by} overflows ' 'max_requests of {max_requests}' @@ -136,11 +146,8 @@ def increment_usage(self, increment_by=1): .format(increment_by=increment_by)) try: - if self.client not in self.whitelisted_clients: - current_usage = self._redis.evalsha( - INCREMENT_SCRIPT_HASH, 1, self._rate_limit_key, self._expire, increment_by) - else: - current_usage = 0 + current_usage = self._redis.evalsha( + INCREMENT_SCRIPT_HASH, 1, self._rate_limit_key, self._expire, increment_by) except NoScriptError: current_usage = self._redis.eval( INCREMENT_SCRIPT, 1, self._rate_limit_key, self._expire, increment_by) @@ -154,6 +161,7 @@ def _is_rate_limit_supported(self): """ Checks if Rate Limit is supported which can basically be found by looking at Redis database version that should be 2.6.0 or greater. + :return: bool """ redis_version = self._redis.info()['redis_version'] @@ -170,19 +178,19 @@ def _reset(self): class RateLimiter(object): - def __init__(self, resource, max_requests, whitelisted_clients=[], expire=None, redis_pool=REDIS_POOL): + def __init__(self, resource, max_requests, ignored_clients=None, expire=None, redis_pool=REDIS_POOL): """ Rate limit factory. Checks if RateLimit is supported when limit is called. :param resource: resource identifier string (i.e. ‘user_pictures’) :param max_requests: integer (i.e. ‘10’) - :param whitelisted_clients: list of ip addresses (i.e. ['127.0.0.1']) + :param ignored_clients: list of ip addresses (i.e. ['127.0.0.1']) :param expire: seconds to wait before resetting counters (i.e. ‘60’) :param redis_pool: instance of redis.ConnectionPool. Default: ConnectionPool(host='127.0.0.1', port=6379, db=0) """ self.resource = resource self.max_requests = max_requests - self.whitelisted_clients = whitelisted_clients + self.ignored_clients = ignored_clients self.expire = expire self.redis_pool = redis_pool @@ -194,7 +202,7 @@ def limit(self, client): resource=self.resource, client=client, max_requests=self.max_requests, - whitelisted_clients=self.whitelisted_clients, + ignored_clients=self.ignored_clients, expire=self.expire, redis_pool=self.redis_pool, ) \ No newline at end of file diff --git a/tests/rate_limit_test.py b/tests/rate_limit_test.py index 2c871f4..f4cde87 100644 --- a/tests/rate_limit_test.py +++ b/tests/rate_limit_test.py @@ -41,11 +41,11 @@ def test_limit_10_max_request(self): self.assertEqual(self.rate_limit.get_usage(), 11) self.assertEqual(self.rate_limit.has_been_reached(), True) - def test_whitelisted_clients(self): + def test_ignored_clients(self): """ - Should not increment counter if client is part of whitelisted_clients list. + Should not increment counter if client is part of ignored_clients list. """ - self.rate_limit = RateLimit(resource='test', client='localhost', whitelisted_clients=['localhost'], + self.rate_limit = RateLimit(resource='test', client='localhost', ignored_clients=['localhost'], max_requests=10, expire=2) self.assertEqual(self.rate_limit.get_usage(), 0) self.assertEqual(self.rate_limit.has_been_reached(), False) From 66c009a14b327ff81aa3ddd9d00239e6b1b94aef Mon Sep 17 00:00:00 2001 From: jay-parikh Date: Fri, 17 Jul 2020 11:41:24 +0530 Subject: [PATCH 3/4] fix: documentation for ignored_clients --- README.rst | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/README.rst b/README.rst index bab78c4..d2e8917 100644 --- a/README.rst +++ b/README.rst @@ -64,6 +64,17 @@ Example: 100 requests per hour except TooManyRequests: return '429 Too Many Requests' +Example: you can also pass an optional list of ignored_clients to bypass Rate Limit + +.. code-block:: python + + from redis_rate_limit import RateLimit, TooManyRequests + try: + with RateLimit(resource='users_list', client='192.168.0.10', max_requests=100, ignored_clients=['192.168.0.10'], expire=3600): + return '200 OK' + except TooManyRequests: + return '429 Too Many Requests' + Example: you can also setup a factory to use it later .. code-block:: python From 0aceb5c816edd6986e57ae96be977317ba5c352c Mon Sep 17 00:00:00 2001 From: jay-parikh Date: Tue, 21 Jul 2020 11:02:02 +0530 Subject: [PATCH 4/4] fix: added blank lines --- redis_rate_limit/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/redis_rate_limit/__init__.py b/redis_rate_limit/__init__.py index 3b79321..4becb75 100644 --- a/redis_rate_limit/__init__.py +++ b/redis_rate_limit/__init__.py @@ -92,6 +92,7 @@ def get_usage(self): """ Returns actual resource usage by client. Note that it could be greater than the maximum number of requests set. + :return: integer: current usage """ return int(self._redis.get(self._rate_limit_key) or 0) @@ -205,4 +206,4 @@ def limit(self, client): ignored_clients=self.ignored_clients, expire=self.expire, redis_pool=self.redis_pool, - ) \ No newline at end of file + )