Fix JWT retry logic (#479)

Author: sujaygarlanka

HISTORY.rst

@@ -8,6 +8,8 @@ Upcoming
 - Fixed bug in get_admin_events function which caused errors when the optional event_types parameter was omitted.
 - Added support for more attribute parameters when uploading new files and new versions of existing files.
 - Combined preflight check and lookup of accelerator URL into a single request for uploads.
+- Fixed JWT retry logic so the a new JTI claim is generated on each retry
+- Fixed retry logic so when a `Retry-After` header is passed back from the API, the SDK waits for the amount of time specified in the header before retrying
 
 2.6.1 (2019-10-24)
 ++++++++++++++++++

boxsdk/auth/jwt_auth.py

@@ -6,13 +6,15 @@
 import json
 import random
 import string
+import time
 
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
 import jwt
 from six import binary_type, string_types, raise_from, text_type
 
+from ..config import API
 from ..exception import BoxOAuthException
 from .oauth2 import OAuth2
 from ..object.user import User
@@ -239,14 +241,30 @@ def _auth_with_jwt(self, sub, sub_type):
         :rtype:
             `unicode`
         """
-        try:
-            return self._construct_and_send_jwt_auth(sub, sub_type)
-        except BoxOAuthException as ex:
-            error_response = ex.network_response
-            box_datetime = self._get_date_header(error_response)
-            if box_datetime is not None and self._was_exp_claim_rejected_due_to_clock_skew(error_response):
-                return self._construct_and_send_jwt_auth(sub, sub_type, box_datetime)
-            raise
+        attempt_number = 0
+        jwt_time = None
+        while True:
+            try:
+                return self._construct_and_send_jwt_auth(sub, sub_type, jwt_time)
+            except BoxOAuthException as ex:
+                network_response = ex.network_response
+                code = network_response.status_code  # pylint: disable=maybe-no-member
+                box_datetime = self._get_date_header(network_response)
+
+                if attempt_number >= API.MAX_RETRY_ATTEMPTS:
+                    raise ex
+
+                if (code == 429 or code >= 500):
+                    jwt_time = None
+                elif box_datetime is not None and self._was_exp_claim_rejected_due_to_clock_skew(network_response):
+                    jwt_time = box_datetime
+                else:
+                    raise ex
+
+                time_delay = self._session.get_retry_after_time(attempt_number, network_response.headers.get('Retry-After', None))  # pylint: disable=maybe-no-member
+                time.sleep(time_delay)
+                attempt_number += 1
+                self._logger.debug('Retrying JWT request')
 
     @staticmethod
     def _get_date_header(network_response):

boxsdk/session/session.py

@@ -24,6 +24,7 @@ class Session(object):
 
     _retry_randomization_factor = 0.5
     _retry_base_interval = 1
+    _JWT_GRANT_TYPE = 'urn:ietf:params:oauth:grant-type:jwt-bearer'
 
     """
     Box API session. Provides automatic retry of failed requests.
@@ -268,7 +269,7 @@ def with_default_network_request_kwargs(self, extra_network_parameters):
     # We updated our retry strategy to use exponential backoff instead of the header returned from the API response.
     # This is something we can remove in latter major bumps.
     # pylint: disable=unused-argument
-    def _get_retry_after_time(self, attempt_number, retry_after_header):
+    def get_retry_after_time(self, attempt_number, retry_after_header):
         """
         Get the amount of time to wait before retrying the API request, using the attempt number that failed to
         calculate the retry time for the next retry attempt.
@@ -285,6 +286,11 @@ def _get_retry_after_time(self, attempt_number, retry_after_header):
         :return:                        Number of seconds to wait before retrying.
         :rtype:                         `Number`
         """
+        if retry_after_header is not None:
+            try:
+                return int(retry_after_header)
+            except (ValueError, TypeError):
+                pass
         min_randomization = 1 - self._retry_randomization_factor
         max_randomization = 1 + self._retry_randomization_factor
         randomization = (random.uniform(0, 1) * (max_randomization - min_randomization)) + min_randomization
@@ -388,7 +394,7 @@ def _prepare_and_send_request(
         network_response = self._send_request(request, **kwargs)
 
         while True:
-            retry = self._get_retry_request_callable(network_response, attempt_number, request)
+            retry = self._get_retry_request_callable(network_response, attempt_number, request, **kwargs)
 
             if retry is None or attempt_number >= API.MAX_RETRY_ATTEMPTS:
                 break
@@ -401,7 +407,7 @@ def _prepare_and_send_request(
 
         return network_response
 
-    def _get_retry_request_callable(self, network_response, attempt_number, request):
+    def _get_retry_request_callable(self, network_response, attempt_number, request, **kwargs):
         """
         Get a callable that retries a request for certain types of failure.
 
@@ -429,11 +435,15 @@ def _get_retry_request_callable(self, network_response, attempt_number, request)
             `callable`
         """
         # pylint:disable=unused-argument
+        data = kwargs.get('data', {})
+        grant_type = None
+        if 'grant_type' in data:
+            grant_type = data['grant_type']
         code = network_response.status_code
-        if code in (202, 429) or code >= 500:
+        if (code in (202, 429) or code >= 500) and grant_type != self._JWT_GRANT_TYPE:
             return partial(
                 self._network_layer.retry_after,
-                self._get_retry_after_time(attempt_number, network_response.headers.get('Retry-After', None)),
+                self.get_retry_after_time(attempt_number, network_response.headers.get('Retry-After', None)),
                 self._send_request,
             )
         return None
@@ -547,7 +557,7 @@ def _renew_session(self, access_token_used):
         new_access_token, _ = self._oauth.refresh(access_token_used)
         return new_access_token
 
-    def _get_retry_request_callable(self, network_response, attempt_number, request):
+    def _get_retry_request_callable(self, network_response, attempt_number, request, **kwargs):
         """
         Get a callable that retries a request for certain types of failure.
 
@@ -581,6 +591,7 @@ def _get_retry_request_callable(self, network_response, attempt_number, request)
             network_response,
             attempt_number,
             request,
+            **kwargs
         )
 
     def _send_request(self, request, **kwargs):

test/unit/auth/test_jwt_auth.py

@@ -491,7 +491,7 @@ def unsuccessful_jwt_response(box_datetime, status_code, error_description, incl
 @pytest.mark.parametrize('jwt_algorithm', ('RS512',))
 @pytest.mark.parametrize('rsa_passphrase', (None,))
 @pytest.mark.parametrize('pass_private_key_by_path', (False,))
-@pytest.mark.parametrize('status_code', (400, 401, 429, 500))
+@pytest.mark.parametrize('status_code', (400, 401))
 @pytest.mark.parametrize('error_description', ('invalid box_sub_type claim', 'invalid kid', "check the 'exp' claim"))
 @pytest.mark.parametrize('error_code', ('invalid_grant', 'bad_request'))
 @pytest.mark.parametrize('include_date_header', (True, False))
@@ -512,8 +512,132 @@ def test_auth_retry_for_invalid_exp_claim(
                     auth.authenticate_instance(enterprise_id)
             else:
                 auth.authenticate_instance(enterprise_id)
-            expected_calls = [call(enterprise_id, 'enterprise')]
+            expected_calls = [call(enterprise_id, 'enterprise', None)]
             if expect_auth_retry:
                 expected_calls.append(call(enterprise_id, 'enterprise', box_datetime.replace(microsecond=0, tzinfo=None)))
             assert len(mock_send_jwt.mock_calls) == len(expected_calls)
             mock_send_jwt.assert_has_calls(expected_calls)
+
+
+@pytest.mark.parametrize('jwt_algorithm', ('RS512',))
+@pytest.mark.parametrize('rsa_passphrase', (None,))
+@pytest.mark.parametrize('pass_private_key_by_path', (False,))
+@pytest.mark.parametrize('status_code', (429,))
+@pytest.mark.parametrize('error_description', ('Request rate limit exceeded',))
+@pytest.mark.parametrize('error_code', ('rate_limit_exceeded',))
+@pytest.mark.parametrize('include_date_header', (False,))
+def test_auth_retry_for_rate_limit_error(
+        jwt_auth_init_mocks,
+        unsuccessful_jwt_response,
+):
+    # pylint:disable=redefined-outer-name
+    enterprise_id = 'fake_enterprise_id'
+    with jwt_auth_init_mocks(assert_authed=False) as params:
+        auth = params[0]
+        with patch.object(auth, '_construct_and_send_jwt_auth') as mock_send_jwt:
+            side_effect = []
+            expected_calls = []
+            # Retries multiple times, but less than max retries. Then succeeds when it gets a token.
+            for _ in range(API.MAX_RETRY_ATTEMPTS - 2):
+                side_effect.append(BoxOAuthException(429, network_response=unsuccessful_jwt_response))
+                expected_calls.append(call(enterprise_id, 'enterprise', None))
+            side_effect.append('jwt_token')
+            expected_calls.append(call(enterprise_id, 'enterprise', None))
+            mock_send_jwt.side_effect = side_effect
+
+            auth.authenticate_instance(enterprise_id)
+            assert len(mock_send_jwt.mock_calls) == len(expected_calls)
+            mock_send_jwt.assert_has_calls(expected_calls)
+
+
+@pytest.mark.parametrize('jwt_algorithm', ('RS512',))
+@pytest.mark.parametrize('rsa_passphrase', (None,))
+@pytest.mark.parametrize('pass_private_key_by_path', (False,))
+@pytest.mark.parametrize('status_code', (429,))
+@pytest.mark.parametrize('error_description', ('Request rate limit exceeded',))
+@pytest.mark.parametrize('error_code', ('rate_limit_exceeded',))
+@pytest.mark.parametrize('include_date_header', (False,))
+def test_auth_max_retries_for_rate_limit_error(
+        jwt_auth_init_mocks,
+        unsuccessful_jwt_response,
+):
+    # pylint:disable=redefined-outer-name
+    enterprise_id = 'fake_enterprise_id'
+    with jwt_auth_init_mocks(assert_authed=False) as params:
+        auth = params[0]
+        with patch.object(auth, '_construct_and_send_jwt_auth') as mock_send_jwt:
+            side_effect = []
+            expected_calls = []
+            # Retries max number of times, then throws the error
+            for _ in range(API.MAX_RETRY_ATTEMPTS + 1):
+                side_effect.append(BoxOAuthException(429, network_response=unsuccessful_jwt_response))
+                expected_calls.append(call(enterprise_id, 'enterprise', None))
+            mock_send_jwt.side_effect = side_effect
+
+            with pytest.raises(BoxOAuthException) as error:
+                auth.authenticate_instance(enterprise_id)
+            assert error.value.status == 429
+            assert len(mock_send_jwt.mock_calls) == len(expected_calls)
+            mock_send_jwt.assert_has_calls(expected_calls)
+
+
+@pytest.mark.parametrize('jwt_algorithm', ('RS512',))
+@pytest.mark.parametrize('rsa_passphrase', (None,))
+@pytest.mark.parametrize('pass_private_key_by_path', (False,))
+@pytest.mark.parametrize('status_code', (500,))
+@pytest.mark.parametrize('error_description', ('Internal Server Error',))
+@pytest.mark.parametrize('error_code', ('internal_server_error',))
+@pytest.mark.parametrize('include_date_header', (False,))
+def test_auth_retry_for_internal_server_error(
+        jwt_auth_init_mocks,
+        unsuccessful_jwt_response,
+):
+    # pylint:disable=redefined-outer-name
+    enterprise_id = 'fake_enterprise_id'
+    with jwt_auth_init_mocks(assert_authed=False) as params:
+        auth = params[0]
+        with patch.object(auth, '_construct_and_send_jwt_auth') as mock_send_jwt:
+            side_effect = []
+            expected_calls = []
+            # Retries multiple times, but less than max retries. Then succeeds when it gets a token.
+            for _ in range(API.MAX_RETRY_ATTEMPTS - 2):
+                side_effect.append(BoxOAuthException(500, network_response=unsuccessful_jwt_response))
+                expected_calls.append(call(enterprise_id, 'enterprise', None))
+            side_effect.append('jwt_token')
+            expected_calls.append(call(enterprise_id, 'enterprise', None))
+            mock_send_jwt.side_effect = side_effect
+
+            auth.authenticate_instance(enterprise_id)
+            assert len(mock_send_jwt.mock_calls) == len(expected_calls)
+            mock_send_jwt.assert_has_calls(expected_calls)
+
+
+@pytest.mark.parametrize('jwt_algorithm', ('RS512',))
+@pytest.mark.parametrize('rsa_passphrase', (None,))
+@pytest.mark.parametrize('pass_private_key_by_path', (False,))
+@pytest.mark.parametrize('status_code', (500,))
+@pytest.mark.parametrize('error_description', ('Internal Server Error',))
+@pytest.mark.parametrize('error_code', ('internal_server_error',))
+@pytest.mark.parametrize('include_date_header', (False,))
+def test_auth_max_retries_for_internal_server_error(
+        jwt_auth_init_mocks,
+        unsuccessful_jwt_response,
+):
+    # pylint:disable=redefined-outer-name
+    enterprise_id = 'fake_enterprise_id'
+    with jwt_auth_init_mocks(assert_authed=False) as params:
+        auth = params[0]
+        with patch.object(auth, '_construct_and_send_jwt_auth') as mock_send_jwt:
+            side_effect = []
+            expected_calls = []
+            # Retries max number of times, then throws the error
+            for _ in range(API.MAX_RETRY_ATTEMPTS + 1):
+                side_effect.append(BoxOAuthException(500, network_response=unsuccessful_jwt_response))
+                expected_calls.append(call(enterprise_id, 'enterprise', None))
+            mock_send_jwt.side_effect = side_effect
+
+            with pytest.raises(BoxOAuthException) as error:
+                auth.authenticate_instance(enterprise_id)
+            assert error.value.status == 500
+            assert len(mock_send_jwt.mock_calls) == len(expected_calls)
+            mock_send_jwt.assert_has_calls(expected_calls)

test/unit/session/test_session.py

@@ -150,7 +150,7 @@ def test_box_session_retries_response_after_retry_after(
     assert box_response.status_code == 200
     assert len(mock_network_layer.retry_after.call_args_list) == 1
     assert isinstance(mock_network_layer.retry_after.call_args[0][0], Number)
-    assert round(mock_network_layer.retry_after.call_args[0][0], 4) == 1.18
+    assert round(mock_network_layer.retry_after.call_args[0][0], 4) == 1
 
 
 @pytest.mark.parametrize('test_method', [
@@ -300,7 +300,7 @@ def test_session_uses_local_config(box_session, mock_network_layer, generic_succ
 )
 def test_get_retry_after_time(box_session, attempt_number, retry_after_header, expected_result):
     with patch('random.uniform', return_value=0.68):
-        retry_time = box_session._get_retry_after_time(attempt_number, retry_after_header)  # pylint: disable=protected-access
+        retry_time = box_session.get_retry_after_time(attempt_number, retry_after_header)  # pylint: disable=protected-access
     retry_time = round(retry_time, 4)
     assert retry_time == expected_result