Periodically refresh JWT public key set (#88)

* Periodic JWT Public Key Refresh Implementation

Author: tmikula-dev

src/handlers/handler_token.py

@@ -21,6 +21,7 @@
 import base64
 import logging
 import os
+from datetime import datetime, timedelta, timezone
 from typing import Dict, Any, cast
 
 import jwt
@@ -41,10 +42,31 @@ class HandlerToken:
     HandlerToken manages token provider URL and public keys for JWT verification.
     """
 
+    _REFRESH_INTERVAL = timedelta(minutes=28)
+
     def __init__(self, config):
         self.provider_url: str = config.get(TOKEN_PROVIDER_URL, "")
         self.public_keys_url: str = config.get(TOKEN_PUBLIC_KEYS_URL) or config.get(TOKEN_PUBLIC_KEY_URL)
         self.public_keys: list[RSAPublicKey] = []
+        self._last_loaded_at: datetime | None = None
+
+    def _refresh_keys_if_needed(self) -> None:
+        """
+        Refresh the public keys if the refresh interval has passed.
+        """
+        logger.debug("Checking if the token public keys need refresh")
+
+        if self._last_loaded_at is None:
+            return
+        now = datetime.now(timezone.utc)
+        if now - self._last_loaded_at < self._REFRESH_INTERVAL:
+            logger.debug("Token public keys are up to date, no refresh needed")
+            return
+        try:
+            logger.debug("Token public keys are stale, refreshing now")
+            self.load_public_keys()
+        except RuntimeError:
+            logger.warning("Token public key refresh failed, using existing keys")
 
     def load_public_keys(self) -> "HandlerToken":
         """
@@ -75,6 +97,7 @@ def load_public_keys(self) -> "HandlerToken":
                 cast(RSAPublicKey, serialization.load_der_public_key(base64.b64decode(raw_key))) for raw_key in raw_keys
             ]
             logger.debug("Loaded %d token public keys", len(self.public_keys))
+            self._last_loaded_at = datetime.now(timezone.utc)
 
             return self
         except (requests.RequestException, ValueError, KeyError, UnsupportedAlgorithm) as exc:
@@ -91,6 +114,8 @@ def decode_jwt(self, token_encoded: str) -> Dict[str, Any]:
         Raises:
             jwt.PyJWTError: If verification fails for all public keys.
         """
+        self._refresh_keys_if_needed()
+
         logger.debug("Decoding JWT")
         for public_key in self.public_keys:
             try:

tests/handlers/test_handler_token.py

@@ -15,13 +15,22 @@
 #
 
 import json
-from unittest.mock import patch
+from datetime import datetime, timedelta, timezone
+from unittest.mock import patch, Mock
 
 import pytest
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
 
 from src.handlers.handler_token import HandlerToken
 
 
+@pytest.fixture
+def token_handler():
+    """Create a HandlerToken instance for testing."""
+    config = {"token_public_keys_url": "https://example.com/keys"}
+    return HandlerToken(config)
+
+
 def test_get_token_endpoint(event_gate_module, make_event):
     event = make_event("/token")
     resp = event_gate_module.lambda_handler(event, None)
@@ -89,3 +98,47 @@ def test_extract_token_empty():
 def test_extract_token_direct_bearer_header():
     token = HandlerToken.extract_token({"Bearer": "  tok123  "})
     assert token == "tok123"
+
+
+## Checking the freshness of public keys
+def test_refresh_keys_not_needed_when_keys_fresh(token_handler):
+    """Keys loaded less than 30 minutes ago should not trigger refresh."""
+    token_handler._last_loaded_at = datetime.now(timezone.utc) - timedelta(minutes=10)
+    token_handler.public_keys = [Mock(spec=RSAPublicKey)]
+
+    with patch.object(token_handler, "load_public_keys") as mock_load:
+        token_handler._refresh_keys_if_needed()
+        mock_load.assert_not_called()
+
+
+def test_refresh_keys_triggered_when_keys_stale(token_handler):
+    """Keys loaded more than 30 minutes ago should trigger refresh."""
+    token_handler._last_loaded_at = datetime.now(timezone.utc) - timedelta(minutes=29)
+    token_handler.public_keys = [Mock(spec=RSAPublicKey)]
+
+    with patch.object(token_handler, "load_public_keys") as mock_load:
+        token_handler._refresh_keys_if_needed()
+        mock_load.assert_called_once()
+
+
+def test_refresh_keys_handles_load_failure_gracefully(token_handler):
+    """If key refresh fails, should log warning and continue with existing keys."""
+    old_key = Mock(spec=RSAPublicKey)
+    token_handler.public_keys = [old_key]
+    token_handler._last_loaded_at = datetime.now(timezone.utc) - timedelta(minutes=29)
+
+    with patch.object(token_handler, "load_public_keys", side_effect=RuntimeError("Network error")):
+        token_handler._refresh_keys_if_needed()
+        assert token_handler.public_keys == [old_key]
+
+
+def test_decode_jwt_triggers_refresh_check(token_handler):
+    """Decoding JWT should check if keys need refresh before decoding."""
+    dummy_key = Mock(spec=RSAPublicKey)
+    token_handler.public_keys = [dummy_key]
+    token_handler._last_loaded_at = datetime.now(timezone.utc) - timedelta(minutes=10)
+
+    with patch.object(token_handler, "_refresh_keys_if_needed") as mock_refresh:
+        with patch("jwt.decode", return_value={"sub": "TestUser"}):
+            token_handler.decode_jwt("dummy-token")
+            mock_refresh.assert_called_once()