[wip]Feature/wb rate limiting 2

dev-requirements.txt

@@ -14,4 +14,3 @@ pytest==5.2.1
 pytest-asyncio==0.10.0
 pytest-cov==2.8.1
 python-coveralls==2.9.3
-redis==3.3.8

requirements.txt

@@ -18,6 +18,7 @@ pyjwt==1.4.0
 python-dateutil==2.5.3
 pytz==2017.2
 sentry-sdk==0.14.4
+redis==3.3.8
 setuptools==37.0.0
 stevedore==1.2.0
 tornado==6.0.3

tests/conftest.py

@@ -6,6 +6,10 @@
     'CELERY_RESULT_BACKEND': 'redis://'
 }
 
+
+from waterbutler.server import settings as server_settings
+server_settings.ENABLE_RATE_LIMITING = False
+
 import aiohttpretty
 
 

tests/core/test_exceptions.py

@@ -38,6 +38,8 @@ class TestExceptionSerialization:
         exceptions.UninitializedRepositoryError,
         exceptions.UnexportableFileTypeError,
         exceptions.InvalidProviderConfigError,
+        exceptions.WaterButlerRedisError,
+        exceptions.TooManyRequests,
     ])
     def test_tolerate_dumb_signature(self, exception_class):
         """In order for WaterButlerError-inheriting exception classes to survive

waterbutler/core/exceptions.py

@@ -3,6 +3,8 @@
 
 from aiohttp.client_exceptions import ContentTypeError
 
+from waterbutler.server import settings
+
 
 DEFAULT_ERROR_MSG = 'An error occurred while making a {response.method} request to {response.url}'
 
@@ -55,6 +57,48 @@ def __str__(self):
         return '{}, {}'.format(self.code, self.message)
 
 
+class TooManyRequests(WaterButlerError):
+    """Indicates the user has sent too many requests in a given amount of time and is being
+    rate-limited. Thrown as HTTP 429, ``Too Many Requests``. Exception response includes headers to
+    inform user when to try again. Headers are:
+
+    * ``Retry-After``: Epoch time after which the rate-limit is reset
+    * ``X-Waterbutler-RateLimiting-Window``: The number of seconds after the first request when\
+    the limit resets
+    * ``X-Waterbutler-RateLimiting-Limit``: Total number of requests that may be sent within the\
+    window
+    * ``X-Waterbutler-RateLimiting-Remaining``: How many more requests can be sent during the window
+    * ``X-Waterbutler-RateLimiting-Reset``: Seconds until the rate-limit is reset
+    """
+    def __init__(self, data):
+        if type(data) != dict:
+            message = ('Too many requests issued, but error lacks necessary data to build proper '
+                       'response. Got:({})'.format(data))
+        else:
+            message = {
+                'error': 'API rate-limiting active due to too many requests',
+                'headers': {
+                    'Retry-After': data['retry_after'],
+                    'X-Waterbutler-RateLimiting-Window': settings.RATE_LIMITING_FIXED_WINDOW_SIZE,
+                    'X-Waterbutler-RateLimiting-Limit': settings.RATE_LIMITING_FIXED_WINDOW_LIMIT,
+                    'X-Waterbutler-RateLimiting-Remaining': data['remaining'],
+                    'X-Waterbutler-RateLimiting-Reset': data['reset'],
+                },
+            }
+        super().__init__(message, code=HTTPStatus.TOO_MANY_REQUESTS, is_user_error=True)
+
+
+class WaterButlerRedisError(WaterButlerError):
+    """Indicates the Redis server has returned an error. Thrown as HTTP 503, ``Service Unavailable``
+    """
+    def __init__(self, redis_command):
+
+        message = {
+            'error': 'The Redis server failed when processing command {}'.format(redis_command),
+        }
+        super().__init__(message, code=HTTPStatus.SERVICE_UNAVAILABLE, is_user_error=False)
+
+
 class InvalidParameters(WaterButlerError):
     """Errors regarding incorrect data being sent to a method should raise either this
     Exception or a subclass thereof.  Defaults status code to 400, Bad Request.

