Cache signing keys

jwt/jwks_client.py

@@ -1,6 +1,6 @@
 import json
 import urllib.request
-from typing import Any, List
+from typing import Any, Dict, List, Tuple
 
 from .api_jwk import PyJWK, PyJWKSet
 from .api_jwt import decode_complete as decode_token
@@ -10,6 +10,8 @@
 class PyJWKClient:
     def __init__(self, uri: str):
         self.uri = uri
+        # Map (uri, kid) to signing keys
+        self._known_signing_keys: Dict[Tuple[str, str], PyJWK] = {}
 
     def fetch_data(self) -> Any:
         with urllib.request.urlopen(self.uri) as response:
@@ -33,10 +35,13 @@ def get_signing_keys(self) -> List[PyJWK]:
         return signing_keys
 
     def get_signing_key(self, kid: str) -> PyJWK:
+        if (self.uri, kid) in self._known_signing_keys:
+            return self._known_signing_keys[(self.uri, kid)]
         signing_keys = self.get_signing_keys()
         signing_key = None
 
         for key in signing_keys:
+            self._known_signing_keys[(self.uri, key.key_id)] = key
             if key.key_id == kid:
                 signing_key = key
                 break

tests/test_jwks_client.py

@@ -37,7 +37,7 @@ def mocked_response(data):
         response.__exit__ = mock.Mock()
         response.read.side_effect = [json.dumps(data)]
         urlopen_mock.return_value = response
-        yield
+        yield urlopen_mock
 
 
 @crypto_required
@@ -88,6 +88,35 @@ def test_get_signing_key(self):
         assert signing_key.key_id == kid
         assert signing_key.public_key_use == "sig"
 
+    def test_get_signing_key_caches_result(self):
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+        kid = "NEE1QURBOTM4MzI5RkFDNTYxOTU1MDg2ODgwQ0UzMTk1QjYyRkRFQw"
+
+        with mocked_response(RESPONSE_DATA) as network_response:
+            jwks_client = PyJWKClient(url)
+            jwks_client.get_signing_key(kid)
+            jwks_client.get_signing_key(kid)
+
+        assert network_response.call_count == 1
+
+    def test_get_signing_key_cache_distinguishes_uris(self):
+        url_1 = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+        url_2 = "https://dev-88evx9ru.auth0.com/.well-known/jwks.json"
+        kid = "NEE1QURBOTM4MzI5RkFDNTYxOTU1MDg2ODgwQ0UzMTk1QjYyRkRFQw"
+
+        with mocked_response(RESPONSE_DATA):
+            jwks_client = PyJWKClient(url_1)
+            jwks_client.get_signing_key(kid)
+
+        # mocked_response does not allow urllib.request.urlopen to be called twice
+        # so a second mock is needed
+        with mocked_response(RESPONSE_DATA):
+            jwks_client.uri = url_2
+            jwks_client.get_signing_key(kid)
+
+        for uri_kid_pair in ((url_1, kid), (url_2, kid)):
+            assert uri_kid_pair in jwks_client._known_signing_keys
+
     def test_get_signing_key_from_jwt(self):
         token = "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsImtpZCI6Ik5FRTFRVVJCT1RNNE16STVSa0ZETlRZeE9UVTFNRGcyT0Rnd1EwVXpNVGsxUWpZeVJrUkZRdyJ9.eyJpc3MiOiJodHRwczovL2Rldi04N2V2eDlydS5hdXRoMC5jb20vIiwic3ViIjoiYVc0Q2NhNzl4UmVMV1V6MGFFMkg2a0QwTzNjWEJWdENAY2xpZW50cyIsImF1ZCI6Imh0dHBzOi8vZXhwZW5zZXMtYXBpIiwiaWF0IjoxNTcyMDA2OTU0LCJleHAiOjE1NzIwMDY5NjQsImF6cCI6ImFXNENjYTc5eFJlTFdVejBhRTJINmtEME8zY1hCVnRDIiwiZ3R5IjoiY2xpZW50LWNyZWRlbnRpYWxzIn0.PUxE7xn52aTCohGiWoSdMBZGiYAHwE5FYie0Y1qUT68IHSTXwXVd6hn02HTah6epvHHVKA2FqcFZ4GGv5VTHEvYpeggiiZMgbxFrmTEY0csL6VNkX1eaJGcuehwQCRBKRLL3zKmA5IKGy5GeUnIbpPHLHDxr-GXvgFzsdsyWlVQvPX2xjeaQ217r2PtxDeqjlf66UYl6oY6AqNS8DH3iryCvIfCcybRZkc_hdy-6ZMoKT6Piijvk_aXdm7-QQqKJFHLuEqrVSOuBqqiNfVrG27QzAPuPOxvfXTVLXL2jek5meH6n-VWgrBdoMFH93QEszEDowDAEhQPHVs0xj7SIzA"
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"