From f28f697310f7cc5cd0cd6354c1c0fa9ade158c67 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Sat, 25 Apr 2026 12:48:23 -0700 Subject: [PATCH] fix(proxy): make third-party extensions opt-in MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, any package registered under the headroom.proxy_extension entry-point group auto-loaded at proxy startup. A user pip-installing a plugin (or pulling one in transitively) would get its middleware running in front of all their LLM traffic with zero opt-in or visibility — the same mechanism that masked the Shield Enterprise streaming bug. Change: install_all() now takes an explicit enabled set (or reads HEADROOM_PROXY_EXTENSIONS). Discovery still runs to enumerate what's available, but only names the operator opted into actually install. The literal '*' is a wildcard for trusted environments. CLI: headroom proxy --proxy-extension shield_enterprise headroom proxy --proxy-extension shield_enterprise,mypkg headroom proxy --proxy-extension '*' Env: HEADROOM_PROXY_EXTENSIONS=shield_enterprise The startup banner now shows discovered + enabled extensions: Extensions: discovered=shield_enterprise (opt-in: --proxy-extension ...) Extensions: ENABLED shield_enterprise (available: shield_enterprise) Extensions: ENABLED (wildcard) shield_enterprise Names that were requested but not found are logged as warnings. Adds proxy_extensions: list[str] | None to ProxyConfig. Plumbs it through CLI -> ProxyConfig -> install_all(enabled=...). This is a behavior change for users who relied on auto-loading. Existing Shield/extension users must add --proxy-extension or set HEADROOM_PROXY_EXTENSIONS to keep their middleware running. --- headroom/cli/proxy.py | 47 +++++++++++++++++++++ headroom/proxy/extensions.py | 82 +++++++++++++++++++++++++++++++++--- headroom/proxy/models.py | 6 +++ headroom/proxy/server.py | 9 ++-- 4 files changed, 135 insertions(+), 9 deletions(-) diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 738d7d8e4..48756832d 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -57,6 +57,17 @@ from .main import main @click.option("--no-optimize", is_flag=True, help="Disable optimization (passthrough mode)") @click.option("--no-cache", is_flag=True, help="Disable semantic caching") @click.option("--no-rate-limit", is_flag=True, help="Disable rate limiting") +@click.option( + "--proxy-extension", + "proxy_extension", + multiple=True, + envvar="HEADROOM_PROXY_EXTENSIONS", + help=( + "Enable a registered proxy extension by entry-point name (opt-in). " + "Repeat the flag or pass a comma-separated list. Use '*' to enable " + "every discovered extension. Env: HEADROOM_PROXY_EXTENSIONS." + ), +) @click.option( "--no-subscription-tracking", is_flag=True, @@ -274,6 +285,7 @@ def proxy( no_optimize: bool, no_cache: bool, no_rate_limit: bool, + proxy_extension: tuple[str, ...], no_subscription_tracking: bool, subscription_poll_interval: int | None, retry_max_attempts: int | None, @@ -403,6 +415,13 @@ def proxy( optimize=not no_optimize, cache_enabled=not no_cache, rate_limit_enabled=not no_rate_limit, + # Flatten repeat-flag tuple AND any comma-separated values inside it. + # `--proxy-extension a,b --proxy-extension c` and `HEADROOM_PROXY_EXTENSIONS=a,b,c` + # both yield ["a", "b", "c"]. None when nothing was supplied. + proxy_extensions=( + [part.strip() for chunk in proxy_extension for part in chunk.split(",") if part.strip()] + or None + ), subscription_tracking_enabled=not no_subscription_tracking, subscription_poll_interval_s=( subscription_poll_interval if subscription_poll_interval is not None else 300 @@ -536,6 +555,33 @@ Memory (Multi-Provider): else: telemetry_line = " Telemetry: DISABLED" + # Discover proxy extensions (third-party packages registered via the + # `headroom.proxy_extension` entry-point group). Surfaced in the banner + # so operators can see what's available + what's currently opted-in. + # Discovery does NOT run extension code; only the explicitly-enabled + # set in config.proxy_extensions actually installs. + try: + from headroom.proxy.extensions import discover as _discover_extensions + + _ext_available = sorted(name for name, _ in _discover_extensions()) + except Exception: # noqa: BLE001 — banner must never crash startup + _ext_available = [] + _ext_enabled = config.proxy_extensions or [] + if not _ext_available: + extensions_line = " Extensions: (none discovered)" + elif not _ext_enabled: + extensions_line = ( + f" Extensions: discovered={','.join(_ext_available)} " + f"(opt-in: --proxy-extension or HEADROOM_PROXY_EXTENSIONS=)" + ) + elif "*" in _ext_enabled: + extensions_line = f" Extensions: ENABLED (wildcard) {','.join(_ext_available)}" + else: + extensions_line = ( + f" Extensions: ENABLED {','.join(sorted(_ext_enabled))} " + f"(available: {','.join(_ext_available)})" + ) + click.echo(f""" ╔═══════════════════════════════════════════════════════════════════════╗ ║ HEADROOM PROXY ║ @@ -551,6 +597,7 @@ Starting proxy server... Rate Limit: {"ENABLED" if config.rate_limit_enabled else "DISABLED"} Memory: {memory_status} License: {license_status} +{extensions_line} {stateless_line}{telemetry_line} {backend_section} Routing: diff --git a/headroom/proxy/extensions.py b/headroom/proxy/extensions.py index 15f2d2044..59b33f12d 100644 --- a/headroom/proxy/extensions.py +++ b/headroom/proxy/extensions.py @@ -17,6 +17,18 @@ Each ``install`` callable is invoked with the FastAPI ``app`` and the OSS makes no assumptions about what extensions do. The interface is deliberately minimal; extensions own the complexity behind it. +**Extensions are opt-in.** Discovery enumerates every registered extension, +but ``install_all`` only invokes those explicitly enabled by the operator. +This protects users from silent behavior changes when a package they didn't +audit gets installed in the same environment (e.g., as a transitive dep). + +Enabling extensions: + + * CLI: ``headroom proxy --proxy-extension shield_enterprise,mypkg`` + * Env: ``HEADROOM_PROXY_EXTENSIONS=shield_enterprise,mypkg`` + * Wildcard: ``--proxy-extension '*'`` enables every discovered extension + (use only when you trust everything in your environment). + Stability contract: this module is load-bearing for the Enterprise build and any third-party extensions. Changes to the signature of ``install(app, config)`` or the entry-point group name require a deprecation cycle. @@ -26,12 +38,14 @@ from __future__ import annotations import importlib.metadata import logging -from collections.abc import Callable, Iterator +import os +from collections.abc import Callable, Iterable, Iterator from typing import Any log = logging.getLogger(__name__) ENTRY_POINT_GROUP = "headroom.proxy_extension" +ENV_VAR = "HEADROOM_PROXY_EXTENSIONS" ProxyExtension = Callable[[Any, Any], None] """Signature: ``install(app: FastAPI, config: ProxyConfig) -> None``.""" @@ -58,17 +72,73 @@ def discover() -> Iterator[tuple[str, ProxyExtension]]: yield entry.name, install -def install_all(app: Any, config: Any) -> list[str]: - """Run every discovered extension's ``install(app, config)``. +def _resolve_enabled(enabled: Iterable[str] | None) -> set[str]: + """Resolve the set of enabled extension names. + + Precedence: explicit ``enabled`` argument > ``HEADROOM_PROXY_EXTENSIONS`` + env var > empty (no extensions). Empty strings and whitespace are + stripped. The literal ``*`` enables all discovered extensions. + """ + raw: Iterable[str] + if enabled is not None: + raw = enabled + else: + raw = (os.environ.get(ENV_VAR) or "").split(",") + out: set[str] = set() + for n in raw: + n = n.strip() + if n: + out.add(n) + return out + + +def install_all( + app: Any, + config: Any, + enabled: Iterable[str] | None = None, +) -> list[str]: + """Run only the explicitly-enabled extensions' ``install(app, config)``. + + Discovery still runs so we can log the universe of available extensions, + but only those whose entry-point ``name`` is in ``enabled`` are invoked. + The literal ``"*"`` in ``enabled`` is a wildcard that enables every + discovered extension. Returns the names of successfully installed extensions. If an extension raises inside ``install()``, the exception propagates — this is the - documented fail-closed signal (e.g., a Shield Enterprise license check - failing should abort startup rather than silently run without protection). + documented fail-closed signal (e.g., a license check failing should + abort startup rather than silently run without protection). """ + enabled_set = _resolve_enabled(enabled) + discovered = list(discover()) + discovered_names = [n for n, _ in discovered] + + if not enabled_set: + if discovered_names: + log.info( + "proxy extensions discovered but disabled (opt-in): %s. " + "Enable with --proxy-extension or %s=.", + ",".join(discovered_names), + ENV_VAR, + ) + return [] + + wildcard = "*" in enabled_set installed: list[str] = [] - for name, install in discover(): + for name, install in discovered: + if not wildcard and name not in enabled_set: + continue install(app, config) installed.append(name) log.info("proxy extension installed: %s", name) + + # Warn about names the user asked for that weren't found. + if not wildcard: + missing = enabled_set - set(discovered_names) + if missing: + log.warning( + "proxy extensions requested but not found: %s (available: %s)", + ",".join(sorted(missing)), + ",".join(discovered_names) or "", + ) return installed diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index 1bbf5f3fd..46c67b575 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -171,6 +171,12 @@ class ProxyConfig: log_file: str | None = None log_full_messages: bool = False + # Third-party proxy extensions (opt-in only). List of entry-point names + # to enable from the `headroom.proxy_extension` group, or `["*"]` for + # wildcard. Empty/None means no extensions run, even if installed. + # CLI: --proxy-extension ; env: HEADROOM_PROXY_EXTENSIONS. + proxy_extensions: list[str] | None = None + # Fallback fallback_enabled: bool = False fallback_provider: str | None = None diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index d16dcd9ca..e8ae00f02 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -1332,11 +1332,14 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: return await call_next(request) # Third-party proxy extensions (Enterprise, custom plugins). Discovered via - # the `headroom.proxy_extension` entry-point group. An extension that raises - # from its install() is a deliberate fail-closed signal and aborts startup. + # the `headroom.proxy_extension` entry-point group, but **opt-in only**: + # only names listed in config.proxy_extensions (CLI: --proxy-extension, + # env: HEADROOM_PROXY_EXTENSIONS) actually get installed. Discovery alone + # never runs third-party code. An extension that raises from its install() + # is a deliberate fail-closed signal and aborts startup. from headroom.proxy.extensions import install_all as _install_extensions - _install_extensions(app, config) + _install_extensions(app, config, enabled=getattr(config, "proxy_extensions", None)) # Health & Metrics @app.get("/livez")