How To Use PyJWT With Django In A Resource Server And Still Keep Parts Of Your Sanity

The OAuth2 spec cleanly separates the role of Authorization Server (AS) from that of Resource Server (RS). The role of the AS, and the whole OAUTH2 dance, is to get an access token that will be accepted by a RS.

Figure 1 from the OAuth2 spec RFC 6749, shows four parties. The Client communicates with the Owner to get Authorization, with the Authorization Server to get an Access Token, and then uses the Access Token with the Resource Server to get access to the service.
Figure 1 from the OAuth2 spec RFC 6749

It’s puzzling. It should be easy, nay, trivial, to implement the Resource Server side in Django, yet it’s not. There are several libraries whose description can be interpreted as “implementing OAuth2”, yet what they all do is implement the Authorization Server side. I want to consume access tokens, not issue them!

Donnie Darko "I made a new friend" meme template.
Patient: "I made a new OAuth2 Implementation for Python", Therapist "OAuth2 implementation, or authorization server", Patient (defeated look): "Authorization server"

(Then of course there’s djangorestframework-simplejwt whose sole, primary, and exclusive functionality is to implement the most stupid JWT pattern known to humankind.)

Now, in theory the access token could be anything. But common Authorization Server implementations (keycloak, authentik, various hosted services) have converged on issuing signed JSON Web Tokens. So what the resource server needs is to be configured with the key to verify the tokens. We could conceivably hard code the key in source or configuration, but that is a bit of a hassle, and anyway, this is the third decade of the third millennium, quite frankly we shouldn’t have to. All server implementations offer a JWKS endpoint where the currently valid keys can be queried (and even a full autodiscovery endpoint, to discover the JWKS endpoint). An implementation of a resource server should, in theory, only need to be pointed at the JWKS endpoint and everything should just work.

The PyJWT documentation has something like this for the purpose:

import jwt

token = "..."
jwks_client = jwt.PyJWKClient(url)
signing_key = jwks_client.get_signing_key_from_jwt(token)
data = jwt.decode(token, signing_key.key, algorithms=["RS256"])

We want to authorize requests. Every time a new request comes in. That’s what a Resource Server does. It uses the provided token to check authorization. And apparently the documentation seems to suggest the correct way to do this is to fetch the keys from the JWKS endpoint on every request.

WTF?

We’ll need some caching. The documentation is mum on the topic. The implementation however is not. Turns out, they have implemented a cache. Only, they have implemented it on the PyJWKClient object itself. And there’s no easy way to hook up a real cache (such as Django’s).

The usual flow for normal Python web frameworks is that no object survives from request to request. Each request gets a clean slate. They may run in the same process sequentially. In different threads in the same process. In multiple processes. Or even async in the same thread in the same process. With the given example code we would be hitting the authorization server JWKS endpoint for every incoming request, adding huge latencies to processing.

In order to retain even a shred of sanity, we have no choice but to turn the JWKClient into a kind of singleton. It looks like this:

import jwt
from django.conf import settings


_jwks_client: Optional[jwt.PyJWKClient] = None


def get_jwks_client() -> jwt.PyJWKClient:
    # PyJWKClient caches responses from the JWKS endpoint *inside* the PyJWKClient object
    global _jwks_client
    if _jwks_client:
        return _jwks_client
    _jwks_client = jwt.PyJWKClient(settings.JWKS_URI)
    return _jwks_client

With this definition in place you can get the signing key as signing_key = get_jwks_client().get_signing_key_from_jwt(token) and will at least get some caching within a process, until the server decides to spawn a new process.

Then, to hook up authentication into Django Rest Framework you’ll do something like this (where User.from_token needs to be something that can turn a verified JWT dict into a User object):

def authenticate_request(request):
    if header := get_authorization_header(request):
        match header.split():
            case b"Bearer", token_bytes:
                token = token_bytes.decode("us-ascii", errors="strict")
                signing_key = get_jwks_client().get_signing_key_from_jwt(token)
                data = jwt.decode(token, signing_key.key, algorithms=["RS256"])

                if data:
                    return User.from_token(data), data