Merge pull request #1488 from PolicyEngine/codex/reapply-thursday-rollback-changes

Reapply audit fixes and migrate Authlib 1.7 JWKS handling

Author: anth-volk

.env-example

@@ -1,2 +1,7 @@
 FLASK_DEBUG=1
-CACHE_REDIS_HOST=redis
+CACHE_REDIS_HOST=redis
+
+# Optional: wipe the local sqlite analytics DB on startup. Only
+# consulted when FLASK_DEBUG=1 and analytics is enabled. Default off
+# so captured debug data is not lost across restarts.
+# RESET_ANALYTICS=1

changelog.d/authlib-1.7-joserfc.fixed.md

@@ -0,0 +1 @@
+Migrate Auth0 JWKS validation to joserfc keys so Authlib 1.7.0 works without the temporary version cap.

changelog.d/reapply-audit-fixes.fixed.md

@@ -0,0 +1 @@
+Reapply selected audit security and correctness fixes that were rolled back in the 0.13.13 restoration.

config/README.md

@@ -270,6 +270,12 @@ The following endpoints remain unprotected:
 - When enabled, all protected endpoints validate JWT tokens against Auth0's JWKS
 - The Auth0 domain and audience must match the configured values
 
+## Analytics reset (debug only)
+
+`RESET_ANALYTICS=1` (or `analytics.reset: true` in YAML) wipes the
+local SQLite analytics DB on startup. This is **only** consulted when
+`FLASK_DEBUG=1`; production never resets the analytics DB.
+
 ## Usage Examples
 
 ### Production Deployment (Current)

policyengine_household_api/api.py

@@ -4,14 +4,10 @@
 
 # Python imports
 import os
-from pathlib import Path
 
 # External imports
 from flask_cors import CORS
 import flask
-from flask_sqlalchemy import SQLAlchemy
-from sqlalchemy.orm import DeclarativeBase
-from dotenv import load_dotenv
 from flask_limiter import Limiter
 from flask_limiter.util import get_remote_address
 from policyengine_household_api.data.analytics_setup import (
@@ -20,7 +16,6 @@
 
 # Internal imports
 from .decorators.auth import create_auth_decorator
-from .constants import VERSION, REPO
 from policyengine_household_api.decorators.analytics import (
     log_analytics_if_enabled,
 )
@@ -40,6 +35,15 @@
 
 app = application = flask.Flask(__name__)
 
+# Reject absurdly large request bodies before any view runs. 10 MiB is
+# well above the largest legitimate household payload we have seen
+# (axes scans push a few hundred KiB) while still capping the memory a
+# single attacker can force us to allocate. Overridable via the
+# ``MAX_CONTENT_LENGTH`` env var (bytes).
+app.config["MAX_CONTENT_LENGTH"] = int(
+    os.getenv("MAX_CONTENT_LENGTH", 10 * 1024 * 1024)
+)
+
 CORS(app)
 
 # Use in-memory storage for rate limiting
@@ -59,6 +63,7 @@
 
 @app.route("/<country_id>/calculate", methods=["POST"])
 @require_auth_if_enabled()
+@limiter.limit("60 per minute")
 @log_analytics_if_enabled
 def calculate(country_id):
     return get_calculate(country_id)
@@ -84,8 +89,11 @@ def readiness_check():
     )
 
 
+# Note: `/calculate_demo` is intentionally public (documented in
+# config/README.md). It is guarded by a conservative rate limit rather
+# than JWT authentication.
 @app.route("/<country_id>/calculate_demo", methods=["POST"])
-@limiter.limit("1 per second")
+@limiter.limit("1 per 10 seconds")
 def calculate_demo(country_id):
     return get_calculate(country_id)
 

policyengine_household_api/auth/validation.py

@@ -1,18 +1,114 @@
 import json
+import logging
+import time
+from threading import Lock
 from urllib.request import urlopen
 
 from authlib.oauth2.rfc7523 import JWTBearerTokenValidator
