jwk_key is reparsed 3x per login (and it's expensive)

[cakemanny]

Describe the bug

The OIDC private key is parsed multiple times on each OpenID Connect login. (3 times by our specific observations)

If we zoom in on the flame graph for a request to the token endpoint, we see that there are 3 calls to JWK.from_pem, twice before signing the ID Token and once again afterwards. In total 81% of the time is wasted parsing and loading the private key from its PEM format.

To Reproduce

Terse version:

• Assume django application and django-oauth-toolkit installed
• Follow steps in the readme to set up / configure Open ID Connect: https://django-oauth-toolkit.readthedocs.io/en/latest/oidc.html#configuration
• Create an application that uses RSA as the algorithm
• Perform an oauth login

Verbose version:

mkdir dot-jwk-bug && cd dot-jwk-bug
python3 -m venv .venv
source .venv/bin/activate
python3 -m pip install django django-oauth-toolkit py-spy
django-admin startproject mysite

Follow DOT installation steps: https://django-oauth-toolkit.readthedocs.io/en/latest/install.html

Edit mysite/mysite/settings.py appending:

OAUTH2_PROVIDER = {
    "OIDC_ENABLED": True,
    "OIDC_RSA_PRIVATE_KEY": os.environ.get("OIDC_RSA_PRIVATE_KEY"),
    "PKCE_REQUIRED": False,  # to simplify repro process
    "SCOPES": {
        "openid": "OpenID Connect scope",
    },
}

openssl genrsa -out oidc.key 4096
export OIDC_RSA_PRIVATE_KEY="$(cat oidc.key)"

(cd mysite
python manage.py migrate
python manage.py createsuperuser
)

Create the account with any uninteresting credentials you will be able to remember

(cd mysite && python manage.py runserver)

• Go to http://127.0.0.1:8000/admin/oauth2_provider/application/ ,
• Login with the super user created earlier
• Create new application
  • Client id: dot_jwk_bug
  • Redirect uris list: http://127.0.0.1:9000/
  • Client type: Public
  • Authorization grant type: Authorization code
  • Copy the Client secret
  • Skip authorization: checked (only to simply reproduction process)
  • Algorithm: RSA with SHA-2 256

export CLIENT_ID=dot_jwk_bug
export CLIENT_SECRET=[client secret here]

Retrieve an oauth access token

We run netcat to listen for the authorisation code:

export CODE=$(echo $'HTTP/1.0 200 OK\r\n\r\nall good' | nc -l 127.0.0.1 9000 | awk '/^GET/ { print $2 }' | awk -F '=' '{ print $2 }')

While that command waits go to the following URL in the the browser where still logged in as the superuser:

http://127.0.0.1:8000/o/authorize/?client_id=dot_jwk_bug&response_type=code&redirect_uri=http://127.0.0.1:9000/&scope=openid

Return to the terminal and finally retrieve an access token:

curl 'http://127.0.0.1:8000/o/token/' -d 'code='"${CODE}"'&grant_type=authorization_code&client_id=dot_jwk_bug&client_secret='"${CLIENT_SECRET}"'&redirect_uri=http://127.0.0.1:9000/'

In order to get the flame chart shown above one can run py-spy while doing this access token retrieval process and ctrl-c once done

sudo py-spy record -o dot-profile.svg -p [PID for the python process]

Expected behaviour

The OIDC private key to be parsed exactly once at application startup, and then never again.

Version

• I have tested with the latest published release and it's still a problem.
• I have tested with the master branch and it's still a problem.

Additional context

In case it's useful to point out, this is sitting behind a property 🤷

https://github.com/jazzband/django-oauth-toolkit/blob/769c0a2fd668f1a0dd3ca80ef8bfd76f8082eb57/oauth2_provider/models.py#L217-L225

And we can already see a couple of the places where it is accessed and thus recalculated:
https://github.com/jazzband/django-oauth-toolkit/blob/769c0a2fd668f1a0dd3ca80ef8bfd76f8082eb57/oauth2_provider/oauth2_validators.py#L838-L844

[n2ygk]

Thanks for the report. It would be great if you could submit a PR to fix it.

[cakemanny]

Happily. I could do with some direction though.

I ruled out caching the jwk_key property because the model instance is instantiated on each request.

I thought about the idea of adding some coercion to the OAuth2ProviderSettings in oauth2_provider/settings.py, a bit like what is being done for these class override settings
https://github.com/jazzband/django-oauth-toolkit/blob/9dd1033f65fa251137ac2d478d2a57028534d964/oauth2_provider/settings.py#L208-L210
I worry it makes the code surprising though.

Another option might be to create an equivalent lower cased version: i.e. accessed like so: oauth2_settings.oidc_rsa_private_key.
I think that solution becomes a bit more complex, since there is already some caching and reloading logic for the settings:
https://github.com/jazzband/django-oauth-toolkit/blob/9dd1033f65fa251137ac2d478d2a57028534d964/oauth2_provider/settings.py#L265-L270
There would need to be a oidc_rsa_private_keys_inactive as well to go with it.

---

I think the simplest however would be to write a wrapper function around the JWK.from_pem using functools.cache to cache the loaded key.
Roughly:

diff --git a/oauth2_provider/views/oidc.py b/oauth2_provider/views/oidc.py
index 38560ae..6019713 100644
--- a/oauth2_provider/views/oidc.py
+++ b/oauth2_provider/views/oidc.py
@@ -1,5 +1,6 @@
 import json
 from urllib.parse import urlparse
+import functools
 
 from django.http import HttpResponse, JsonResponse
 from django.urls import reverse
@@ -74,6 +75,10 @@ class ConnectDiscoveryInfoView(OIDCOnlyMixin, View):
         return response
 
 
+@functools.cache
+def jwk_from_pem(pem_string):
+    return jwk.JWK.from_pem(pem_string.encode('utf-8'))
+
 class JwksInfoView(OIDCOnlyMixin, View):
     """
     View used to show oidc json web key set document
@@ -86,7 +91,7 @@ class JwksInfoView(OIDCOnlyMixin, View):
                 oauth2_settings.OIDC_RSA_PRIVATE_KEY,
                 *oauth2_settings.OIDC_RSA_PRIVATE_KEYS_INACTIVE,
             ]:
-                key = jwk.JWK.from_pem(pem.encode("utf8"))
+                key = jwk_from_pem(pem)
                 data = {"alg": "RS256", "use": "sig", "kid": key.thumbprint()}
                 data.update(json.loads(key.export_public()))
                 keys.append(data)

I'm not quite sure which module it ought to live in though. It needs to be imported and used by both the JwksInfoView as well as by the jwk_key property on the model.

[tonial]

Wouldn't using Django's django.utils.functional.cached_property instead of property do the trick ?

[cakemanny]

Hi @tonial, I don't think it goes far enough, since the Application model instance is instantiated for each request.
It would eliminate one PEM load per request, leave 2 too many still from my perspective.

On my machine, that means the request time goes from around 1.5s to 1.2s

The options I mentioned that I started investigating reduce request time to around 100ms (most of which then becomes of the client secret using pbkdf2). It seems the signing gets faster once the key has already been used.

Flame graph for memoized PEM load

As I write this, I am now just noticing that most of the time spent signing for the first time after creating the JWK object is extracting the private key out (_get_private_key (jwcrypto/jwk.py:896)) . Which I think is an interesting observation.