From 5a2c48443049e080ce0604dee9ca46e8423b7c7c Mon Sep 17 00:00:00 2001 From: Garm Date: Fri, 17 Apr 2026 18:40:38 +0200 Subject: [PATCH] refactor(telemetry): centralize stack slug validation + cardinality cap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review feedback on PR #194: prior code validated X-Headroom-Stack slugs differently at each ingress — the Prometheus counter capped length but the env-var path and beacon didn't. Under a misconfigured or malicious client sending arbitrary X-Headroom-Stack values, requests_by_stack could grow unbounded in memory, in the Prometheus scrape, and in the JSONB telemetry payload. - New normalize_stack(raw) in headroom/telemetry/context.py: strips, lowercases, enforces ^[a-z][a-z0-9_]{0,63}$. Single chokepoint. - detect_stack now routes HEADROOM_STACK and stats dominant-slug through it; invalid env values fall through to the agent-type path or default "proxy". - PrometheusMetrics.record_stack routes through normalize_stack and rejects new slugs once the dict hits MAX_DISTINCT_STACKS (32); existing slugs still increment so valid callers aren't starved. - 10 new unit tests covering normalize_stack charset/length/empty cases, the cardinality cap, and invalid-env fallback paths. Co-Authored-By: Claude Opus 4.7 (1M context) --- headroom/proxy/prometheus_metrics.py | 14 +++-- headroom/telemetry/context.py | 36 +++++++++++-- tests/test_telemetry_context.py | 77 +++++++++++++++++++++++++++- 3 files changed, 119 insertions(+), 8 deletions(-) diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index 6752f1f6b..40e49ecd6 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -199,13 +199,19 @@ class PrometheusMetrics: ``stack`` is the ``X-Headroom-Stack`` header value (e.g. ``adapter_ts_openai``). Called once per inbound request from the - proxy's stack middleware; a no-op when the header is absent. + proxy's stack middleware; a no-op when the header is absent, fails + validation, or would exceed the cardinality cap. """ - if not stack: + from headroom.telemetry.context import MAX_DISTINCT_STACKS, normalize_stack + + slug = normalize_stack(stack) + if not slug: return - slug = stack.strip().lower() - if not slug or len(slug) > 64: + if ( + slug not in self.requests_by_stack + and len(self.requests_by_stack) >= MAX_DISTINCT_STACKS + ): return self.requests_by_stack[slug] += 1 diff --git a/headroom/telemetry/context.py b/headroom/telemetry/context.py index a9546bb2a..f3494a2b8 100644 --- a/headroom/telemetry/context.py +++ b/headroom/telemetry/context.py @@ -15,6 +15,7 @@ from __future__ import annotations import logging import os +import re from typing import Any logger = logging.getLogger(__name__) @@ -22,6 +23,35 @@ logger = logging.getLogger(__name__) _KNOWN_WRAP_AGENTS = frozenset({"claude", "copilot", "codex", "aider", "cursor", "openclaw"}) +# Stack slugs must start with a letter and contain only [a-z0-9_], max 64 chars. +# Applied at every ingress (env var, HTTP header, stats aggregation) so downstream +# sinks (Prometheus labels, Supabase column, JSONB payload) see a bounded vocabulary. +_STACK_SLUG_RE = re.compile(r"^[a-z][a-z0-9_]{0,63}$") + +# Cardinality cap on the per-process requests_by_stack dict. Protects the +# Prometheus scrape, the in-memory counter, and the JSONB telemetry payload +# from unbounded label explosion when clients send arbitrary X-Headroom-Stack +# header values. +MAX_DISTINCT_STACKS = 32 + + +def normalize_stack(raw: str | None) -> str | None: + """Validate and normalize a stack slug. + + Returns the lowercased/stripped slug if it matches ``^[a-z][a-z0-9_]{0,63}$``, + else ``None``. All external stack identifiers (env var, HTTP header, stats + keys) must pass through this function — it is the single chokepoint that + bounds cardinality and rejects garbage before it reaches Prometheus or the + Supabase telemetry row. + """ + + if not raw: + return None + slug = raw.strip().lower() + if not _STACK_SLUG_RE.match(slug): + return None + return slug + def _slug_from_agent_type(agent_type: str) -> str: """Return ``wrap_`` for known agents, otherwise ``unknown``.""" @@ -82,9 +112,9 @@ def detect_stack(stats: dict[str, Any] | None = None) -> str: """ try: - explicit = os.environ.get("HEADROOM_STACK") + explicit = normalize_stack(os.environ.get("HEADROOM_STACK")) if explicit: - return explicit.strip().lower() + return explicit agent_type = os.environ.get("HEADROOM_AGENT_TYPE") if agent_type: @@ -97,7 +127,7 @@ def detect_stack(stats: dict[str, Any] | None = None) -> str: if total > 0: dominant, count = max(by_stack.items(), key=lambda kv: kv[1]) if count / total >= 0.8: - return str(dominant) + return normalize_stack(str(dominant)) or "unknown" return "mixed" return "proxy" diff --git a/tests/test_telemetry_context.py b/tests/test_telemetry_context.py index 1d055a8c0..9658abdf7 100644 --- a/tests/test_telemetry_context.py +++ b/tests/test_telemetry_context.py @@ -6,7 +6,12 @@ from types import SimpleNamespace import pytest -from headroom.telemetry.context import detect_install_mode, detect_stack +from headroom.telemetry.context import ( + MAX_DISTINCT_STACKS, + detect_install_mode, + detect_stack, + normalize_stack, +) @pytest.fixture(autouse=True) @@ -92,3 +97,73 @@ class TestDetectStack: monkeypatch.setenv("HEADROOM_STACK", "wrap_claude") stats = {"requests": {"by_stack": {"adapter_ts_openai": 100}}} assert detect_stack(stats) == "wrap_claude" + + def test_invalid_env_falls_through_to_proxy(self, monkeypatch): + # Garbage env var → normalize_stack rejects → falls back to default + monkeypatch.setenv("HEADROOM_STACK", "bad slug with spaces!") + assert detect_stack() == "proxy" + + def test_invalid_env_allows_agent_type_fallback(self, monkeypatch): + monkeypatch.setenv("HEADROOM_STACK", "Has-Dashes-And-Caps") + monkeypatch.setenv("HEADROOM_AGENT_TYPE", "claude") + assert detect_stack() == "wrap_claude" + + +class TestNormalizeStack: + def test_empty_and_none(self): + assert normalize_stack(None) is None + assert normalize_stack("") is None + assert normalize_stack(" ") is None + + def test_lowercases_and_strips(self): + assert normalize_stack(" Wrap_Claude ") == "wrap_claude" + + def test_rejects_invalid_charset(self): + assert normalize_stack("has-dashes") is None + assert normalize_stack("has spaces") is None + assert normalize_stack("has.dots") is None + assert normalize_stack("has/slashes") is None + assert normalize_stack("1_starts_with_digit") is None + + def test_accepts_valid_slugs(self): + for slug in ("proxy", "wrap_claude", "adapter_ts_openai", "a", "a1_2_3"): + assert normalize_stack(slug) == slug + + def test_rejects_over_64_chars(self): + assert normalize_stack("a" * 64) == "a" * 64 + assert normalize_stack("a" * 65) is None + + +class TestRecordStackValidation: + """PrometheusMetrics.record_stack must route through normalize_stack and + respect the cardinality cap.""" + + def _metrics(self): + from headroom.proxy.prometheus_metrics import PrometheusMetrics + + return PrometheusMetrics() + + def test_ignores_invalid_slug(self): + m = self._metrics() + m.record_stack("bad slug!") + m.record_stack("has-dashes") + m.record_stack("") + m.record_stack(None) + assert dict(m.requests_by_stack) == {} + + def test_counts_valid_slug(self): + m = self._metrics() + m.record_stack("wrap_claude") + m.record_stack("WRAP_CLAUDE") + assert m.requests_by_stack["wrap_claude"] == 2 + + def test_cardinality_cap_rejects_new_slugs(self): + m = self._metrics() + for i in range(MAX_DISTINCT_STACKS): + m.record_stack(f"slug_{i}") + assert len(m.requests_by_stack) == MAX_DISTINCT_STACKS + m.record_stack("slug_overflow") + assert "slug_overflow" not in m.requests_by_stack + # but existing slugs still increment + m.record_stack("slug_0") + assert m.requests_by_stack["slug_0"] == 2