Cache signing keys (jpadilla#611)

* Cache the result of get_signing_key

* Include URI in key for getting known signing keys

* Add test_get_signing_key_caches_result test

* Add test to make sure multiple uris are being distinguished in key caching

* Ignore URI in caching

* Use functools.lru_cache to cache signing keys

* Allow opting out of key caching

* Allow adjusting max cached keys

* Add jpadilla#611 change to CHANGELOG.rst

* Update CHANGELOG.rst

Co-authored-by: José Padilla <jpadilla@webapplicate.com>

Co-authored-by: José Padilla <jpadilla@webapplicate.com>

Author: stevenpitts

CHANGELOG.rst

@@ -18,6 +18,8 @@ Fixed
 Added
 ~~~~~
 
+- Add caching by default to PyJWKClient `#611 <https://github.com/jpadilla/pyjwt/pull/611>`__
+
 `v2.0.1 <https://github.com/jpadilla/pyjwt/compare/2.0.0...2.0.1>`__
 --------------------------------------------------------------------
 

jwt/jwks_client.py

@@ -1,5 +1,6 @@
 import json
 import urllib.request
+from functools import lru_cache
 from typing import Any, List
 
 from .api_jwk import PyJWK, PyJWKSet
@@ -8,8 +9,12 @@
 
 
 class PyJWKClient:
-    def __init__(self, uri: str):
+    def __init__(self, uri: str, cache_keys: bool = True, max_cached_keys: int = 16):
         self.uri = uri
+        if cache_keys:
+            # Cache signing keys
+            # Ignore mypy (https://github.com/python/mypy/issues/2427)
+            self.get_signing_key = lru_cache(maxsize=max_cached_keys)(self.get_signing_key)  # type: ignore
 
     def fetch_data(self) -> Any:
         with urllib.request.urlopen(self.uri) as response:

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,38 @@ 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"
+
+        jwks_client = PyJWKClient(url)
+
+        with mocked_response(RESPONSE_DATA):
+            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) as repeated_call:
+            jwks_client.get_signing_key(kid)
+
+        assert repeated_call.call_count == 0
+
+    def test_get_signing_key_does_not_cache_opt_out(self):
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+        kid = "NEE1QURBOTM4MzI5RkFDNTYxOTU1MDg2ODgwQ0UzMTk1QjYyRkRFQw"
+
+        jwks_client = PyJWKClient(url, cache_keys=False)
+
+        with mocked_response(RESPONSE_DATA):
+            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) as repeated_call:
+            jwks_client.get_signing_key(kid)
+
+        assert repeated_call.call_count == 1
+
     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"