waterbutler/server/api/v1/core.py

@@ -1,3 +1,5 @@
+import logging
+
 import tornado.web
 import tornado.gen
 import tornado.iostream
@@ -8,6 +10,8 @@
 from waterbutler.server import utils
 from waterbutler.core import exceptions
 
+logger = logging.getLogger(__name__)
+
 
 class BaseHandler(utils.CORsMixin, utils.UtilMixin, tornado.web.RequestHandler):
 
@@ -25,7 +29,23 @@ def write_error(self, status_code, exc_info):
                     scope.level = 'info'
 
                 self.set_status(int(exc.code))
-                finish_args = [exc.data] if exc.data else [{'code': exc.code, 'message': exc.message}]
+
+                # If the exception has a `data` property then we need to handle that with care.
+                # The expectation is that we need to return a structured response.  For now, assume
+                # that involves setting the response headers to the value of the `headers`
+                # attribute of the `data`, while also serializing the entire `data` data structure.
+                if exc.data:
+                    self.set_header('Content-Type', 'application/json')
+                    headers = exc.data.get('headers', None)
+                    if headers:
+                        for key, value in headers.items():
+                            self.set_header(key, value)
+                    finish_args = [exc.data]
+                    self.write(exc.data)
+                else:
+                    finish_args = [{'code': exc.code, 'message': exc.message}]
+                    self.write(exc.message)
+
             elif issubclass(etype, tasks.WaitTimeOutError):
                 self.set_status(202)
                 scope.level = 'info'

waterbutler/server/api/v1/provider/__init__.py

@@ -14,10 +14,13 @@
 from waterbutler.core import remote_logging
 from waterbutler.server.auth import AuthHandler
 from waterbutler.core.log_payload import LogPayload
+from waterbutler.core.exceptions import TooManyRequests
 from waterbutler.core.streams import RequestStreamReader
+from waterbutler.server.settings import ENABLE_RATE_LIMITING
 from waterbutler.server.api.v1.provider.create import CreateMixin
 from waterbutler.server.api.v1.provider.metadata import MetadataMixin
 from waterbutler.server.api.v1.provider.movecopy import MoveCopyMixin
+from waterbutler.server.api.v1.provider.ratelimiting import RateLimitingMixin
 
 logger = logging.getLogger(__name__)
 auth_handler = AuthHandler(settings.AUTH_HANDLERS)
@@ -32,13 +35,22 @@ def list_or_value(value):
     return [item.decode('utf-8') for item in value]
 
 
+# TODO: the order should be reverted though it doesn't have any functional effect for this class.
 @tornado.web.stream_request_body
-class ProviderHandler(core.BaseHandler, CreateMixin, MetadataMixin, MoveCopyMixin):
+class ProviderHandler(core.BaseHandler, CreateMixin, MetadataMixin, MoveCopyMixin, RateLimitingMixin):
     PRE_VALIDATORS = {'put': 'prevalidate_put', 'post': 'prevalidate_post'}
     POST_VALIDATORS = {'put': 'postvalidate_put'}
     PATTERN = r'/resources/(?P<resource>(?:\w|\d)+)/providers/(?P<provider>(?:\w|\d)+)(?P<path>/.*/?)'
 
     async def prepare(self, *args, **kwargs):
+
+        if ENABLE_RATE_LIMITING:
+            logger.debug('>>> checking for rate-limiting')
+            limit_hit, data = self.rate_limit()
+            if limit_hit:
+                raise TooManyRequests(data=data)
+            logger.debug('>>> rate limiting check passed ...')
+
         method = self.request.method.lower()
 
         # TODO Find a nicer way to handle this

waterbutler/server/api/v1/provider/ratelimiting.py

