mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(proxy): pilot hardening — inbound auth, security headers, audit log, air-gap switch (#1537)
## Description Tier-2 pilot security hardening from the engineering plan (Tier-1 landed in #1515). Four operator-facing controls, each verified open on `main` and grounded in a real exposure or enterprise requirement rather than a form answer: - **Optional inbound auth token (`HEADROOM_PROXY_TOKEN`).** When set, non-loopback callers to the data plane must present it (`Authorization: Bearer <token>` or `X-Headroom-Proxy-Token`); loopback callers and health probes are exempt. Constant-time (bytes) comparison. Closes the gap where the Docker image binds `0.0.0.0:8787` and exposes unauthenticated `/v1/*` routes to the pod network. A loud startup warning fires when binding a non-loopback host with no token set. - **Response security headers** (`X-Content-Type-Options: nosniff`, `X-Frame-Options: DENY`, `Referrer-Policy: no-referrer`, HSTS) on every response, including 401s. - **Audit log for state-mutating admin endpoints.** A structured `headroom.audit` JSON event (source IP, method, path, status) for `/admin/*`, `/cache/clear`, `/stats/reset`; `/admin/runtime-env` additionally records the changed key names (values omitted so secrets are never logged). Logger-only — safe under `HEADROOM_STATELESS` (no new file writes). - **Air-gap master switch (`HEADROOM_OFFLINE=1`).** Hard-disables all outbound egress in one flag — telemetry beacon, update check, license/usage reporter, and HuggingFace model downloads (forces `HF_HUB_OFFLINE`/`TRANSFORMERS_OFFLINE`) — and logs an offline banner. The first three live in one outermost security middleware that wraps every inbound request; the offline switch is centralized in a new top-level `headroom/offline.py` predicate the egress paths consult. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) ## Changes Made - `headroom/offline.py` (new): `is_offline()` predicate + `apply_offline_env()`; consulted by `beacon.is_telemetry_enabled`, `update_check.is_update_check_enabled`, and the license-reporter gate. - `headroom/proxy/audit.py` (new): `headroom.audit` structured logger + `record_admin_action` / `is_auditable_path`. - `headroom/proxy/server.py`: outermost `_security_gate` middleware (token enforcement + security headers + admin audit), offline activation + banner in `create_app`, non-loopback-no-token startup warning, `runtime-env` change auditing, env wiring in `_proxy_config_from_env`. - `headroom/proxy/models.py`: `ProxyConfig.proxy_token` and `ProxyConfig.offline`. - `headroom/cli/proxy.py`: env wiring + a Security banner line (flags the open-bind case). - `headroom/telemetry/beacon.py`, `headroom/update_check.py`: offline short-circuit. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ ruff check <changed files> All checks passed! $ mypy headroom/proxy/server.py headroom/offline.py headroom/proxy/audit.py headroom/proxy/models.py \ headroom/telemetry/beacon.py headroom/update_check.py Success: no issues found $ pytest tests/test_proxy_hardening.py -q 15 passed $ pytest tests/ -q (full suite, model/eval-dependent dirs ignored) 32 failed, 7131 passed, 126 skipped in 646s ``` The 32 failures are pre-existing/environmental, not introduced by this change — verified by running the same tests on `main` (they fail identically there). They are all `...Real` / `...live` / `real_api` integration tests that make live backend calls: AWS Bedrock returns "model is Legacy, access denied" on this host's credentials, plus a local tree-sitter version that requires `bytes`. On CI (no AWS/API creds) these tests skip. None touch the hardening code paths. ## Real Behavior Proof - Environment: macOS, Python 3.12, repo `.venv`; tests via FastAPI `TestClient` against `create_app`. - Exact command / steps: configure `ProxyConfig(proxy_token="...")`, then issue requests from a non-loopback client (`client=("203.0.113.5", ...)`) and a loopback client (`client=("127.0.0.1", ...)`). - Observed result: non-loopback request with no/!wrong token → 401; with correct `Authorization: Bearer` or `X-Headroom-Proxy-Token` → not 401; loopback and `/livez`/`/readyz` → never challenged. Every response (incl. the 401) carries `X-Content-Type-Options: nosniff` / `X-Frame-Options: DENY`. `POST /cache/clear` emits a `headroom.audit` JSON line with the source IP, path, and status. `HEADROOM_OFFLINE=1` makes `is_telemetry_enabled()` and `is_update_check_enabled()` return False and sets `HF_HUB_OFFLINE`/`TRANSFORMERS_OFFLINE`. - Not tested: WebSocket routes (see limitations); live upstream proxying of `/v1/*` (covered by existing integration tests / CI). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes Known limitations (by design / scope, documented in code): - WebSocket routes are not covered by the HTTP token middleware (`@app.middleware("http")` does not run for WS). The HTTP data plane is the main surface; the open-bind warning still applies. Follow-up. - The token keys off the direct peer IP — behind a same-host reverse proxy all requests appear loopback, so enforce auth at the reverse proxy in that topology (same property as the existing loopback guard). - OTEL metrics export is intentionally left on under offline mode — it targets the customer's own sink, not a Headroom phone-home. CHANGELOG not updated (handled by release tooling).
This commit is contained in:
Tejas Chopra
GitHub
parent
840871cb96
commit
546ab553dc
8 changed files with 437 additions and 3 deletions
|
|
@ -1121,6 +1121,9 @@ def proxy(
|
|||
disable_kompress_fallback=disable_kompress_fallback,
|
||||
disable_kompress_anthropic=disable_kompress_anthropic,
|
||||
disable_kompress_openai=disable_kompress_openai,
|
||||
# Optional inbound auth token + air-gap switch (env-driven).
|
||||
proxy_token=os.environ.get("HEADROOM_PROXY_TOKEN") or None,
|
||||
offline=_get_env_bool("HEADROOM_OFFLINE", False),
|
||||
# Code graph: live file watcher for incremental reindexing
|
||||
code_graph_watcher=code_graph,
|
||||
# Read lifecycle: ON by default (use --no-read-lifecycle to disable)
|
||||
|
|
@ -1281,6 +1284,26 @@ Memory (Multi-Provider):
|
|||
f"(available: {','.join(_ext_available)})"
|
||||
)
|
||||
|
||||
# Security posture line: inbound auth token + air-gap mode, and a loud
|
||||
# flag for the open-bind case (non-loopback host with no token).
|
||||
from headroom.proxy.loopback_guard import is_loopback_host
|
||||
|
||||
_auth_on = bool(config.proxy_token or os.environ.get("HEADROOM_PROXY_TOKEN"))
|
||||
if config.offline:
|
||||
_security_status = "OFFLINE (all egress disabled)" + (
|
||||
" · inbound token REQUIRED (non-loopback)" if _auth_on else ""
|
||||
)
|
||||
elif _auth_on:
|
||||
_security_status = "inbound token REQUIRED for non-loopback callers"
|
||||
elif not is_loopback_host(config.host):
|
||||
_security_status = (
|
||||
"WARNING non-loopback bind with NO token — /v1/* is UNAUTHENTICATED "
|
||||
"(set HEADROOM_PROXY_TOKEN)"
|
||||
)
|
||||
else:
|
||||
_security_status = "loopback-only (no inbound token)"
|
||||
security_line = f" Security: {_security_status}"
|
||||
|
||||
# Code-aware status line — same logic the inner banner uses, surfaced here
|
||||
# so the click-CLI banner is a complete picture (avoids the dual-banner
|
||||
# confusion this branch retired).
|
||||
|
|
@ -1334,6 +1357,7 @@ Starting proxy server...
|
|||
{code_aware_line}
|
||||
{context_tool_line}
|
||||
{extensions_line}
|
||||
{security_line}
|
||||
{stateless_line}{telemetry_line}
|
||||
{backend_section}{tuning_section}
|
||||
|
||||
|
|
|
|||
36
headroom/offline.py
Normal file
36
headroom/offline.py
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
"""Air-gap / no-egress master switch (``HEADROOM_OFFLINE``).
|
||||
|
||||
A single predicate the individual egress paths consult so a regulated or
|
||||
air-gapped deployment can disable **all** outbound network access with one
|
||||
flag: the telemetry beacon, the update check, the license/usage reporter, and
|
||||
HuggingFace model downloads. Each of those already had its own opt-out; this
|
||||
is the one switch that turns them all off together and fails closed.
|
||||
|
||||
Kept at the top level (depends only on the stdlib) so any layer — telemetry,
|
||||
proxy, model code — can import it without creating a package cycle.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
_TRUE_VALUES = frozenset({"1", "true", "yes", "on"})
|
||||
|
||||
OFFLINE_ENV = "HEADROOM_OFFLINE"
|
||||
|
||||
|
||||
def is_offline() -> bool:
|
||||
"""Return True when ``HEADROOM_OFFLINE`` selects fully-offline operation."""
|
||||
return os.environ.get(OFFLINE_ENV, "").strip().lower() in _TRUE_VALUES
|
||||
|
||||
|
||||
def apply_offline_env() -> None:
|
||||
"""Force HuggingFace/Transformers offline so model code uses only locally
|
||||
cached artifacts and never reaches the Hub.
|
||||
|
||||
Idempotent and uses ``setdefault`` so an explicit operator override (e.g.
|
||||
``HF_HUB_OFFLINE=0``) still wins. Call once early in startup.
|
||||
"""
|
||||
if is_offline():
|
||||
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
||||
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
||||
59
headroom/proxy/audit.py
Normal file
59
headroom/proxy/audit.py
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
"""Lightweight audit log for administrative / state-mutating proxy actions.
|
||||
|
||||
Emits one structured JSON line per sensitive action (``/admin/*`` runtime
|
||||
changes, cache clears, stats resets) to the dedicated ``headroom.audit``
|
||||
logger. Operators capture this stream the same way they capture the proxy
|
||||
log; routing/retention is theirs to configure.
|
||||
|
||||
Logger-only by design — it writes **no** dedicated file, so it is safe under
|
||||
``HEADROOM_STATELESS`` (no filesystem writes) and respects whatever log sink
|
||||
the deployment already uses.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
audit_logger = logging.getLogger("headroom.audit")
|
||||
|
||||
# Paths whose requests mutate runtime state or expose stored content and so
|
||||
# warrant an audit trail. Matched by exact value or, for ``/admin/``, prefix.
|
||||
_ADMIN_PREFIX = "/admin/"
|
||||
_SENSITIVE_EXACT = frozenset({"/cache/clear", "/stats/reset"})
|
||||
|
||||
|
||||
def is_auditable_path(path: str) -> bool:
|
||||
"""Return True when requests to ``path`` should be audited."""
|
||||
return path.startswith(_ADMIN_PREFIX) or path in _SENSITIVE_EXACT
|
||||
|
||||
|
||||
def _client_ip(request: Any) -> str | None:
|
||||
client = getattr(request, "client", None)
|
||||
return getattr(client, "host", None) if client is not None else None
|
||||
|
||||
|
||||
def record_admin_action(
|
||||
*,
|
||||
request: Any,
|
||||
action: str,
|
||||
status_code: int,
|
||||
details: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Emit a structured audit event. Never raises (audit must not break a
|
||||
request); logs at WARNING on its own failure."""
|
||||
try:
|
||||
event: dict[str, Any] = {
|
||||
"event": "headroom_admin_audit",
|
||||
"action": action,
|
||||
"method": getattr(request, "method", None),
|
||||
"path": getattr(getattr(request, "url", None), "path", None),
|
||||
"source_ip": _client_ip(request),
|
||||
"status_code": status_code,
|
||||
}
|
||||
if details:
|
||||
event["details"] = details
|
||||
audit_logger.info(json.dumps(event, ensure_ascii=False, default=str))
|
||||
except Exception: # noqa: BLE001 — auditing must never break the request
|
||||
audit_logger.warning("audit event emission failed", exc_info=True)
|
||||
|
|
@ -342,6 +342,18 @@ class ProxyConfig:
|
|||
# Stateless mode — disable all filesystem writes for read-only / container deployments
|
||||
stateless: bool = False
|
||||
|
||||
# Optional inbound auth. When set, non-loopback requests to the data-plane
|
||||
# routes must present this token (``Authorization: Bearer <token>`` or the
|
||||
# ``X-Headroom-Proxy-Token`` header). Loopback callers are exempt. Closes the
|
||||
# gap where a container bound to 0.0.0.0 exposes unauthenticated /v1/* routes
|
||||
# to the pod network. Env: HEADROOM_PROXY_TOKEN.
|
||||
proxy_token: str | None = None
|
||||
|
||||
# Air-gap master switch — hard-disable ALL outbound network egress
|
||||
# (telemetry beacon, update check, license/usage reporter, HuggingFace model
|
||||
# downloads) for fully offline / regulated deployments. Env: HEADROOM_OFFLINE=1.
|
||||
offline: bool = False
|
||||
|
||||
# Unit 4: Bounded pre-upstream concurrency for Anthropic replay storms.
|
||||
#
|
||||
# Caps the number of simultaneous requests allowed to run the
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ import argparse
|
|||
import asyncio
|
||||
import concurrent.futures
|
||||
import contextlib
|
||||
import hmac
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
|
@ -90,6 +91,7 @@ from headroom.observability import (
|
|||
shutdown_headroom_tracing,
|
||||
shutdown_otel_metrics,
|
||||
)
|
||||
from headroom.offline import apply_offline_env, is_offline
|
||||
from headroom.pipeline import PipelineExtensionManager, PipelineStage
|
||||
from headroom.providers.proxy_routes import register_provider_routes
|
||||
from headroom.providers.registry import (
|
||||
|
|
@ -104,6 +106,7 @@ from headroom.providers.registry import (
|
|||
resolve_api_targets,
|
||||
)
|
||||
from headroom.proxy import runtime_env
|
||||
from headroom.proxy.audit import is_auditable_path, record_admin_action
|
||||
from headroom.proxy.auth_mode import should_stamp_codex_client
|
||||
from headroom.proxy.background_compression import BackgroundCompressor
|
||||
|
||||
|
|
@ -134,6 +137,7 @@ from headroom.proxy.helpers import (
|
|||
jitter_delay_ms,
|
||||
retry_after_ms,
|
||||
)
|
||||
from headroom.proxy.loopback_guard import is_loopback_host
|
||||
from headroom.proxy.memory_handler import MemoryConfig, MemoryHandler
|
||||
|
||||
# Data models (extracted to headroom/proxy/models.py for maintainability)
|
||||
|
|
@ -1073,9 +1077,11 @@ class HeadroomProxy(
|
|||
"hint=bridge_syncs_only_the_legacy_DB_today_per-project_bridge_follow-up_planned"
|
||||
)
|
||||
|
||||
# Usage Reporter (license validation + phone-home for managed/enterprise)
|
||||
# Usage Reporter (license validation + phone-home for managed/enterprise).
|
||||
# Suppressed entirely in offline mode — the air-gap switch must stop all
|
||||
# egress, including license phone-home, even when a key is configured.
|
||||
self.usage_reporter: UsageReporter | None = None
|
||||
if config.license_key:
|
||||
if config.license_key and not (config.offline or is_offline()):
|
||||
from headroom.telemetry.reporter import UsageReporter
|
||||
|
||||
self.usage_reporter = UsageReporter(
|
||||
|
|
@ -1876,6 +1882,20 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
_setup_file_logging()
|
||||
|
||||
config = config or ProxyConfig()
|
||||
|
||||
# Air-gap master switch. Propagate config.offline to the env so the
|
||||
# env-based egress predicates (telemetry, update check, license) all honor
|
||||
# it, force HF/transformers offline before any model code loads, and
|
||||
# announce that every outbound path is disabled.
|
||||
if config.offline:
|
||||
os.environ.setdefault("HEADROOM_OFFLINE", "1")
|
||||
if is_offline():
|
||||
apply_offline_env()
|
||||
logger.warning(
|
||||
"event=proxy_offline_mode air-gap active — all outbound egress disabled "
|
||||
"(telemetry, update check, license reporter, HuggingFace downloads)"
|
||||
)
|
||||
|
||||
proxy = HeadroomProxy(config)
|
||||
|
||||
# cc-switch reconciler (opt-in: HEADROOM_CC_SWITCH_RECONCILE=1).
|
||||
|
|
@ -2462,6 +2482,86 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
)
|
||||
return response
|
||||
|
||||
# ── Security gate (registered last → runs outermost) ──────────────────
|
||||
# Three concerns, kept together because they all wrap every inbound
|
||||
# request: optional inbound auth on the data plane, response security
|
||||
# headers, and an audit trail for state-mutating admin endpoints.
|
||||
_proxy_token = config.proxy_token or os.environ.get("HEADROOM_PROXY_TOKEN") or None
|
||||
# Pre-encode once for constant-time comparison (compare_digest on str raises
|
||||
# TypeError for non-ASCII input, which would turn a 401 into a 500).
|
||||
_proxy_token_bytes = _proxy_token.encode("utf-8") if _proxy_token else b""
|
||||
# Health/readiness probes must stay reachable without a token so
|
||||
# orchestrators can check a container that binds non-loopback.
|
||||
_AUTH_EXEMPT_PATHS = frozenset({"/health", "/healthz", "/livez", "/readyz"})
|
||||
|
||||
# Loud warning when a non-loopback bind has no token configured: that is the
|
||||
# exact shape (e.g. the Docker 0.0.0.0 image) that exposes unauthenticated
|
||||
# /v1/* routes to the surrounding network.
|
||||
if not _proxy_token and not is_loopback_host(getattr(config, "host", None)):
|
||||
logger.warning(
|
||||
"event=proxy_open_bind host=%s — proxy is bound to a non-loopback "
|
||||
"interface with no HEADROOM_PROXY_TOKEN set; the /v1/* data-plane "
|
||||
"routes are reachable WITHOUT authentication. Set HEADROOM_PROXY_TOKEN "
|
||||
"to require a bearer token from non-loopback callers.",
|
||||
getattr(config, "host", None),
|
||||
)
|
||||
|
||||
def _apply_security_headers(response) -> None:
|
||||
# setdefault: never clobber a header an upstream/handler already set.
|
||||
response.headers.setdefault("X-Content-Type-Options", "nosniff")
|
||||
response.headers.setdefault("X-Frame-Options", "DENY")
|
||||
response.headers.setdefault("Referrer-Policy", "no-referrer")
|
||||
response.headers.setdefault(
|
||||
"Strict-Transport-Security", "max-age=31536000; includeSubDomains"
|
||||
)
|
||||
|
||||
def _extract_proxy_token(headers) -> str | None:
|
||||
auth = str(headers.get("authorization") or "")
|
||||
if auth.lower().startswith("bearer "):
|
||||
return auth[7:].strip() or None
|
||||
raw = headers.get("x-headroom-proxy-token")
|
||||
return str(raw) if raw else None
|
||||
|
||||
@app.middleware("http")
|
||||
async def _security_gate(request, call_next):
|
||||
# 1) Optional inbound auth. When a token is configured, require it on
|
||||
# non-loopback requests; loopback callers and health probes are
|
||||
# exempt. Loopback is the same trust boundary the admin/debug
|
||||
# endpoints already use (see loopback_guard).
|
||||
if _proxy_token:
|
||||
path = request.url.path
|
||||
client = getattr(request, "client", None)
|
||||
client_host = getattr(client, "host", None) if client is not None else None
|
||||
if path not in _AUTH_EXEMPT_PATHS and not is_loopback_host(client_host):
|
||||
provided = _extract_proxy_token(request.headers)
|
||||
if provided is None or not hmac.compare_digest(
|
||||
provided.encode("utf-8", "replace"), _proxy_token_bytes
|
||||
):
|
||||
logger.warning(
|
||||
"event=proxy_auth_rejected path=%s client=%s reason=%s",
|
||||
path,
|
||||
client_host,
|
||||
"missing_token" if provided is None else "bad_token",
|
||||
)
|
||||
rejection = JSONResponse(status_code=401, content={"error": "unauthorized"})
|
||||
_apply_security_headers(rejection)
|
||||
return rejection
|
||||
|
||||
response = await call_next(request)
|
||||
_apply_security_headers(response)
|
||||
|
||||
# 2) Audit trail for admin / state-mutating endpoints.
|
||||
try:
|
||||
if is_auditable_path(request.url.path):
|
||||
record_admin_action(
|
||||
request=request,
|
||||
action="admin_request",
|
||||
status_code=response.status_code,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("admin audit emission failed", exc_info=True)
|
||||
return response
|
||||
|
||||
# Third-party proxy extensions (Enterprise, custom plugins). Discovered via
|
||||
# the `headroom.proxy_extension` entry-point group, but **opt-in only**:
|
||||
# only names listed in config.proxy_extensions (CLI: --proxy-extension,
|
||||
|
|
@ -2582,6 +2682,16 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
applied = runtime_env.set_overrides(body)
|
||||
if applied:
|
||||
logger.info("runtime-env hot-reload applied: %s", sorted(applied))
|
||||
# Record which runtime-env keys changed (the "what" of a config
|
||||
# change) in addition to the generic admin-request audit emitted by
|
||||
# the security middleware. Values are intentionally omitted — keys
|
||||
# alone avoid logging any secret values that were set.
|
||||
record_admin_action(
|
||||
request=request,
|
||||
action="runtime_env_update",
|
||||
status_code=200,
|
||||
details={"changed_keys": sorted(applied)},
|
||||
)
|
||||
return JSONResponse(
|
||||
status_code=200,
|
||||
content={"applied": applied, "runtime_env": runtime_env.effective_runtime_env()},
|
||||
|
|
@ -3931,6 +4041,8 @@ def _proxy_config_from_env() -> ProxyConfig:
|
|||
keepalive_expiry=_get_env_float("HEADROOM_KEEPALIVE_EXPIRY", 90.0),
|
||||
http2=_get_env_bool("HEADROOM_HTTP2", True),
|
||||
periodic_toin_stats_enabled=_get_env_bool("HEADROOM_PERIODIC_TOIN_STATS", True),
|
||||
proxy_token=os.environ.get("HEADROOM_PROXY_TOKEN") or None,
|
||||
offline=_get_env_bool("HEADROOM_OFFLINE", False),
|
||||
mode=normalize_proxy_mode(_get_env_str("HEADROOM_MODE", PROXY_MODE_TOKEN)),
|
||||
read_maturation=_get_env_bool("HEADROOM_READ_MATURATION", False),
|
||||
read_maturation_quiesce_turns=_get_env_int("HEADROOM_READ_MATURATION_QUIESCE_TURNS", 5),
|
||||
|
|
|
|||
|
|
@ -28,6 +28,10 @@ def is_telemetry_enabled() -> bool:
|
|||
feeds the in-process collector and the ``/stats`` endpoint; nothing is
|
||||
transmitted to Headroom Labs.
|
||||
"""
|
||||
from headroom.offline import is_offline
|
||||
|
||||
if is_offline():
|
||||
return False
|
||||
val = os.environ.get("HEADROOM_TELEMETRY", "").lower().strip()
|
||||
return val in _ON_VALUES
|
||||
|
||||
|
|
|
|||
|
|
@ -54,10 +54,15 @@ def _env_on(name: str) -> bool:
|
|||
def is_update_check_enabled() -> bool:
|
||||
"""Whether the update check / banner should run at all.
|
||||
|
||||
Disabled by ``HEADROOM_UPDATE_CHECK=off``, stateless mode
|
||||
Disabled by ``HEADROOM_UPDATE_CHECK=off``, offline mode
|
||||
(``HEADROOM_OFFLINE``), stateless mode
|
||||
(``HEADROOM_STATELESS=true``/``1``/``yes``/``on``, matching the proxy's own
|
||||
parsing), or any CI environment (``CI`` set).
|
||||
"""
|
||||
from headroom.offline import is_offline
|
||||
|
||||
if is_offline():
|
||||
return False
|
||||
if _env_off("HEADROOM_UPDATE_CHECK"):
|
||||
return False
|
||||
if _env_on("HEADROOM_STATELESS"):
|
||||
|
|
|
|||
182
tests/test_proxy_hardening.py
Normal file
182
tests/test_proxy_hardening.py
Normal file
|
|
@ -0,0 +1,182 @@
|
|||
"""Tests for the Tier-2 pilot hardening features:
|
||||
|
||||
- 2.1 optional inbound auth token (HEADROOM_PROXY_TOKEN) on the data plane
|
||||
- 3.1 response security headers
|
||||
- 2.4 admin/state-mutating audit log
|
||||
- 2.2 air-gap master switch (HEADROOM_OFFLINE)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("fastapi")
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from headroom.cache.compression_store import reset_compression_store
|
||||
from headroom.offline import apply_offline_env, is_offline
|
||||
from headroom.proxy.audit import is_auditable_path
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
NONLOOPBACK = ("203.0.113.5", 44444) # TEST-NET-3, never loopback
|
||||
LOOPBACK = ("127.0.0.1", 12345)
|
||||
|
||||
|
||||
def _make_app(**overrides):
|
||||
reset_compression_store()
|
||||
config = ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
**overrides,
|
||||
)
|
||||
return create_app(config)
|
||||
|
||||
|
||||
# ───────────────────────────── 2.1 inbound auth token ─────────────────────
|
||||
|
||||
|
||||
class TestInboundAuthToken:
|
||||
def test_no_token_configured_leaves_data_plane_open(self):
|
||||
"""Default (no token): non-loopback callers are not challenged."""
|
||||
app = _make_app()
|
||||
with TestClient(app, base_url="http://testserver", client=NONLOOPBACK) as c:
|
||||
assert c.get("/livez").status_code == 200
|
||||
|
||||
def test_token_set_rejects_nonloopback_without_credential(self):
|
||||
app = _make_app(proxy_token="s3cr3t-token")
|
||||
with TestClient(app, base_url="http://testserver", client=NONLOOPBACK) as c:
|
||||
resp = c.get("/stats")
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_token_set_accepts_correct_bearer(self):
|
||||
app = _make_app(proxy_token="s3cr3t-token")
|
||||
with TestClient(app, base_url="http://testserver", client=NONLOOPBACK) as c:
|
||||
resp = c.get("/stats", headers={"Authorization": "Bearer s3cr3t-token"})
|
||||
assert resp.status_code != 401
|
||||
|
||||
def test_token_set_accepts_custom_header(self):
|
||||
app = _make_app(proxy_token="s3cr3t-token")
|
||||
with TestClient(app, base_url="http://testserver", client=NONLOOPBACK) as c:
|
||||
resp = c.get("/stats", headers={"X-Headroom-Proxy-Token": "s3cr3t-token"})
|
||||
assert resp.status_code != 401
|
||||
|
||||
def test_token_set_rejects_wrong_token(self):
|
||||
app = _make_app(proxy_token="s3cr3t-token")
|
||||
with TestClient(app, base_url="http://testserver", client=NONLOOPBACK) as c:
|
||||
resp = c.get("/stats", headers={"Authorization": "Bearer wrong"})
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_loopback_is_exempt_from_token(self):
|
||||
"""Loopback callers (same trust boundary as admin routes) skip the token."""
|
||||
app = _make_app(proxy_token="s3cr3t-token")
|
||||
with TestClient(app, base_url="http://127.0.0.1", client=LOOPBACK) as c:
|
||||
assert c.get("/stats").status_code != 401
|
||||
|
||||
def test_health_endpoints_exempt_even_nonloopback(self):
|
||||
"""Orchestrator health probes must work without the token."""
|
||||
app = _make_app(proxy_token="s3cr3t-token")
|
||||
with TestClient(app, base_url="http://testserver", client=NONLOOPBACK) as c:
|
||||
assert c.get("/livez").status_code == 200
|
||||
assert c.get("/readyz").status_code in (200, 503) # ready/not-ready, never 401
|
||||
|
||||
|
||||
# ───────────────────────────── 3.1 security headers ───────────────────────
|
||||
|
||||
|
||||
class TestSecurityHeaders:
|
||||
def test_headers_present_on_responses(self):
|
||||
app = _make_app()
|
||||
with TestClient(app, base_url="http://127.0.0.1", client=LOOPBACK) as c:
|
||||
h = c.get("/livez").headers
|
||||
assert h.get("X-Content-Type-Options") == "nosniff"
|
||||
assert h.get("X-Frame-Options") == "DENY"
|
||||
assert h.get("Referrer-Policy") == "no-referrer"
|
||||
assert "max-age=" in h.get("Strict-Transport-Security", "")
|
||||
|
||||
def test_headers_present_on_401(self):
|
||||
app = _make_app(proxy_token="s3cr3t-token")
|
||||
with TestClient(app, base_url="http://testserver", client=NONLOOPBACK) as c:
|
||||
resp = c.get("/stats")
|
||||
assert resp.status_code == 401
|
||||
assert resp.headers.get("X-Content-Type-Options") == "nosniff"
|
||||
|
||||
|
||||
# ───────────────────────────── 2.4 admin audit log ────────────────────────
|
||||
|
||||
|
||||
class TestAdminAuditLog:
|
||||
def test_auditable_path_classification(self):
|
||||
assert is_auditable_path("/admin/runtime-env")
|
||||
assert is_auditable_path("/cache/clear")
|
||||
assert is_auditable_path("/stats/reset")
|
||||
assert not is_auditable_path("/v1/messages")
|
||||
assert not is_auditable_path("/livez")
|
||||
|
||||
def test_cache_clear_emits_audit_event(self):
|
||||
# Capture the dedicated audit logger directly (the proxy's logging setup
|
||||
# configures propagation, so attach to the logger rather than rely on
|
||||
# caplog's root handler).
|
||||
messages: list[str] = []
|
||||
|
||||
class _Capture(logging.Handler):
|
||||
def emit(self, record: logging.LogRecord) -> None:
|
||||
messages.append(record.getMessage())
|
||||
|
||||
handler = _Capture()
|
||||
audit_logger = logging.getLogger("headroom.audit")
|
||||
audit_logger.setLevel(logging.INFO)
|
||||
audit_logger.addHandler(handler)
|
||||
try:
|
||||
app = _make_app()
|
||||
with TestClient(app, base_url="http://127.0.0.1", client=LOOPBACK) as c:
|
||||
assert c.post("/cache/clear").status_code == 200
|
||||
finally:
|
||||
audit_logger.removeHandler(handler)
|
||||
|
||||
assert messages, "expected an audit record for /cache/clear"
|
||||
assert any("/cache/clear" in m for m in messages)
|
||||
assert any("headroom_admin_audit" in m for m in messages)
|
||||
assert any('"source_ip": "127.0.0.1"' in m for m in messages)
|
||||
|
||||
|
||||
# ───────────────────────────── 2.2 air-gap switch ─────────────────────────
|
||||
|
||||
|
||||
class TestOfflineSwitch:
|
||||
def test_is_offline_reads_env(self, monkeypatch):
|
||||
monkeypatch.delenv("HEADROOM_OFFLINE", raising=False)
|
||||
assert is_offline() is False
|
||||
monkeypatch.setenv("HEADROOM_OFFLINE", "1")
|
||||
assert is_offline() is True
|
||||
monkeypatch.setenv("HEADROOM_OFFLINE", "off")
|
||||
assert is_offline() is False
|
||||
|
||||
def test_offline_disables_telemetry(self, monkeypatch):
|
||||
from headroom.telemetry.beacon import is_telemetry_enabled
|
||||
|
||||
monkeypatch.setenv("HEADROOM_TELEMETRY", "on")
|
||||
monkeypatch.setenv("HEADROOM_OFFLINE", "1")
|
||||
assert is_telemetry_enabled() is False # offline overrides the opt-in
|
||||
|
||||
def test_offline_disables_update_check(self, monkeypatch):
|
||||
from headroom.update_check import is_update_check_enabled
|
||||
|
||||
monkeypatch.delenv("CI", raising=False)
|
||||
monkeypatch.delenv("HEADROOM_STATELESS", raising=False)
|
||||
monkeypatch.setenv("HEADROOM_OFFLINE", "1")
|
||||
assert is_update_check_enabled() is False
|
||||
|
||||
def test_apply_offline_env_sets_hf_offline(self, monkeypatch):
|
||||
monkeypatch.delenv("HF_HUB_OFFLINE", raising=False)
|
||||
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False)
|
||||
monkeypatch.setenv("HEADROOM_OFFLINE", "1")
|
||||
apply_offline_env()
|
||||
import os
|
||||
|
||||
assert os.environ.get("HF_HUB_OFFLINE") == "1"
|
||||
assert os.environ.get("TRANSFORMERS_OFFLINE") == "1"
|
||||
Loading…
Add table
Add a link
Reference in a new issue