Refresh token with complex object

[kush99993s]

Hello everyone,

I am trying to do is refresh token with complex object. My code is as follows:

from flask import Blueprint, jsonify, request
from flask_jwt_extended import jwt_required, create_access_token, \
    jwt_refresh_token_required, create_refresh_token,\
    get_jwt_identity, get_jwt_claims
from utility import to_dict
from my_app.auth.models import User
from my_app import app
from flask_jwt_extended import JWTManager


auth = Blueprint("auth", __name__)
jwt = JWTManager(app)


# This is an example of a complex object that we could build
# a JWT from. In practice, this will likely be something
# like a SQLAlchemy instance.
class UserObject:
    def __init__(self, username, roles):
        self.username = username
        self.roles = roles


# Create a function that will be called whenever create_access_token
# is used. It will take whatever object is passed into the
# create_access_token method, and lets us define what custom claims
# should be added to the access token.
@jwt.user_claims_loader
def add_claims_to_access_token(user):
    return {'roles': user.roles}


# Create a function that will be called whenever create_access_token
# is used. It will take whatever object is passed into the
# create_access_token method, and lets us define what the identity
# of the access token should be.
@jwt.user_identity_loader
def user_identity_lookup(user):
    return user.username


@auth.route('/login', methods=['POST'])
def login():
    data = to_dict(request.form, request.args, request.get_json())
    username = data['username']
    password = data['password']

    if (username is None) or (password is None):
        return jsonify({"msg": "Bad username or password"}), 401

    log_user = User(username, password)

    if not log_user.IsUser:
        return jsonify({"msg": "Bad username or password"}), 401

    if len(log_user.dashboard_groups) == 0:
        return jsonify({"msg": "Forbidden"}), 403

    user = UserObject(log_user.username, log_user.roles)

    ret = {
        'access_token': create_access_token(identity=user),
        'refresh_token': create_refresh_token(identity=user)
    }
    return jsonify(ret), 200


# The jwt_refresh_token_required decorator insures a valid refresh
# token is present in the request before calling this endpoint. We
# can use the get_jwt_identity() function to get the identity of
# the refresh token, and use the create_access_token() function again
# to make a new access token for this identity.
@auth.route('/refresh', methods=['POST'])
@jwt_refresh_token_required
def refresh():
    user = UserObject(get_jwt_identity(), get_jwt_claims())
    ret = {
        'access_token': create_access_token(identity=user),
        'refresh_token': create_refresh_token(identity=user)
    }
    return jsonify(ret), 200


@auth.route('/protected', methods=['GET'])
@jwt_required
def protected():
    username = get_jwt_identity()
    roles = get_jwt_claims()
    print(roles)
    return jsonify(logged_in_as=username), 200

If I log in using my username and password, login function works just fine, I get my refresh token and access token. then If I hit protected endpoint, it works fine, it prints me all roles that I have store.

However, if I hit refresh endpoint, then hit protected endpoint, I don't see any roles printed. I believe when I try to connect to refresh endpoint. I loose all role related data. when I print(get_jwt_claims()), it doesn't print anything in refresh end point.

If I am not wrong, when I hist refresh endpoint then it is looking for roles base upon refresh token. which is null and that is creating issue.

[kush99993s]

if we look at _create_access_token function, which create new access token, it add user_claims=self._user_claims_callback(identity)

access_token = encode_access_token(
            identity=self._user_identity_callback(identity),
            secret=config.encode_key,
            algorithm=config.algorithm,
            expires_delta=expires_delta,
            fresh=fresh,
            user_claims=self._user_claims_callback(identity),
            csrf=config.csrf_protect,
            identity_claim=config.identity_claim
        )

Where if we look at _create_refresh_token, we are not doing that

 refresh_token = encode_refresh_token(
            identity=self._user_identity_callback(identity),
            secret=config.encode_key,
            algorithm=config.algorithm,
            expires_delta=expires_delta,
            csrf=config.csrf_protect,
            identity_claim=config.identity_claim
        )

if we add user_claims=self._user_claims_callback(identity), then we can pass user claim from refresh token to new token, otherwise we loose that information.

[vimalloc]

In most cases, people would solve this by pulling their whole user object out from the underlying datastore, and using that to re-create the access token when using the refresh token endpoint. Because the refresh token should be used so seldomly, in most cases the hit to your underlying datastore shouldn't be a performance concern. As a side note, you probably don't want to build another refresh token in your refresh endpoint, the same refresh token can be used for as long as it is not expired. In your example, this might look like:

@auth.route('/refresh', methods=['POST'])
@jwt_refresh_token_required
def refresh():
    refresh_user = User(get_jwt_identity(), None)
    user = UserObject(refresh_user.username, refresh_user.roles)
    ret = {
        'access_token': create_access_token(identity=user),
    }
    return jsonify(ret), 200

I could potentially look at adding a flag that would keep the claims on both the access and refresh tokens, but unless I am missing something I imagine that wouldn't be very useful for the common case. Would the above example work for your case?

[roubaeli]

hi everyone,

would my example be a valid use case for this feature? our API authentication is using third-party service and storing some needed information in the token because the call is quite heavy...we wanted to use refreshing of that token to avoid the need to make that call every few minutes when the access token expires which is obviously not possible if the refresh token needs to get all info again...

[vimalloc]

It sounds like this is a feature people want and would actually use. I'm in the middle of moving so it will likely be a couple weeks before I could get to this, but if someone else wanted to take a stab at it I would be happy to help out and get it merged. I'm thinking it should be an option so that we don't break backwords compatibility (app.config['JWT_CLAIM_IN_REFRESH_TOKEN'] = True or something along those lines).

[vimalloc]

Released in version 3.10.0. Thanks for the PR @roubaeli! 👍