Adds callback methods for verifying the user claims of an access token

refs #64 #70

Author: vimalloc

docs/changing_default_behavior.rst

@@ -43,6 +43,12 @@ Possible loader functions are:
     * - **user_loader_error_loader**
       - Function that is called when the user_loader callback function returns **None**
       - Takes one argument - The identity of the user who failed to load
+    * - **claims_verification_loader**
+      - Function that is called to verify the custom **user_claims** data. Must return True or False
+      - Takes one argument - The custom user_claims dict in an access token
+    * - **claims_verification_failed_loader**
+      - Function that is called when the user claims verification callback returns False
+      - None
 
 Dynamic token expires time
 ~~~~~~~~~~~~~~~~~~~~~~~~~~

flask_jwt_extended/default_callbacks.py

@@ -83,3 +83,18 @@ def default_user_loader_error_callback(identity):
     status code
     """
     return jsonify({'msg': "Error loading the user {}".format(identity)}), 401
+
+
+def default_claims_verification_callback(user_claims):
+    """
+    By default, we do not do any verification of the user claims.
+    """
+    return True
+
+
+def default_claims_verification_failed_callback():
+    """
+    By default, if the user claims verification failed, we return a generic
+    error message with a 400 status code
+    """
+    return jsonify({'msg': 'User claims verification failed'}), 400

flask_jwt_extended/exceptions.py

@@ -62,3 +62,11 @@ class UserLoadError(JWTExtendedException):
     that it cannot or will not load a user for the given identity.
     """
     pass
+
+
+class UserClaimsVerificationError(JWTExtendedException):
+    """
+    Error raised when the claims_verification_callback function returns False,
+    indicating that the expected user claims are invalid
+    """
+    pass

flask_jwt_extended/jwt_manager.py