@@ -0,0 +1,129 @@
+import hashlib
+import logging
+from datetime import datetime, timedelta
+
+from redis import Redis
+from redis.exceptions import RedisError
+
+from waterbutler.server import settings
+from waterbutler.core.exceptions import WaterButlerRedisError
+
+logger = logging.getLogger(__name__)
+
+
+class RateLimitingMixin:
+    """ Rate-limiting WB API with Redis using the "Fixed Window" algorithm.
+    """
+
+    def __init__(self):
+
+        self.WINDOW_SIZE = settings.RATE_LIMITING_FIXED_WINDOW_SIZE
+        self.WINDOW_LIMIT = settings.RATE_LIMITING_FIXED_WINDOW_LIMIT
+        self.redis_conn = Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT,
+                                password=settings.REDIS_PASSWORD)
+
+    def rate_limit(self):
+        """ Check with the WB Redis server on whether to rate-limit a request.  Returns a tuple.
+        First value is `True` if the limit is reached, `False` otherwise.  Second value is the
+        rate-limiting metadata (nbr of requests remaining, time to reset, etc.) if the request was
+        rate-limited.
+        """
+
+        limit_check, redis_key = self.get_auth_naive()
+        logger.debug('>>> RATE LIMITING >>> check={} key={}'.format(limit_check, redis_key))
+        if not limit_check:
+            return False, None
+
+        try:
+            counter = self.redis_conn.incr(redis_key)
+        except RedisError:
+            raise WaterButlerRedisError('INCR {}'.format(redis_key))
+
+        if counter > self.WINDOW_LIMIT:
+            # The key exists and the limit has been reached.
+            try:
+                retry_after = self.redis_conn.ttl(redis_key)
+            except RedisError:
+                raise WaterButlerRedisError('TTL {}'.format(redis_key))
+            logger.debug('>>> RATE LIMITING >>> FAIL >>> key={} '
+                         'counter={} url={}'.format(redis_key, counter, self.request.full_url()))
+            data = {
+                'retry_after': int(retry_after),
+                'remaining': 0,
+                'reset': str(datetime.now() + timedelta(seconds=int(retry_after))),
+            }
+            return True, data
+        elif counter == 1:
+            # The key does not exist and `.incr()` returns 1 by default.
+            try:
+                self.redis_conn.expire(redis_key, self.WINDOW_SIZE)
+            except RedisError:
+                raise WaterButlerRedisError('EXPIRE {} {}'.format(redis_key, self.WINDOW_SIZE))
+            logger.debug('>>> RATE LIMITING >>> NEW >>> key={} '
+                         'counter={} url={}'.format(redis_key, counter, self.request.full_url()))
+        else:
+            # The key exists and the limit has not been reached.
+            logger.debug('>>> RATE LIMITING >>> PASS >>> key={} '
+                         'counter={} url={}'.format(redis_key, counter, self.request.full_url()))
+
+        return False, None
+
+    def get_auth_naive(self):
+        """ Get the obfuscated authentication / authorization credentials from the request.  Return
+        a tuple ``(limit_check, auth_key)`` that tells the rate-limiter 1) whether to rate-limit,
+        and 2) if so, limit by what key.
+
+        Refer to ``tornado.httputil.HTTPServerRequest`` for more info on tornado's request object:
+            https://www.tornadoweb.org/en/stable/httputil.html#tornado.httputil.HTTPServerRequest
+
+        This is a NAIVE implementation in which WaterButler rate-limiter only checks the existence
+        of auth creds in the requests without further verifying them with the OSF.  Invalid creds
+        will fail the next OSF auth part anyway even if it passes the rate-limiter.
+
+        There are four types of auth: 1) OAuth access token, 2) basic auth w/ base64-encoded
+        username/password, 3) OSF cookie, and 4) no auth.  The naive implementation checks each
+        method in this order.  Only cookie-based auth is permitted to bypass the rate-limiter.
+        This order does not care about the validity of the auth mechanism.  An invalid Basic auth
+        header + an OSF cookie will be rate-limited according to the Basic auth header.
+
+        TODO: check with OSF API auth to see how it deals with multiple auth options.
+        """
+
+        auth_hdrs = self.request.headers.get('Authorization', None)
+
+        # CASE 1: Requests with a bearer token (PAT or OAuth)
+        if auth_hdrs and auth_hdrs.startswith('Bearer '):  # Bearer token
+            bearer_token = auth_hdrs.split(' ')[1] if auth_hdrs.startswith('Bearer ') else None
+            logger.debug('>>> RATE LIMITING >>> AUTH:TOKEN >>> {}'.format(bearer_token))
+            return True, 'TOKEN__{}'.format(self._obfuscate_creds(bearer_token))
+
+        # CASE 2: Requests with basic auth using username and password
+        if auth_hdrs and auth_hdrs.startswith('Basic '):  # Basic auth
+            basic_creds = auth_hdrs.split(' ')[1] if auth_hdrs.startswith('Basic ') else None
+            logger.debug('>>> RATE LIMITING >>> AUTH:BASIC >>> {}'.format(basic_creds))
+            return True, 'BASIC__{}'.format(self._obfuscate_creds(basic_creds))
+
+        # CASE 3: Requests with OSF cookies
+        # SECURITY WARNING: Must check cookie last since it can only be allowed when used alone!
+        cookies = self.request.cookies or None
+        if cookies and cookies.get('osf'):
+            osf_cookie = cookies.get('osf').value
+            logger.debug('>>> RATE LIMITING >>> AUTH:COOKIE >>> {}'.format(osf_cookie))
+            return False, 'COOKIE_{}'.format(self._obfuscate_creds(osf_cookie))
+
+        # TODO: Work with DevOps to make sure that the `remote_ip` is the real IP instead of our
+        #       load balancers.  In addition, check relevatn HTTP headers as well.
+        # CASE 4: Requests without any expected auth (case 1, 2 or 3 above).
+        remote_ip = self.request.remote_ip or 'NOI.PNO.IPN.OIP'
+        logger.debug('>>> RATE LIMITING >>> AUTH:NONE >>> {}'.format(remote_ip))
+        return True, 'NOAUTH_{}'.format(self._obfuscate_creds(remote_ip))
+
+    @staticmethod
+    def _obfuscate_creds(creds):
+        """Obfuscate authentication/authorization credentials: cookie, access token and password.
+
+        It is not recommended to store the plain OSF cookie or the OAuth bearer token as key and it
+        is evil to store the base64-encoded username and password as key since it is reversible.
+        """
+
+        return hashlib.sha256(creds.encode('utf-8')).hexdigest().upper()

