From 882cafb4ba5e6cc1afbf904c6c3dc358945e92ee Mon Sep 17 00:00:00 2001 From: Basraah Date: Mon, 5 Jun 2017 08:36:25 +1000 Subject: [PATCH] Discord API rate limiting (#799) Added discord too many requests handling decorator Added tests for core Discord manager functions Added discord backoff retry for celery Added tests for update groups backoff Support per-route and global rate limiting --- alliance_auth/tests/test_settings.py | 12 +- services/modules/discord/manager.py | 117 ++++++++++++ services/modules/discord/tasks.py | 6 +- services/modules/discord/tests.py | 255 +++++++++++++++++++++++++++ 4 files changed, 383 insertions(+), 7 deletions(-) diff --git a/alliance_auth/tests/test_settings.py b/alliance_auth/tests/test_settings.py index 118b4a9f..52b09fa5 100644 --- a/alliance_auth/tests/test_settings.py +++ b/alliance_auth/tests/test_settings.py @@ -398,12 +398,12 @@ TEAMSPEAK3_PUBLIC_URL = os.environ.get('AA_TEAMSPEAK3_PUBLIC_URL', 'example.com' # DISCORD_CALLBACK_URL - oauth callback url # DISCORD_SYNC_NAMES - enable to force discord nicknames to be set to eve char name (bot needs Manage Nicknames permission) ###################################### -DISCORD_GUILD_ID = os.environ.get('AA_DISCORD_GUILD_ID', '') -DISCORD_BOT_TOKEN = os.environ.get('AA_DISCORD_BOT_TOKEN', '') -DISCORD_INVITE_CODE = os.environ.get('AA_DISCORD_INVITE_CODE', '') -DISCORD_APP_ID = os.environ.get('AA_DISCORD_APP_ID', '') -DISCORD_APP_SECRET = os.environ.get('AA_DISCORD_APP_SECRET', '') -DISCORD_CALLBACK_URL = os.environ.get('AA_DISCORD_CALLBACK_URL', 'http://example.com/discord_callback') +DISCORD_GUILD_ID = os.environ.get('AA_DISCORD_GUILD_ID', '0118999') +DISCORD_BOT_TOKEN = os.environ.get('AA_DISCORD_BOT_TOKEN', 'bottoken') +DISCORD_INVITE_CODE = os.environ.get('AA_DISCORD_INVITE_CODE', 'invitecode') +DISCORD_APP_ID = os.environ.get('AA_DISCORD_APP_ID', 'appid') +DISCORD_APP_SECRET = os.environ.get('AA_DISCORD_APP_SECRET', 'secret') +DISCORD_CALLBACK_URL = os.environ.get('AA_DISCORD_CALLBACK_URL', 'http://example.com/discord/callback') DISCORD_SYNC_NAMES = 'True' == os.environ.get('AA_DISCORD_SYNC_NAMES', 'False') ###################################### diff --git a/services/modules/discord/manager.py b/services/modules/discord/manager.py index e9b5547d..e0b369dc 100644 --- a/services/modules/discord/manager.py +++ b/services/modules/discord/manager.py @@ -5,9 +5,12 @@ import re from django.conf import settings from services.models import GroupCache from requests_oauthlib import OAuth2Session +from functools import wraps import logging import datetime +import time from django.utils import timezone +from django.core.cache import cache logger = logging.getLogger(__name__) @@ -31,6 +34,118 @@ SCOPES = [ GROUP_CACHE_MAX_AGE = datetime.timedelta(minutes=30) +class DiscordApiException(Exception): + def __init__(self): + super(Exception, self).__init__() + + +class DiscordApiTooBusy(DiscordApiException): + def __init__(self): + super(DiscordApiException, self).__init__() + self.message = "The Discord API is too busy to process this request now, please try again later." + + +class DiscordApiBackoff(DiscordApiException): + def __init__(self, retry_after, global_ratelimit): + super(DiscordApiException, self).__init__() + self.retry_after = retry_after + self.global_ratelimit = global_ratelimit + + +cache_time_format = '%Y-%m-%d %H:%M:%S' + + +def api_backoff(func): + """ + Decorator, Handles HTTP 429 "Too Many Requests" messages from the Discord API + If blocking=True is specified, this function will block and retry + the function up to max_retries=n times, or 3 if retries is not specified. + If the API call still recieves a backoff timer this function will raise + a exception. + If the caller chooses blocking=False, the decorator will raise a DiscordApiBackoff + exception and the caller can choose to retry after the given timespan available in + the retry_after property in seconds. + """ + + class PerformBackoff(Exception): + def __init__(self, retry_after, retry_datetime, global_ratelimit): + super(Exception, self).__init__() + self.retry_after = int(retry_after) + self.retry_datetime = retry_datetime + self.global_ratelimit = global_ratelimit + + @wraps(func) + def decorated(*args, **kwargs): + blocking = kwargs.get('blocking', False) + retries = kwargs.get('max_retries', 3) + + # Strip our parameters + if 'max_retries' in kwargs: + del kwargs['max_retries'] + if 'blocking' in kwargs: + del kwargs['blocking'] + + cache_key = 'DISCORD_BACKOFF_' + func.__name__ + cache_global_key = 'DISCORD_BACKOFF_GLOBAL' + + while retries > 0: + try: + try: + # Check global backoff first, then route backoff + existing_global_backoff = cache.get(cache_global_key) + existing_backoff = existing_global_backoff or cache.get(cache_key) + if existing_backoff: + backoff_timer = datetime.datetime.strptime(existing_backoff, cache_time_format) + if backoff_timer > datetime.datetime.utcnow(): + backoff_seconds = (backoff_timer - datetime.datetime.utcnow()).total_seconds() + logger.debug("Still under backoff for {} seconds, backing off" % backoff_seconds) + # Still under backoff + raise PerformBackoff( + retry_after=backoff_seconds, + retry_datetime=backoff_timer, + global_ratelimit=bool(existing_global_backoff) + ) + logger.debug("Calling API calling function") + func(*args, **kwargs) + break + except requests.HTTPError as e: + if e.response.status_code == 429: + if 'Retry-After' in e.response.headers: + retry_after = e.response.headers['Retry-After'] + else: + # Pick some random time + retry_after = 5 + + logger.info("Received backoff from API of %s seconds, handling" % retry_after) + # Store value in redis + backoff_until = (datetime.datetime.utcnow() + + datetime.timedelta(seconds=int(retry_after))) + global_backoff = bool(e.response.headers.get('X-RateLimit-Global', False)) + if global_backoff: + logger.info("Global backoff!!") + cache.set(cache_global_key, backoff_until.strftime(cache_time_format), retry_after) + else: + cache.set(cache_key, backoff_until.strftime(cache_time_format), retry_after) + raise PerformBackoff(retry_after=retry_after, retry_datetime=backoff_until, + global_ratelimit=global_backoff) + else: + # Not 429, re-raise + raise e + except PerformBackoff as bo: + # Sleep if we're blocking + if blocking: + logger.info("Blocking Back off from API calls for %s seconds" % bo.retry_after) + time.sleep(10 if bo.retry_after > 10 else bo.retry_after) + else: + # Otherwise raise exception and let caller handle the backoff + raise DiscordApiBackoff(retry_after=bo.retry_after, global_ratelimit=bo.global_ratelimit) + finally: + retries -= 1 + if retries == 0: + raise DiscordApiTooBusy() + return decorated + + class DiscordOAuthManager: def __init__(self): pass @@ -191,6 +306,7 @@ class DiscordOAuthManager: DiscordOAuthManager.__update_group_cache() @staticmethod + @api_backoff def update_groups(user_id, groups): custom_headers = {'content-type': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN} group_ids = [DiscordOAuthManager.__group_name_to_id(DiscordOAuthManager._sanitize_groupname(g)) for g in groups] @@ -199,3 +315,4 @@ class DiscordOAuthManager: r = requests.patch(path, headers=custom_headers, json=data) logger.debug("Received status code %s after setting user roles" % r.status_code) r.raise_for_status() + diff --git a/services/modules/discord/tasks.py b/services/modules/discord/tasks.py index 4f0ff61c..a791b494 100644 --- a/services/modules/discord/tasks.py +++ b/services/modules/discord/tasks.py @@ -9,7 +9,7 @@ from django.core.exceptions import ObjectDoesNotExist from eveonline.managers import EveManager from notifications import notify -from services.modules.discord.manager import DiscordOAuthManager +from services.modules.discord.manager import DiscordOAuthManager, DiscordApiBackoff from services.tasks import only_one from .models import DiscordUser @@ -74,6 +74,10 @@ class DiscordTasks: logger.debug("Updating user %s discord groups to %s" % (user, groups)) try: DiscordOAuthManager.update_groups(user.discord.uid, groups) + except DiscordApiBackoff as bo: + logger.info("Discord group sync API back off for %s, " + "retrying in %s seconds" % (user, bo.retry_after)) + raise task_self.retry(countdown=bo.retry_after) except Exception as e: if task_self: logger.exception("Discord group sync failed for %s, retrying in 10 mins" % user) diff --git a/services/modules/discord/tests.py b/services/modules/discord/tests.py index a6775ede..8dd34191 100644 --- a/services/modules/discord/tests.py +++ b/services/modules/discord/tests.py @@ -17,6 +17,10 @@ from alliance_auth.tests.auth_utils import AuthUtils from .auth_hooks import DiscordService from .models import DiscordUser from .tasks import DiscordTasks +from .manager import DiscordOAuthManager + +import requests_mock +import datetime MODULE_PATH = 'services.modules.discord' @@ -198,3 +202,254 @@ class DiscordViewsTestCase(TestCase): self.assertRedirects(response, expected_url='/en/services/', target_status_code=200) with self.assertRaises(ObjectDoesNotExist): discord_user = User.objects.get(pk=self.member.pk).discord + + +class DiscordManagerTestCase(TestCase): + def setUp(self): + pass + + def test__sanitize_groupname(self): + test_group_name = ' Group Name_Test_' + group_name = DiscordOAuthManager._sanitize_groupname(test_group_name) + + self.assertEqual(group_name, 'GroupName_Test') + + def test_generate_Bot_add_url(self): + from . import manager + bot_add_url = DiscordOAuthManager.generate_bot_add_url() + + auth_url = manager.AUTH_URL + real_bot_add_url = '{}?client_id=appid&scope=bot&permissions={}'.format(auth_url, manager.BOT_PERMISSIONS) + self.assertEqual(bot_add_url, real_bot_add_url) + + def test_generate_oauth_redirect_url(self): + from . import manager + import urllib + import sys + oauth_url = DiscordOAuthManager.generate_oauth_redirect_url() + + self.assertIn(manager.AUTH_URL, oauth_url) + self.assertIn('+'.join(manager.SCOPES), oauth_url) + self.assertIn(settings.DISCORD_APP_ID, oauth_url) + if sys.version_info[0] < 3: + # Py2 + self.assertIn(urllib.quote_plus(settings.DISCORD_CALLBACK_URL), oauth_url) + else: # Py3 + self.assertIn(urllib.parse.quote_plus(settings.DISCORD_CALLBACK_URL), oauth_url) + + @mock.patch(MODULE_PATH + '.manager.OAuth2Session') + def test__process_callback_code(self, oauth): + from . import manager + instance = oauth.return_value + instance.fetch_token.return_value = {'access_token': 'mywonderfultoken'} + + token = DiscordOAuthManager._process_callback_code('12345') + + self.assertTrue(oauth.called) + args, kwargs = oauth.call_args + self.assertEqual(args[0], settings.DISCORD_APP_ID) + self.assertEqual(kwargs['redirect_uri'], settings.DISCORD_CALLBACK_URL) + self.assertTrue(instance.fetch_token.called) + args, kwargs = instance.fetch_token.call_args + self.assertEqual(args[0], manager.TOKEN_URL) + self.assertEqual(kwargs['client_secret'], settings.DISCORD_APP_SECRET) + self.assertEqual(kwargs['code'], '12345') + self.assertEqual(token['access_token'], 'mywonderfultoken') + + @mock.patch(MODULE_PATH + '.manager.DiscordOAuthManager._process_callback_code') + @requests_mock.Mocker() + def test_add_user(self, oauth_token, m): + from . import manager + import json + + # Arrange + oauth_token.return_value = {'access_token': 'accesstoken'} + + headers = {'accept': 'application/json', 'authorization': 'Bearer accesstoken'} + + m.register_uri('POST', + manager.DISCORD_URL + '/invites/'+str(settings.DISCORD_INVITE_CODE), + request_headers=headers, + text='{}') + + m.register_uri('GET', + manager.DISCORD_URL + "/users/@me", + request_headers=headers, + text=json.dumps({'id': "123456"})) + + # Act + return_value = DiscordOAuthManager.add_user('abcdef') + + # Assert + self.assertEqual(return_value, '123456') + self.assertEqual(m.call_count, 2) + + @requests_mock.Mocker() + def test_delete_user(self, m): + from . import manager + import json + + # Arrange + headers = {'accept': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN} + + user_id = 12345 + request_url = '{}/guilds/{}/members/{}'.format(manager.DISCORD_URL, settings.DISCORD_GUILD_ID, user_id) + m.register_uri('DELETE', + request_url, + request_headers=headers, + text=json.dumps({})) + + # Act + result = DiscordOAuthManager.delete_user(user_id) + + # Assert + self.assertTrue(result) + + ### + # Test 404 (already deleted) + # Arrange + m.register_uri('DELETE', + request_url, + request_headers=headers, + status_code=404) + + # Act + result = DiscordOAuthManager.delete_user(user_id) + + # Assert + self.assertTrue(result) + + ### + # Test 500 (some random API error) + # Arrange + m.register_uri('DELETE', + request_url, + request_headers=headers, + status_code=500) + + # Act + result = DiscordOAuthManager.delete_user(user_id) + + # Assert + self.assertFalse(result) + + @requests_mock.Mocker() + def test_update_nickname(self, m): + from . import manager + import json + # Arrange + headers = {'content-type': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN} + + user_id = 12345 + request_url = '{}/guilds/{}/members/{}'.format(manager.DISCORD_URL, settings.DISCORD_GUILD_ID, user_id) + m.patch(request_url, + request_headers=headers) + + # Act + result = DiscordOAuthManager.update_nickname(user_id, 'somenick') + + # Assert + self.assertTrue(result) + + @mock.patch(MODULE_PATH + '.manager.DiscordOAuthManager._DiscordOAuthManager__get_group_cache') + @requests_mock.Mocker() + def test_update_groups(self, group_cache, m): + from . import manager + import json + + # Arrange + groups = ['Member', 'Blue', 'Special Group'] + + group_cache.return_value = [{'id': 111, 'name': 'Member'}, + {'id': 222, 'name': 'Blue'}, + {'id': 333, 'name': 'SpecialGroup'}, + {'id': 444, 'name': 'NotYourGroup'}] + + headers = {'content-type': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN} + user_id = 12345 + request_url = '{}/guilds/{}/members/{}'.format(manager.DISCORD_URL, settings.DISCORD_GUILD_ID, user_id) + + m.patch(request_url, + request_headers=headers) + + # Act + DiscordOAuthManager.update_groups(user_id, groups) + + # Assert + self.assertEqual(len(m.request_history), 1, 'Must be one HTTP call made') + history = json.loads(m.request_history[0].text) + self.assertIn('roles', history, "'The request must send JSON object with the 'roles' key") + self.assertIn(111, history['roles'], 'The group id 111 must be added to the request') + self.assertIn(222, history['roles'], 'The group id 222 must be added to the request') + self.assertIn(333, history['roles'], 'The group id 333 must be added to the request') + self.assertNotIn(444, history['roles'], 'The group id 444 must NOT be added to the request') + + @mock.patch(MODULE_PATH + '.manager.cache') + @mock.patch(MODULE_PATH + '.manager.DiscordOAuthManager._DiscordOAuthManager__get_group_cache') + @requests_mock.Mocker() + def test_update_groups_backoff(self, group_cache, djcache, m): + from . import manager + + # Arrange + groups = ['Member'] + group_cache.return_value = [{'id': 111, 'name': 'Member'}] + + headers = {'content-type': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN} + user_id = 12345 + request_url = '{}/guilds/{}/members/{}'.format(manager.DISCORD_URL, settings.DISCORD_GUILD_ID, user_id) + + djcache.get.return_value = None # No existing backoffs in cache + + m.patch(request_url, + request_headers=headers, + headers={'Retry-After': '200'}, + status_code=429) + + # Act & Assert + with self.assertRaises(manager.DiscordApiBackoff) as bo: + try: + DiscordOAuthManager.update_groups(user_id, groups, blocking=False) + except manager.DiscordApiBackoff as bo: + self.assertEqual(bo.retry_after, 200, 'Retry-After time must be equal to Retry-After set in header') + self.assertFalse(bo.global_ratelimit, 'global_ratelimit must be False') + raise bo + + self.assertTrue(djcache.set.called) + args, kwargs = djcache.set.call_args + self.assertEqual(args[0], 'DISCORD_BACKOFF_update_groups') + self.assertTrue(datetime.datetime.strptime(args[1], manager.cache_time_format) > datetime.datetime.now()) + + @mock.patch(MODULE_PATH + '.manager.cache') + @mock.patch(MODULE_PATH + '.manager.DiscordOAuthManager._DiscordOAuthManager__get_group_cache') + @requests_mock.Mocker() + def test_update_groups_global_backoff(self, group_cache, djcache, m): + from . import manager + + # Arrange + groups = ['Member'] + group_cache.return_value = [{'id': 111, 'name': 'Member'}] + + headers = {'content-type': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN} + user_id = 12345 + request_url = '{}/guilds/{}/members/{}'.format(manager.DISCORD_URL, settings.DISCORD_GUILD_ID, user_id) + + djcache.get.return_value = None # No existing backoffs in cache + + m.patch(request_url, + request_headers=headers, + headers={'Retry-After': '200', 'X-RateLimit-Global': 'true'}, + status_code=429) + + # Act & Assert + with self.assertRaises(manager.DiscordApiBackoff) as bo: + try: + DiscordOAuthManager.update_groups(user_id, groups, blocking=False) + except manager.DiscordApiBackoff as bo: + self.assertEqual(bo.retry_after, 200, 'Retry-After time must be equal to Retry-After set in header') + self.assertTrue(bo.global_ratelimit, 'global_ratelimit must be True') + raise bo + + self.assertTrue(djcache.set.called) + args, kwargs = djcache.set.call_args + self.assertEqual(args[0], 'DISCORD_BACKOFF_GLOBAL') + self.assertTrue(datetime.datetime.strptime(args[1], manager.cache_time_format) > datetime.datetime.now())