mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
refactor(proxy): isolate forwarded header policy (#1942)
## Description Extract the trusted forwarded-header trust policy into `headroom.proxy.forwarded_policy`, leaving `forwarded_headers` as the FastAPI/request-state adapter. This makes CIDR parsing, peer trust, leftmost forwarded-for handling, and rejection decisions deterministic and directly testable without request/logging side effects. Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [x] Code refactoring (no functional changes) ## Changes Made - Added `ForwardedHeaderInputs` and `ForwardedHeaderResolution` as pure policy value objects. - Moved CIDR parsing, IP normalization, trust membership, header splitting, and forwarded-header resolution into `headroom.proxy.forwarded_policy`. - Kept `headroom.proxy.forwarded_headers` as the request adapter with the same public API and compatibility helper names. - Added direct tests for trusted, rejected, direct-client, IPv4-mapped IPv6, and leftmost `X-Forwarded-For` policy behavior. - Included the LiteLLM callback hook compatibility shim needed for repo-wide mypy on branches based on `main`. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text python -m pytest tests/test_forwarded_policy.py tests/test_forwarded_headers.py tests/test_litellm_callback.py tests/test_compress_api.py::TestLiteLLMCallback -q 51 passed in 6.36s python -m ruff check . All checks passed! python -m ruff format --check . 1095 files already formatted python -m mypy headroom --ignore-missing-imports Success: no issues found in 409 source files ``` ## Real Behavior Proof - Environment: Windows, Python 3.13.13, local worktree `C:\git\headroom-pr-slice8`. - Exact command / steps: Ran the focused pytest suite plus repo-wide Ruff, format check, and mypy commands listed above. - Observed result: The existing request-facing forwarded-header behavior remains covered by `tests/test_forwarded_headers.py`, while the extracted pure policy is covered by `tests/test_forwarded_policy.py`. - Not tested: Full test suite locally; CI will run the full matrix. ## 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 - [ ] 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 ## Screenshots (if applicable) N/A ## Additional Notes Documentation and changelog updates are not applicable for this internal refactor. The LiteLLM shim is repeated here because this branch is intentionally independent from the other open architecture slices and must stay green against current `main`.
This commit is contained in:
parent
0ce09fb63f
commit
cb38f79377
3 changed files with 210 additions and 78 deletions
|
|
@ -57,6 +57,15 @@ import logging
|
||||||
import os
|
import os
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
from headroom.proxy.forwarded_policy import (
|
||||||
|
ForwardedHeaderInputs,
|
||||||
|
header_first,
|
||||||
|
normalize_ip,
|
||||||
|
parse_cidr_list,
|
||||||
|
peer_is_trusted_gateway,
|
||||||
|
resolve_forwarded_headers,
|
||||||
|
)
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from fastapi import Request
|
from fastapi import Request
|
||||||
|
|
||||||
|
|
@ -87,20 +96,7 @@ def _parse_cidr_list(
|
||||||
would silently downgrade the proxy from "strict" to "more strict",
|
would silently downgrade the proxy from "strict" to "more strict",
|
||||||
masking the operator's intent.
|
masking the operator's intent.
|
||||||
"""
|
"""
|
||||||
if not raw or not raw.strip():
|
return parse_cidr_list(raw)
|
||||||
return ()
|
|
||||||
nets: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = []
|
|
||||||
for chunk in raw.split(","):
|
|
||||||
entry = chunk.strip()
|
|
||||||
if not entry:
|
|
||||||
# Tolerate `"10.0.0.0/8,"` — trailing comma is benign.
|
|
||||||
continue
|
|
||||||
# ``strict=False`` so ``10.0.0.1/8`` is accepted as the network
|
|
||||||
# ``10.0.0.0/8`` instead of rejecting host bits — operators
|
|
||||||
# routinely paste a sample IP with a netmask and expect it to
|
|
||||||
# mean the network.
|
|
||||||
nets.append(ipaddress.ip_network(entry, strict=False))
|
|
||||||
return tuple(nets)
|
|
||||||
|
|
||||||
|
|
||||||
def load_trusted_gateway_cidrs(
|
def load_trusted_gateway_cidrs(
|
||||||
|
|
@ -130,43 +126,7 @@ def _normalize_ip(
|
||||||
Returns ``None`` on malformed input; callers treat that as "not a
|
Returns ``None`` on malformed input; callers treat that as "not a
|
||||||
trusted gateway".
|
trusted gateway".
|
||||||
"""
|
"""
|
||||||
try:
|
return normalize_ip(host)
|
||||||
addr = ipaddress.ip_address(host)
|
|
||||||
except ValueError:
|
|
||||||
return None
|
|
||||||
if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None:
|
|
||||||
return addr.ipv4_mapped
|
|
||||||
return addr
|
|
||||||
|
|
||||||
|
|
||||||
def peer_is_trusted_gateway(
|
|
||||||
peer_host: str | None,
|
|
||||||
cidrs: tuple[ipaddress.IPv4Network | ipaddress.IPv6Network, ...],
|
|
||||||
) -> bool:
|
|
||||||
"""Return True iff ``peer_host`` is inside any of the allow-list CIDRs.
|
|
||||||
|
|
||||||
Empty allow-list → always False (the strict-secure default).
|
|
||||||
``None`` peer (TestClient / UDS) → False; we never trust a peer we
|
|
||||||
can't even identify.
|
|
||||||
"""
|
|
||||||
if not cidrs:
|
|
||||||
return False
|
|
||||||
if peer_host is None:
|
|
||||||
return False
|
|
||||||
addr = _normalize_ip(peer_host)
|
|
||||||
if addr is None:
|
|
||||||
return False
|
|
||||||
for net in cidrs:
|
|
||||||
# IP/network family must match; ipaddress raises TypeError if
|
|
||||||
# we pass an IPv4 address into an IPv6 network membership test
|
|
||||||
# in some versions, so we family-gate first.
|
|
||||||
if isinstance(addr, ipaddress.IPv4Address) and isinstance(net, ipaddress.IPv6Network):
|
|
||||||
continue
|
|
||||||
if isinstance(addr, ipaddress.IPv6Address) and isinstance(net, ipaddress.IPv4Network):
|
|
||||||
continue
|
|
||||||
if addr in net:
|
|
||||||
return True
|
|
||||||
return False
|
|
||||||
|
|
||||||
|
|
||||||
def _peer_host(request: Any) -> str | None:
|
def _peer_host(request: Any) -> str | None:
|
||||||
|
|
@ -185,10 +145,7 @@ def _header_first(value: str) -> str:
|
||||||
leftmost hop is the only one whose authenticity the immediate
|
leftmost hop is the only one whose authenticity the immediate
|
||||||
gateway can vouch for, and beyond that we have no trust signal.
|
gateway can vouch for, and beyond that we have no trust signal.
|
||||||
"""
|
"""
|
||||||
if not value:
|
return header_first(value)
|
||||||
return ""
|
|
||||||
head, _, _ = value.partition(",")
|
|
||||||
return head.strip()
|
|
||||||
|
|
||||||
|
|
||||||
def _read_header(request: Any, name: str) -> str:
|
def _read_header(request: Any, name: str) -> str:
|
||||||
|
|
@ -249,29 +206,22 @@ def _resolve(request: Any) -> tuple[str, dict[str, str]]:
|
||||||
if cached_ip is not None and cached_fwd is not None:
|
if cached_ip is not None and cached_fwd is not None:
|
||||||
return cached_ip, cached_fwd
|
return cached_ip, cached_fwd
|
||||||
|
|
||||||
peer_host = _peer_host(request) or ""
|
inputs = ForwardedHeaderInputs(
|
||||||
fwd_for_raw = _read_header(request, "x-forwarded-for")
|
peer_host=_peer_host(request) or "",
|
||||||
fwd_proto_raw = _read_header(request, "x-forwarded-proto")
|
forwarded_for=_read_header(request, "x-forwarded-for"),
|
||||||
fwd_host_raw = _read_header(request, "x-forwarded-host")
|
forwarded_proto=_read_header(request, "x-forwarded-proto"),
|
||||||
|
forwarded_host=_read_header(request, "x-forwarded-host"),
|
||||||
cidrs = load_trusted_gateway_cidrs()
|
)
|
||||||
trusted = peer_is_trusted_gateway(peer_host or None, cidrs)
|
resolution = resolve_forwarded_headers(inputs, load_trusted_gateway_cidrs())
|
||||||
|
if resolution.rejected:
|
||||||
if trusted:
|
_emit_rejection_event(
|
||||||
client_ip = _header_first(fwd_for_raw) or peer_host
|
inputs.peer_host or None,
|
||||||
forwarded = {
|
inputs.forwarded_for,
|
||||||
"for": _header_first(fwd_for_raw),
|
inputs.forwarded_proto,
|
||||||
"proto": fwd_proto_raw.strip(),
|
inputs.forwarded_host,
|
||||||
"host": fwd_host_raw.strip(),
|
)
|
||||||
}
|
client_ip = resolution.client_ip
|
||||||
else:
|
forwarded = resolution.forwarded
|
||||||
# Headers may be absent (legitimate direct client). Only emit
|
|
||||||
# the rejection event if the peer ACTUALLY tried to set one —
|
|
||||||
# otherwise we'd spam logs for every direct request.
|
|
||||||
if fwd_for_raw or fwd_proto_raw or fwd_host_raw:
|
|
||||||
_emit_rejection_event(peer_host or None, fwd_for_raw, fwd_proto_raw, fwd_host_raw)
|
|
||||||
client_ip = peer_host
|
|
||||||
forwarded = {"for": "", "proto": "", "host": ""}
|
|
||||||
|
|
||||||
if state is not None:
|
if state is not None:
|
||||||
try:
|
try:
|
||||||
|
|
|
||||||
112
headroom/proxy/forwarded_policy.py
Normal file
112
headroom/proxy/forwarded_policy.py
Normal file
|
|
@ -0,0 +1,112 @@
|
||||||
|
"""Pure policy for trusted forwarded headers."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import ipaddress
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
Network = ipaddress.IPv4Network | ipaddress.IPv6Network
|
||||||
|
Address = ipaddress.IPv4Address | ipaddress.IPv6Address
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class ForwardedHeaderInputs:
|
||||||
|
"""Raw connection and forwarded-header values before trust evaluation."""
|
||||||
|
|
||||||
|
peer_host: str
|
||||||
|
forwarded_for: str = ""
|
||||||
|
forwarded_proto: str = ""
|
||||||
|
forwarded_host: str = ""
|
||||||
|
|
||||||
|
@property
|
||||||
|
def has_forwarded_headers(self) -> bool:
|
||||||
|
return bool(self.forwarded_for or self.forwarded_proto or self.forwarded_host)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class ForwardedHeaderResolution:
|
||||||
|
"""Deterministic forwarded-header trust decision."""
|
||||||
|
|
||||||
|
client_ip: str
|
||||||
|
forwarded: dict[str, str]
|
||||||
|
trusted: bool
|
||||||
|
rejected: bool
|
||||||
|
|
||||||
|
|
||||||
|
def parse_cidr_list(raw: str) -> tuple[Network, ...]:
|
||||||
|
"""Parse a comma-separated CIDR list. Empty / whitespace returns empty."""
|
||||||
|
if not raw or not raw.strip():
|
||||||
|
return ()
|
||||||
|
nets: list[Network] = []
|
||||||
|
for chunk in raw.split(","):
|
||||||
|
entry = chunk.strip()
|
||||||
|
if not entry:
|
||||||
|
continue
|
||||||
|
nets.append(ipaddress.ip_network(entry, strict=False))
|
||||||
|
return tuple(nets)
|
||||||
|
|
||||||
|
|
||||||
|
def normalize_ip(host: str) -> Address | None:
|
||||||
|
"""Parse ``host`` as IP, normalizing IPv4-mapped IPv6 to IPv4."""
|
||||||
|
try:
|
||||||
|
addr = ipaddress.ip_address(host)
|
||||||
|
except ValueError:
|
||||||
|
return None
|
||||||
|
if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None:
|
||||||
|
return addr.ipv4_mapped
|
||||||
|
return addr
|
||||||
|
|
||||||
|
|
||||||
|
def peer_is_trusted_gateway(
|
||||||
|
peer_host: str | None,
|
||||||
|
cidrs: tuple[Network, ...],
|
||||||
|
) -> bool:
|
||||||
|
"""Return True iff ``peer_host`` is inside any allow-listed CIDR."""
|
||||||
|
if not cidrs or peer_host is None:
|
||||||
|
return False
|
||||||
|
addr = normalize_ip(peer_host)
|
||||||
|
if addr is None:
|
||||||
|
return False
|
||||||
|
for net in cidrs:
|
||||||
|
if isinstance(addr, ipaddress.IPv4Address) and isinstance(net, ipaddress.IPv6Network):
|
||||||
|
continue
|
||||||
|
if isinstance(addr, ipaddress.IPv6Address) and isinstance(net, ipaddress.IPv4Network):
|
||||||
|
continue
|
||||||
|
if addr in net:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def header_first(value: str) -> str:
|
||||||
|
"""Return the leftmost element of a comma-separated header value."""
|
||||||
|
if not value:
|
||||||
|
return ""
|
||||||
|
head, _, _ = value.partition(",")
|
||||||
|
return head.strip()
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_forwarded_headers(
|
||||||
|
inputs: ForwardedHeaderInputs,
|
||||||
|
cidrs: tuple[Network, ...],
|
||||||
|
) -> ForwardedHeaderResolution:
|
||||||
|
"""Resolve client IP and sanitized forwarded headers from pure inputs."""
|
||||||
|
trusted = peer_is_trusted_gateway(inputs.peer_host or None, cidrs)
|
||||||
|
if trusted:
|
||||||
|
forwarded_for = header_first(inputs.forwarded_for)
|
||||||
|
return ForwardedHeaderResolution(
|
||||||
|
client_ip=forwarded_for or inputs.peer_host,
|
||||||
|
forwarded={
|
||||||
|
"for": forwarded_for,
|
||||||
|
"proto": inputs.forwarded_proto.strip(),
|
||||||
|
"host": inputs.forwarded_host.strip(),
|
||||||
|
},
|
||||||
|
trusted=True,
|
||||||
|
rejected=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
return ForwardedHeaderResolution(
|
||||||
|
client_ip=inputs.peer_host,
|
||||||
|
forwarded={"for": "", "proto": "", "host": ""},
|
||||||
|
trusted=False,
|
||||||
|
rejected=inputs.has_forwarded_headers,
|
||||||
|
)
|
||||||
70
tests/test_forwarded_policy.py
Normal file
70
tests/test_forwarded_policy.py
Normal file
|
|
@ -0,0 +1,70 @@
|
||||||
|
"""Tests for pure trusted-forwarded-header policy."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from headroom.proxy.forwarded_policy import (
|
||||||
|
ForwardedHeaderInputs,
|
||||||
|
header_first,
|
||||||
|
parse_cidr_list,
|
||||||
|
peer_is_trusted_gateway,
|
||||||
|
resolve_forwarded_headers,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolve_trusted_peer_honors_sanitized_forwarded_values() -> None:
|
||||||
|
cidrs = parse_cidr_list("10.0.0.0/8")
|
||||||
|
result = resolve_forwarded_headers(
|
||||||
|
ForwardedHeaderInputs(
|
||||||
|
peer_host="10.0.0.5",
|
||||||
|
forwarded_for="203.0.113.7, 10.0.0.99",
|
||||||
|
forwarded_proto=" https ",
|
||||||
|
forwarded_host=" api.example.com ",
|
||||||
|
),
|
||||||
|
cidrs,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.trusted is True
|
||||||
|
assert result.rejected is False
|
||||||
|
assert result.client_ip == "203.0.113.7"
|
||||||
|
assert result.forwarded == {
|
||||||
|
"for": "203.0.113.7",
|
||||||
|
"proto": "https",
|
||||||
|
"host": "api.example.com",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolve_untrusted_peer_rejects_forwarded_values() -> None:
|
||||||
|
cidrs = parse_cidr_list("10.0.0.0/8")
|
||||||
|
result = resolve_forwarded_headers(
|
||||||
|
ForwardedHeaderInputs(
|
||||||
|
peer_host="8.8.8.8",
|
||||||
|
forwarded_for="203.0.113.7",
|
||||||
|
forwarded_proto="https",
|
||||||
|
forwarded_host="api.example.com",
|
||||||
|
),
|
||||||
|
cidrs,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.trusted is False
|
||||||
|
assert result.rejected is True
|
||||||
|
assert result.client_ip == "8.8.8.8"
|
||||||
|
assert result.forwarded == {"for": "", "proto": "", "host": ""}
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolve_direct_untrusted_peer_without_forwarded_values_is_not_rejection() -> None:
|
||||||
|
result = resolve_forwarded_headers(ForwardedHeaderInputs(peer_host="8.8.8.8"), ())
|
||||||
|
|
||||||
|
assert result.trusted is False
|
||||||
|
assert result.rejected is False
|
||||||
|
assert result.client_ip == "8.8.8.8"
|
||||||
|
assert result.forwarded == {"for": "", "proto": "", "host": ""}
|
||||||
|
|
||||||
|
|
||||||
|
def test_peer_trust_handles_ipv4_mapped_ipv6() -> None:
|
||||||
|
cidrs = parse_cidr_list("10.0.0.0/8")
|
||||||
|
|
||||||
|
assert peer_is_trusted_gateway("::ffff:10.0.0.1", cidrs) is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_header_first_uses_leftmost_forwarded_for_hop() -> None:
|
||||||
|
assert header_first("203.0.113.7, 10.0.0.99, 10.0.0.5") == "203.0.113.7"
|
||||||
Loading…
Add table
Add a link
Reference in a new issue