waterbutler/server/settings.py

@@ -30,3 +30,20 @@
 if not settings.DEBUG:
     assert HMAC_SECRET, 'HMAC_SECRET must be specified when not in debug mode'
 HMAC_SECRET = (HMAC_SECRET or 'changeme').encode('utf-8')
+
+
+# Configs for WB API Rate-limiting with Redis
+ENABLE_RATE_LIMITING = config.get_bool('ENABLE_RATE_LIMITING', False)
+REDIS_HOST = config.get('REDIS_HOST', '192.168.168.167')
+REDIS_PORT = config.get('REDIS_PORT', '6379')
+REDIS_PASSWORD = config.get('REDIS_PASSWORD', None)
+
+# Number of seconds until the redis key expires
+RATE_LIMITING_FIXED_WINDOW_SIZE = int(config.get('RATE_LIMITING_FIXED_WINDOW_SIZE', 3600))
+
+# number of reqests permitted while the redis key is active
+RATE_LIMITING_FIXED_WINDOW_LIMIT = int(config.get('RATE_LIMITING_FIXED_WINDOW_LIMIT', 3600))
+
+# test env envar mapping
+RATE_LIMITING_FIXED_WINDOW_SIZE = 360
+RATE_LIMITING_FIXED_WINDOW_LIMIT = 30