@@ -5,13 +5,16 @@
 from flask_jwt_extended.config import config
 from flask_jwt_extended.exceptions import (
     JWTDecodeError, NoAuthorizationError, InvalidHeaderError, WrongTokenError,
-    RevokedTokenError, FreshTokenRequired, CSRFError, UserLoadError
+    RevokedTokenError, FreshTokenRequired, CSRFError, UserLoadError,
+    UserClaimsVerificationError
 )
 from flask_jwt_extended.default_callbacks import (
     default_expired_token_callback, default_user_claims_callback,
     default_user_identity_callback, default_invalid_token_callback,
     default_unauthorized_callback, default_needs_fresh_token_callback,
-    default_revoked_token_callback, default_user_loader_error_callback
+    default_revoked_token_callback, default_user_loader_error_callback,
+    default_claims_verification_callback,
+    default_claims_verification_failed_callback
 )
 from flask_jwt_extended.tokens import (
     encode_refresh_token, encode_access_token
@@ -40,6 +43,8 @@ def __init__(self, app=None):
         self._user_loader_callback = None
         self._user_loader_error_callback = default_user_loader_error_callback
         self._token_in_blacklist_callback = None
+        self._claims_verification_callback = default_claims_verification_callback
+        self._claims_verification_failed_callback = default_claims_verification_failed_callback
 
         # Register this extension with the flask app now (if it is provided)
         if app is not None:
@@ -110,6 +115,10 @@ def handler_user_load_error(e):
             identity = get_jwt_identity()
             return self._user_loader_error_callback(identity)
 
+        @app.errorhandler(UserClaimsVerificationError)
+        def handle_failed_user_claims_verification(e):
+            return self._claims_verification_failed_callback()
+
     @staticmethod
     def _set_default_configuration_options(app):
         """
@@ -296,6 +305,31 @@ def token_in_blacklist_loader(self, callback):
         self._token_in_blacklist_callback = callback
         return callback
 
+    def claims_verification_loader(self, callback):
+        """
+        Sets the callback function for checking if the custom user claims are
+        valid for this access token.
+
+        This callback function must take one parameter, which is the custom
+        user claims present in the access token. This callback function should
+        return True if the user claims are valid, False otherwise.
+        """
+        self._claims_verification_callback = callback
+        return callback
+
+    def claims_verification_failed_loader(self, callback):
+        """
+        Sets the callback method to be called if the user claims verification
+        method returns False, indicating that the user claims are not valid.
+
+        The default implementation will return the json:
+        '{"msg": "User claims verification failed"})' with a 400 status code
+
+        Callback must be a function that takes no arguments.
+        """
+        self._claims_verification_failed_callback = callback
+        return callback
+
     def create_refresh_token(self, identity, expires_delta=None):
         """
         Creates a new refresh token

flask_jwt_extended/utils.py

@@ -106,6 +106,11 @@ def token_in_blacklist(*args, **kwargs):
     return jwt_manager._token_in_blacklist_callback(*args, **kwargs)
 
 
+def verify_token_claims(*args, **kwargs):
+    jwt_manager = _get_jwt_manager()
+    return jwt_manager._claims_verification_callback(*args, **kwargs)
+
+
 def get_csrf_token(encoded_token):
     token = decode_jwt(
         encoded_token,

flask_jwt_extended/view_decorators.py

@@ -10,12 +10,13 @@
 from flask_jwt_extended.config import config
 from flask_jwt_extended.exceptions import (
     InvalidHeaderError, NoAuthorizationError, WrongTokenError,
-    FreshTokenRequired, CSRFError, UserLoadError, RevokedTokenError
+    FreshTokenRequired, CSRFError, UserLoadError, RevokedTokenError,
+    UserClaimsVerificationError
 )
 from flask_jwt_extended.tokens import decode_jwt
 from flask_jwt_extended.utils import (
     has_user_loader, user_loader, token_in_blacklist,
-    has_token_in_blacklist_callback
+    has_token_in_blacklist_callback, verify_token_claims
 )
 
 
@@ -207,6 +208,11 @@ def _decode_jwt_from_request(request_type):
     if decoded_token['type'] != request_type:
         raise WrongTokenError('Only {} tokens can access this endpoint'.format(request_type))
 
+    # Check if the custom claims in access tokens are valid
+    if request_type == 'access':
+        if not verify_token_claims(decoded_token['user_claims']):
+            raise UserClaimsVerificationError('user_claims verification failed')
+
     # If blacklisting is enabled, see if this token has been revoked
     if _token_blacklisted(decoded_token, request_type):
         raise RevokedTokenError('Token has been revoked')

tests/test_jwt_manager.py

@@ -214,3 +214,24 @@ def custom_user_loader_error(identity):
 
             self.assertEqual(status_code, 404)
             self.assertEqual(data, {'msg': 'Not found'})
+
+    def test_claims_verification(self):
+        with self.app.test_request_context():
+            m = JWTManager(self.app)
+
+            @m.claims_verification_loader
+            def user_claims_verification(claims):
+                return 'foo' in claims
+
+            @m.claims_verification_failed_loader
+            def user_claims_verification_failed():
+                return jsonify({'msg': 'Test'}), 404
+
+            result = m._claims_verification_callback({'bar': 'baz'})
+            self.assertEqual(result, False)
+
+            result = m._claims_verification_failed_callback()
+            status_code, data = self._parse_callback_result(result)
+
+            self.assertEqual(status_code, 404)
+            self.assertEqual(data, {'msg': 'Test'})

tests/test_protected_endpoints.py

@@ -22,7 +22,6 @@ def setUp(self):
         self.app.config['JWT_ALGORITHM'] = 'HS256'
         self.app.config['JWT_ACCESS_TOKEN_EXPIRES'] = timedelta(seconds=1)
         self.app.config['JWT_REFRESH_TOKEN_EXPIRES'] = timedelta(seconds=1)
-        self.app.config['JWT_IDENTITY_CLAIM'] = 'sub'
         self.jwt_manager = JWTManager(self.app)
         self.client = self.app.test_client()
 

tests/test_user_claims_verification.py

@@ -0,0 +1,87 @@
+import unittest
+
+from flask import Flask, jsonify, json
+
+from flask_jwt_extended import JWTManager, create_access_token, jwt_required
+
+
+class TestUserClaimsVerification(unittest.TestCase):
+
+    def setUp(self):
+        self.app = Flask(__name__)
+        self.app.secret_key = 'super=secret'
+        self.jwt_manager = JWTManager(self.app)
+        self.client = self.app.test_client()
+
+        @self.jwt_manager.claims_verification_loader
+        def claims_verification(user_claims):
+            expected_keys = ['foo', 'bar']
+            for key in expected_keys:
+                if key not in user_claims:
+                    return False
+            return True
+
+        @self.app.route('/auth/login', methods=['POST'])
+        def login():
+            ret = {'access_token': create_access_token('test')}
+            return jsonify(ret), 200
+
+        @self.app.route('/protected')
+        @jwt_required
+        def protected():
+            return jsonify({'msg': "hello world"})
+
+    def _jwt_get(self, url, jwt, header_name='Authorization', header_type='Bearer'):
+        header_type = '{} {}'.format(header_type, jwt).strip()
+        response = self.client.get(url, headers={header_name: header_type})
+        status_code = response.status_code
+        data = json.loads(response.get_data(as_text=True))
+        return status_code, data
+
+    def test_valid_user_claims(self):
+        @self.jwt_manager.user_claims_loader
+        def user_claims_callback(identity):
+            return {'foo': 'baz', 'bar': 'boom'}
+
+        response = self.client.post('/auth/login')
+        data = json.loads(response.get_data(as_text=True))
+        access_token = data['access_token']
+
+        status, data = self._jwt_get('/protected', access_token)
+        self.assertEqual(data, {'msg': 'hello world'})
+        self.assertEqual(status, 200)
+
+    def test_empty_claims_verification_error(self):
+        response = self.client.post('/auth/login')
+        data = json.loads(response.get_data(as_text=True))
+        access_token = data['access_token']
+
+        status, data = self._jwt_get('/protected', access_token)
+        self.assertEqual(data, {'msg': 'User claims verification failed'})
+        self.assertEqual(status, 400)
+
+    def test_bad_claims_verification_error(self):
+        @self.jwt_manager.user_claims_loader
+        def user_claims_callback(identity):
+            return {'super': 'banana'}
+
+        response = self.client.post('/auth/login')
+        data = json.loads(response.get_data(as_text=True))
+        access_token = data['access_token']
+
+        status, data = self._jwt_get('/protected', access_token)
+        self.assertEqual(data, {'msg': 'User claims verification failed'})
+        self.assertEqual(status, 400)
+
+    def test_bad_claims_custom_error_callback(self):
+        @self.jwt_manager.claims_verification_failed_loader
+        def user_claims_callback():
+            return jsonify({'foo': 'bar'}), 404
+
+        response = self.client.post('/auth/login')
+        data = json.loads(response.get_data(as_text=True))
+        access_token = data['access_token']
+
+        status, data = self._jwt_get('/protected', access_token)
+        self.assertEqual(data, {'foo': 'bar'})
+        self.assertEqual(status, 404)