Merge pull request #1518 from roboflow/fix/issue-with-legacy-request-handler

Fix bug with improper parsing multipart-request causing reading the request body stream twice

Author: PawelPeczek-Roboflow

inference/core/interfaces/http/dependencies.py

@@ -1,17 +1,26 @@
-from typing import Optional
+from typing import Optional, Union
 
 from fastapi import Request
+from starlette.datastructures import UploadFile
 
+from inference.core.exceptions import InputImageLoadError
 
-async def request_body_content_only_to_be_used_in_legacy_request_handler(
+
+async def parse_body_content_for_legacy_request_handler(
     request: Request,
-) -> Optional[bytes]:
+) -> Optional[Union[bytes, UploadFile]]:
     content_type = request.headers.get("Content-Type")
+    if content_type is None:
+        return None
     image_reference_in_query = request.query_params.get("image")
-    if (
-        content_type is None
-        or "multipart/form-data" in content_type
-        or image_reference_in_query
-    ):
+    if "multipart/form-data" in content_type:
+        form_data = await request.form()
+        if "file" not in form_data:
+            raise InputImageLoadError(
+                message="Expected image to be send in part named 'file' of multipart/form-data request",
+                public_message="Expected image to be send in part named 'file' of multipart/form-data request",
+            )
+        return form_data["file"]
+    if content_type is None or image_reference_in_query:
         return None
     return await request.body()

inference/core/interfaces/http/http_api.py

@@ -18,11 +18,11 @@
     Path,
     Query,
     Request,
-    UploadFile,
 )
 from fastapi.responses import JSONResponse, RedirectResponse, Response
 from fastapi.staticfiles import StaticFiles
 from fastapi_cprofile.profiler import CProfileMiddleware
+from starlette.datastructures import UploadFile
 from starlette.middleware.base import BaseHTTPMiddleware
 
 from inference.core import logger
@@ -161,12 +161,13 @@
 from inference.core.exceptions import (
     ContentTypeInvalid,
     ContentTypeMissing,
+    InputImageLoadError,
     MissingServiceSecretError,
     RoboflowAPINotAuthorizedError,
 )
 from inference.core.interfaces.base import BaseInterface
 from inference.core.interfaces.http.dependencies import (
-    request_body_content_only_to_be_used_in_legacy_request_handler,
+    parse_body_content_for_legacy_request_handler,
 )
 from inference.core.interfaces.http.error_handlers import (
     with_route_exceptions,
@@ -2395,10 +2396,8 @@ def legacy_infer_from_request(
                 background_tasks: BackgroundTasks,
                 request: Request,
                 request_body: Annotated[
-                    Optional[bytes],
-                    Depends(
-                        request_body_content_only_to_be_used_in_legacy_request_handler
-                    ),
+                    Optional[Union[bytes, UploadFile]],
+                    Depends(parse_body_content_for_legacy_request_handler),
                 ],
                 dataset_id: str = Path(
                     description="ID of a Roboflow dataset corresponding to the model to use for inference OR workspace ID"
@@ -2492,7 +2491,6 @@ def legacy_infer_from_request(
                     "external",
                     description="The detailed source information of the inference request",
                 ),
-                file: Optional[UploadFile] = None,
             ):
                 """
                 Legacy inference endpoint for object detection, instance segmentation, and classification.
@@ -2511,36 +2509,35 @@ def legacy_infer_from_request(
                     f"Reached legacy route /:dataset_id/:version_id with {dataset_id}/{version_id}"
                 )
                 model_id = f"{dataset_id}/{version_id}"
-
                 if confidence >= 1:
                     confidence /= 100
                 elif confidence < 0.01:
                     confidence = 0.01
 
                 if overlap >= 1:
                     overlap /= 100
-
                 if image is not None:
                     request_image = InferenceRequestImage(type="url", value=image)
                 else:
                     if "Content-Type" not in request.headers:
                         raise ContentTypeMissing(
                             f"Request must include a Content-Type header"
                         )
-                    if file is not None:
-                        base64_image_str = file.file.read()
+                    if isinstance(request_body, UploadFile):
+                        base64_image_str = request_body.file.read()
                         base64_image_str = base64.b64encode(base64_image_str)
                         request_image = InferenceRequestImage(
                             type="base64", value=base64_image_str.decode("ascii")
                         )
-                    elif (
-                        "application/x-www-form-urlencoded"
-                        in request.headers["Content-Type"]
-                        or "application/json" in request.headers["Content-Type"]
-                    ):
+                    elif isinstance(request_body, bytes):
                         request_image = InferenceRequestImage(
                             type=image_type, value=request_body
                         )
+                    elif request_body is None:
+                        raise InputImageLoadError(
+                            message="Image not found in request body.",
+                            public_message="Image not found in request body.",
+                        )
                     else:
                         raise ContentTypeInvalid(
                             f"Invalid Content-Type: {request.headers['Content-Type']}"
@@ -2701,6 +2698,15 @@ def model_add_legacy(
             async def dashboard_guard():
                 return Response(status_code=404)
 
+        @app.exception_handler(InputImageLoadError)
+        async def unicorn_exception_handler(request: Request, exc: InputImageLoadError):
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "message": f"Could not load input image. Cause: {exc.get_public_error_details()}"
+                },
+            )
+
         app.mount(
             "/",
             StaticFiles(directory="./inference/landing/out", html=True),

inference/core/version.py

@@ -1,4 +1,4 @@
-__version__ = "0.54.0"
+__version__ = "0.54.1"
 
 
 if __name__ == "__main__":