From 0056a23e737b281b71e135ebd3096d733b651cc1 Mon Sep 17 00:00:00 2001 From: parthrohit22 Date: Fri, 28 Aug 2026 21:35:53 +0100 Subject: [PATCH 1/3] fix(api): rate-limit /ready and /metrics, expose db pool telemetry Closes the remaining scope of #296. The connection-checkout leak itself was already fixed by #306 (g.db + the teardown handler, with a real PostgreSQL-backed integration test proving connections balance across repeated readiness probes). What was still open: /ready and /metrics are unauthenticated by design (probe/scrape endpoints must never require a token), which also made them the one place an unauthenticated caller could trigger repeated pooled-connection work with no rate limiting at all, and there was no visibility into how close the pool was to exhaustion before it happened. - Add api.observability.probe_rate_limit: an in-memory, per-process, per-source-IP rate limiter for probe/scrape endpoints. Deliberately not the existing Postgres-backed api.rate_limit.rate_limit, which would add a database round trip (and a second, separately-tracked pooled connection under its own g.db) to the exact endpoint whose job is to protect the database from overload. The connection pool it guards is itself process-local under Gunicorn's multi-worker model, so a per-process budget is the matching granularity, not a weaker substitute for a shared one. Wired onto /ready (budget of 5 per 10s per source IP, half the default DB_POOL_MAX_CONN) and /metrics (20 per 10s, generous for normal Prometheus scrape intervals). The check runs before the view body, so a rejected request never reaches the database work it would otherwise trigger. - Add api.models.finding.get_pool_stats(): a point-in-time snapshot of the shared pool's in-use/idle/max-connection counts and utilization percentage. Reports only counts - never the DSN, host, or credentials - so it's safe on a public surface. Returns zeroed stats before any connection has been made instead of raising. - Wire get_pool_stats() into three new Prometheus gauges (openshield_db_pool_connections_in_use/idle/max), refreshed lazily on every /metrics scrape via a provider callback registered from api/app.py - api/observability.py stays free of project imports (its own documented constraint, since the worker reuses it too) by never importing api.models.finding directly. - Document in docs/deployment/render.md that Render's Blueprint format has no path-based access control, so the in-app rate limiter is a defense-in-depth backstop, not a substitute for restricting network reachability to /ready and /metrics at whatever reverse proxy/CDN/ WAF fronts a real deployment - that configuration is operational, outside what render.yaml can express. New tests in tests/test_readiness_hardening.py cover: the rate limiter's budget/window/per-IP-isolation/testing-bypass/key-pruning behavior in isolation, /ready actually rejecting a source once its budget is spent without touching the database for the rejected request, get_pool_stats()'s zero/nonzero/never-leaks-the-dsn behavior, and /metrics rendering the three new gauges (and surviving a broken stats provider without failing the whole scrape). Verified: full backend suite (796 passed, 3 skipped - pre-existing, unrelated), including tests/test_observability.py's real PostgreSQL-backed readiness-leak test run against a local Postgres instance to confirm the new decorator doesn't disturb #306's fix; ruff check and format --check clean. Signed-off-by: Parth J Rohit Signed-off-by: parthrohit22 --- CHANGELOG.md | 2 + api/app.py | 21 ++- api/models/finding.py | 38 +++++ api/observability.py | 129 ++++++++++++++- docs/deployment/render.md | 27 ++++ tests/test_readiness_hardening.py | 261 ++++++++++++++++++++++++++++++ 6 files changed, 475 insertions(+), 3 deletions(-) create mode 100644 tests/test_readiness_hardening.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ad685417..8d48dee3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,12 +21,14 @@ OpenShield uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - Harvest Now Decrypt Later exposure window calculation per cryptographic asset - Deterministic Render deployment workflow for separate API and worker services - Terraform configuration for Render, Vercel, and GitHub OIDC +- Database connection pool utilization and exhaustion telemetry (`openshield_db_pool_connections_*`) on `/metrics` ### Fixed - High-severity CodeQL findings in Python and JavaScript code - Security findings identified during Semgrep analysis - Sensitive identity metadata removed from scanner debug logging +- `/ready` and `/metrics` rate-limited per source IP so the unauthenticated probe/scrape surface can no longer be used to exhaust the database connection pool ### Security diff --git a/api/app.py b/api/app.py index 765682bb..2e57002f 100644 --- a/api/app.py +++ b/api/app.py @@ -10,8 +10,15 @@ from flask_cors import CORS from werkzeug.middleware.proxy_fix import ProxyFix -from api.models.finding import DatabaseManager -from api.observability import configure_logging, get_request_id, init_app, init_sentry +from api.models.finding import DatabaseManager, get_pool_stats +from api.observability import ( + configure_logging, + get_request_id, + init_app, + init_sentry, + probe_rate_limit, + set_pool_stats_provider, +) load_dotenv() @@ -32,6 +39,14 @@ _MAX_AUTHORIZATION_HEADER_LENGTH = 8192 _GENERATE_CMD = 'python -c "import secrets; print(secrets.token_urlsafe(32))"' +# Generous enough for legitimate manual or automated readiness checks from +# one source, but bounded well under the default pool size +# (DB_POOL_MAX_CONN=10) so a single caller can never claim more than half +# the pool's capacity by itself, even if every allowed request in the +# window lands at once. See probe_rate_limit's docstring for why this is +# in-memory rather than the shared Postgres-backed rate_limit(). +_READY_MAX_REQUESTS_PER_WINDOW = 5 + def _is_production() -> bool: return ( @@ -111,6 +126,7 @@ def create_app() -> Flask: # every later before_request handler (including JWT auth) and to the # error handlers. Also mounts the public /metrics endpoint. init_app(app) + set_pool_stats_provider(get_pool_stats) # ------------------------------------------------------------------ # # Configuration & Security # @@ -233,6 +249,7 @@ def health(): return jsonify({"status": "ok"}) @app.get("/ready") + @probe_rate_limit(_READY_MAX_REQUESTS_PER_WINDOW) def ready(): """Readiness probe: 200 when the database is reachable, else 503.""" try: diff --git a/api/models/finding.py b/api/models/finding.py index 0f366924..cc5506b2 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -46,6 +46,44 @@ def _get_pool(dsn: str) -> "psycopg2.pool.ThreadedConnectionPool": return pool +def get_pool_stats(dsn: Optional[str] = None) -> Dict[str, Any]: + """Return a point-in-time snapshot of the shared connection pool's utilization. + + Reports only counts (in-use / idle / configured maximum) - never the DSN, + host, credentials, or anything else that could describe the database + deployment. Safe to expose on a public surface (Prometheus /metrics, the + /ready probe) precisely because it carries no operational secrets, only + capacity numbers an operator needs to see the pool approaching exhaustion + before it happens. + + Returns zeroed stats with no pool created yet (no request has connected + since process start) rather than raising, so callers on the request path + (like /ready) never fail because of a stats lookup. + """ + dsn = dsn or os.environ.get("DATABASE_URL", "") + with _POOLS_LOCK: + pool = _POOLS.get(dsn) + + if pool is None: + return {"max_connections": _POOL_MAX_CONN, "in_use": 0, "idle": 0, "utilization_percent": 0.0} + + # psycopg2's pool has no public stats API; _used/_pool/maxconn are the + # same attributes getconn()/putconn() themselves mutate under this same + # lock, so a snapshot taken while holding it can't land mid-mutation. + with pool._lock: + in_use = len(pool._used) + idle = len(pool._pool) + max_connections = pool.maxconn + + utilization_percent = round((in_use / max_connections) * 100, 1) if max_connections else 0.0 + return { + "max_connections": max_connections, + "in_use": in_use, + "idle": idle, + "utilization_percent": utilization_percent, + } + + FRAMEWORK_FILE_MAP = { "cis": "cis_azure_benchmark.json", "nist": "nist_csf.json", diff --git a/api/observability.py b/api/observability.py index 4b7bcfa6..ddc6286b 100644 --- a/api/observability.py +++ b/api/observability.py @@ -14,10 +14,14 @@ import logging import os import re +import threading import time import uuid +from collections import deque +from functools import wraps +from typing import Callable, Deque, Dict, Optional, Tuple -from flask import Flask, Response, g, request +from flask import Flask, Response, current_app, g, jsonify, request from prometheus_client import ( CONTENT_TYPE_LATEST, Counter, @@ -82,6 +86,49 @@ "Latency of outbound LLM provider requests in seconds.", ["provider"], ) +DB_POOL_CONNECTIONS_IN_USE = Gauge( + "openshield_db_pool_connections_in_use", + "Connections currently checked out from the shared database pool.", +) +DB_POOL_CONNECTIONS_IDLE = Gauge( + "openshield_db_pool_connections_idle", + "Connections currently idle in the shared database pool.", +) +DB_POOL_CONNECTIONS_MAX = Gauge( + "openshield_db_pool_connections_max", + "Configured maximum size of the shared database pool.", +) + +# Set by api.app at startup (see set_pool_stats_provider). Kept as a plain +# module-level callable rather than importing api.models.finding directly - +# this module stays free of project imports (see module docstring) so it +# can be reused from the worker without pulling in the API's DB layer. +_pool_stats_provider: Optional[Callable[[], Dict[str, int]]] = None + + +def set_pool_stats_provider(fn: Callable[[], Dict[str, int]]) -> None: + """Register the callable /metrics uses to refresh the DB pool gauges. + + Called lazily on every scrape (see ``metrics()`` below) rather than on a + timer, so the numbers are exact at scrape time instead of aging between + scrapes. + """ + global _pool_stats_provider + _pool_stats_provider = fn + + +def _refresh_pool_metrics() -> None: + if _pool_stats_provider is None: + return + try: + stats = _pool_stats_provider() + DB_POOL_CONNECTIONS_IN_USE.set(stats["in_use"]) + DB_POOL_CONNECTIONS_IDLE.set(stats["idle"]) + DB_POOL_CONNECTIONS_MAX.set(stats["max_connections"]) + except Exception: + # A stats lookup must never take /metrics down - the rest of the + # scrape (HTTP/scan counters etc.) is still valid without it. + logger.exception("Failed to refresh database pool metrics") # --------------------------------------------------------------------------- # @@ -150,6 +197,84 @@ def get_request_id() -> str: return rid +# --------------------------------------------------------------------------- # +# Probe/scrape rate limiting # +# --------------------------------------------------------------------------- # +# /health, /ready, and /metrics are exempt from JWT auth by design (see +# api.app._ALWAYS_PUBLIC) so uptime checkers and Prometheus scrapers can +# reach them without a token. That also makes /ready and /metrics the one +# place an unauthenticated caller can trigger repeated backend work - a +# pooled database checkout for /ready - without ever presenting a token. +# +# api.rate_limit.rate_limit (used elsewhere in the API) is the wrong tool +# here: it does its own Postgres round trip per check, which would add +# database load to /ready - the exact endpoint whose job is to protect the +# database pool from overload - and risks checking out a second, separate +# pooled connection under its own `g.db` alongside the one /ready's handler +# already manages, only one of which the teardown handler would return. +# +# This limiter is pure in-memory and per-process instead. The connection +# pool it protects (api.models.finding._POOLS) is itself process-local under +# Gunicorn's multi-worker model, so a per-process budget is the matching +# granularity, not a weaker substitute for a shared one: it bounds exactly +# the pool a single worker process can exhaust. It is a defense-in-depth +# backstop, not a replacement for restricting these paths at whatever +# reverse proxy/CDN/WAF fronts the deployment - see the "Restricting +# probe/scrape endpoints at the edge" section of docs/deployment/render.md. +_PROBE_WINDOW_SECONDS = 10.0 +# Prometheus scrapes on a fixed interval (typically 15-30s) from a small, +# stable set of scraper IPs, so this stays generous; it exists to blunt a +# single caller hammering the endpoint, not to constrain normal scraping. +_METRICS_MAX_REQUESTS_PER_WINDOW = 20 +_probe_lock = threading.Lock() +_probe_hits: Dict[Tuple[str, str], Deque[float]] = {} + + +def probe_rate_limit(max_requests: int, window_seconds: float = _PROBE_WINDOW_SECONDS): + """Limit a probe/scrape view to ``max_requests`` per ``window_seconds`` per client IP. + + The check runs before the wrapped view body, so a request rejected here + never reaches the database work it would otherwise trigger. Disabled in + testing mode, matching the convention in api.rate_limit.rate_limit. + """ + + def decorator(fn): + @wraps(fn) + def wrapped(*args, **kwargs): + if current_app.testing: + return fn(*args, **kwargs) + + key = (request.remote_addr or "unknown", request.path) + now = time.monotonic() + with _probe_lock: + hits = _probe_hits.get(key) + if hits is not None: + cutoff = now - window_seconds + while hits and hits[0] < cutoff: + hits.popleft() + if not hits: + # Prune empty entries immediately rather than letting + # every distinct caller's key accumulate forever - + # a key only exists while it has a hit in-window. + del _probe_hits[key] + hits = None + if hits is None: + hits = deque() + allowed = len(hits) < max_requests + if allowed: + hits.append(now) + _probe_hits[key] = hits + + if not allowed: + return jsonify({"status": "rate_limited"}), 429 + + return fn(*args, **kwargs) + + return wrapped + + return decorator + + # --------------------------------------------------------------------------- # # Flask wiring # # --------------------------------------------------------------------------- # @@ -179,5 +304,7 @@ def _record_observability(response: Response) -> Response: return response @app.get("/metrics") + @probe_rate_limit(_METRICS_MAX_REQUESTS_PER_WINDOW) def metrics() -> Response: + _refresh_pool_metrics() return Response(generate_latest(), content_type=CONTENT_TYPE_LATEST) diff --git a/docs/deployment/render.md b/docs/deployment/render.md index 6a53c2f7..9e21e100 100644 --- a/docs/deployment/render.md +++ b/docs/deployment/render.md @@ -69,6 +69,33 @@ the Flask application. Leaving this value unset enables wildcard CORS and is not acceptable for a real staging or production environment. Do not add this setting to worker services. +## Restricting probe/scrape endpoints at the edge + +`/health`, `/ready`, and `/metrics` are intentionally exempt from JWT auth (see +`api.app._ALWAYS_PUBLIC`) so uptime checkers and Prometheus scrapers can reach them +without a token. `/health` is a pure liveness check with no backend dependency and is +meant to stay reachable from anywhere - it is what `render.yaml`'s +`healthCheckPath` uses. `/ready` and `/metrics` are different: `/ready` checks out a +pooled database connection on every call, and `/metrics` returns operational counters +(including database pool utilization from `openshield_db_pool_connections_*` - +counts only, never the DSN, host, or credentials - see `get_pool_stats()` in +`api/models/finding.py`). Neither should be reachable by arbitrary internet clients. + +Render's Blueprint format has no path-based access control, so this repository +cannot restrict `/ready`/`/metrics` from `render.yaml` itself. The application +carries its own in-process, in-memory rate limit on both paths +(`api.observability.probe_rate_limit`, wired in `api/app.py`) as a defense-in-depth +backstop - deliberately not the shared Postgres-backed `api.rate_limit.rate_limit` +used elsewhere, since that would add database load to the exact endpoint meant to +protect the database from overload. That backstop bounds a single source hammering +one process; it is not a substitute for restricting network reachability. + +Whoever operates the reverse proxy, CDN, or WAF in front of a Render deployment +(Cloudflare or similar, if one is configured) should restrict `/ready` and `/metrics` +to known monitoring/scraping source IPs or an internal network path, the same way +they would for any other backend-only operational endpoint. This is an operational +configuration step outside this repository, not something `render.yaml` can express. + ## Coordinated deterministic deployment For the selected environment, the workflow: diff --git a/tests/test_readiness_hardening.py b/tests/test_readiness_hardening.py new file mode 100644 index 00000000..60d05aef --- /dev/null +++ b/tests/test_readiness_hardening.py @@ -0,0 +1,261 @@ +"""Tests for issue #296's remaining scope: in-memory rate limiting for the +/ready and /metrics probe/scrape endpoints, and database pool utilization +telemetry. + +The connection-leak fix itself (returning the pool connection via `g.db` + +the teardown handler) and its regression coverage are #306 - see +test_observability.py::test_ready_returns_every_real_connection_to_the_pool. +This file only covers what was still open after that: bounding how much +pooled-connection work an unauthenticated caller can trigger per window, and +exposing pool in-use/idle/max counts without leaking anything about the +underlying database (DSN, host, credentials). +""" + +import threading +from unittest.mock import MagicMock + +import flask +import pytest + +import api.models.finding as finding_module +import api.observability as observability +from api.app import _READY_MAX_REQUESTS_PER_WINDOW, create_app +from api.models.finding import get_pool_stats +from api.observability import probe_rate_limit + + +@pytest.fixture(autouse=True) +def _clear_probe_rate_limit_state(): + """Isolate each test's rate-limit budget - the tracked-hits dict is a + module-level global shared across every app instance in this process.""" + observability._probe_hits.clear() + yield + observability._probe_hits.clear() + + +def _build_ready_app(db_mock, monkeypatch): + """Build a real app with /ready's rate limiter active (TESTING left + False - probe_rate_limit is a no-op under app.testing, and these tests + exist specifically to exercise it).""" + monkeypatch.setattr("api.app.DatabaseManager", MagicMock(return_value=db_mock)) + return create_app() + + +# --------------------------------------------------------------------------- # +# probe_rate_limit() - decorator behavior in isolation # +# --------------------------------------------------------------------------- # + + +def _tiny_app(max_requests=2, window_seconds=5.0): + app = flask.Flask(__name__) + + @app.get("/probe") + @probe_rate_limit(max_requests, window_seconds=window_seconds) + def probe(): + return {"ok": True} + + return app + + +def test_probe_rate_limit_allows_up_to_the_budget_then_429s(): + client = _tiny_app(max_requests=2).test_client() + + assert client.get("/probe").status_code == 200 + assert client.get("/probe").status_code == 200 + resp = client.get("/probe") + + assert resp.status_code == 429 + + +def test_probe_rate_limit_budget_resets_after_the_window(monkeypatch): + client = _tiny_app(max_requests=2, window_seconds=5.0).test_client() + fake_now = [1000.0] + monkeypatch.setattr(observability.time, "monotonic", lambda: fake_now[0]) + + assert client.get("/probe").status_code == 200 + assert client.get("/probe").status_code == 200 + assert client.get("/probe").status_code == 429 + + fake_now[0] += 5.1 # past the window + assert client.get("/probe").status_code == 200 + + +def test_probe_rate_limit_budget_is_independent_per_source_ip(): + client = _tiny_app(max_requests=2).test_client() + + assert client.get("/probe", environ_overrides={"REMOTE_ADDR": "203.0.113.5"}).status_code == 200 + assert client.get("/probe", environ_overrides={"REMOTE_ADDR": "203.0.113.5"}).status_code == 200 + assert client.get("/probe", environ_overrides={"REMOTE_ADDR": "203.0.113.5"}).status_code == 429 + + # A different source IP has its own, untouched budget. + assert client.get("/probe", environ_overrides={"REMOTE_ADDR": "198.51.100.9"}).status_code == 200 + + +def test_probe_rate_limit_is_a_noop_under_testing_mode(): + app = _tiny_app(max_requests=2) + app.config["TESTING"] = True + client = app.test_client() + + for _ in range(5): + assert client.get("/probe").status_code == 200 + + +def test_probe_rate_limit_prunes_expired_keys_instead_of_growing_forever(monkeypatch): + """A key with no hits inside the window must not linger in the tracking + dict - otherwise every distinct one-off caller would accumulate state + forever, which is its own memory-growth DoS vector.""" + client = _tiny_app(max_requests=1, window_seconds=1.0).test_client() + fake_now = [1000.0] + monkeypatch.setattr(observability.time, "monotonic", lambda: fake_now[0]) + + assert client.get("/probe", environ_overrides={"REMOTE_ADDR": "203.0.113.5"}).status_code == 200 + assert ("203.0.113.5", "/probe") in observability._probe_hits + + fake_now[0] += 1.1 + assert client.get("/probe", environ_overrides={"REMOTE_ADDR": "203.0.113.5"}).status_code == 200 + # The stale hit expired and only the new one remains tracked - the key + # was pruned and recreated, not left growing with dead timestamps. + assert len(observability._probe_hits[("203.0.113.5", "/probe")]) == 1 + + +# --------------------------------------------------------------------------- # +# /ready - rate limiting wired into the real route # +# --------------------------------------------------------------------------- # + + +def test_ready_rate_limits_a_single_source_after_the_window_budget(monkeypatch): + healthy_db = MagicMock() + healthy_db.ping.return_value = True + app = _build_ready_app(healthy_db, monkeypatch) + client = app.test_client() + + for _ in range(_READY_MAX_REQUESTS_PER_WINDOW): + assert client.get("/ready", environ_overrides={"REMOTE_ADDR": "203.0.113.5"}).status_code == 200 + + limited = client.get("/ready", environ_overrides={"REMOTE_ADDR": "203.0.113.5"}) + assert limited.status_code == 429 + # A rejected-before-checkout request never touches the database at all. + assert healthy_db.ping.call_count == _READY_MAX_REQUESTS_PER_WINDOW + + +def test_ready_rate_limit_does_not_affect_other_source_ips(monkeypatch): + healthy_db = MagicMock() + healthy_db.ping.return_value = True + app = _build_ready_app(healthy_db, monkeypatch) + client = app.test_client() + + for _ in range(_READY_MAX_REQUESTS_PER_WINDOW): + client.get("/ready", environ_overrides={"REMOTE_ADDR": "203.0.113.5"}) + assert client.get("/ready", environ_overrides={"REMOTE_ADDR": "203.0.113.5"}).status_code == 429 + + fresh_source = client.get("/ready", environ_overrides={"REMOTE_ADDR": "198.51.100.9"}) + assert fresh_source.status_code == 200 + + +# --------------------------------------------------------------------------- # +# get_pool_stats() # +# --------------------------------------------------------------------------- # + + +def test_pool_stats_report_zero_before_any_connection_is_made(): + stats = get_pool_stats("postgresql://unused-test-dsn/db") + + assert stats == { + "max_connections": finding_module._POOL_MAX_CONN, + "in_use": 0, + "idle": 0, + "utilization_percent": 0.0, + } + + +def test_pool_stats_reflect_checked_out_and_idle_connections(monkeypatch): + dsn = "postgresql://stats-test-dsn/db" + monkeypatch.setattr(finding_module, "_POOL_MAX_CONN", 4) + + fake_pool = MagicMock() + fake_pool.maxconn = 4 + fake_pool._used = {1: object(), 2: object()} + fake_pool._pool = [object()] + fake_pool._lock = threading.Lock() + + with finding_module._POOLS_LOCK: + finding_module._POOLS[dsn] = fake_pool + try: + stats = get_pool_stats(dsn) + finally: + with finding_module._POOLS_LOCK: + finding_module._POOLS.pop(dsn, None) + + assert stats == {"max_connections": 4, "in_use": 2, "idle": 1, "utilization_percent": 50.0} + + +def test_pool_stats_utilization_percent_handles_a_zero_size_pool(): + """A pool that hasn't been created yet reports the configured ceiling + with zero utilization, never a division by zero.""" + dsn = "postgresql://never-connected-test-dsn/db" + stats = get_pool_stats(dsn) + + assert stats["utilization_percent"] == 0.0 + + +def test_pool_stats_never_reveal_the_dsn_host_or_credentials(): + dsn = "postgresql://produser:s3cr3t-password@internal-db-host.example:5432/openshield" + stats = get_pool_stats(dsn) + + serialized = repr(stats) + assert "s3cr3t-password" not in serialized + assert "internal-db-host" not in serialized + assert "produser" not in serialized + assert set(stats.keys()) == {"max_connections", "in_use", "idle", "utilization_percent"} + + +# --------------------------------------------------------------------------- # +# /metrics - pool gauges reflect get_pool_stats() at scrape time # +# --------------------------------------------------------------------------- # + + +def test_metrics_endpoint_reports_pool_utilization_gauges(monkeypatch): + dsn = "postgresql://metrics-test-dsn/db" + monkeypatch.setenv("DATABASE_URL", dsn) + monkeypatch.setattr(finding_module, "_POOL_MAX_CONN", 4) + + fake_pool = MagicMock() + fake_pool.maxconn = 4 + fake_pool._used = {1: object(), 2: object(), 3: object()} + fake_pool._pool = [object()] + fake_pool._lock = threading.Lock() + + with finding_module._POOLS_LOCK: + finding_module._POOLS[dsn] = fake_pool + try: + app = create_app() + app.config["TESTING"] = True + resp = app.test_client().get("/metrics") + body = resp.get_data(as_text=True) + finally: + with finding_module._POOLS_LOCK: + finding_module._POOLS.pop(dsn, None) + + assert resp.status_code == 200 + assert "openshield_db_pool_connections_in_use 3.0" in body + assert "openshield_db_pool_connections_idle 1.0" in body + assert "openshield_db_pool_connections_max 4.0" in body + # The DSN itself must never appear in a public scrape endpoint's body. + assert dsn not in body + + +def test_metrics_endpoint_survives_a_pool_stats_failure(monkeypatch): + """A broken stats provider must not take the whole /metrics scrape down - + the HTTP/scan counters are still valid without the pool gauges.""" + + def _boom(): + raise RuntimeError("stats provider exploded") + + monkeypatch.setattr(observability, "_pool_stats_provider", _boom) + app = create_app() + app.config["TESTING"] = True + + resp = app.test_client().get("/metrics") + + assert resp.status_code == 200 + assert "openshield_scan_duration_seconds" in resp.get_data(as_text=True) From 5854316e8836bae759b1612ec9aebb9b5d5df006 Mon Sep 17 00:00:00 2001 From: parthrohit22 Date: Sat, 29 Aug 2026 14:05:56 +0100 Subject: [PATCH 2/3] fix(api): restore dropped constant, bound probe-limiter memory Two fixes on top of the dev merge that landed on this branch: 1. The merge into dev (which now carries #294/#320's own changes to api/app.py) dropped this branch's _READY_MAX_REQUESTS_PER_WINDOW constant definition while keeping its usage on the /ready route, leaving api/app.py with an undefined name that only surfaced at create_app() call time (ruff's F821 caught it as CI's first failure; Backend Tests failed for the same underlying reason). Restored the constant and its comment. 2. m-khan-97's review: probe_rate_limit's per-key cleanup only ever prunes the exact (address, path) key the current request touches. A caller that continuously rotates its source address - or a spoofed forwarded address wherever the trusted-proxy boundary is misconfigured - creates a new one-shot dictionary entry per address that's never revisited and therefore never pruned, making the limiter's own tracking dict an unbounded memory sink. Fixed with the two things asked for: - A periodic global sweep (every _PROBE_SWEEP_INTERVAL calls, not every call - a full-dict scan per request would defeat the point of a cheap in-memory limiter) that prunes every key whose hits have all expired, not just the current request's key. - A hard cap (_PROBE_MAX_TRACKED_KEYS) on distinct tracked keys, with deterministic least-recently-touched eviction via an OrderedDict instead of the previous plain dict - every hit (including one that just survives a sweep) moves its key to the end, so eviction always drops the coldest entry first. Also added the suggested (non-blocking) Retry-After header on 429. New tests in tests/test_readiness_hardening.py cover exactly the scenario m-khan-97 described: many distinct one-shot addresses, advance past the window, trigger cleanup through different addresses, and prove the stale keys are gone and the map stays bounded. Plus direct hard-cap/LRU-eviction tests and the Retry-After header. Verified: full backend suite (862 passed, 5 skipped - pre-existing/ environment-only), including all 17 tests in tests/test_readiness_hardening.py and the pre-existing tests/test_auth.py / tests/test_observability.py suites unaffected. ruff check and format --check clean. Signed-off-by: Parth J Rohit Signed-off-by: parthrohit22 --- api/app.py | 8 +++ api/observability.py | 63 +++++++++++++++++++--- tests/test_readiness_hardening.py | 88 +++++++++++++++++++++++++++++++ 3 files changed, 152 insertions(+), 7 deletions(-) diff --git a/api/app.py b/api/app.py index 5de9df14..75c30521 100644 --- a/api/app.py +++ b/api/app.py @@ -49,6 +49,14 @@ _KNOWN_ROLES = {"viewer", "operator", "admin"} _WRITE_ROLES = {"operator", "admin"} +# Generous enough for legitimate manual or automated readiness checks from +# one source, but bounded well under the default pool size +# (DB_POOL_MAX_CONN=10) so a single caller can never claim more than half +# the pool's capacity by itself, even if every allowed request in the +# window lands at once. See probe_rate_limit's docstring for why this is +# in-memory rather than the shared Postgres-backed rate_limit(). +_READY_MAX_REQUESTS_PER_WINDOW = 5 + def _is_production() -> bool: return ( diff --git a/api/observability.py b/api/observability.py index ddc6286b..ab0d4a1a 100644 --- a/api/observability.py +++ b/api/observability.py @@ -17,7 +17,7 @@ import threading import time import uuid -from collections import deque +from collections import OrderedDict, deque from functools import wraps from typing import Callable, Deque, Dict, Optional, Tuple @@ -226,8 +226,45 @@ def get_request_id() -> str: # stable set of scraper IPs, so this stays generous; it exists to blunt a # single caller hammering the endpoint, not to constrain normal scraping. _METRICS_MAX_REQUESTS_PER_WINDOW = 20 +# Hard ceiling on distinct (address, path) keys tracked at once. Same-key +# cleanup alone isn't enough: a caller that continuously rotates its source +# address (or a spoofed forwarded address wherever the trusted-proxy +# boundary is misconfigured) creates a new one-shot dict entry per address, +# and a key that's never revisited is never pruned by the per-call cleanup +# below - making the tracking dict itself an unbounded memory sink. This +# caps it regardless of how many distinct addresses show up. +_PROBE_MAX_TRACKED_KEYS = 10_000 +# A full scan over every tracked key on every single request would undercut +# the point of a cheap in-memory limiter, so the sweep below runs every Nth +# call instead of every call. This bounds how long a key whose owner never +# returns can survive to at most this many requests' worth of accumulation +# - the hard cap above is what actually bounds worst-case memory regardless +# of traffic shape or how the sweep is scheduled. +_PROBE_SWEEP_INTERVAL = 200 _probe_lock = threading.Lock() -_probe_hits: Dict[Tuple[str, str], Deque[float]] = {} +# An OrderedDict, not a plain dict: every touch (new hit or a sweep pruning +# it survives) moves a key to the end, so the front is always the least +# recently touched key - the correct, deterministic thing to evict first +# when the hard cap above is reached. +_probe_hits: "OrderedDict[Tuple[str, str], Deque[float]]" = OrderedDict() +_probe_call_count = 0 + + +def _sweep_expired_probe_hits(now: float, window_seconds: float) -> None: + """Remove every tracked key whose hits have all expired. + + Must be called with _probe_lock already held. This is the global + cleanup pass: the per-call logic in probe_rate_limit only ever prunes + the one key the current request touched, which does nothing for a + key that's never hit again. + """ + cutoff = now - window_seconds + for key in list(_probe_hits.keys()): + hits = _probe_hits[key] + while hits and hits[0] < cutoff: + hits.popleft() + if not hits: + del _probe_hits[key] def probe_rate_limit(max_requests: int, window_seconds: float = _PROBE_WINDOW_SECONDS): @@ -244,29 +281,41 @@ def wrapped(*args, **kwargs): if current_app.testing: return fn(*args, **kwargs) + global _probe_call_count key = (request.remote_addr or "unknown", request.path) now = time.monotonic() with _probe_lock: + _probe_call_count += 1 + if _probe_call_count % _PROBE_SWEEP_INTERVAL == 0: + _sweep_expired_probe_hits(now, window_seconds) + hits = _probe_hits.get(key) if hits is not None: cutoff = now - window_seconds while hits and hits[0] < cutoff: hits.popleft() if not hits: - # Prune empty entries immediately rather than letting - # every distinct caller's key accumulate forever - - # a key only exists while it has a hit in-window. del _probe_hits[key] hits = None + if hits is None: + if len(_probe_hits) >= _PROBE_MAX_TRACKED_KEYS: + # Deterministic eviction: drop the least recently + # touched key, not an arbitrary/insertion-order one. + _probe_hits.popitem(last=False) hits = deque() + allowed = len(hits) < max_requests if allowed: hits.append(now) - _probe_hits[key] = hits + # Re-inserting (or just touching) moves this key to the end - + # it's the most recently active key, so it's the last thing + # eviction should reach for. + _probe_hits[key] = hits + _probe_hits.move_to_end(key) if not allowed: - return jsonify({"status": "rate_limited"}), 429 + return jsonify({"status": "rate_limited"}), 429, {"Retry-After": str(int(window_seconds))} return fn(*args, **kwargs) diff --git a/tests/test_readiness_hardening.py b/tests/test_readiness_hardening.py index 60d05aef..9d8f85fc 100644 --- a/tests/test_readiness_hardening.py +++ b/tests/test_readiness_hardening.py @@ -118,6 +118,94 @@ def test_probe_rate_limit_prunes_expired_keys_instead_of_growing_forever(monkeyp assert len(observability._probe_hits[("203.0.113.5", "/probe")]) == 1 +def test_probe_rate_limit_sweeps_stale_one_shot_addresses_globally(monkeypatch): + """A caller that rotates its source address to defeat same-key cleanup + (issue: each address is a one-shot key that's never revisited, so it's + never pruned) must still be bounded - the periodic global sweep has to + catch keys nobody ever hits a second time, not just the one the current + request touched.""" + monkeypatch.setattr(observability, "_PROBE_SWEEP_INTERVAL", 5) + client = _tiny_app(max_requests=1, window_seconds=1.0).test_client() + fake_now = [1000.0] + monkeypatch.setattr(observability.time, "monotonic", lambda: fake_now[0]) + + # 20 distinct one-shot source addresses, each hit exactly once - none of + # them is ever revisited, so per-key cleanup alone would never touch them. + for i in range(20): + addr = f"203.0.113.{i}" + assert client.get("/probe", environ_overrides={"REMOTE_ADDR": addr}).status_code == 200 + assert len(observability._probe_hits) == 20 + + # Advance well past the window, then trigger enough calls (from yet more + # different addresses) for the sweep interval to fire. + fake_now[0] += 5.0 + for i in range(20, 25): + client.get("/probe", environ_overrides={"REMOTE_ADDR": f"203.0.113.{i}"}) + + # Every one of the original 20 stale one-shot keys is gone. Only recent + # activity (some subset of the newest addresses) remains tracked. + for i in range(20): + assert ("203.0.113.{}".format(i), "/probe") not in observability._probe_hits + assert len(observability._probe_hits) <= 5 + + +def test_probe_rate_limit_enforces_a_hard_cap_with_lru_eviction(monkeypatch): + """Even within the window (nothing has expired yet), the number of + tracked keys must never exceed the hard cap - bounds worst-case memory + against a burst of many distinct addresses regardless of timing.""" + monkeypatch.setattr(observability, "_PROBE_MAX_TRACKED_KEYS", 3) + monkeypatch.setattr(observability, "_PROBE_SWEEP_INTERVAL", 1_000_000) # don't let the sweep interfere + client = _tiny_app(max_requests=5, window_seconds=100.0).test_client() + + client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.1"}) + client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.2"}) + client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.3"}) + assert list(observability._probe_hits.keys()) == [ + ("10.0.0.1", "/probe"), + ("10.0.0.2", "/probe"), + ("10.0.0.3", "/probe"), + ] + + # A 4th distinct address at the cap evicts the least-recently-touched + # key (10.0.0.1, never touched again) rather than growing past the cap. + client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.4"}) + assert len(observability._probe_hits) == 3 + assert ("10.0.0.1", "/probe") not in observability._probe_hits + assert ("10.0.0.4", "/probe") in observability._probe_hits + + +def test_probe_rate_limit_touching_a_key_protects_it_from_eviction(monkeypatch): + """Re-hitting an existing key must move it to the back of the eviction + order - an active caller must not be evicted just because other callers + showed up afterward.""" + monkeypatch.setattr(observability, "_PROBE_MAX_TRACKED_KEYS", 2) + monkeypatch.setattr(observability, "_PROBE_SWEEP_INTERVAL", 1_000_000) + client = _tiny_app(max_requests=5, window_seconds=100.0).test_client() + + client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.1"}) + client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.2"}) + # Touch 10.0.0.1 again - it's now the most recently active, not the + # least, even though it was inserted first. + client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.1"}) + + # A 3rd distinct address should evict 10.0.0.2 (untouched since its + # first hit), not 10.0.0.1 (touched most recently). + client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.3"}) + assert ("10.0.0.1", "/probe") in observability._probe_hits + assert ("10.0.0.2", "/probe") not in observability._probe_hits + assert ("10.0.0.3", "/probe") in observability._probe_hits + + +def test_probe_rate_limit_429_includes_retry_after_header(): + client = _tiny_app(max_requests=1, window_seconds=7.0).test_client() + + client.get("/probe") + resp = client.get("/probe") + + assert resp.status_code == 429 + assert resp.headers.get("Retry-After") == "7" + + # --------------------------------------------------------------------------- # # /ready - rate limiting wired into the real route # # --------------------------------------------------------------------------- # From 30249c01f3480ece60a8eeb6461af07d19c4d242 Mon Sep 17 00:00:00 2001 From: parthrohit22 Date: Tue, 1 Sep 2026 15:37:56 +0100 Subject: [PATCH 3/3] fix(api): address TFT444's 3 probe-limiter findings 1. move_to_end(key) ran on every request including rejected ones, so an attacker hammering an already-exhausted key kept it permanently at the back of the eviction order while quiet, legitimate keys drifted toward the front and got evicted instead - defeating the memory cap's actual purpose for exactly the caller it exists to bound. Now only a request that counts against the budget (allowed) refreshes a key's position; a rejected request leaves it wherever it already was. 2. _probe_hits is one dict shared by every probe_rate_limit-decorated endpoint, but the periodic sweep applied whichever endpoint's request happened to trigger it - its own window_seconds - to every tracked key regardless of which endpoint's window actually applies to it. Dormant today since /ready and /metrics both default to the same 10s window, but the first endpoint added with a different one would have caused premature resets or lingering stale entries across every other endpoint's keys. Added _ProbeEntry to carry window_seconds alongside each key's own hits deque, so the sweep (and the per-call prune) always uses the window that key was actually registered under. 3. Documented the ProxyFix(x_for=1) trust boundary explicitly in api/app.py: what it assumes (Render's edge is the only thing able to append to X-Forwarded-For before this process sees it) and what breaks if that's violated (a directly-reachable origin lets a caller set their own forwarded IP per request, which is equivalent to no per-IP rate limiting at all for every control that depends on request.remote_addr). New regression tests for both behavioral fixes: one proves a caller hammering an exhausted key doesn't stay artificially warm while a quiet legitimate key gets evicted in its place; the other runs two endpoints with different window_seconds sharing the tracked-key dict and proves a sweep triggered by one doesn't misapply its window to the other's entries. Verified: full backend suite (864 passed, 5 skipped - pre-existing/ environment-only), all 19 tests in tests/test_readiness_hardening.py including the 2 new ones, ruff check and format --check clean. Signed-off-by: Parth J Rohit Signed-off-by: parthrohit22 --- api/app.py | 16 +++++- api/observability.py | 91 +++++++++++++++++++++---------- tests/test_readiness_hardening.py | 75 ++++++++++++++++++++++++- 3 files changed, 150 insertions(+), 32 deletions(-) diff --git a/api/app.py b/api/app.py index 75c30521..0afc9033 100644 --- a/api/app.py +++ b/api/app.py @@ -126,7 +126,21 @@ def create_app() -> Flask: # Trust exactly one reverse-proxy hop (Render's edge) for the client IP # and scheme, so request.remote_addr reflects the real caller instead of # collapsing every client onto Render's proxy address. Rate limiting and - # any other per-IP logic depend on this being accurate. + # any other per-IP logic (api.observability.probe_rate_limit, + # api.rate_limit.rate_limit) depend on this being accurate. + # + # This is a trust boundary, not just a convenience setting: x_for=1 makes + # Flask take the *last* entry of an inbound X-Forwarded-For header as the + # real client IP, on the assumption that Render's edge is the only thing + # capable of appending to it before the request reaches this process. If + # the origin were ever reachable directly - bypassing Render's edge, e.g. + # a misconfigured DNS record or a leaked origin IP - a direct caller's own + # X-Forwarded-For header would be trusted as-is, and they could set it to + # a fresh IP on every request. Every per-IP control in this file (the + # probe-endpoint limiter, the Postgres-backed rate limiter) would then + # bucket each request as a "new" caller, which is equivalent to no rate + # limiting for that path at all. Keeping the origin unreachable except + # through Render's edge is what this setting's correctness depends on. app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1) # ------------------------------------------------------------------ # diff --git a/api/observability.py b/api/observability.py index ab0d4a1a..0d00f06c 100644 --- a/api/observability.py +++ b/api/observability.py @@ -242,28 +242,52 @@ def get_request_id() -> str: # of traffic shape or how the sweep is scheduled. _PROBE_SWEEP_INTERVAL = 200 _probe_lock = threading.Lock() -# An OrderedDict, not a plain dict: every touch (new hit or a sweep pruning -# it survives) moves a key to the end, so the front is always the least -# recently touched key - the correct, deterministic thing to evict first -# when the hard cap above is reached. -_probe_hits: "OrderedDict[Tuple[str, str], Deque[float]]" = OrderedDict() + + +class _ProbeEntry: + """One tracked (address, path) key's rate-limit state. + + window_seconds is stored per entry, not taken from whichever request + happens to trigger the periodic sweep: _probe_hits is one dict shared + across every probe_rate_limit-decorated view, so a sweep triggered by + /metrics's decorator (its own window_seconds closure) must not apply + that window to keys tracked for /ready, or vice versa. Both endpoints + currently default to the same 10s window, which is exactly why this + was dormant rather than visibly broken - the first endpoint added with + a different window would have hit premature resets or lingering stale + entries for every other endpoint's keys. + """ + + __slots__ = ("hits", "window_seconds") + + def __init__(self, window_seconds: float) -> None: + self.hits: Deque[float] = deque() + self.window_seconds = window_seconds + + +# An OrderedDict, not a plain dict: every touch (a request actually within +# budget, or a sweep pruning it survives) moves a key to the end, so the +# front is always the least recently touched key - the correct, +# deterministic thing to evict first when the hard cap above is reached. +_probe_hits: "OrderedDict[Tuple[str, str], _ProbeEntry]" = OrderedDict() _probe_call_count = 0 -def _sweep_expired_probe_hits(now: float, window_seconds: float) -> None: +def _sweep_expired_probe_hits(now: float) -> None: """Remove every tracked key whose hits have all expired. Must be called with _probe_lock already held. This is the global cleanup pass: the per-call logic in probe_rate_limit only ever prunes the one key the current request touched, which does nothing for a - key that's never hit again. + key that's never hit again. Each entry's own window_seconds is used, + not the sweep caller's - see _ProbeEntry. """ - cutoff = now - window_seconds for key in list(_probe_hits.keys()): - hits = _probe_hits[key] - while hits and hits[0] < cutoff: - hits.popleft() - if not hits: + entry = _probe_hits[key] + cutoff = now - entry.window_seconds + while entry.hits and entry.hits[0] < cutoff: + entry.hits.popleft() + if not entry.hits: del _probe_hits[key] @@ -287,32 +311,39 @@ def wrapped(*args, **kwargs): with _probe_lock: _probe_call_count += 1 if _probe_call_count % _PROBE_SWEEP_INTERVAL == 0: - _sweep_expired_probe_hits(now, window_seconds) - - hits = _probe_hits.get(key) - if hits is not None: - cutoff = now - window_seconds - while hits and hits[0] < cutoff: - hits.popleft() - if not hits: + _sweep_expired_probe_hits(now) + + entry = _probe_hits.get(key) + if entry is not None: + cutoff = now - entry.window_seconds + while entry.hits and entry.hits[0] < cutoff: + entry.hits.popleft() + if not entry.hits: del _probe_hits[key] - hits = None + entry = None - if hits is None: + if entry is None: if len(_probe_hits) >= _PROBE_MAX_TRACKED_KEYS: # Deterministic eviction: drop the least recently # touched key, not an arbitrary/insertion-order one. _probe_hits.popitem(last=False) - hits = deque() + entry = _ProbeEntry(window_seconds) + # OrderedDict places a newly-inserted key at the end, + # so a brand-new entry doesn't need an explicit + # move_to_end - it's already the most recent. + _probe_hits[key] = entry - allowed = len(hits) < max_requests + allowed = len(entry.hits) < max_requests if allowed: - hits.append(now) - # Re-inserting (or just touching) moves this key to the end - - # it's the most recently active key, so it's the last thing - # eviction should reach for. - _probe_hits[key] = hits - _probe_hits.move_to_end(key) + entry.hits.append(now) + # Only a request that actually counted against the + # budget refreshes this key's position. A rejected + # request must not keep an over-budget key artificially + # warm - otherwise a caller that never lets its own + # budget recover (an attacker) stays permanently + # protected from eviction while quiet, legitimate + # keys drift toward the front and get evicted instead. + _probe_hits.move_to_end(key) if not allowed: return jsonify({"status": "rate_limited"}), 429, {"Retry-After": str(int(window_seconds))} diff --git a/tests/test_readiness_hardening.py b/tests/test_readiness_hardening.py index 9d8f85fc..75c9fde1 100644 --- a/tests/test_readiness_hardening.py +++ b/tests/test_readiness_hardening.py @@ -115,7 +115,7 @@ def test_probe_rate_limit_prunes_expired_keys_instead_of_growing_forever(monkeyp assert client.get("/probe", environ_overrides={"REMOTE_ADDR": "203.0.113.5"}).status_code == 200 # The stale hit expired and only the new one remains tracked - the key # was pruned and recreated, not left growing with dead timestamps. - assert len(observability._probe_hits[("203.0.113.5", "/probe")]) == 1 + assert len(observability._probe_hits[("203.0.113.5", "/probe")].hits) == 1 def test_probe_rate_limit_sweeps_stale_one_shot_addresses_globally(monkeypatch): @@ -196,6 +196,79 @@ def test_probe_rate_limit_touching_a_key_protects_it_from_eviction(monkeypatch): assert ("10.0.0.3", "/probe") in observability._probe_hits +def test_probe_rate_limit_rejected_requests_do_not_refresh_eviction_order(monkeypatch): + """A caller already over budget (rejected with 429) must not get its key + moved to the back of the eviction order just by hammering the endpoint - + otherwise an attacker who never lets their own budget recover stays + permanently protected from eviction, while quiet legitimate keys drift + toward the front and get evicted in its place. Only a request that + actually counts against the budget may refresh a key's position.""" + monkeypatch.setattr(observability, "_PROBE_MAX_TRACKED_KEYS", 2) + monkeypatch.setattr(observability, "_PROBE_SWEEP_INTERVAL", 1_000_000) + client = _tiny_app(max_requests=1, window_seconds=100.0).test_client() + + # 10.0.0.1 spends its one-request budget, then keeps hammering - every + # further call is rejected (429), not counted toward the budget again. + assert client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.1"}).status_code == 200 + for _ in range(5): + assert client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.1"}).status_code == 429 + + # A quiet second address takes the other slot at the cap. + client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.2"}) + + # A 3rd distinct address arrives. If the repeated 429s had kept + # 10.0.0.1 "warm", 10.0.0.2 (touched once, then quiet) would be the + # one evicted instead - exactly backwards from what should happen. + client.get("/probe", environ_overrides={"REMOTE_ADDR": "10.0.0.3"}) + assert ("10.0.0.1", "/probe") not in observability._probe_hits + assert ("10.0.0.2", "/probe") in observability._probe_hits + assert ("10.0.0.3", "/probe") in observability._probe_hits + + +def test_probe_rate_limit_sweep_uses_each_keys_own_window_not_the_callers(monkeypatch): + """_probe_hits is one dict shared by every probe_rate_limit-decorated + endpoint. A sweep triggered by one endpoint's request must apply each + tracked key's own window_seconds, not the window of whichever endpoint + happened to trigger the sweep - otherwise a short-window endpoint's + sweep could prematurely wipe a long-window endpoint's still-valid + entries, or a long-window endpoint's sweep could leave a short-window + endpoint's genuinely-expired entries lingering.""" + monkeypatch.setattr(observability, "_PROBE_SWEEP_INTERVAL", 1) + fake_now = [1000.0] + monkeypatch.setattr(observability.time, "monotonic", lambda: fake_now[0]) + + app = flask.Flask(__name__) + + @app.get("/short") + @probe_rate_limit(5, window_seconds=1.0) + def short_window(): + return {"ok": True} + + @app.get("/long") + @probe_rate_limit(5, window_seconds=1000.0) + def long_window(): + return {"ok": True} + + client = app.test_client() + + # Seed one key per endpoint. + client.get("/short", environ_overrides={"REMOTE_ADDR": "10.0.0.1"}) + client.get("/long", environ_overrides={"REMOTE_ADDR": "10.0.0.2"}) + assert ("10.0.0.1", "/short") in observability._probe_hits + assert ("10.0.0.2", "/long") in observability._probe_hits + + # Advance past the short window but nowhere near the long one, then + # trigger a sweep via a request to the *long*-window endpoint. A sweep + # that used the triggering request's own window (1000s) would wrongly + # treat /short's entry as unexpired; using each entry's own window + # correctly expires it while leaving /long's entry alone. + fake_now[0] += 2.0 + client.get("/long", environ_overrides={"REMOTE_ADDR": "10.0.0.3"}) # sweep interval is 1, fires every call + + assert ("10.0.0.1", "/short") not in observability._probe_hits + assert ("10.0.0.2", "/long") in observability._probe_hits + + def test_probe_rate_limit_429_includes_retry_after_header(): client = _tiny_app(max_requests=1, window_seconds=7.0).test_client()