84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
 
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from packetstation.proxy import TrustedProxies, TrustedProxyMiddleware, forwarded
|
||
|
|
|
||
|
|
PROXY = "172.18.0.5"
|
||
|
|
|
||
|
|
|
||
|
|
def trusted(addresses: set[str] | None = None) -> TrustedProxies:
|
||
|
|
t = TrustedProxies("npm_app_1", resolver=lambda _h: set(addresses or {PROXY}))
|
||
|
|
t.refresh()
|
||
|
|
return t
|
||
|
|
|
||
|
|
|
||
|
|
def H(**kw: str) -> list[tuple[bytes, bytes]]:
|
||
|
|
return [(k.replace("_", "-").encode(), v.encode()) for k, v in kw.items()]
|
||
|
|
|
||
|
|
|
||
|
|
def test_trusted_peer_uses_first_xff_and_proto() -> None:
|
||
|
|
ip, scheme = forwarded(PROXY, H(x_forwarded_for="10.1.2.3, 172.18.0.9", x_forwarded_proto="https"), trusted())
|
||
|
|
assert ip == "10.1.2.3" and scheme == "https"
|
||
|
|
|
||
|
|
|
||
|
|
def test_untrusted_peer_ignores_headers() -> None:
|
||
|
|
ip, scheme = forwarded("10.9.9.9", H(x_forwarded_for="1.2.3.4", x_forwarded_proto="https"), trusted())
|
||
|
|
assert ip == "10.9.9.9" and scheme is None
|
||
|
|
|
||
|
|
|
||
|
|
def test_empty_host_trusts_nothing() -> None:
|
||
|
|
t = TrustedProxies("")
|
||
|
|
t.refresh()
|
||
|
|
assert forwarded(PROXY, H(x_forwarded_for="1.2.3.4"), t) == (PROXY, None)
|
||
|
|
|
||
|
|
|
||
|
|
def test_garbage_xff_and_proto_ignored() -> None:
|
||
|
|
ip, scheme = forwarded(PROXY, H(x_forwarded_for="not-an-ip", x_forwarded_proto="gopher"), trusted())
|
||
|
|
assert ip == PROXY and scheme is None
|
||
|
|
|
||
|
|
|
||
|
|
def test_resolution_failure_keeps_last_good_set() -> None:
|
||
|
|
calls = {"n": 0}
|
||
|
|
|
||
|
|
def resolver(_h: str) -> set[str]:
|
||
|
|
calls["n"] += 1
|
||
|
|
if calls["n"] > 1:
|
||
|
|
raise OSError("no such host")
|
||
|
|
return {PROXY}
|
||
|
|
|
||
|
|
t = TrustedProxies("npm_app_1", resolver=resolver)
|
||
|
|
t.refresh()
|
||
|
|
t.refresh()
|
||
|
|
assert t.addresses == frozenset({PROXY})
|
||
|
|
|
||
|
|
|
||
|
|
def test_re_resolve_picks_up_new_address() -> None:
|
||
|
|
addrs = [{PROXY}, {"172.18.0.77"}]
|
||
|
|
t = TrustedProxies("npm_app_1", resolver=lambda _h: addrs.pop(0))
|
||
|
|
t.refresh()
|
||
|
|
t.refresh()
|
||
|
|
assert t.trusts("172.18.0.77") and not t.trusts(PROXY)
|
||
|
|
|
||
|
|
|
||
|
|
def _run_mw(peer: str, headers: list[tuple[bytes, bytes]], t: TrustedProxies) -> dict[str, Any]:
|
||
|
|
seen: dict[str, Any] = {}
|
||
|
|
|
||
|
|
async def app(scope: Any, receive: Any, send: Any) -> None:
|
||
|
|
seen.update(scope)
|
||
|
|
|
||
|
|
scope = {"type": "http", "client": (peer, 5555), "scheme": "http", "headers": headers, "path": "/"}
|
||
|
|
asyncio.run(TrustedProxyMiddleware(app, t)(scope, None, None)) # type: ignore[arg-type]
|
||
|
|
return seen
|
||
|
|
|
||
|
|
|
||
|
|
def test_middleware_rewrites_for_trusted_peer() -> None:
|
||
|
|
s = _run_mw(PROXY, H(x_forwarded_for="10.1.2.3", x_forwarded_proto="https"), trusted())
|
||
|
|
assert s["client"] == ("10.1.2.3", 5555) and s["scheme"] == "https"
|
||
|
|
|
||
|
|
|
||
|
|
def test_middleware_leaves_untrusted_peer() -> None:
|
||
|
|
s = _run_mw("10.5.5.5", H(x_forwarded_for="10.1.2.3", x_forwarded_proto="https"), trusted())
|
||
|
|
assert s["client"] == ("10.5.5.5", 5555) and s["scheme"] == "http"
|