From 627f72832d5982aea7da72e4076a1bf9b7c59263 Mon Sep 17 00:00:00 2001 From: Saifeddine ALOUI Date: Thu, 12 Sep 2024 09:08:44 +0200 Subject: [PATCH] Update security.py --- lollms/security.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/lollms/security.py b/lollms/security.py index 708ceed..a33e169 100644 --- a/lollms/security.py +++ b/lollms/security.py @@ -9,6 +9,7 @@ import re import platform import string from lollms.utilities import PackageManager +from starlette.middleware.base import BaseHTTPMiddleware if not PackageManager.check_package_installed("lxml"): PackageManager.install_package("lxml") @@ -334,3 +335,11 @@ if __name__=="__main__": print(f"Original: {path}, Sanitized: {sanitized}") except HTTPException as e: print(f"Original: {path}, Exception: {e.detail}") + +class MultipartBoundaryCheck(BaseHTTPMiddleware): + async def dispatch(self, request, call_next): + if request.headers.get("content-type", "").startswith("multipart/form-data"): + boundary = request.headers.get("content-type").split("boundary=")[-1] + if len(boundary) > 70: # Adjust this limit as needed + return JSONResponse(status_code=400, content={"detail": "Invalid boundary length"}) + return await call_next(request)