mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
from fastapi import Request
|
|
|
|
from litellm.proxy.auth.network import (
|
|
TrustedProxyConfig,
|
|
resolve_client_ip,
|
|
resolve_network_context,
|
|
)
|
|
|
|
TRUSTED = TrustedProxyConfig(use_forwarded_for=True, trusted_proxy_cidrs=["10.0.0.0/8"])
|
|
|
|
|
|
def make_request(
|
|
*,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
client: Optional[Tuple[str, int]] = ("203.0.113.7", 5555),
|
|
) -> Request:
|
|
raw_headers: List[Tuple[bytes, bytes]] = [
|
|
(key.lower().encode(), value.encode()) for key, value in (headers or {}).items()
|
|
]
|
|
scope: Dict[str, Any] = {
|
|
"type": "http",
|
|
"http_version": "1.1",
|
|
"method": "GET",
|
|
"path": "/",
|
|
"raw_path": b"/",
|
|
"query_string": b"",
|
|
"headers": raw_headers,
|
|
"client": client,
|
|
"server": ("testserver", 80),
|
|
"scheme": "http",
|
|
}
|
|
return Request(scope)
|
|
|
|
|
|
def test_xff_ignored_when_forwarding_disabled():
|
|
config = TrustedProxyConfig(
|
|
use_forwarded_for=False, trusted_proxy_cidrs=["10.0.0.0/8"]
|
|
)
|
|
request = make_request(
|
|
headers={"x-forwarded-for": "203.0.113.9"}, client=("10.0.0.1", 1)
|
|
)
|
|
ip, via_proxy = resolve_client_ip(request, config)
|
|
assert ip == "10.0.0.1"
|
|
assert via_proxy is False
|
|
|
|
|
|
def test_xff_honored_from_trusted_peer():
|
|
request = make_request(
|
|
headers={"x-forwarded-for": "203.0.113.9, 10.0.0.5"}, client=("10.0.0.1", 1)
|
|
)
|
|
ip, via_proxy = resolve_client_ip(request, TRUSTED)
|
|
assert ip == "203.0.113.9"
|
|
assert via_proxy is True
|
|
|
|
|
|
def test_ipv4_mapped_peer_and_hop_match_ipv4_trusted_ranges():
|
|
request = make_request(headers={"x-forwarded-for": "203.0.113.9, ::ffff:10.0.0.5"}, client=("::ffff:10.0.0.1", 1))
|
|
ip, via_proxy = resolve_client_ip(request, TRUSTED)
|
|
assert ip == "203.0.113.9"
|
|
assert via_proxy is True
|
|
|
|
|
|
def test_ipv4_mapped_peer_still_matches_mapped_notation_trusted_range():
|
|
config = TrustedProxyConfig(use_forwarded_for=True, trusted_proxy_cidrs=["::ffff:10.0.0.0/104"])
|
|
request = make_request(headers={"x-forwarded-for": "203.0.113.9"}, client=("::ffff:10.0.0.1", 1))
|
|
ip, via_proxy = resolve_client_ip(request, config)
|
|
assert ip == "203.0.113.9"
|
|
assert via_proxy is True
|
|
|
|
|
|
def test_spoofed_xff_from_untrusted_peer_is_ignored():
|
|
request = make_request(
|
|
headers={"x-forwarded-for": "203.0.113.9"}, client=("8.8.8.8", 1)
|
|
)
|
|
ip, via_proxy = resolve_client_ip(request, TRUSTED)
|
|
assert ip == "8.8.8.8"
|
|
assert via_proxy is False
|
|
|
|
|
|
def test_right_to_left_parse_skips_chained_trusted_proxies():
|
|
request = make_request(
|
|
headers={"x-forwarded-for": "198.51.100.4, 10.1.1.1, 10.0.0.9"},
|
|
client=("10.0.0.1", 1),
|
|
)
|
|
ip, via_proxy = resolve_client_ip(request, TRUSTED)
|
|
assert ip == "198.51.100.4"
|
|
assert via_proxy is True
|
|
|
|
|
|
def test_all_trusted_hops_fall_back_to_peer():
|
|
request = make_request(
|
|
headers={"x-forwarded-for": "10.1.1.1, 10.0.0.9"}, client=("10.0.0.1", 1)
|
|
)
|
|
ip, via_proxy = resolve_client_ip(request, TRUSTED)
|
|
assert ip == "10.0.0.1"
|
|
assert via_proxy is True
|
|
|
|
|
|
def test_invalid_xff_token_is_skipped():
|
|
request = make_request(
|
|
headers={"x-forwarded-for": "not-an-ip, 203.0.113.50"}, client=("10.0.0.1", 1)
|
|
)
|
|
ip, _ = resolve_client_ip(request, TRUSTED)
|
|
assert ip == "203.0.113.50"
|
|
|
|
|
|
def test_network_context_captures_host_and_proxy_flag():
|
|
request = make_request(
|
|
headers={"x-forwarded-for": "203.0.113.9", "host": "proxy.litellm.ai"},
|
|
client=("10.0.0.1", 1),
|
|
)
|
|
ctx = resolve_network_context(request, TRUSTED)
|
|
assert ctx.client_ip == "203.0.113.9"
|
|
assert ctx.host == "proxy.litellm.ai"
|
|
assert ctx.via_trusted_proxy is True
|