diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 48c8acee4a..9bbbc1a430 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/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: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 13f9475c5c..990cb52337 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/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") diff --git a/litellm/tests/test_user_api_key_auth.py b/litellm/tests/test_user_api_key_auth.py index 33f055b37d..ad057ee572 100644 --- a/litellm/tests/test_user_api_key_auth.py +++ b/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(): """