24 Jun, 2020

Flask-JWT-Extended × Flask-Login

Apparently I do webshit now

For the past few months, I’ve been working on building a backend for $STARTUP, with a bunch of friends. I’ll probably write in detail about it when we launch our beta. The backend is your bog standard REST API, built on Flask—if you didn’t guess from the title already.

Our existing codebase heavily relies on Flask-Login; it offers some pretty neat interfaces for dealing with users and their states. However, its default mode of operation—sessions—don’t really fit into a Flask app that’s really just an API. It’s not optimal. Besides, this is what JWTs were built for.

I won’t bother delving deep into JSON web tokens, but the general flow is like so:

The neat thing about tokens is you can store stuff in them—“claims”, as they’re called.

returning an access_token to the client

The access_token is sent to the client upon login. The idea is simple, perform your usual checks (username / password etc.) and login the user via flask_login.login_user. Generate an access token using flask_jwt_extended.create_access_token, store your user identity in it (and other claims) and return it to the user in your 200 response.

Here’s the excerpt from our codebase.

access_token = create_access_token(identity=email)
login_user(user, remember=request.json["remember"])
return good("Logged in successfully!", access_token=access_token)

But, for login_user to work, we need to setup a custom user loader to pull out the identity from the request and return the user object.

defining a custom user loader in Flask-Login

By default, Flask-Login handles user loading via the user_loader decorator, which should return a user object. However, since we want to pull a user object from the incoming request (the token contains it), we’ll have to write a custom user loader via the request_loader decorator.

# Checks the 'Authorization' header by default.
app.config["JWT_TOKEN_LOCATION"] = ["json"]

# Defaults to 'identity', but the spec prefers 'sub'.
app.config["JWT_IDENTITY_CLAIM"] = "sub"

@login.request_loader
def load_person_from_request(request):
    try:
        token = request.json["access_token"]
    except Exception:
        return None
    data = decode_token(token)
    # this can be your 'User' class
    person = PersonSignup.query.filter_by(email=data["sub"]).first()
    if person:
        return person
    return None

There’s just one mildly annoying thing to deal with, though. Flask-Login insists on setting a session cookie. We will have to disable this behaviour ourselves. And the best part? There’s no documentation for this—well there is, but it’s incomplete and points to deprecated functions.

To do this, we define a custom session interface, like so:

from flask.sessions import SecureCookieSessionInterface
from flask import g
from flask_login import user_loaded_from_request

@user_loaded_from_request.connect
def user_loaded_from_request(app, user=None):
    g.login_via_request = True


class CustomSessionInterface(SecureCookieSessionInterface):
    def should_set_cookie(self, *args, **kwargs):
        return False

    def save_session(self, *args, **kwargs):
        if g.get("login_via_request"):
            return
        return super(CustomSessionInterface, self).save_session(*args, **kwargs)


app.session_interface = CustomSessionInterface()

In essence, this checks the global store g for login_via_request and and doesn’t set a cookie in that case. I’ve submitted a PR upstream for this to be included in the docs (#514).

Questions or comments? Send an email.