Skip to content

Commit

Permalink
Implement rate-limiting in WB with redis
Browse files Browse the repository at this point in the history
 * Add rate-limiting support via redis to WaterButler. This is
   implementation uses the fixed window algorithm.

   METHOD: Users are distinguished first by their credentials, and
   second by their IP address. The rate limiter recognizes different
   types of auth and rate-limits each type separately. The 4 types are
   OSF cookie, OAuth bearer token, basic auth w/ base64- encoded
   username/password, and un-authed.

   OSF cookies, OAuth access tokens, and base64-encoded
   usernames/passwords are used as redis keys during rate-limiting. We
   must obfuscate them for the same reason we store only password
   hashes in database. `hashlib.hash256().hexdigest()` is used in this
   case. In addition, a prefix is added to the digest to identifty
   which type it is. The "No Auth" case is hashed as well (though
   unnecessarily) so that the keys all have the same look and length.

   Auth by OSF cookie currently bypasses the rate limiter to avoid
   throttling web users.

   CONFIGURATION: Rate limiting settings are found in
   `waterbutler.server.settings`.  By default, WB allows 3600 requests
   per hour.  Rate-limiting is turned OFF by default; set
   `ENABLE_RATE_LIMITING` to `True` turn it on.

   BEHAVIOR: Return the Retry-After header in the 429 response if the limit is
   hit

   * Add X-WB-RL headers to responses for blocked requests

   * Use Redis TTL command to obtain the time until reset.

   * Remove redundant Redis `.get()` for key existence check

   * In Redis, the INCR command increments the number stored at key by
     1. If the key does not exist, it is set to 0 before performing
     the operation. The python redis implementation makes this even
     simpler by setting the value to 1 without doing the increment.

   * The rate limiter now throws 503 if WB Redis fails

 * grab-bag of wb-rl updates

   * Bump redis dep to 3.3.8.  No consequential changes.

   * Don't throw errors in the error handling: provide a fallback for
     the `resource` attribute if rate-limiting kicks in before that
     has been initialized.

   * Update some docstrings to clarify return values and process.

   * Refactor test rate-limiting auth testing to only extract data as
     needed.

   * Add docs to settings; use `config.get_bool()` on booleans.
  • Loading branch information
cslzchen authored and felliott committed Oct 12, 2021
1 parent 47ad15e commit 52cb391
Show file tree
Hide file tree
Showing 8 changed files with 217 additions and 3 deletions.
1 change: 0 additions & 1 deletion dev-requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions tests/core/test_exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
39 changes: 39 additions & 0 deletions waterbutler/core/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}'

Expand Down Expand Up @@ -55,6 +57,43 @@ 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 ("rate limiting").
TODO - Optional: Return a few customized Waterbutler headers with rate-limiting details:
* X-Waterbutler-RateLimiting-Window
* X-Waterbutler-RateLimiting-Limit
* X-Waterbutler-RateLimiting-Remaining
* X-Waterbutler-RateLimiting-Reset
"""
def __init__(self, data):
if type(data) != dict:
message = 'This is a thing: {}'.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. Returns 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.
Expand Down
22 changes: 21 additions & 1 deletion waterbutler/server/api/v1/core.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import logging

import tornado.web
import tornado.gen
import tornado.iostream
Expand All @@ -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):

Expand All @@ -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'
Expand Down
14 changes: 13 additions & 1 deletion waterbutler/server/api/v1/provider/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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.info('>>> preparing request ...')
limit_hit, data = self.rate_limit()
if limit_hit:
raise TooManyRequests(data=data)
logger.info('>>> rate limiting check passed ...')

method = self.request.method.lower()

# TODO Find a nicer way to handle this
Expand Down
129 changes: 129 additions & 0 deletions waterbutler/server/api/v1/provider/ratelimiting.py
Original file line number Diff line number Diff line change
@@ -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):

# TODO: set different parameters for different types of auth
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_DOMAIN, port=settings.REDIS_PORT)

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.info('>>> 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.info('>>> 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.info('>>> 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.info('>>> 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.info('>>> 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.info('>>> 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.info('>>> 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.info('>>> 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()
12 changes: 12 additions & 0 deletions waterbutler/server/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,3 +30,15 @@
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_DOMAIN = config.get('REDIS_DOMAIN', '192.168.168.167')
REDIS_PORT = config.get('REDIS_PORT', '6379')

# Number of seconds until the redis key expires
RATE_LIMITING_FIXED_WINDOW_SIZE = int(config.get('RATE_LIMITING_FIXED_WINDOW_SIZE', 360))

# number of reqests permitted while the redis key is active
RATE_LIMITING_FIXED_WINDOW_LIMIT = int(config.get('RATE_LIMITING_FIXED_WINDOW_LIMIT', 36))

0 comments on commit 52cb391

Please sign in to comment.