Add user_loader feature (refs #49)

Author: vimalloc

flask_jwt_extended/__init__.py

@@ -6,7 +6,7 @@
 from .utils import (
     create_refresh_token, create_access_token, get_jwt_identity,
     get_jwt_claims, set_access_cookies, set_refresh_cookies,
-    unset_jwt_cookies, get_raw_jwt
+    unset_jwt_cookies, get_raw_jwt, get_current_user, current_user
 )
 from .blacklist import (
     revoke_token, unrevoke_token, get_stored_tokens, get_all_stored_tokens,

flask_jwt_extended/default_callbacks.py

@@ -74,3 +74,12 @@ def default_revoked_token_callback():
     return a general error message with a 401 status code
     """
     return jsonify({'msg': 'Token has been revoked'}), 401
+
+
+def default_user_loader_error_callback(identity):
+    """
+    By default, if a user_loader callback is defined and the callback
+    function returns None, we return a general error message with a 401
+    status code
+    """
+    return jsonify({'msg': "Error loading the user {}".format(identity)}), 401

flask_jwt_extended/exceptions.py

@@ -54,3 +54,11 @@ class FreshTokenRequired(JWTExtendedException):
     protected by fresh_jwt_required
     """
     pass
+
+
+class UserLoadError(JWTExtendedException):
+    """
+    Error raised when a user_loader callback function returns None, indicating
+    that it cannot or will not load a user for the given identity.
+    """
+    pass

flask_jwt_extended/jwt_manager.py

@@ -6,19 +6,18 @@
 from flask_jwt_extended.config import config
 from flask_jwt_extended.exceptions import (
     JWTDecodeError, NoAuthorizationError, InvalidHeaderError, WrongTokenError,
-    RevokedTokenError, FreshTokenRequired, CSRFError
+    RevokedTokenError, FreshTokenRequired, CSRFError, UserLoadError
 )
 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_unauthorized_callback, default_needs_fresh_token_callback,
+    default_revoked_token_callback, default_user_loader_error_callback
 )
 from flask_jwt_extended.tokens import (
-    encode_refresh_token, decode_jwt,
-    encode_access_token
+    encode_refresh_token, decode_jwt, encode_access_token
 )
+from flask_jwt_extended.utils import get_jwt_identity
 
 
 class JWTManager(object):
@@ -39,6 +38,8 @@ def __init__(self, app=None):
         self._unauthorized_callback = default_unauthorized_callback
         self._needs_fresh_token_callback = default_needs_fresh_token_callback
         self._revoked_token_callback = default_revoked_token_callback
+        self._user_loader_callback = None
+        self._user_loader_error_callback = default_user_loader_error_callback
 
         # Register this extension with the flask app now (if it is provided)
         if app is not None:
@@ -101,6 +102,14 @@ def handle_revoked_token_error(e):
         def handle_fresh_token_required(e):
             return self._needs_fresh_token_callback()
 
+        @app.errorhandler(UserLoadError)
+        def handler_user_load_error(e):
+            # The identity is already saved before this exception was raised,
+            # otherwise a different exception would be raised, which is why we
+            # can safely call get_jwt_identity() here
+            identity = get_jwt_identity()
+            return self._user_loader_error_callback(identity)
+
     @staticmethod
     def _set_default_configuration_options(app):
         """
@@ -244,6 +253,50 @@ def revoked_token_loader(self, callback):
         self._revoked_token_callback = callback
         return callback
 
+    def user_loader_callback_loader(self, callback):
+        """
+        Sets the callback method to be called to load a user on a protected
+        endpoint.
+
+        By default this is not is not used.
+
+        If a callback method is passed in here, it must take one argument,
+        which is the identity of the user to load. It must return the user
+        object, or None in the case of an error (which will cause the TODO
+        error handler to be hit)
+        """
+        self._user_loader_callback = callback
+        return callback
+
+    def user_loader_error_loader(self, callback):
+        """
+        Sets the callback method to be called if a user fails or is refused
+        to load when calling the _user_loader_callback function (indicated by
+        that function returning None)
+
+        The default implementation will return json:
+        '{"msg": "Error loading the user <identity>"}' with a 400 status code.
+
+        Callback must be a function that takes one argument, the identity of the
+        user who failed to load.
+        """
+        self._user_loader_error_callback = callback
+        return callback
+
+    def has_user_loader(self):
+        """
+        Returns True if a user_loader_callback has been defined in this
+        application, False otherwise
+        """
+        return self._user_loader_callback is not None
+
+    def user_loader(self, identity):
+        """
+        Calls the _user_loader_callback function (if it is defined) and returns
+        the resulting user from this callback.
+        """
+        return self._user_loader_callback(identity)
+
     def create_refresh_token(self, identity):
         """
         Creates a new refresh token

flask_jwt_extended/utils.py

@@ -1,4 +1,6 @@
 from flask import current_app
+from werkzeug.local import LocalProxy
+
 try:
     from flask import _app_ctx_stack as ctx_stack
 except ImportError:  # pragma: no cover
@@ -8,6 +10,10 @@
 from flask_jwt_extended.tokens import decode_jwt
 
 
+# Proxy to access the current user
+current_user = LocalProxy(lambda: get_current_user())
+
+
 def get_raw_jwt():
     """
     Returns the python dictionary which has all of the data in this JWT. If no
@@ -32,6 +38,15 @@ def get_jwt_claims():
     return get_raw_jwt().get('user_claims', {})
 
 
+def get_current_user():
+    """
+    Returns the loaded user from a user_loader callback in a protected endpoint.
+    If no user was loaded, or if no user_loader callback was defined, this will
+    return None
+    """
+    return getattr(ctx_stack.top, 'jwt_user', None)
+
+
 def get_jti(encoded_token):
     """
     Returns the JTI given the JWT encoded token
@@ -60,6 +75,16 @@ def create_refresh_token(*args, **kwargs):
     return jwt_manager.create_refresh_token(*args, **kwargs)
 
 
+def user_loader(*args, **kwargs):
+    jwt_manager = _get_jwt_manager()
+    return jwt_manager.user_loader(*args, **kwargs)
+
+
+def has_user_loader(*args, **kwargs):
+    jwt_manager = _get_jwt_manager()
+    return jwt_manager.has_user_loader(*args, **kwargs)
+
+
 def get_csrf_token(encoded_token):
     token = decode_jwt(encoded_token, config.decode_key, config.algorithm, csrf=True)
     return token['csrf']

flask_jwt_extended/view_decorators.py

@@ -11,9 +11,10 @@
 from flask_jwt_extended.config import config
 from flask_jwt_extended.exceptions import (
     InvalidHeaderError, NoAuthorizationError, WrongTokenError,
-    FreshTokenRequired, CSRFError
+    FreshTokenRequired, CSRFError, UserLoadError
 )
 from flask_jwt_extended.tokens import decode_jwt
+from flask_jwt_extended.utils import has_user_loader, user_loader
 
 
 def jwt_required(fn):
@@ -28,10 +29,9 @@ def jwt_required(fn):
     """
     @wraps(fn)
     def wrapper(*args, **kwargs):
-        # Save the jwt in the context so that it can be accessed later by
-        # the various endpoints that is using this decorator
         jwt_data = _decode_jwt_from_request(request_type='access')
         ctx_stack.top.jwt = jwt_data
+        _load_user(jwt_data['identity'])
         return fn(*args, **kwargs)
     return wrapper
 
@@ -49,15 +49,11 @@ def jwt_optional(fn):
     @wraps(fn)
     def wrapper(*args, **kwargs):
         try:
-            # If an acceptable JWT is found in the request, put it into
-            # the application context
             jwt_data = _decode_jwt_from_request(request_type='access')
             ctx_stack.top.jwt = jwt_data
+            _load_user(jwt_data['identity'])
         except NoAuthorizationError:
-            # Allow request to proceed if no authorization header is present
-            # in the request, but don't modify application context
             pass
-        # Return the decorated function in either case
         return fn(*args, **kwargs)
     return wrapper
 
@@ -78,9 +74,8 @@ def wrapper(*args, **kwargs):
         if not jwt_data['fresh']:
             raise FreshTokenRequired('Fresh token required')
 
-        # Save the jwt in the context so that it can be accessed later by
-        # the various endpoints that is using this decorator
         ctx_stack.top.jwt = jwt_data
+        _load_user(jwt_data['identity'])
         return fn(*args, **kwargs)
     return wrapper
 
@@ -93,14 +88,22 @@ def jwt_refresh_token_required(fn):
     """
     @wraps(fn)
     def wrapper(*args, **kwargs):
-        # Save the jwt in the context so that it can be accessed later by
-        # the various endpoints that is using this decorator
         jwt_data = _decode_jwt_from_request(request_type='refresh')
         ctx_stack.top.jwt = jwt_data
+        _load_user(jwt_data['identity'])
         return fn(*args, **kwargs)
     return wrapper
 
 
+def _load_user(identity):
+    if has_user_loader():
+        user = user_loader(identity)
+        if user is None:
+            raise UserLoadError("user_loader returned None for {}".format(identity))
+        else:
+            ctx_stack.top.jwt_user = user
+
+
 def _decode_jwt_from_headers():
     header_name = config.header_name
     header_type = config.header_type

tests/test_jwt_manager.py

@@ -32,7 +32,12 @@ def test_class_init(self):
     def test_default_user_claims_callback(self):
         identity = 'foobar'
         m = JWTManager(self.app)
-        assert m._user_claims_callback(identity) == {}
+        self.assertEqual(m._user_claims_callback(identity), {})
+
+    def test_default_user_identity_callback(self):
+        identity = 'foobar'
+        m = JWTManager(self.app)
+        self.assertEqual(m._user_identity_callback(identity), identity)
 
     def test_default_expired_token_callback(self):
         with self.app.test_request_context():
@@ -80,6 +85,24 @@ def test_default_revoked_token_callback(self):
             self.assertEqual(status_code, 401)
             self.assertEqual(data, {'msg': 'Token has been revoked'})
 
+    def test_default_user_loader_callback(self):
+        m = JWTManager(self.app)
+        self.assertEqual(m._user_loader_callback, None)
+
+    def test_default_user_loader_error_callback(self):
+        with self.app.test_request_context():
+            identity = 'foobar'
+            m = JWTManager(self.app)
+            result = m._user_loader_error_callback(identity)
+            status_code, data = self._parse_callback_result(result)
+
+            self.assertEqual(status_code, 401)
+            self.assertEqual(data, {'msg': 'Error loading the user foobar'})
+
+    def test_default_has_user_loader(self):
+        m = JWTManager(self.app)
+        self.assertEqual(m.has_user_loader(), False)
+
     def test_custom_user_claims_callback(self):
         identity = 'foobar'
         m = JWTManager(self.app)
@@ -159,3 +182,33 @@ def custom_revoken_token():
 
             self.assertEqual(status_code, 422)
             self.assertEqual(data, {'err': 'Nice knowing you!'})
+
+    def test_custom_user_loader(self):
+        with self.app.test_request_context():
+            m = JWTManager(self.app)
+
+            @m.user_loader_callback_loader
+            def custom_user_loader(identity):
+                if identity == 'foo':
+                    return None
+                return identity
+
+            identity = 'foobar'
+            result = m._user_loader_callback(identity)
+            self.assertEqual(result, identity)
+            self.assertEqual(m.has_user_loader(), True)
+
+    def test_custom_user_loader_error_callback(self):
+        with self.app.test_request_context():
+            m = JWTManager(self.app)
+
+            @m.user_loader_error_loader
+            def custom_user_loader_error(identity):
+                return jsonify({'msg': 'Not found'}), 404
+
+            identity = 'foobar'
+            result = m._user_loader_error_callback(identity)
+            status_code, data = self._parse_callback_result(result)
+
+            self.assertEqual(status_code, 404)
+            self.assertEqual(data, {'msg': 'Not found'})

tests/test_protected_endpoints.py

@@ -62,7 +62,6 @@ def partially_protected():
                 return jsonify({'msg': "protected hello world"})
             return jsonify({'msg': "unprotected hello world"})
 
-
     def _jwt_post(self, url, jwt):
         response = self.client.post(url, content_type='application/json',
                                     headers={'Authorization': 'Bearer {}'.format(jwt)})