-from authlib.jose.rfc7517.jwk import JsonWebKey
+from joserfc.jwk import KeySet
+
+logger = logging.getLogger(__name__)
+
+JWKS_FETCH_TIMEOUT = 10  # seconds
+# Minimum wait between back-to-back lazy retries after a failure.
+# Keeps us from hammering Auth0 when it is actively degraded.
+JWKS_RETRY_INTERVAL_SECONDS = 30
+
+
+# Module-level cache of successful JWKS fetches, keyed by issuer. Only
+# successes are cached so that a transient failure is retried on the
+# next authenticated request (``lru_cache`` would have memoised the
+# ``None`` return, making the "lazy retry" dead code).
+_jwks_cache: dict = {}
+# Records the monotonic timestamp of the most recent *failed* fetch
+# per-issuer so we can rate-limit retries without caching the failure
+# itself.
+_jwks_last_failure: dict = {}
+_jwks_lock = Lock()
+
+
+def _fetch_jwks_uncached(issuer: str):
+    """Fetch the JWKS for an Auth0 issuer, bypassing the cache.
+
+    Returns a joserfc key set on success, ``None`` on failure. Errors
+    are logged rather than raised so that a transient Auth0 outage
+    doesn't crash the process at import time.
+    """
+    jwks_url = f"{issuer}.well-known/jwks.json"
+    try:
+        with urlopen(jwks_url, timeout=JWKS_FETCH_TIMEOUT) as response:
+            return KeySet.import_key_set(json.loads(response.read()))
+    except Exception as e:
+        logger.warning(f"Failed to fetch JWKS from {jwks_url}: {e}")
+        return None
+
+
+def _fetch_jwks(issuer: str):
+    """Fetch JWKS, caching only successful results.
+
+    On failure we record the time but do not memoise the ``None`` — a
+    later call will retry (subject to ``JWKS_RETRY_INTERVAL_SECONDS``
+    backoff) so that the validator self-heals once Auth0 recovers.
+    """
+    with _jwks_lock:
+        cached = _jwks_cache.get(issuer)
+        if cached is not None:
+            return cached
+        last_failure = _jwks_last_failure.get(issuer)
+        if (
+            last_failure is not None
+            and time.monotonic() - last_failure < JWKS_RETRY_INTERVAL_SECONDS
+        ):
+            # Too soon after the last failure — don't hammer Auth0.
+            return None
+
+    # Fetch outside the lock so a slow network call doesn't block
+    # other threads that might be serving requests with a cached key.
+    key_set = _fetch_jwks_uncached(issuer)
+
+    with _jwks_lock:
+        if key_set is not None:
+            _jwks_cache[issuer] = key_set
+            _jwks_last_failure.pop(issuer, None)
+        else:
+            _jwks_last_failure[issuer] = time.monotonic()
+    return key_set
+
+
+def _clear_jwks_cache():
+    """Test helper: wipe the success/failure caches."""
+    with _jwks_lock:
+        _jwks_cache.clear()
+        _jwks_last_failure.clear()
 
 
 class Auth0JWTBearerTokenValidator(JWTBearerTokenValidator):
     def __init__(self, domain, audience):
         issuer = f"https://{domain}/"
-        jsonurl = urlopen(f"{issuer}.well-known/jwks.json")
-        public_key = JsonWebKey.import_key_set(json.loads(jsonurl.read()))
+
+        public_key = _fetch_jwks(issuer)
+        if public_key is None:
+            # Retry on next token validation rather than failing hard
+            # at construction time. A missing key set means token
+            # validation will fail cleanly inside authlib.
+            logger.warning(
+                "JWKS unavailable at construction; will retry on first "
+                "token validation."
+            )
+
         super(Auth0JWTBearerTokenValidator, self).__init__(public_key)
+        self._issuer = issuer
         self.claims_options = {
             "exp": {"essential": True},
             "aud": {"essential": True, "value": audience},
             "iss": {"essential": True, "value": issuer},
         }
+
+    def authenticate_token(self, token_string):
+        # Lazy-refresh the JWKS if the initial fetch failed. Because
+        # ``_fetch_jwks`` only caches successes, this call will retry
+        # the network fetch (subject to a short backoff) until Auth0
+        # responds.
+        if self.public_key is None:
+            self.public_key = _fetch_jwks(self._issuer)
+        return super().authenticate_token(token_string)

policyengine_household_api/country.py

@@ -1,4 +1,5 @@
 import importlib
+import logging
 from flask import Response
 import json
 from policyengine_core.taxbenefitsystems import TaxBenefitSystem
