diff --git a/plugins/cron_providers/chronos/verify.py b/plugins/cron_providers/chronos/verify.py index 99c8db93e4bd..831ff3d0a0a7 100644 --- a/plugins/cron_providers/chronos/verify.py +++ b/plugins/cron_providers/chronos/verify.py @@ -19,6 +19,7 @@ hand-roll JWT verification. from __future__ import annotations import logging +import threading from typing import Any, Callable, Dict, Optional logger = logging.getLogger("cron.chronos.verify") @@ -27,6 +28,44 @@ logger = logging.getLogger("cron.chronos.verify") # JWT (without this claim) must NOT be replayable against /api/cron/fire. _FIRE_PURPOSE = "cron_fire" +# Process-wide cache of PyJWKClient instances, keyed by JWKS URL. +# +# WHY THIS EXISTS: a PyJWKClient caches the fetched JWKS (signing keys) on the +# INSTANCE. Constructing a fresh client per fire therefore threw that cache +# away and forced a synchronous JWKS HTTP GET to the portal on EVERY fire. Under +# a burst of concurrent fires (an instance with several cron jobs firing in the +# same window) that fanned out into N simultaneous JWKS fetches, which the +# portal rate-limited (HTTP 403) — verification then failed and the agent +# answered 401. When a fetch was merely slow rather than rate-limited, it blocked +# the event loop long enough that the fire webhook could not return its 202 +# before the relay's 30s timeout (observed in prod as relay 504s concentrated on +# high-job-count instances). Reusing one client per URL keeps the signing keys +# cached (NAS keys rotate rarely), so the steady state is zero JWKS fetches per +# fire. See docs/chronos-managed-cron-contract.md and the betterstack triage. +_JWK_CLIENTS: Dict[str, Any] = {} +_JWK_CLIENTS_LOCK = threading.Lock() + + +def _get_jwk_client(jwks_url: str) -> Any: + """Return a process-cached PyJWKClient for ``jwks_url`` (one per URL). + + PyJWKClient does its own key caching internally (``cache_keys``/``lifespan``); + the whole point here is to reuse the SAME instance across fires so that cache + is actually hit instead of discarded. Double-checked-locked so concurrent + fires resolve to a single shared client without racing. + """ + client = _JWK_CLIENTS.get(jwks_url) + if client is not None: + return client + with _JWK_CLIENTS_LOCK: + client = _JWK_CLIENTS.get(jwks_url) + if client is None: + from jwt import PyJWKClient + + client = PyJWKClient(jwks_url) + _JWK_CLIENTS[jwks_url] = client + return client + def verify_nas_fire_token( *, @@ -60,12 +99,15 @@ def verify_nas_fire_token( try: import jwt - from jwt import PyJWKClient # Resolve the signing key from the JWKS endpoint by the token's kid. signing_key = None if jwks_or_key.startswith("http://") or jwks_or_key.startswith("https://"): - jwk_client = PyJWKClient(jwks_or_key) + # Reuse a process-cached client so the JWKS fetch is amortised across + # fires (a fresh client per fire re-fetched the JWKS every time and, + # under concurrent fires, tripped the portal's rate limit → 403 → + # 401, or blocked the event loop past the relay's 30s timeout → 504). + jwk_client = _get_jwk_client(jwks_or_key) signing_key = jwk_client.get_signing_key_from_jwt(token).key else: # A PEM public key passed inline (test / pinned-key deployments). diff --git a/tests/plugins/test_chronos_verify.py b/tests/plugins/test_chronos_verify.py index 5747a81721ef..3f8b64ab1929 100644 --- a/tests/plugins/test_chronos_verify.py +++ b/tests/plugins/test_chronos_verify.py @@ -152,6 +152,7 @@ def test_empty_token_refused(rsa_keys): def test_jwks_url_path_resolves_key(rsa_keys, monkeypatch): """The JWKS-URL branch resolves the signing key via PyJWKClient.""" + from plugins.cron_providers.chronos import verify as verify_mod from plugins.cron_providers.chronos.verify import verify_nas_fire_token priv, pub = rsa_keys @@ -168,6 +169,8 @@ def test_jwks_url_path_resolves_key(rsa_keys, monkeypatch): return FakeKey() monkeypatch.setattr("jwt.PyJWKClient", FakeJWKClient) + # Isolate from the process-wide client cache (other tests may have populated it). + monkeypatch.setattr(verify_mod, "_JWK_CLIENTS", {}) claims = verify_nas_fire_token( token=token, expected_audience=AUD, jwks_or_key="https://portal.nousresearch.com/.well-known/jwks.json", @@ -176,6 +179,56 @@ def test_jwks_url_path_resolves_key(rsa_keys, monkeypatch): assert claims is not None and claims["purpose"] == "cron_fire" +def test_jwks_client_is_cached_across_calls(rsa_keys, monkeypatch): + """Regression (Chronos relay 403/401 + 504 storm): the JWKS client must be + constructed ONCE per URL and reused across fires, NOT rebuilt per call. + + A fresh PyJWKClient per fire discards its key cache and forces a synchronous + JWKS HTTP GET on every fire; a burst of concurrent fires then fans out into N + simultaneous fetches that the portal rate-limits (403 → agent 401) or that + block the event loop past the relay's 30s timeout (504). Reusing one cached + client keeps the steady state at zero fetches per fire. This test fails + against the pre-fix code (construct_count == N) and passes with the cache + (construct_count == 1). + """ + from plugins.cron_providers.chronos import verify as verify_mod + from plugins.cron_providers.chronos.verify import verify_nas_fire_token + + priv, pub = rsa_keys + url = "https://portal.nousresearch.com/.well-known/jwks.json" + + counters = {"construct": 0, "fetch": 0} + + class FakeKey: + key = pub + + class CountingJWKClient: + def __init__(self, u): + counters["construct"] += 1 + + def get_signing_key_from_jwt(self, tok): + counters["fetch"] += 1 + return FakeKey() + + monkeypatch.setattr("jwt.PyJWKClient", CountingJWKClient) + # Start from an empty cache so this test's count is deterministic. + monkeypatch.setattr(verify_mod, "_JWK_CLIENTS", {}) + + for _ in range(5): + token = _mint(priv, _base_claims()) + claims = verify_nas_fire_token( + token=token, expected_audience=AUD, jwks_or_key=url, issuer=ISS, + ) + assert claims is not None + + # The client is built once and reused; only the fetch (served from the + # client's own cache in production) is per-call. + assert counters["construct"] == 1, ( + f"expected 1 PyJWKClient construction, got {counters['construct']} " + "(a fresh client per fire is the bug this guards against)" + ) + + def test_get_fire_verifier_returns_nas_verifier(): from plugins.cron_providers.chronos.verify import get_fire_verifier, verify_nas_fire_token