address comments

Signed-off-by: Lucas <lyoon@redhat.com>

Author: JslYoon

src/app/endpoints/vector_stores.py

@@ -3,7 +3,6 @@
 import asyncio
 import os
 import traceback
-from io import BytesIO
 from typing import Annotated, Any
 
 from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile, status
@@ -14,6 +13,7 @@
 from authorization.middleware import authorize
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
+from constants import DEFAULT_MAX_FILE_UPLOAD_SIZE
 from log import get_logger
 from models.config import Action
 from models.requests import (
@@ -22,6 +22,7 @@
     VectorStoreUpdateRequest,
 )
 from models.responses import (
+    AbstractErrorResponse,
     FileResponse,
     ForbiddenResponse,
     InternalServerErrorResponse,
@@ -63,6 +64,7 @@
 
 file_responses: dict[int | str, dict[str, Any]] = {
     200: FileResponse.openapi_response(),
+    400: {"description": "Bad Request - Invalid file upload"},
     401: UnauthorizedResponse.openapi_response(
         examples=["missing header", "missing token"]
     ),
@@ -434,22 +436,69 @@ async def create_file(
 
     Raises:
         HTTPException:
-            - 400: Bad request (e.g., file too large)
+            - 400: Bad request (e.g., file too large, invalid format)
             - 401: Authentication failed
             - 403: Authorization failed
             - 500: Lightspeed Stack configuration not loaded
             - 503: Unable to connect to Llama Stack
     """
     _ = auth
-    _ = request
 
     check_configuration_loaded(configuration)
 
+    # Check Content-Length header BEFORE reading to prevent DoS via memory exhaustion
+    content_length = request.headers.get("content-length")
+    if content_length:
+        try:
+            size = int(content_length)
+            if size > DEFAULT_MAX_FILE_UPLOAD_SIZE:
+                response = AbstractErrorResponse(
+                    response="File too large",
+                    cause=(
+                        f"File size {size} bytes exceeds maximum allowed "
+                        f"size of {DEFAULT_MAX_FILE_UPLOAD_SIZE} bytes "
+                        f"({DEFAULT_MAX_FILE_UPLOAD_SIZE // (1024 * 1024)} MB)"
+                    ),
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                )
+                raise HTTPException(**response.model_dump())
+        except ValueError:
+            # Invalid Content-Length header, continue and validate after reading
+            pass
+
+    # file.size attribute if available
+    if hasattr(file, "size") and file.size is not None:
+        if file.size > DEFAULT_MAX_FILE_UPLOAD_SIZE:
+            response = AbstractErrorResponse(
+                response="File too large",
+                cause=(
+                    f"File size {file.size} bytes exceeds maximum allowed "
+                    f"size of {DEFAULT_MAX_FILE_UPLOAD_SIZE} bytes "
+                    f"({DEFAULT_MAX_FILE_UPLOAD_SIZE // (1024 * 1024)} MB)"
+                ),
+                status_code=status.HTTP_400_BAD_REQUEST,
+            )
+            raise HTTPException(**response.model_dump())
+
     try:
         client = AsyncLlamaStackClientHolder().get_client()
 
-        # Read file content
+        # Read file content once
         content = await file.read()
+
+        # Verify actual size after reading
+        if len(content) > DEFAULT_MAX_FILE_UPLOAD_SIZE:
+            response = AbstractErrorResponse(
+                response="File too large",
+                cause=(
+                    f"File content size {len(content)} bytes exceeds maximum "
+                    f"allowed size of {DEFAULT_MAX_FILE_UPLOAD_SIZE} bytes "
+                    f"({DEFAULT_MAX_FILE_UPLOAD_SIZE // (1024 * 1024)} MB)"
+                ),
+                status_code=status.HTTP_400_BAD_REQUEST,
+            )
+            raise HTTPException(**response.model_dump())
+
         filename = file.filename or "uploaded_file"
 
         # Add .txt extension if no extension present
@@ -463,10 +512,12 @@ async def create_file(
             len(content),
         )
 
-        # Convert to BytesIO for Llama Stack client
-        # The client expects bytes, io.IOBase, PathLike, or a tuple
+        # BytesIO wraps the bytes object for library client compatibility
+        # Note: BytesIO doesn't copy the data, it creates a file-like view
+        from io import BytesIO
+
         file_bytes = BytesIO(content)
-        file_bytes.name = filename  # Set the filename attribute
+        file_bytes.name = filename
 
         file_obj = await client.files.create(
             file=file_bytes,
@@ -487,7 +538,13 @@ async def create_file(
         raise HTTPException(**response.model_dump()) from e
     except BadRequestError as e:
         logger.error("Bad request for file upload: %s", e)
-        response = ServiceUnavailableResponse(backend_name="Llama Stack", cause=str(e))
+        # BadRequestError from Llama Stack indicates client error (e.g., file too large)
+        # Map to 400 Bad Request, not 503 Service Unavailable
+        response = AbstractErrorResponse(
+            response="Invalid file upload",
+            cause=f"File upload rejected by Llama Stack: {str(e)}",
+            status_code=status.HTTP_400_BAD_REQUEST,
+        )
         raise HTTPException(**response.model_dump()) from e
     except Exception as e:
         full_trace = traceback.format_exc()
@@ -572,9 +629,12 @@ async def add_file_to_vector_store(
                 else:
                     raise  # Re-raise if not a lock error or max retries reached
         if not vs_file:
-            raise HTTPException(
-                status_code=500, detail="Failed to create vector store file"
+            # Use standard error response model for consistency
+            response = InternalServerErrorResponse(
+                response="Failed to create vector store file",
+                cause="All retry attempts failed to create the vector store file",
             )
+            raise HTTPException(**response.model_dump())
         logger.info(
             "Vector store file created - ID: %s, status: %s, last_error: %s",
             vs_file.id,

src/constants.py

@@ -128,6 +128,10 @@
 DEFAULT_AUTHENTICATION_MODULE = AUTH_MOD_NOOP
 # Maximum allowed size for base64-encoded x-rh-identity header (bytes)
 DEFAULT_RH_IDENTITY_MAX_HEADER_SIZE = 8192
+
+# Maximum allowed file upload size (bytes) - 100MB default
+# Protects against DoS attacks via large file uploads
+DEFAULT_MAX_FILE_UPLOAD_SIZE = 100 * 1024 * 1024  # 100 MB
 DEFAULT_JWT_UID_CLAIM = "user_id"
 DEFAULT_JWT_USER_NAME_CLAIM = "username"
 

tests/unit/app/endpoints/test_vector_stores.py

@@ -402,6 +402,7 @@ async def test_create_file_success(mocker: MockerFixture) -> None:
     # Mock UploadFile
     mock_file = mocker.AsyncMock()
     mock_file.filename = "test.txt"
+    mock_file.size = 12  # Size of "test content"
     mock_file.read.return_value = b"test content"
 
     response = await create_file(request=request, auth=auth, file=mock_file)
@@ -821,6 +822,7 @@ async def test_create_file_connection_error(mocker: MockerFixture) -> None:
 
     mock_file = mocker.AsyncMock()
     mock_file.filename = "test.txt"
+    mock_file.size = 12  # Size of "test content"
     mock_file.read.return_value = b"test content"
 
     with pytest.raises(HTTPException) as e:
@@ -854,11 +856,75 @@ async def test_create_file_bad_request(mocker: MockerFixture) -> None:
 
     mock_file = mocker.AsyncMock()
     mock_file.filename = "test.txt"
+    mock_file.size = 12  # Size of "test content"
     mock_file.read.return_value = b"test content"
 
     with pytest.raises(HTTPException) as e:
         await create_file(request=request, auth=auth, file=mock_file)
-    assert e.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE
+
+    assert e.value.status_code == status.HTTP_400_BAD_REQUEST
+
+
+@pytest.mark.asyncio
+async def test_create_file_too_large(mocker: MockerFixture) -> None:
+    """Test create file with file size exceeding limit."""
+    mock_authorization_resolvers(mocker)
+
+    config_dict = get_test_config()
+    cfg = AppConfig()
+    cfg.init_from_dict(config_dict)
+
+    mocker.patch("app.endpoints.vector_stores.configuration", cfg)
+
+    request = get_test_request()
+    auth = get_test_auth()
+
+    # Create a mock file that exceeds the size limit
+    mock_file = mocker.AsyncMock()
+    mock_file.filename = "large_file.pdf"
+    mock_file.size = 200 * 1024 * 1024  # 200 MB (exceeds 100 MB limit)
+    mock_file.read.return_value = b"x" * (200 * 1024 * 1024)
+
+    with pytest.raises(HTTPException) as e:
+        await create_file(request=request, auth=auth, file=mock_file)
+
+    assert e.value.status_code == status.HTTP_400_BAD_REQUEST
+    assert "too large" in str(e.value.detail).lower()
+
+
+@pytest.mark.asyncio
+async def test_create_file_content_length_too_large(mocker: MockerFixture) -> None:
+    """Test create file with Content-Length header exceeding limit."""
+    mock_authorization_resolvers(mocker)
+
+    config_dict = get_test_config()
+    cfg = AppConfig()
+    cfg.init_from_dict(config_dict)
+
+    mocker.patch("app.endpoints.vector_stores.configuration", cfg)
+
+    # Create request with large Content-Length header
+    request = Request(
+        scope={
+            "type": "http",
+            "headers": [
+                (b"authorization", b"Bearer test-token"),
+                (b"content-length", b"209715200"),  # 200 MB
+            ],
+        }
+    )
+    auth = get_test_auth()
+
+    # Create a mock file
+    mock_file = mocker.AsyncMock()
+    mock_file.filename = "large_file.pdf"
+    mock_file.size = None  # No size attribute
+
+    with pytest.raises(HTTPException) as e:
+        await create_file(request=request, auth=auth, file=mock_file)
+
+    assert e.value.status_code == status.HTTP_400_BAD_REQUEST
+    assert "too large" in str(e.value.detail).lower()
 
 
 @pytest.mark.asyncio