refactor(telemetry): centralize stack slug validation + cardinality cap

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) <noreply@anthropic.com>
This commit is contained in:
Garm 2026-04-17 18:40:38 +02:00
parent d81b677f80
commit 5a2c484430
3 changed files with 119 additions and 8 deletions

View file

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

View file

@ -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_<agent>`` 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"

View file

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