@@ -431,8 +432,12 @@ def calculate(
 
             return household, None
 
-        except Exception as e:
-            print(f"Error computing tracer output: {e}")
+        except Exception:
+            # Re-raise so endpoints/household.py (which unpacks
+            # ``(result, computation_tree_uuid)``) can surface a real
+            # 500 instead of a TypeError on ``None`` unpacking.
+            logging.exception("Tracer failed while computing household")
+            raise
 
 
 def create_policy_reform(policy_data: dict) -> dict:
@@ -477,7 +482,7 @@ def apply(self):
 
 
 def get_requested_computations(household: dict):
-    requested_computations = dpath.util.search(
+    requested_computations = dpath.search(
         household,
         "*/*/*/*",
         afilter=lambda t: t is None,

policyengine_household_api/data/analytics_setup.py

@@ -47,11 +47,21 @@ def initialize_analytics_db_if_enabled(app):
         db_url = (
             REPO / "policyengine_household_api" / "data" / "policyengine.db"
         )
-        if Path(db_url).exists():
+        # Only wipe the analytics DB when explicitly requested via
+        # RESET_ANALYTICS=1 (or the ``analytics.reset`` config flag).
+        should_reset = os.getenv("RESET_ANALYTICS", "").lower() in (
+            "1",
+            "true",
+            "yes",
+        ) or get_config_value("analytics.reset", False)
+        if should_reset and Path(db_url).exists():
             Path(db_url).unlink()
         if not Path(db_url).exists():
             Path(db_url).touch()
-        app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:////" + str(db_url)
+        # sqlite: absolute paths require exactly three slashes plus the
+        # leading "/" from the absolute path (=> "sqlite:////tmp/x.db").
+        # db_url here is already absolute, so use an f-string.
+        app.config["SQLALCHEMY_DATABASE_URI"] = f"sqlite:///{db_url}"
     else:
         app.config["SQLALCHEMY_DATABASE_URI"] = "mysql+pymysql://"
         app.config["SQLALCHEMY_ENGINE_OPTIONS"] = {"creator": getconn}

policyengine_household_api/decorators/analytics.py

@@ -5,7 +5,7 @@
 
 from functools import wraps
 from flask import request
-from datetime import datetime
+from datetime import datetime, timezone
 import jwt
 import logging
 from policyengine_household_api.constants import VERSION
@@ -14,10 +14,56 @@
 )
 from policyengine_household_api.data.analytics_setup import db
 from policyengine_household_api.data.models import Visit
+from policyengine_household_api.utils.config_loader import get_config_value
 
 logger = logging.getLogger(__name__)
 
 
+# Cache the JWKS client so we don't re-fetch keys on every request.
+_jwks_client_cache: dict = {}
+
+
+def _get_jwks_client(auth0_address: str):
+    """Return a cached PyJWKClient for the given Auth0 domain."""
+    if auth0_address not in _jwks_client_cache:
+        jwks_url = f"https://{auth0_address}/.well-known/jwks.json"
+        _jwks_client_cache[auth0_address] = jwt.PyJWKClient(jwks_url)
+    return _jwks_client_cache[auth0_address]
+
+
+def _verified_sub_claim(token: str) -> str | None:
+    """
+    Return the token's ``sub`` claim if signature verification succeeds
+    against the configured Auth0 JWKS, else ``None``.
+
+    If Auth0 configuration is missing (e.g. in a dev environment) the
+    claim cannot be trusted and we return ``None`` so the caller can
+    store a null client_id rather than an attacker-controlled value.
+    """
+    auth0_address = get_config_value("auth.auth0.address", "")
+    auth0_audience = get_config_value("auth.auth0.audience", "")
+    if not auth0_address or not auth0_audience:
+        return None
+
+    try:
+        signing_key = _get_jwks_client(auth0_address).get_signing_key_from_jwt(
+            token
+        )
+        claims = jwt.decode(
+            token,
+            signing_key.key,
+            algorithms=["RS256"],
+            audience=auth0_audience,
+            issuer=f"https://{auth0_address}/",
+            options={"verify_signature": True},
+        )
+    except Exception as e:
+        logger.debug(f"JWT signature verification failed: {e}")
+        return None
+
+    return claims.get("sub")
+
+
 def log_analytics_if_enabled(func):
     """
     Decorator that logs analytics only if analytics is enabled in configuration.
@@ -45,25 +91,34 @@ def decorated_function(*args, **kwargs):
             # Create a record that will be emitted to the db
             new_visit = Visit()
 
-            # Pull client_id from JWT
+            # Pull client_id from JWT. We only trust the `sub` claim
+            # when the token signature has been verified against the
+            # Auth0 JWKS. If verification fails (bad signature, JWKS
+            # unreachable, etc.) we still record the visit but drop
+            # the client_id so that attackers cannot spoof analytics
+            # identities simply by crafting an unsigned JWT.
             try:
                 auth_header = str(request.authorization)
                 token = auth_header.split(" ")[1]
-                decoded_token = jwt.decode(
-                    token, options={"verify_signature": False}
-                )
-                client_id = decoded_token["sub"]
-
-                suffix_to_slice = "@clients"
-                if (
-                    len(client_id) >= len(suffix_to_slice)
-                    and client_id[-len(suffix_to_slice) :] == suffix_to_slice
-                ):
-                    client_id = client_id[: -len(suffix_to_slice)]
-                new_visit.client_id = client_id
+                client_id = _verified_sub_claim(token)
+
+                if client_id is None:
+                    new_visit.client_id = None
+                else:
+                    suffix_to_slice = "@clients"
+                    if (
+                        len(client_id) >= len(suffix_to_slice)
+                        and client_id[-len(suffix_to_slice) :]
+                        == suffix_to_slice
+                    ):
+                        client_id = client_id[: -len(suffix_to_slice)]
+                    new_visit.client_id = client_id
             except Exception as e:
                 logger.debug(f"Could not extract client_id from JWT: {e}")
-                new_visit.client_id = "unknown"
+                # Match the verified-fail path: a missing/unparseable
+                # header must also be stored as NULL, never as a
+                # sentinel string we'd have to filter out downstream.
+                new_visit.client_id = None
 
             # Set API version
             new_visit.api_version = VERSION
@@ -75,8 +130,9 @@ def decorated_function(*args, **kwargs):
             # Set content_length_bytes
             new_visit.content_length_bytes = request.content_length
 
-            # Set the date and time
-            now = datetime.utcnow()
+            # Set the date and time (timezone-aware; utcnow() is
+            # deprecated in Python 3.12+)
+            now = datetime.now(timezone.utc)
             new_visit.datetime = now
 
             # Emit the new record to the db