Prevent too large requests (#3048)

* prevent too large requests

* reject file upload attempts

* lint

Author: avishniakov

src/zenml/config/server_config.py

@@ -29,6 +29,7 @@
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY,
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE,
     DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS,
+    DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES,
     DEFAULT_ZENML_SERVER_NAME,
     DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE,
@@ -231,6 +232,8 @@ class ServerConfiguration(BaseModel):
         auto_activate: Whether to automatically activate the server and create a
             default admin user account with an empty password during the initial
             deployment.
+        max_request_body_size_in_bytes: The maximum size of the request body in
+            bytes. If not specified, the default value of 256 Kb will be used.
     """
 
     deployment_type: ServerDeploymentType = ServerDeploymentType.OTHER
@@ -319,6 +322,10 @@ class ServerConfiguration(BaseModel):
 
     thread_pool_size: int = DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE
 
+    max_request_body_size_in_bytes: int = (
+        DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES
+    )
+
     _deployment_id: Optional[UUID] = None
 
     @model_validator(mode="before")

src/zenml/constants.py

@@ -317,6 +317,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int:
 DEFAULT_ZENML_SERVER_SECURE_HEADERS_REPORT_TO = "default"
 DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD = False
 DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS = 30
+DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES = 256 * 1024 * 1024
 
 # Configurations to decide which resources report their usage and check for
 # entitlement in the case of a cloud deployment. Expected Format is this:

src/zenml/zen_server/zen_server_api.py

@@ -24,16 +24,21 @@
 from asyncio.log import logger
 from datetime import datetime, timedelta, timezone
 from genericpath import isfile
-from typing import Any, List
+from typing import Any, List, Set
 
 from anyio import to_thread
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import ORJSONResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
+from starlette.middleware.base import (
+    BaseHTTPMiddleware,
+    RequestResponseEndpoint,
+)
 from starlette.middleware.cors import CORSMiddleware
-from starlette.responses import FileResponse
+from starlette.responses import FileResponse, JSONResponse, Response
+from starlette.types import ASGIApp
 
 import zenml
 from zenml.analytics import source_context
@@ -143,6 +148,79 @@ def validation_exception_handler(
     return ORJSONResponse(error_detail(exc, ValueError), status_code=422)
 
 
+class RequestBodyLimit(BaseHTTPMiddleware):
+    """Limits the size of the request body."""
+
+    def __init__(self, app: ASGIApp, max_bytes: int) -> None:
+        """Limits the size of the request body.
+
+        Args:
+            app: The FastAPI app.
+            max_bytes: The maximum size of the request body.
+        """
+        super().__init__(app)
+        self.max_bytes = max_bytes
+
+    async def dispatch(
+        self, request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
+        """Limits the size of the request body.
+
+        Args:
+            request: The incoming request.
+            call_next: The next function to be called.
+
+        Returns:
+            The response to the request.
+        """
+        if content_length := request.headers.get("content-length"):
+            if int(content_length) > self.max_bytes:
+                return Response(status_code=413)  # Request Entity Too Large
+        return await call_next(request)
+
+
+class RestrictFileUploadsMiddleware(BaseHTTPMiddleware):
+    """Restrict file uploads to certain paths."""
+
+    def __init__(self, app: FastAPI, allowed_paths: Set[str]):
+        """Restrict file uploads to certain paths.
+
+        Args:
+            app: The FastAPI app.
+            allowed_paths: The allowed paths.
+        """
+        super().__init__(app)
+        self.allowed_paths = allowed_paths
+
+    async def dispatch(
+        self, request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
+        """Restrict file uploads to certain paths.
+
+        Args:
+            request: The incoming request.
+            call_next: The next function to be called.
+
+        Returns:
+            The response to the request.
+        """
+        if request.method == "POST":
+            content_type = request.headers.get("content-type", "")
+            if (
+                "multipart/form-data" in content_type
+                and request.url.path not in self.allowed_paths
+            ):
+                return JSONResponse(
+                    status_code=403,
+                    content={
+                        "detail": "File uploads are not allowed on this endpoint."
+                    },
+                )
+        return await call_next(request)
+
+
+ALLOWED_FOR_FILE_UPLOAD: Set[str] = set()
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=server_config().cors_allow_origins,
@@ -151,6 +229,13 @@ def validation_exception_handler(
     allow_headers=["*"],
 )
 
+app.add_middleware(
+    RequestBodyLimit, max_bytes=server_config().max_request_body_size_in_bytes
+)
+app.add_middleware(
+    RestrictFileUploadsMiddleware, allowed_paths=ALLOWED_FOR_FILE_UPLOAD
+)
+
 
 @app.middleware("http")
 async def set_secure_headers(request: Request, call_next: Any) -> Any: