allow setting allowed routes on proxy

litellm/proxy/auth/auth_utils.py

@@ -1,13 +1,123 @@
 import re
 import sys
 import traceback
+from typing import List, Optional, Tuple
 
-from fastapi import Request
+from fastapi import HTTPException, Request, status
 
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import *
 
 
+def _get_request_ip_address(
+    request: Request, use_x_forwarded_for: Optional[bool] = False
+) -> Optional[str]:
+
+    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 = ""
+
+    return client_ip
+
+
+def _check_valid_ip(
+    allowed_ips: Optional[List[str]],
+    request: Request,
+    use_x_forwarded_for: Optional[bool] = False,
+) -> Tuple[bool, Optional[str]]:
+    """
+    Returns if ip is allowed or not
+    """
+    if allowed_ips is None:  # if not set, assume true
+        return True, None
+
+    # if general_settings.get("use_x_forwarded_for") is True then use x-forwarded-for
+    client_ip = _get_request_ip_address(
+        request=request, use_x_forwarded_for=use_x_forwarded_for
+    )
+
+    # Check if IP address is allowed
+    if client_ip not in allowed_ips:
+        return False, client_ip
+
+    return True, client_ip
+
+
+def is_request_body_safe(request_body: dict) -> bool:
+    """
+    Check if the request body is safe.
+
+    A malicious user can set the api_base to their own domain and invoke POST /chat/completions to intercept and steal the OpenAI API key.
+    Relevant issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997
+    """
+    banned_params = ["api_base", "base_url"]
+
+    for param in banned_params:
+        if param in request_body:
+            raise ValueError(f"BadRequest: {param} is not allowed in request body")
+
+    return True
+
+
+async def pre_db_read_auth_checks(
+    request: Request,
+    request_data: dict,
+    route: str,
+):
+    """
+    1. Checks if request size is under max_request_size_mb (if set)
+    2. Check if request body is safe (example user has not set api_base in request body)
+    3. Check if IP address is allowed (if set)
+    4. Check if request route is an allowed route on the proxy (if set)
+
+    Returns:
+    - True
+
+    Raises:
+    - HTTPException if request fails initial auth checks
+    """
+    from litellm.proxy.proxy_server import general_settings, premium_user
+
+    # Check 1. request size
+    await check_if_request_size_is_safe(request=request)
+
+    # Check 2. Request body is safe
+    is_request_body_safe(request_body=request_data)
+
+    # Check 3. Check if IP address is allowed
+    is_valid_ip, passed_in_ip = _check_valid_ip(
+        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:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=f"Access forbidden: IP address {passed_in_ip} not allowed.",
+        )
+
+    # Check 4. Check if request route is an allowed route on the proxy
+    if "allowed_routes" in general_settings:
+        _allowed_routes = general_settings["allowed_routes"]
+        if premium_user is not True:
+            verbose_proxy_logger.error(
+                f"Trying to set allowed_routes. This is an Enterprise feature. {CommonProxyErrors.not_premium_user.value}"
+            )
+        if route not in _allowed_routes:
+            verbose_proxy_logger.error(
+                f"Route {route} not in allowed_routes={_allowed_routes}"
+            )
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail=f"Access forbidden: Route {route} not allowed",
+            )
+
+
 def route_in_additonal_public_routes(current_route: str):
     """
     Helper to check if the user defined public_routes on config.yaml