feat(user_api_key_auth.py): allow restricting calls by IP address

Allows admin to restrict which IP addresses can make calls to the proxy
This commit is contained in:
Krrish Dholakia 2024-07-08 15:58:15 -07:00
parent 95739c3778
commit f3cc57bc6f
2 changed files with 78 additions and 0 deletions

View file

@ -136,6 +136,19 @@ async def user_api_key_auth(
enable_jwt_auth: true enable_jwt_auth: true
``` ```
""" """
### FILTER IP ADDRESS ###
is_valid_ip = _check_valid_ip(
allowed_ips=general_settings.get("allowed_ips", None), request=request
)
if not is_valid_ip:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Access forbidden: IP address not allowed.",
)
route: str = request.url.path route: str = request.url.path
if ( if (
@ -1208,3 +1221,22 @@ def _get_user_role(user_id_information: Optional[list]):
_user = user_id_information[0] _user = user_id_information[0]
return _user.get("user_role") return _user.get("user_role")
def _check_valid_ip(allowed_ips: Optional[List[str]], request: Request) -> 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:
client_ip = request.client.host
else:
client_ip = None
# Check if IP address is allowed
if client_ip not in allowed_ips:
return False
return True

View file

@ -0,0 +1,46 @@
# What is this?
## Unit tests for user_api_key_auth helper functions
import os
import sys
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
from typing import List, Optional
from unittest.mock import MagicMock
import pytest
import litellm
class Request:
def __init__(self, client_ip: Optional[str] = None):
self.client = MagicMock()
self.client.host = client_ip
@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(
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)
assert _check_valid_ip(allowed_ips, request) == expected_result # type: ignore