Merge pull request #5181 from BerriAI/litellm_check_x_fwded_for

[Feat-Proxy Security] Allow Using `x-forwarded-for` for enforcing + tracking  ip address

litellm/proxy/auth/user_api_key_auth.py

@@ -124,7 +124,9 @@ async def user_api_key_auth(
         await check_if_request_size_is_safe(request=request)
 
         is_valid_ip = _check_valid_ip(
-            allowed_ips=general_settings.get("allowed_ips", None), request=request
+            allowed_ips=general_settings.get("allowed_ips", None),
+            use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False),
+            request=request,
         )
 
         if not is_valid_ip:
@@ -1206,17 +1208,23 @@ def _get_user_role(
     return role
 
 
-def _check_valid_ip(allowed_ips: Optional[List[str]], request: Request) -> bool:
+def _check_valid_ip(
+    allowed_ips: Optional[List[str]],
+    request: Request,
+    use_x_forwarded_for: Optional[bool] = False,
+) -> bool:
     """
     Returns if ip is allowed or not
     """
     if allowed_ips is None:  # if not set, assume true
         return True
 
-    if request.client is not None:
+    # if general_settings.get("use_x_forwarded_for") is True then use x-forwarded-for
+    client_ip = None
+    if use_x_forwarded_for is True and "x-forwarded-for" in request.headers:
+        client_ip = request.headers["x-forwarded-for"]
+    elif request.client is not None:
         client_ip = request.client.host
-    else:
-        client_ip = None
 
     # Check if IP address is allowed
     if client_ip not in allowed_ips:

litellm/proxy/litellm_pre_call_utils.py

@@ -177,7 +177,17 @@ async def add_litellm_data_to_request(
     requester_ip_address = ""
     if premium_user is True:
         # Only set the IP Address for Enterprise Users
+
+        # logic for tracking IP Address
         if (
+            general_settings is not None
+            and general_settings.get("use_x_forwarded_for") is True
+            and request is not None
+            and hasattr(request, "headers")
+            and "x-forwarded-for" in request.headers
+        ):
+            requester_ip_address = request.headers["x-forwarded-for"]
+        elif (
             request is not None
             and hasattr(request, "client")
             and hasattr(request.client, "host")

litellm/tests/test_user_api_key_auth.py

@@ -7,7 +7,7 @@ import sys
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
-from typing import List, Optional
+from typing import Dict, List, Optional
 from unittest.mock import MagicMock
 
 import pytest
@@ -16,9 +16,10 @@ import litellm
 
 
 class Request:
-    def __init__(self, client_ip: Optional[str] = None):
+    def __init__(self, client_ip: Optional[str] = None, headers: Optional[dict] = None):
         self.client = MagicMock()
         self.client.host = client_ip
+        self.headers: Dict[str, str] = {}
 
 
 @pytest.mark.parametrize(
@@ -46,6 +47,34 @@ def test_check_valid_ip(
     assert _check_valid_ip(allowed_ips, request) == expected_result  # type: ignore
 
 
+# test x-forwarder for is used when user has opted in
+
+
+@pytest.mark.parametrize(
+    "allowed_ips, client_ip, expected_result",
+    [
+        (None, "127.0.0.1", True),  # No IP restrictions, should be allowed
+        (["127.0.0.1"], "127.0.0.1", True),  # IP in allowed list
+        (["192.168.1.1"], "127.0.0.1", False),  # IP not in allowed list
+        ([], "127.0.0.1", False),  # Empty allowed list, no IP should be allowed
+        (["192.168.1.1", "10.0.0.1"], "10.0.0.1", True),  # IP in allowed list
+        (
+            ["192.168.1.1"],
+            None,
+            False,
+        ),  # Request with no client IP should not be allowed
+    ],
+)
+def test_check_valid_ip_sent_with_x_forwarded_for(
+    allowed_ips: Optional[List[str]], client_ip: Optional[str], expected_result: bool
+):
+    from litellm.proxy.auth.user_api_key_auth import _check_valid_ip
+
+    request = Request(client_ip, headers={"X-Forwarded-For": client_ip})
+
+    assert _check_valid_ip(allowed_ips, request, use_x_forwarded_for=True) == expected_result  # type: ignore
+
+
 @pytest.mark.asyncio
 async def test_check_blocked_team():
     """