Repository: meta-llama/llama-stack
File: src/llama_stack_api/common/upload_safety.py (Python)

Code before:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

"""Upload safety utilities for enforcing file size limits.

Provides a bounded-read helper that reads UploadFile objects in chunks,
aborting early when the configured maximum size is exceeded. This avoids
buffering arbitrarily large uploads into memory before rejecting them.
"""

import io

from fastapi import UploadFile
from starlette.datastructures import Headers

from llama_stack_api.common.errors import FileTooLargeError

# Default maximum upload size: 100 MB
DEFAULT_MAX_UPLOAD_SIZE_BYTES: int = 100 * 1024 * 1024

# Chunk size for incremental reads (1 MB)
_READ_CHUNK_SIZE: int = 1 * 1024 * 1024


async def read_upload_with_size_limit(
    file: UploadFile,
    max_upload_bytes: int = DEFAULT_MAX_UPLOAD_SIZE_BYTES,
) -> bytes:
    """Read an uploaded file with a size limit, aborting early if exceeded.

    Performs two levels of checking:
    1. Pre-check: if ``file.size`` is populated (from Content-Length), reject
       immediately without reading any bytes.
    2. Chunked read: reads the file in fixed-size chunks, tracking total bytes.
       Stops as soon as the cumulative size exceeds *max_upload_bytes*.

    Args:
        file: The FastAPI UploadFile to read.
        max_upload_bytes: Maximum allowed file size in bytes.

    Returns:
        The full file content as bytes.

    Raises:
        FileTooLargeError: If the file exceeds *max_upload_bytes*.
    """

Review comment (suggestion):
```suggestion
# Chunk size for incremental reads (1 MiB)
```
small nit this is technically a mebibyte (MiB) since this is powers of two (1024=2^10).


Code after:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

"""Upload safety utilities for enforcing file size limits.

Provides a bounded-read helper that reads UploadFile objects in chunks,
aborting early when the configured maximum size is exceeded. This avoids
buffering arbitrarily large uploads into memory before rejecting them.
"""

import io

from fastapi import UploadFile
from starlette.datastructures import Headers

from llama_stack_api.common.errors import FileTooLargeError

# Default maximum upload size: 100 MiB
DEFAULT_MAX_UPLOAD_SIZE_BYTES: int = 100 * 1024 * 1024

# Chunk size for incremental reads (1 MiB)
_READ_CHUNK_SIZE: int = 1 * 1024 * 1024


async def read_upload_with_size_limit(
    file: UploadFile,
    max_upload_bytes: int = DEFAULT_MAX_UPLOAD_SIZE_BYTES,
) -> bytes:
    """Read an uploaded file with a size limit, aborting early if exceeded.

    Performs two levels of checking:
    1. Pre-check: if ``file.size`` is populated (from Content-Length), reject
       immediately without reading any bytes.
    2. Chunked read: reads the file in fixed-size chunks, tracking total bytes.
       Stops as soon as the cumulative size exceeds *max_upload_bytes*.

    Args:
        file: The FastAPI UploadFile to read.
        max_upload_bytes: Maximum allowed file size in bytes.

    Returns:
        The full file content as bytes.

    Raises:
        FileTooLargeError: If the file exceeds *max_upload_bytes*.
    """