fix(cron/chronos): cache PyJWKClient across fires to stop JWKS fetch storm (#64641)

The inbound cron-fire verifier constructed a fresh PyJWKClient on every
fire, discarding the client's key cache and forcing a synchronous JWKS
HTTP GET to the portal on each fire. Under a burst of concurrent fires
(a hosted instance with several cron jobs firing in the same window) this
fanned out into N simultaneous JWKS fetches that the portal rate-limited
(HTTP 403 -> verification fails -> agent 401), or that 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).

Cache one PyJWKClient per JWKS URL at module scope (double-checked lock)
so the signing keys are reused across fires; NAS keys rotate rarely, so
the steady state is zero JWKS fetches per fire.

Regression test proves 5 fires -> 1 client construction (was 5).
This commit is contained in:
Ben Barclay 2026-07-15 09:27:35 +10:00 committed by GitHub
parent df5700ebe3
commit 9884b4faad
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 97 additions and 2 deletions

View file

@ -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).

View file

@ -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