diff --git a/indexnow.py b/indexnow.py index 84bbc40..1bf3422 100644 --- a/indexnow.py +++ b/indexnow.py @@ -36,10 +36,58 @@ _STATE_PATH = os.environ.get("THERMOGRAPH_INDEXNOW_STATE_FILE") or os.path.join( _ENDPOINT = "https://api.indexnow.org/indexnow" # a shared endpoint; it fans out to all participants _BATCH = 10000 # IndexNow's max URLs per request +# One shared client instead of a bare httpx.post per batch — same reused-client +# pattern as web/app.py's _frontend_client and notifications/discord.py's _client. +# 30s default matches submit()'s own prior per-call default; still overridable +# per call (submit_all's fan-out passes its own `timeout` through unchanged). +_client = httpx.Client(timeout=30) + _lock = threading.Lock() _key = None +def _claim_file(value: str) -> str: + """Persist a freshly generated key as `_KEY_PATH`'s content, atomic + first-writer-wins — same race, same fix as push.py's VAPID `_claim_file` + (see its docstring): on a cold /state volume several workers can hit the + missing-file branch at once, and without this each would generate and cache + its OWN key, diverging from the file and from each other (the key served at + /{key}.txt would then randomly mismatch whichever key a given worker signed + a submission with). Write to a private temp file, then claim the real path + with a hard link (atomic; a loser gets EEXIST with no half-written window); + a loser reads back the winner's file instead of keeping its own generation.""" + tmp_path = f"{_KEY_PATH}.{os.getpid()}.tmp" + try: + os.makedirs(_DATA_DIR, exist_ok=True) + with open(tmp_path, "w", encoding="utf-8") as f: + f.write(value) + os.chmod(tmp_path, 0o600) + except OSError: + log.warning("could not persist IndexNow key to %s; using an in-memory key", _KEY_PATH) + return value + try: + try: + os.link(tmp_path, _KEY_PATH) + return value + except FileExistsError: + try: + with open(_KEY_PATH, encoding="utf-8") as f: + winner = f.read().strip() + if winner: + return winner + except OSError: + pass + return value + except OSError: + log.warning("could not persist IndexNow key to %s; using an in-memory key", _KEY_PATH) + return value + finally: + try: + os.unlink(tmp_path) + except OSError: + pass + + def key() -> str: """The IndexNow key (env → file → generate-and-persist), cached for the process.""" global _key @@ -60,14 +108,10 @@ def key() -> str: return _key except OSError: pass - _key = secrets.token_hex(16) # 32 hex chars — within IndexNow's 8–128 range - try: - os.makedirs(_DATA_DIR, exist_ok=True) - with open(_KEY_PATH, "w", encoding="utf-8") as f: - f.write(_key) - os.chmod(_KEY_PATH, 0o600) - except OSError: - log.warning("could not persist IndexNow key to %s; using an in-memory key", _KEY_PATH) + # Cache whatever _claim_file settles on — see its docstring: our own + # generation if we won the race to create the file, the winner's key + # read back from disk if we lost it. + _key = _claim_file(secrets.token_hex(16)) # 32 hex chars — within IndexNow's 8–128 range return _key @@ -83,8 +127,8 @@ def submit(urls, host: str, key_location: str, scheme: str = "https", timeout: f batch = urls[i:i + _BATCH] payload = {"host": host, "key": k, "keyLocation": key_location, "urlList": batch} try: - r = httpx.post(_ENDPOINT, json=payload, timeout=timeout, - headers={"Content-Type": "application/json; charset=utf-8"}) + r = _client.post(_ENDPOINT, json=payload, timeout=timeout, + headers={"Content-Type": "application/json; charset=utf-8"}) except httpx.HTTPError as e: log.warning("IndexNow request failed: %s", e) statuses.append("error") diff --git a/notifications/discord.py b/notifications/discord.py index 776fa2f..2eb0e24 100644 --- a/notifications/discord.py +++ b/notifications/discord.py @@ -47,6 +47,12 @@ _COLD = 0x4393C3 _TIMEOUT = httpx.Timeout(10.0) +# One shared client instead of a bare httpx.post per call — same reused-client +# pattern as web/app.py's _frontend_client, so every webhook/bot-REST call in +# this module reuses one connection pool instead of paying a fresh TCP+TLS +# handshake per notification. +_client = httpx.Client(timeout=_TIMEOUT) + def enabled() -> bool: return bool(WEBHOOK_URL) or weather_enabled() @@ -131,7 +137,7 @@ def post_daily_feed(feed: dict | None = None) -> bool: "embeds": [embed], } try: - resp = httpx.post(WEBHOOK_URL, json=payload, timeout=_TIMEOUT) + resp = _client.post(WEBHOOK_URL, json=payload, timeout=_TIMEOUT) ok = 200 <= resp.status_code < 300 except Exception: # noqa: BLE001 - best-effort side channel pass @@ -159,10 +165,10 @@ def _bot_post(path: str, json_body: dict) -> httpx.Response | None: """POST to the bot REST API with one 429 retry. None on a transport error.""" headers = {"Authorization": f"Bot {BOT_TOKEN}"} try: - resp = httpx.post(f"{_API}{path}", json=json_body, headers=headers, timeout=_TIMEOUT) + resp = _client.post(f"{_API}{path}", json=json_body, headers=headers, timeout=_TIMEOUT) if resp.status_code == 429: time.sleep(_retry_after(resp)) - resp = httpx.post(f"{_API}{path}", json=json_body, headers=headers, timeout=_TIMEOUT) + resp = _client.post(f"{_API}{path}", json=json_body, headers=headers, timeout=_TIMEOUT) return resp except Exception: # noqa: BLE001 - best-effort side channel return None diff --git a/notifications/notify.py b/notifications/notify.py index d7c82d9..a26702f 100644 --- a/notifications/notify.py +++ b/notifications/notify.py @@ -23,6 +23,7 @@ A forecast-fetch failure just skips the cell, so a pass never dies on a rate lim """ import concurrent.futures import datetime +import logging import os import threading import time @@ -42,6 +43,8 @@ from notifications import push from accounts.db import sync_session_maker from accounts.models import AccessToken, Notification, PushSubscription, Subscription, User +log = logging.getLogger("thermograph.notify") + _BASE = os.environ.get("THERMOGRAPH_BASE", "/thermograph").strip("/") BASE = f"/{_BASE}" if _BASE else "" @@ -57,6 +60,14 @@ FORECAST_HORIZON_DAYS = 7 # — the rest are picked up on later passes. MAX_ARCHIVE_FETCHES_PER_PASS = int(os.environ.get("THERMOGRAPH_NOTIFY_ARCHIVE_FETCHES", "4")) +# Soft wall-clock budget for one whole pass, same discipline as the archive-fetch +# cap above: checked between cells (never mid-cell) so one pathological pass — a +# slow upstream, a DB hiccup, hundreds of cells each doing real work — can't run +# unbounded and push the next pass's wake past its own interval. 2x INTERVAL by +# default: generous enough that an ordinarily-slow pass still finishes, but a +# genuinely runaway one still yields before it would overlap the next wake. +PASS_DEADLINE_SECONDS = int(os.environ.get("THERMOGRAPH_NOTIFY_PASS_DEADLINE", str(INTERVAL * 2))) + # Metrics whose low tail also counts as unusual when a subscription is two-sided # (cold snaps, unusually calm/dry). Precipitation is one-directional (high only). TWO_SIDED_METRICS = set(grading.TEMP_METRICS) @@ -171,19 +182,40 @@ def _send_discord_job(job: tuple[str, str, str, str]) -> None: # hundreds of subscribers no longer stretches a pass out send-by-send. SEND_WORKERS = int(os.environ.get("THERMOGRAPH_NOTIFY_SEND_WORKERS", "8")) +# Belt-and-suspenders on top of push.py's own send timeout: f.result() below has +# no timeout by default and would otherwise wait forever for a future that never +# resolves, which — same as an untimed webpush() call — wedges the notifier for +# the process's life under the singleton flock. Generous on purpose; this only +# needs to fire if a per-send timeout was somehow bypassed, not to race it. +SEND_RESULT_TIMEOUT = int(os.environ.get("THERMOGRAPH_NOTIFY_SEND_RESULT_TIMEOUT", "60")) + def _flush_sends(push_jobs, discord_jobs) -> list[int]: """Deliver every gathered push/Discord job concurrently. Returns the push-subscription row ids the push service reported gone, for the caller - to prune. Blocks until every job has finished (or errored) — a pass still - waits for delivery to complete, it just no longer does so serially.""" + to prune. Blocks until every job has finished or errored, and gives up + waiting on (but does not lose track of) any single job past + SEND_RESULT_TIMEOUT — a pass still waits for delivery to complete, it just + no longer does so serially, and no misbehaving future can hang the result + loop past a bounded wait.""" with concurrent.futures.ThreadPoolExecutor(max_workers=SEND_WORKERS) as pool: futures = [pool.submit(_send_push_job, job) for job in push_jobs] futures += [pool.submit(_send_discord_job, job) for job in discord_jobs] gone = [] for f in futures: try: - result = f.result() + result = f.result(timeout=SEND_RESULT_TIMEOUT) + except concurrent.futures.TimeoutError: + # f.result() gave up waiting, but the worker thread itself can't be + # killed — ThreadPoolExecutor.__exit__ still calls shutdown(wait=True), + # so the pool as a whole won't return until that thread actually + # finishes. The real backstop is the per-call socket timeouts this + # commit adds (push.py's _SEND_TIMEOUT, discord.py's _TIMEOUT): they + # bound how long a "stuck" send can truly run. This catch just stops + # THIS result from being reported/counted as gone, and keeps checking + # the remaining futures instead of getting stuck on one. + log.warning("push/Discord send exceeded %ss timeout; treating as failed", SEND_RESULT_TIMEOUT) + continue except Exception: # noqa: BLE001 - one bad send must not lose the rest continue if result is not None: @@ -370,7 +402,18 @@ def run_pass() -> int: by_cell: dict[str, list] = {} for sub in subs: by_cell.setdefault(sub.cell_id, []).append(sub) - for cell_id, cell_subs in by_cell.items(): + cells_skipped = 0 + for i, (cell_id, cell_subs) in enumerate(by_cell.items()): + # Checked BETWEEN cells, not inside _process_cell — a soft budget, not + # a hard per-item timeout. Once tripped, every remaining cell this pass + # is skipped outright (not attempted-and-abandoned mid-flight) so a + # skipped cell's subscriptions are simply picked up whole on the next + # pass, same as a rate-limited one already is. + if time.perf_counter() - t0 > PASS_DEADLINE_SECONDS: + cells_skipped = len(by_cell) - i + log.warning("notifier pass exceeded %ss deadline; skipping %d/%d remaining cell(s)", + PASS_DEADLINE_SECONDS, cells_skipped, len(by_cell)) + break try: created += _process_cell(session, cell_id, cell_subs, today, now, archive_budget) except Exception: # noqa: BLE001 - one bad cell must not abort the pass @@ -382,6 +425,7 @@ def run_pass() -> int: audit.log_activity("notify.pass", { "subs_evaluated": len(subs), "cells": len(by_cell), "created": created, "archives_fetched": MAX_ARCHIVE_FETCHES_PER_PASS - archive_budget[0], + "cells_skipped_deadline": cells_skipped, "duration_ms": round((time.perf_counter() - t0) * 1000.0, 1)}) return created diff --git a/notifications/push.py b/notifications/push.py index 0c40a51..594607f 100644 --- a/notifications/push.py +++ b/notifications/push.py @@ -39,6 +39,15 @@ _VAPID_PATH = os.environ.get("THERMOGRAPH_VAPID_FILE") or os.path.join(_DATA_DIR # The VAPID "sub" claim — a contact the push service can reach about our traffic. _CONTACT = os.environ.get("THERMOGRAPH_VAPID_CONTACT", "mailto:admin@thermograph.app") +# pywebpush 2.0.0 forwards `timeout` straight to requests.post with no default of +# its own — omit the kwarg here and the send blocks with NO timeout at all (the +# "10s default" people expect only applies if requests itself is called bare). +# One hung push endpoint must never wedge the notifier: the singleton flock +# (core/singleton.py) is held for the process's whole life, so no other worker +# can take over while this thread is stuck, and /healthz keeps reporting green +# on a notifier that's actually frozen. +_SEND_TIMEOUT = 15 + _lock = threading.Lock() _keys = None # cached {"private_key": str, "public_key": str} (base64url-raw) @@ -60,6 +69,59 @@ def _generate() -> dict: return {"private_key": _b64url(scalar), "public_key": _b64url(raw_pub)} +def _claim_file(data: dict) -> dict: + """Persist a freshly generated keypair as `_VAPID_PATH`'s content, atomic + first-writer-wins. `_lock` only keeps this process's own threads from racing + each other — on a cold /state volume every uvicorn *worker* (a separate + process) hits the missing-file branch at boot together, and without this each + would generate its OWN keypair and cache it in-process, silently diverging + from the file and from each other (a subscription signed against one worker's + public key fails to verify under another's private key — the exact incident + class deploy/entrypoint.sh's comments name). + + Write the full keypair to a private temp file first, then claim the real path + with a hard link: `os.link` is atomic and the loser gets EEXIST immediately, + with no window where `_VAPID_PATH` exists but is only half-written (unlike + O_CREAT|O_EXCL directly on the destination followed by a separate write). A + loser discards its own generation and reads back the winner's file instead, + so every process ends up caching the SAME keys the file actually holds.""" + tmp_path = f"{_VAPID_PATH}.{os.getpid()}.tmp" + try: + os.makedirs(_DATA_DIR, exist_ok=True) + with open(tmp_path, "w", encoding="utf-8") as f: + json.dump(data, f) + os.chmod(tmp_path, 0o600) + except OSError: + log.warning("could not persist VAPID keys to %s; using an in-memory pair", _VAPID_PATH) + return data + try: + try: + os.link(tmp_path, _VAPID_PATH) + return data + except FileExistsError: + # Lost the race — someone else's file is now the truth. Read IT back + # rather than keep our own now-orphaned generation. + try: + with open(_VAPID_PATH, encoding="utf-8") as f: + winner = json.load(f) + if winner.get("private_key") and winner.get("public_key"): + return winner + except (OSError, ValueError): + pass + # Winner's file was unreadable/corrupt (very rare) — fall back to our + # own in-memory generation rather than crash; the next _load() call + # (a new process, or after the file heals) retries the file. + return data + except OSError: + log.warning("could not persist VAPID keys to %s; using an in-memory pair", _VAPID_PATH) + return data + finally: + try: + os.unlink(tmp_path) + except OSError: + pass + + def _load() -> dict: """Resolve the keypair once (env → file → generate) and cache it.""" global _keys @@ -81,14 +143,10 @@ def _load() -> dict: return _keys except (OSError, ValueError): pass - _keys = _generate() - try: - os.makedirs(_DATA_DIR, exist_ok=True) - with open(_VAPID_PATH, "w", encoding="utf-8") as f: - json.dump(_keys, f) - os.chmod(_VAPID_PATH, 0o600) - except OSError: - log.warning("could not persist VAPID keys to %s; using an in-memory pair", _VAPID_PATH) + # Cache whatever _claim_file settles on — our own generation if we won + # the race to create the file, or the winner's keys read back from disk + # if we lost it. Either way this process's cache matches the file. + _keys = _claim_file(_generate()) return _keys @@ -111,6 +169,7 @@ def send(subscription_info: dict, payload: dict) -> str: vapid_private_key=keys["private_key"], vapid_claims={"sub": _CONTACT}, ttl=86400, + timeout=_SEND_TIMEOUT, ) return "ok" except WebPushException as e: diff --git a/tests/notifications/test_discord.py b/tests/notifications/test_discord.py index f62763e..64a50a6 100644 --- a/tests/notifications/test_discord.py +++ b/tests/notifications/test_discord.py @@ -1,5 +1,5 @@ """Discord daily-feed post: embed shape, webhook delivery, and the once-a-day guard. -No network — httpx.post is stubbed, like the IndexNow tests.""" +No network — the shared client's .post is stubbed, like the IndexNow tests.""" import datetime import time @@ -63,7 +63,7 @@ def test_build_embed_rejects_empty_feed(): def test_post_disabled_without_webhook(monkeypatch): monkeypatch.setattr(discord, "WEBHOOK_URL", "") posted = [] - monkeypatch.setattr(discord.httpx, "post", lambda *a, **k: posted.append(a) or None) + monkeypatch.setattr(discord._client, "post", lambda *a, **k: posted.append(a) or None) assert discord.post_daily_feed(_feed([_CARD_HOT])) is False assert posted == [] # never touched the network @@ -75,7 +75,7 @@ class _Resp: def test_post_sends_embed_to_the_webhook(monkeypatch): monkeypatch.setattr(discord, "WEBHOOK_URL", "https://discord.test/webhook/abc") calls = [] - monkeypatch.setattr(discord.httpx, "post", + monkeypatch.setattr(discord._client, "post", lambda url, json=None, **k: (calls.append((url, json)), _Resp(204))[1]) assert discord.post_daily_feed(_feed([_CARD_HOT])) is True (url, payload), = calls @@ -91,7 +91,7 @@ def test_post_also_broadcasts_to_weather_channel(monkeypatch): monkeypatch.setattr(discord, "BOT_TOKEN", "bot-abc") monkeypatch.setattr(discord, "WEATHER_CHANNEL_ID", "chan-weather") calls = [] - monkeypatch.setattr(discord.httpx, "post", + monkeypatch.setattr(discord._client, "post", lambda url, json=None, **k: (calls.append((url, json)), _Resp(204))[1]) assert discord.post_daily_feed(_feed([_CARD_HOT])) is True (url, body), = calls @@ -104,14 +104,14 @@ def test_post_is_best_effort_on_http_error(monkeypatch): def _boom(*a, **k): raise RuntimeError("network down") - monkeypatch.setattr(discord.httpx, "post", _boom) + monkeypatch.setattr(discord._client, "post", _boom) # Never raises; just reports failure. assert discord.post_daily_feed(_feed([_CARD_HOT])) is False def test_post_skips_stale_or_empty_feed(monkeypatch): monkeypatch.setattr(discord, "WEBHOOK_URL", "https://discord.test/webhook/abc") - monkeypatch.setattr(discord.httpx, "post", lambda *a, **k: _Resp(204)) + monkeypatch.setattr(discord._client, "post", lambda *a, **k: _Resp(204)) # Yesterday's feed is stale -> no post. yesterday = (datetime.date.today() - datetime.timedelta(days=1)).isoformat() assert discord.post_daily_feed(_feed([_CARD_HOT], date=yesterday)) is False diff --git a/tests/notifications/test_discord_dm.py b/tests/notifications/test_discord_dm.py index 1e1b34e..aa9cc56 100644 --- a/tests/notifications/test_discord_dm.py +++ b/tests/notifications/test_discord_dm.py @@ -28,7 +28,7 @@ class _Resp: def _mock_posts(monkeypatch, script): - """Route discord.httpx.post by URL; `script` maps a URL substring -> _Resp.""" + """Route discord._client.post by URL; `script` maps a URL substring -> _Resp.""" calls = [] def _post(url, json=None, headers=None, timeout=None): @@ -37,7 +37,7 @@ def _mock_posts(monkeypatch, script): if frag in url: return resp return _Resp(404) - monkeypatch.setattr(discord.httpx, "post", _post) + monkeypatch.setattr(discord._client, "post", _post) return calls @@ -97,7 +97,7 @@ def test_bot_post_retries_once_on_429(monkeypatch): monkeypatch.setattr(discord, "BOT_TOKEN", "bot-abc") monkeypatch.setattr(discord.time, "sleep", lambda s: None) # don't actually wait seq = iter([_Resp(429, {"retry_after": 0.01}), _Resp(200, {"id": "c"})]) - monkeypatch.setattr(discord.httpx, "post", lambda *a, **k: next(seq)) + monkeypatch.setattr(discord._client, "post", lambda *a, **k: next(seq)) resp = discord._bot_post("/users/@me/channels", {"recipient_id": "d"}) assert resp.status_code == 200 # the retry succeeded diff --git a/tests/notifications/test_notify.py b/tests/notifications/test_notify.py index 8c22038..74399f1 100644 --- a/tests/notifications/test_notify.py +++ b/tests/notifications/test_notify.py @@ -2,12 +2,14 @@ wording (pure logic), plus an integration check that a full pass fetches a MISSING archive once but never re-fetches a cached one (against conftest's throwaway DB).""" import datetime +import time import types import uuid import numpy as np import polars as pl +from core import audit from data import climate from accounts import db from notifications import notify @@ -211,3 +213,72 @@ def test_inapp_notification_survives_push_error(monkeypatch): assert len(_user_notifications(uid)) == 1 # the in-app write is unaffected assert _push_count() == 1 # a mere error doesn't prune + + +# --- send-result timeout: one hung future must not wedge the pass ----------- +def test_flush_sends_treats_timeout_as_a_failed_send(monkeypatch): + monkeypatch.setattr(notify, "SEND_RESULT_TIMEOUT", 0.05) + + def _hang(job): + time.sleep(0.3) # longer than the timeout above + return None + monkeypatch.setattr(notify, "_send_push_job", _hang) + + # Must not raise, and a timed-out send is never reported as a "gone" + # endpoint to prune (it never actually got a definitive answer). + gone = notify._flush_sends([object()], []) + assert gone == [] + + +def test_flush_sends_still_returns_gone_ids_within_the_timeout(monkeypatch): + monkeypatch.setattr(notify, "SEND_RESULT_TIMEOUT", 5) + monkeypatch.setattr(notify, "_send_push_job", lambda job: job) # echoes the id straight back + gone = notify._flush_sends([7, 8], []) + assert sorted(gone) == [7, 8] + + +# --- pass wall-clock deadline: one pathological pass can't run unbounded ---- +def _seed_two_cell_subscriptions(): + db.Base.metadata.create_all(db.sync_engine) + with db.sync_session_maker() as s: + s.execute(delete(User)) + s.commit() + uid1, uid2 = uuid.uuid4(), uuid.uuid4() + s.add(User(id=uid1, email="deadline1@example.com", hashed_password="x", is_active=True)) + s.add(User(id=uid2, email="deadline2@example.com", hashed_password="x", is_active=True)) + s.commit() + s.add(Subscription(user_id=uid1, cell_id="500_600", label="A", lat=1.0, lon=2.0, + threshold=95, metrics=["tmax"], kind="observed", two_sided=False)) + s.add(Subscription(user_id=uid2, cell_id="700_800", label="B", lat=3.0, lon=4.0, + threshold=95, metrics=["tmax"], kind="observed", two_sided=False)) + s.commit() + return uid1, uid2 + + +def test_run_pass_stops_at_wall_clock_deadline(monkeypatch): + uid1, uid2 = _seed_two_cell_subscriptions() + _cached_extreme(monkeypatch) # every cell would otherwise trigger a notification + monkeypatch.setattr(notify, "PASS_DEADLINE_SECONDS", 0) # trips before the first cell + + logged = [] + monkeypatch.setattr(audit, "log_activity", + lambda kind, data: logged.append((kind, data))) + + notify.run_pass() + + # A deadline of 0 must skip every cell rather than process any of them. + assert _user_notifications(uid1) == [] + assert _user_notifications(uid2) == [] + pass_log = next(data for kind, data in logged if kind == "notify.pass") + assert pass_log["cells_skipped_deadline"] == 2 + + +def test_run_pass_processes_normally_within_the_deadline(monkeypatch): + uid1, uid2 = _seed_two_cell_subscriptions() + _cached_extreme(monkeypatch) + monkeypatch.setattr(notify, "PASS_DEADLINE_SECONDS", 300) # generous — shouldn't trip + + notify.run_pass() + + assert len(_user_notifications(uid1)) == 1 + assert len(_user_notifications(uid2)) == 1 diff --git a/tests/notifications/test_push.py b/tests/notifications/test_push.py new file mode 100644 index 0000000..422583f --- /dev/null +++ b/tests/notifications/test_push.py @@ -0,0 +1,96 @@ +"""VAPID key resolution: the atomic first-writer-wins claim of the keypair file +(push.py's `_claim_file`), and that `_load()` caches whatever that settles on +rather than its own local generation when it loses the race. + +No network — `send()`'s pywebpush call isn't exercised here (that's notify.py's +`test_push_dispatched_on_new_notification` etc., which stub `push.send` itself).""" +import json + +from notifications import push + + +def _reset_state(monkeypatch, path): + monkeypatch.setattr(push, "_VAPID_PATH", str(path)) + monkeypatch.setattr(push, "_DATA_DIR", str(path.parent)) + monkeypatch.setattr(push, "_keys", None) + monkeypatch.delenv("THERMOGRAPH_VAPID_PRIVATE_KEY", raising=False) + monkeypatch.delenv("THERMOGRAPH_VAPID_PUBLIC_KEY", raising=False) + + +def test_claim_file_first_writer_wins(monkeypatch, tmp_path): + path = tmp_path / "vapid.json" + _reset_state(monkeypatch, path) + data = {"private_key": "priv-a", "public_key": "pub-a"} + + result = push._claim_file(data) + + assert result == data + assert json.loads(path.read_text()) == data + # No leftover temp file. + assert list(tmp_path.iterdir()) == [path] + + +def test_claim_file_reads_back_winner_on_race(monkeypatch, tmp_path): + """A process that loses the os.link() race must discard its own freshly + generated keypair and use whatever the winner actually persisted.""" + path = tmp_path / "vapid.json" + _reset_state(monkeypatch, path) + winner = {"private_key": "priv-winner", "public_key": "pub-winner"} + path.write_text(json.dumps(winner)) # simulates another process winning first + loser = {"private_key": "priv-loser", "public_key": "pub-loser"} + + result = push._claim_file(loser) + + assert result == winner + assert json.loads(path.read_text()) == winner # the loser never touched the file + + +def test_load_caches_winner_keys_not_its_own_generation(monkeypatch, tmp_path): + """End-to-end through _load(): on a cold /state volume, a worker that hits the + generate branch but loses the write race must end up with the SAME in-process + cached keys the file actually holds — not the keys it generated locally, + which would silently diverge from every other worker (a subscription signed + against one worker's public key then fails to verify under another's + private key).""" + path = tmp_path / "vapid.json" + _reset_state(monkeypatch, path) + winner = {"private_key": "priv-winner", "public_key": "pub-winner"} + + def _generate_and_lose_the_race(): + # Between our cache-miss file read and our own _claim_file() call, a + # concurrent worker wins and writes the real file first. + path.write_text(json.dumps(winner)) + return {"private_key": "priv-mine", "public_key": "pub-mine"} + + monkeypatch.setattr(push, "_generate", _generate_and_lose_the_race) + + result = push._load() + + assert result == winner + assert push._keys == winner # the process-wide cache matches the file + assert push.public_key() == "pub-winner" + + +def test_load_reads_existing_file_without_generating(monkeypatch, tmp_path): + path = tmp_path / "vapid.json" + _reset_state(monkeypatch, path) + existing = {"private_key": "priv-x", "public_key": "pub-x"} + path.write_text(json.dumps(existing)) + + def boom(): + raise AssertionError("must not generate when a valid file already exists") + monkeypatch.setattr(push, "_generate", boom) + + assert push._load() == existing + + +def test_load_prefers_env_over_file(monkeypatch, tmp_path): + path = tmp_path / "vapid.json" + _reset_state(monkeypatch, path) + path.write_text(json.dumps({"private_key": "priv-file", "public_key": "pub-file"})) + monkeypatch.setenv("THERMOGRAPH_VAPID_PRIVATE_KEY", " priv-env ") + monkeypatch.setenv("THERMOGRAPH_VAPID_PUBLIC_KEY", " pub-env ") + + result = push._load() + + assert result == {"private_key": "priv-env", "public_key": "pub-env"} diff --git a/tests/test_indexnow.py b/tests/test_indexnow.py index d8a7c42..9045ad1 100644 --- a/tests/test_indexnow.py +++ b/tests/test_indexnow.py @@ -69,7 +69,7 @@ def test_indexnow_submit_all_builds_payload(monkeypatch): status_code = 200 text = "ok" - monkeypatch.setattr(indexnow.httpx, "post", + monkeypatch.setattr(indexnow._client, "post", lambda url, json=None, **kw: (calls.append((url, json)), _Resp())[1]) res = indexnow.submit_all("https://thermograph.org") assert res["submitted"] == res["total"] > 100 @@ -87,3 +87,53 @@ def test_indexnow_if_changed_state(monkeypatch, tmp_path): assert indexnow._read_state() == "" # first deploy → would submit indexnow._write_state(sig) assert indexnow._read_state() == sig # unchanged → deploy would skip + + +# --- key() race: atomic first-writer-wins on a cold /state volume ----------- +# Same race, same fix as push.py's VAPID keypair (see tests/notifications/ +# test_push.py) — several workers can hit the missing-key-file branch at once; +# only one may actually create the file, everyone else must read IT back. + +def _reset_key_state(monkeypatch, path): + monkeypatch.setattr(indexnow, "_KEY_PATH", str(path)) + monkeypatch.setattr(indexnow, "_DATA_DIR", str(path.parent)) + monkeypatch.setattr(indexnow, "_key", None) + monkeypatch.delenv("THERMOGRAPH_INDEXNOW_KEY", raising=False) + + +def test_claim_file_first_writer_wins(monkeypatch, tmp_path): + path = tmp_path / "indexnow_key.txt" + _reset_key_state(monkeypatch, path) + + result = indexnow._claim_file("key-a") + + assert result == "key-a" + assert path.read_text() == "key-a" + assert list(tmp_path.iterdir()) == [path] # no leftover temp file + + +def test_claim_file_reads_back_winner_on_race(monkeypatch, tmp_path): + path = tmp_path / "indexnow_key.txt" + _reset_key_state(monkeypatch, path) + path.write_text("key-winner") # simulates another process winning first + + result = indexnow._claim_file("key-loser") + + assert result == "key-winner" + assert path.read_text() == "key-winner" # the loser never touched the file + + +def test_key_caches_winner_not_its_own_generation(monkeypatch, tmp_path): + path = tmp_path / "indexnow_key.txt" + _reset_key_state(monkeypatch, path) + + def _lose_the_race(*a, **k): + path.write_text("key-winner") + return "a" * 32 # what secrets.token_hex(16) would have produced locally + + monkeypatch.setattr(indexnow.secrets, "token_hex", _lose_the_race) + + result = indexnow.key() + + assert result == "key-winner" + assert indexnow._key == "key-winner" # the process-wide cache matches the file