Cache loading of JWK object from OIDC private key

Author: cakemanny

oauth2_provider/models.py

@@ -19,6 +19,7 @@
 from .generators import generate_client_id, generate_client_secret
 from .scopes import get_scopes_backend
 from .settings import oauth2_settings
+from .utils import jwk_from_pem
 from .validators import RedirectURIValidator, WildcardSet
 
 
@@ -234,7 +235,7 @@ def jwk_key(self):
         if self.algorithm == AbstractApplication.RS256_ALGORITHM:
             if not oauth2_settings.OIDC_RSA_PRIVATE_KEY:
                 raise ImproperlyConfigured("You must set OIDC_RSA_PRIVATE_KEY to use RSA algorithm")
-            return jwk.JWK.from_pem(oauth2_settings.OIDC_RSA_PRIVATE_KEY.encode("utf8"))
+            return jwk_from_pem(oauth2_settings.OIDC_RSA_PRIVATE_KEY)
         elif self.algorithm == AbstractApplication.HS256_ALGORITHM:
             return jwk.JWK(kty="oct", k=base64url_encode(self.client_secret))
         raise ImproperlyConfigured("This application does not support signed tokens")

oauth2_provider/utils.py

@@ -0,0 +1,12 @@
+import functools
+
+from jwcrypto import jwk
+
+
+@functools.lru_cache()
+def jwk_from_pem(pem_string):
+    """
+    A cached version of jwcrypto.JWK.from_pem.
+    Converting from PEM is expensive for large keys such as those using RSA.
+    """
+    return jwk.JWK.from_pem(pem_string.encode("utf-8"))

oauth2_provider/views/oidc.py

@@ -7,7 +7,7 @@
 from django.utils.decorators import method_decorator
 from django.views.decorators.csrf import csrf_exempt
 from django.views.generic import FormView, View
-from jwcrypto import jwk, jwt
+from jwcrypto import jwt
 from jwcrypto.common import JWException
 from jwcrypto.jws import InvalidJWSObject
 from jwcrypto.jwt import JWTExpired
@@ -30,6 +30,7 @@
     get_refresh_token_model,
 )
 from ..settings import oauth2_settings
+from ..utils import jwk_from_pem
 from .mixins import OAuthLibMixin, OIDCLogoutOnlyMixin, OIDCOnlyMixin
 
 
@@ -114,7 +115,7 @@ def get(self, request, *args, **kwargs):
                 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)