security - check max request size

litellm/proxy/auth/auth_utils.py

@@ -1,5 +1,7 @@
 import re
 
+from fastapi import Request
+
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import *
 
@@ -75,3 +77,69 @@ def is_llm_api_route(route: str) -> bool:
                 return True
 
     return False
+
+
+async def check_if_request_size_is_safe(request: Request) -> bool:
+    """
+    Enterprise Only:
+        - Checks if the request size is within the limit
+
+    Args:
+        request (Request): The incoming request.
+
+    Returns:
+        bool: True if the request size is within the limit, False otherwise.
+    """
+    from litellm.proxy.proxy_server import general_settings, premium_user
+
+    max_request_size_mb = general_settings.get("max_request_size_mb", None)
+    if max_request_size_mb is not None:
+        # Check if premium user
+        if premium_user is not True:
+            verbose_proxy_logger.warning(
+                f"using max_request_size_mb - not checking -  this is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}"
+            )
+            return True
+
+        # Get the request body
+        content_length = request.headers.get("content-length")
+
+        if content_length:
+            header_size = int(content_length)
+            header_size_mb = bytes_to_mb(bytes_value=header_size)
+            verbose_proxy_logger.debug(
+                f"content_length request size in MB={header_size_mb}"
+            )
+
+            if header_size_mb > max_request_size_mb:
+                raise ProxyException(
+                    message=f"Request size is too large. Request size is {header_size_mb} MB. Max size is {max_request_size_mb} MB",
+                    type=ProxyErrorTypes.bad_request_error.value,
+                    code=400,
+                    param="content-length",
+                )
+        else:
+            # If Content-Length is not available, read the body
+            body = await request.body()
+            body_size = len(body)
+            request_size_mb = bytes_to_mb(bytes_value=body_size)
+
+            verbose_proxy_logger.debug(
+                f"request body request size in MB={request_size_mb}"
+            )
+            if request_size_mb > max_request_size_mb:
+                raise ProxyException(
+                    message=f"Request size is too large. Request size is {request_size_mb} MB. Max size is {max_request_size_mb} MB",
+                    type=ProxyErrorTypes.bad_request_error.value,
+                    code=400,
+                    param="content-length",
+                )
+
+    return True
+
+
+def bytes_to_mb(bytes_value: int):
+    """
+    Helper to convert bytes to MB
+    """
+    return bytes_value / (1024 * 1024)