111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
 
|
"""Trust X-Forwarded-For / X-Forwarded-Proto only from the proxy's resolved addresses."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import ipaddress
|
||
|
|
import logging
|
||
|
|
import socket
|
||
|
|
from collections.abc import Callable, Iterable
|
||
|
|
|
||
|
|
from starlette.requests import Request
|
||
|
|
from starlette.types import ASGIApp, Receive, Scope, Send
|
||
|
|
|
||
|
|
log = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
REFRESH_SECONDS = 60
|
||
|
|
|
||
|
|
|
||
|
|
def _ip(text: str) -> str | None:
|
||
|
|
try:
|
||
|
|
return str(ipaddress.ip_address(text.strip()))
|
||
|
|
except ValueError:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def resolve_host(host: str) -> set[str]:
|
||
|
|
"""Every address the name resolves to (IPv4 and IPv6)."""
|
||
|
|
infos = socket.getaddrinfo(host, None, proto=socket.IPPROTO_TCP)
|
||
|
|
out = {ip for info in infos if (ip := _ip(str(info[4][0]).split("%")[0]))}
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
class TrustedProxies:
|
||
|
|
"""The current set of trusted peer addresses. A failed refresh keeps the last good set."""
|
||
|
|
|
||
|
|
def __init__(self, host: str, resolver: Callable[[str], set[str]] = resolve_host):
|
||
|
|
self.host = host.strip()
|
||
|
|
self.addresses: frozenset[str] = frozenset()
|
||
|
|
self._resolver = resolver
|
||
|
|
|
||
|
|
def refresh(self) -> None:
|
||
|
|
if not self.host:
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
found = self._resolver(self.host)
|
||
|
|
except OSError as e:
|
||
|
|
log.warning(
|
||
|
|
"can't resolve TRUSTED_PROXY_HOST %s (%s); keeping %s",
|
||
|
|
self.host,
|
||
|
|
e,
|
||
|
|
sorted(self.addresses) or "nothing",
|
||
|
|
)
|
||
|
|
return
|
||
|
|
if found != self.addresses:
|
||
|
|
log.info("trusted proxy %s is %s", self.host, sorted(found))
|
||
|
|
self.addresses = frozenset(found)
|
||
|
|
|
||
|
|
def trusts(self, peer: str | None) -> bool:
|
||
|
|
return bool(peer) and _ip(peer or "") in self.addresses
|
||
|
|
|
||
|
|
async def run(self) -> None:
|
||
|
|
"""Background loop: re-resolve every 60 s."""
|
||
|
|
while True:
|
||
|
|
await asyncio.sleep(REFRESH_SECONDS)
|
||
|
|
await asyncio.to_thread(self.refresh)
|
||
|
|
|
||
|
|
|
||
|
|
def forwarded(
|
||
|
|
peer: str | None, headers: Iterable[tuple[bytes, bytes]], trusted: TrustedProxies
|
||
|
|
) -> tuple[str | None, str | None]:
|
||
|
|
"""(client ip, scheme) to use; (peer, None) when the peer is not trusted or headers are unusable."""
|
||
|
|
if not trusted.trusts(peer):
|
||
|
|
return peer, None
|
||
|
|
xff = proto = None
|
||
|
|
for k, v in headers:
|
||
|
|
name = k.lower()
|
||
|
|
if name == b"x-forwarded-for" and xff is None:
|
||
|
|
xff = v.decode("latin-1")
|
||
|
|
elif name == b"x-forwarded-proto" and proto is None:
|
||
|
|
proto = v.decode("latin-1").split(",")[0].strip().lower()
|
||
|
|
client = peer
|
||
|
|
if xff:
|
||
|
|
first = _ip(xff.split(",")[0])
|
||
|
|
if first:
|
||
|
|
client = first
|
||
|
|
scheme = proto if proto in ("http", "https") else None
|
||
|
|
return client, scheme
|
||
|
|
|
||
|
|
|
||
|
|
class TrustedProxyMiddleware:
|
||
|
|
def __init__(self, app: ASGIApp, trusted: TrustedProxies):
|
||
|
|
self.app = app
|
||
|
|
self.trusted = trusted
|
||
|
|
|
||
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||
|
|
if scope["type"] in ("http", "websocket"):
|
||
|
|
client = scope.get("client")
|
||
|
|
peer = client[0] if client else None
|
||
|
|
ip, scheme = forwarded(peer, scope.get("headers", []), self.trusted)
|
||
|
|
if ip != peer or scheme:
|
||
|
|
scope = dict(scope)
|
||
|
|
scope["client"] = (ip, client[1] if client else 0)
|
||
|
|
if scheme:
|
||
|
|
scope["scheme"] = scheme if scope["type"] == "http" else ("wss" if scheme == "https" else "ws")
|
||
|
|
await self.app(scope, receive, send)
|
||
|
|
|
||
|
|
|
||
|
|
def client_ip(request: Request) -> str | None:
|
||
|
|
"""The real client address (after trusted-proxy handling)."""
|
||
|
|
return request.client.host if request.client else None
|