Allow to specify user_claims at token creation (#229)

Author: jeanphix

flask_jwt_extended/jwt_manager.py

@@ -436,14 +436,12 @@ def encode_key_loader(self, callback):
         self._encode_key_callback = callback
         return callback
 
-    def _create_refresh_token(self, identity, expires_delta=None):
+    def _create_refresh_token(self, identity, expires_delta=None, user_claims=None):
         if expires_delta is None:
             expires_delta = config.refresh_expires
 
-        if config.user_claims_in_refresh_token:
+        if user_claims is None and config.user_claims_in_refresh_token:
             user_claims = self._user_claims_callback(identity)
-        else:
-            user_claims = None
 
         refresh_token = encode_refresh_token(
             identity=self._user_identity_callback(identity),
@@ -458,17 +456,20 @@ def _create_refresh_token(self, identity, expires_delta=None):
         )
         return refresh_token
 
-    def _create_access_token(self, identity, fresh=False, expires_delta=None):
+    def _create_access_token(self, identity, fresh=False, expires_delta=None, user_claims=None):
         if expires_delta is None:
             expires_delta = config.access_expires
 
+        if user_claims is None:
+            user_claims = self._user_claims_callback(identity)
+
         access_token = encode_access_token(
             identity=self._user_identity_callback(identity),
             secret=self._encode_key_callback(identity),
             algorithm=config.algorithm,
             expires_delta=expires_delta,
             fresh=fresh,
-            user_claims=self._user_claims_callback(identity),
+            user_claims=user_claims,
             csrf=config.csrf_protect,
             identity_claim_key=config.identity_claim_key,
             user_claims_key=config.user_claims_key,

flask_jwt_extended/utils.py

@@ -114,7 +114,7 @@ def _get_jwt_manager():
                            "application before using this method")
 
 
-def create_access_token(identity, fresh=False, expires_delta=None):
+def create_access_token(identity, fresh=False, expires_delta=None, user_claims=None):
     """
     Create a new access token.
 
@@ -134,13 +134,14 @@ def create_access_token(identity, fresh=False, expires_delta=None):
                           expiration. If this is None, it will use the
                           'JWT_ACCESS_TOKEN_EXPIRES` config value
                           (see :ref:`Configuration Options`)
+    :param user_claims: Optionnal JSON serializable to override user claims.
     :return: An encoded access token
     """
     jwt_manager = _get_jwt_manager()
-    return jwt_manager._create_access_token(identity, fresh, expires_delta)
+    return jwt_manager._create_access_token(identity, fresh, expires_delta, user_claims)
 
 
-def create_refresh_token(identity, expires_delta=None):
+def create_refresh_token(identity, expires_delta=None, user_claims=None):
     """
     Creates a new refresh token.
 
@@ -155,10 +156,11 @@ def create_refresh_token(identity, expires_delta=None):
                           expiration. If this is None, it will use the
                           'JWT_REFRESH_TOKEN_EXPIRES` config value
                           (see :ref:`Configuration Options`)
+    :param user_claims: Optionnal JSON serializable to override user claims.
     :return: An encoded refresh token
     """
     jwt_manager = _get_jwt_manager()
-    return jwt_manager._create_refresh_token(identity, expires_delta)
+    return jwt_manager._create_refresh_token(identity, expires_delta, user_claims)
 
 
 def has_user_loader():

tests/test_user_claims_loader.py

@@ -137,3 +137,41 @@ def add_claims(identity):
     response = test_client.get('/protected2', headers=make_headers(refresh_token))
     assert response.get_json() == {'foo': 'bar'}
     assert response.status_code == 200
+
+
+def test_user_claim_in_refresh_token_specified_at_creation(app):
+    app.config['JWT_CLAIMS_IN_REFRESH_TOKEN'] = True
+
+    with app.test_request_context():
+        refresh_token = create_refresh_token('username', user_claims={'foo': 'bar'})
+
+    test_client = app.test_client()
+    response = test_client.get('/protected2', headers=make_headers(refresh_token))
+    assert response.get_json() == {'foo': 'bar'}
+    assert response.status_code == 200
+
+
+def test_user_claims_in_access_token_specified_at_creation(app):
+    with app.test_request_context():
+        access_token = create_access_token('username', user_claims={'foo': 'bar'})
+
+    test_client = app.test_client()
+    response = test_client.get('/protected', headers=make_headers(access_token))
+    assert response.get_json() == {'foo': 'bar'}
+    assert response.status_code == 200
+
+
+def test_user_claims_in_access_token_specified_at_creation_override(app):
+    jwt = get_jwt_manager(app)
+
+    @jwt.user_claims_loader
+    def add_claims(identity):
+        return {'default': 'value'}
+
+    with app.test_request_context():
+        access_token = create_access_token('username', user_claims={'foo': 'bar'})
+
+    test_client = app.test_client()
+    response = test_client.get('/protected', headers=make_headers(access_token))
+    assert response.get_json() == {'foo': 'bar'}
+    assert response.status_code == 200