fix(auth/cognito): cache JWKS and offload fetch via asyncio.to_thread

Author: fgogolli

src/orb/providers/aws/auth/cognito_strategy.py

@@ -2,11 +2,14 @@
 
 from __future__ import annotations
 
+import asyncio
+import time
 from base64 import urlsafe_b64decode
 from typing import TYPE_CHECKING, Any, Optional
 
 import boto3
 import jwt
+import requests
 from botocore.config import Config
 from botocore.exceptions import ClientError
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers
@@ -34,6 +37,11 @@
 class CognitoAuthStrategy(AuthPort):
     """Authentication strategy using AWS Cognito User Pools."""
 
+    # Class-level JWKS cache: maps jwks_url → {"fetched_at": float, "keys": list[dict]}
+    # Shared across all instances so pool restarts benefit from warm entries.
+    _jwks_cache: dict[str, dict[str, Any]] = {}
+    _cache_ttl_seconds: int = 3600
+
     def __init__(
         self,
         logger: LoggingPort,
@@ -311,8 +319,15 @@ async def _get_public_key(self, kid: str) -> Any:
         """
         Get public key from Cognito JWKS endpoint.
 
-        Fetches the JWKS, locates the entry matching ``kid``, and converts it
-        to a ``cryptography`` RSAPublicKey object suitable for use with PyJWT.
+        Fetches the JWKS (using the class-level cache; re-fetches after TTL),
+        locates the entry matching ``kid``, and converts it to a ``cryptography``
+        RSAPublicKey object suitable for use with PyJWT.
+
+        The synchronous ``requests.get`` call is offloaded via
+        ``asyncio.to_thread`` so the event loop is not blocked during the fetch.
+
+        On ``jwt.InvalidTokenError`` (non-RSA key) the cache entry is evicted so
+        that a key rotation does not lock out valid tokens indefinitely.
 
         Args:
             kid: Key ID from token header
@@ -324,18 +339,13 @@ async def _get_public_key(self, kid: str) -> Any:
             jwt.InvalidTokenError: If the matched key is not an RSA key
         """
         try:
-            # In production, you would cache JWKS and implement appropriate key rotation
-            # This is a simplified implementation
-            import requests
-
-            # Add timeout to prevent hanging connections (security best practice)
-            response = requests.get(self.jwks_url, timeout=30)
-            response.raise_for_status()
-            jwks = response.json()
+            jwks = await self._fetch_jwks_cached()
 
             for key in jwks.get("keys", []):
                 if key.get("kid") == kid:
                     if key.get("kty") != "RSA":
+                        # Evict the cache entry so a subsequent rotation is picked up.
+                        self._jwks_cache.pop(self.jwks_url, None)
                         raise jwt.InvalidTokenError(
                             f"Unsupported key type: {key.get('kty')!r}. Only RSA keys are supported."
                         )
@@ -351,6 +361,28 @@ async def _get_public_key(self, kid: str) -> Any:
             self._logger.error("Failed to get public key: %s", e)
             return None
 
+    async def _fetch_jwks_cached(self) -> dict[str, Any]:
+        """Return the JWKS document for this strategy's endpoint.
+
+        Serves from the class-level cache when the entry is younger than
+        ``_cache_ttl_seconds``; otherwise fetches fresh via
+        ``asyncio.to_thread`` and repopulates the cache.
+        """
+        cached = self._jwks_cache.get(self.jwks_url)
+        if cached is not None:
+            age = time.monotonic() - cached["fetched_at"]
+            if age < self._cache_ttl_seconds:
+                return cached["jwks"]  # type: ignore[return-value]
+
+        def _do_fetch() -> dict[str, Any]:
+            response = requests.get(self.jwks_url, timeout=30)
+            response.raise_for_status()
+            return response.json()  # type: ignore[no-any-return]
+
+        jwks: dict[str, Any] = await asyncio.to_thread(_do_fetch)
+        self._jwks_cache[self.jwks_url] = {"fetched_at": time.monotonic(), "jwks": jwks}
+        return jwks
+
     def _map_groups_to_roles(self, groups: list[str]) -> list[str]:
         """
         Map Cognito groups to application roles.

tests/providers/aws/unit/test_cognito_strategy.py

@@ -37,6 +37,9 @@ def _make_strategy(enabled: bool = True) -> Any:
 
     logger = _make_logger()
 
+    # Clear the class-level JWKS cache so each test starts from a clean state.
+    CognitoAuthStrategy._jwks_cache.clear()
+
     with patch("boto3.client") as _mock_boto3:
         _mock_boto3.return_value = MagicMock()
         strategy = CognitoAuthStrategy(
@@ -384,3 +387,67 @@ async def test_get_public_key_non_2xx_propagates_as_invalid_token():
 
     assert result.status == AuthStatus.INVALID
     assert "Unable to verify token signature" in (result.error_message or "")
+
+
+# ---------------------------------------------------------------------------
+# JWKS caching behaviour
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+@pytest.mark.unit
+async def test_jwks_fetch_is_cached():
+    """_get_public_key called twice with the same kid only fetches JWKS once."""
+    strategy = _make_strategy()
+
+    mock_response = MagicMock()
+    mock_response.json.return_value = {"keys": []}  # kid not found → returns None both times
+
+    with patch("requests.get", return_value=mock_response) as mock_get:
+        await strategy._get_public_key("some-kid")
+        await strategy._get_public_key("some-kid")
+
+    # The HTTP fetch should happen only once; second call should hit the cache.
+    mock_get.assert_called_once()
+
+
+@pytest.mark.asyncio
+@pytest.mark.unit
+async def test_jwks_cache_evicted_on_invalid_key_type():
+    """Cache entry is evicted when _get_public_key raises InvalidTokenError for a non-RSA key.
+
+    A subsequent call should re-fetch from the endpoint so that a key rotation
+    is picked up without requiring a process restart.
+    """
+    import jwt as pyjwt
+
+    strategy = _make_strategy()
+
+    ec_jwks = {
+        "keys": [
+            {
+                "kid": "ec-kid",
+                "kty": "EC",
+                "crv": "P-256",
+                "x": "abc",
+                "y": "def",
+            }
+        ]
+    }
+
+    mock_response = MagicMock()
+    mock_response.json.return_value = ec_jwks
+
+    with patch("requests.get", return_value=mock_response) as mock_get:
+        # First call: fetches JWKS and raises because key type is not RSA.
+        with pytest.raises(pyjwt.InvalidTokenError):
+            await strategy._get_public_key("ec-kid")
+
+        # Cache entry must have been evicted on the error.
+        assert strategy.jwks_url not in strategy._jwks_cache
+
+        # Second call: must re-fetch (not serve from cache).
+        with pytest.raises(pyjwt.InvalidTokenError):
+            await strategy._get_public_key("ec-kid")
+
+    assert mock_get.call_count == 2