Sem-Ver: feature Reduce the time taken to generate a jwt by caching loaded private key instances

* Sem-Ver: feature Reduce the time taken to generate a jwt by passing through a load private key instance to jwt.encode.

Signed-off-by: David Black <dblack@atlassian.com>

* Sem-Ver: feature Introduce a private keys cache so as to support PrivateKeyRetriever that can/may provide different load method results.

Signed-off-by: David Black <dblack@atlassian.com>

Author: dbaxa

atlassian_jwt_auth/signer.py

@@ -3,6 +3,8 @@
 import random
 
 import jwt
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
 
 from atlassian_jwt_auth import algorithms
 from atlassian_jwt_auth import key
@@ -16,6 +18,7 @@ def __init__(self, issuer, private_key_retriever, **kwargs):
         self.lifetime = kwargs.get('lifetime', datetime.timedelta(hours=1))
         self.algorithm = kwargs.get('algorithm', 'RS256')
         self.subject = kwargs.get('subject', None)
+        self._private_keys_cache = dict()
 
         if self.algorithm not in set(
                 algorithms.get_permitted_algorithm_names()):
@@ -25,6 +28,25 @@ def __init__(self, issuer, private_key_retriever, **kwargs):
             raise ValueError("lifetime, '%s',exceeds the allowed 1 hour max" %
                              (self.lifetime))
 
+    def _obtain_private_key(self, key_identifier, private_key_pem):
+        """ returns a loaded instance of the given private key either from
+            cache or from the given private_key_pem.
+        """
+        priv_key = self._private_keys_cache.get(key_identifier.key_id, None)
+        if priv_key is not None:
+            return priv_key
+        if not isinstance(private_key_pem, bytes):
+            private_key_pem = private_key_pem.encode()
+        priv_key = serialization.load_pem_private_key(
+            private_key_pem,
+            password=None,
+            backend=default_backend()
+        )
+        if len(self._private_keys_cache) > 10:
+            self._private_keys_cache = dict()
+        self._private_keys_cache[key_identifier.key_id] = priv_key
+        return priv_key
+
     def _generate_claims(self, audience, **kwargs):
         """ returns a new dictionary of claims. """
         now = self._now()
@@ -48,9 +70,11 @@ def generate_jwt(self, audience, **kwargs):
         """ returns a new signed jwt for use. """
         key_identifier, private_key_pem = self.private_key_retriever.load(
             self.issuer)
+        private_key = self._obtain_private_key(
+            key_identifier, private_key_pem)
         return jwt.encode(
             self._generate_claims(audience, **kwargs),
-            key=private_key_pem,
+            key=private_key,
             algorithm=self.algorithm,
             headers={'kid': key_identifier.key_id})
 

atlassian_jwt_auth/tests/test_signer.py

@@ -2,6 +2,7 @@
 import unittest
 
 import mock
+from cryptography.hazmat.primitives import serialization
 
 import atlassian_jwt_auth
 from atlassian_jwt_auth.tests import utils
@@ -72,9 +73,16 @@ def test_generate_jwt(self, m_jwt_encode):
         jwt_auth_signer.generate_jwt(expected_aud)
         m_jwt_encode.assert_called_with(
             expected_claims,
-            key=self._private_key_pem,
+            key=mock.ANY,
             algorithm=self.algorithm,
             headers={'kid': expected_key_id})
+        for name, args, kwargs in m_jwt_encode.mock_calls:
+            call_private_key = kwargs['key'].private_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PrivateFormat.TraditionalOpenSSL,
+                encryption_algorithm=serialization.NoEncryption()
+            )
+            self.assertEqual(call_private_key, self._private_key_pem)
 
 
 class JWTAuthSignerRS256Test(