Vulnerable code sample

import asyncio
from typing import ClassVar, Dict, List, Optional, Set, Type

import httpx
from PIL import Image
from pydantic import AnyHttpUrl, BaseModel, Field, PrivateAttr
from aiohttp import ClientSession
from urllib.parse import urlparse

from vllm.logger import init_logger
from vllm.multimodal.image import ImagePixelData

logger = init_logger(__name__)


class MultiModalData(BaseModel):
    """
    Data class for holding multimodal data.
    """
    # Each field should be a list of data of a specific modality
    image: Optional[List[ImagePixelData]] = None


class MediaConnector(BaseModel):
    """
    A class that loads data from a URL and converts it to a specific
    data model.
    """
    _SUPPORTED_SCHEMES: ClassVar[Set[str]] = {"http", "https"}

    # The data model to convert the loaded data to.
    _data_model: ClassVar[Type[MultiModalData]]

    # A client session to be shared by all instances of the class.
    _client_session: ClassVar[Optional[ClientSession]] = PrivateAttr(None)

    # A httpx session to be shared by all instances of the class.
    _httpx_client: ClassVar[Optional[httpx.Client]] = PrivateAttr(None)
    _httpx_async_client: ClassVar[Optional[httpx.AsyncClient]] = PrivateAttr(
        None)

    # The URL to load the data from.
    url: AnyHttpUrl

    # A set of allowed hosts for the URL. If not specified, all hosts are
    # allowed.
    allowed_hosts: Optional[Set[str]] = Field(default=None, repr=False)

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._is_url_allowed()

    def _is_url_allowed(self) -> bool:
        """
        Check if the URL is allowed by the allowed_hosts set.
        """
        if self.allowed_hosts is None:
            return True
        return urlparse(str(self.url)).hostname in self.allowed_hosts

    async def load_from_url_async(self) -> MultiModalData:
        """
        Load data from the URL and convert it to the data model.
        This method is asynchronous.
        """
        if self._httpx_async_client is None:
            self._httpx_async_client = httpx.AsyncClient()
        response = await self._httpx_async_client.get(str(self.url))
        response.raise_for_status()

        # Get the content type of the response
        content_type = response.headers.get("content-type", "")
        if "image" in content_type:
            image_data = Image.open(response.content).convert("RGB")
            return self._data_model(image=[image_data])
        else:
            raise NotImplementedError(
                f"Content type {content_type} is not supported")

    def load_from_url(self) -> MultiModalData:
        """
        Load data from the URL and convert it to the data model.
        """
        if self._httpx_client is None:
            self._httpx_client = httpx.Client()
        response = self._httpx_client.get(str(self.url))
        response.raise_for_status()

        # Get the content type of the response
        content_type = response.headers.get("content-type", "")
        if "image" in content_type:
            image_data = Image.open(response.content).convert("RGB")
            return self._data_model(image=[image_data])
        else:
            raise NotImplementedError(
                f"Content type {content_type} is not supported")


class ImageConnector(MediaConnector):
    """
    A class that loads an image from a URL and converts it to an ImagePixelData
    object.
    """
    _data_model: ClassVar[Type[MultiModalData]] = MultiModalData


# A registry of media connectors.
MEDIA_CONNECTORS: Dict[str, Type[MediaConnector]] = {
    "image": ImageConnector
}


def process_multimodal_inputs(http_body: Dict) -> Dict:
    """
    Process the multimodal inputs from the HTTP request body.
    It downloads the media from the URLs and replaces the URLs
    with the downloaded media.
    """
    multi_modal_data = http_body.pop("multi_modal_data", None)
    if multi_modal_data:
        # Create a new dictionary to avoid modifying the original
        processed_multi_modal_data: Dict[str, list] = {}
        for media_type, media_sources in multi_modal_data.items():
            if media_type not in MEDIA_CONNECTORS:
                continue

            connector_cls = MEDIA_CONNECTORS[media_type]

            # Gather all media downloads
            tasks = [
                connector_cls(url=media_source).load_from_url_async()
                for media_source in media_sources
            ]

            # Run all downloads concurrently
            results = asyncio.run(asyncio.gather(*tasks))

            # Merge the results
            for result in results:
                result_dict = result.dict()
                for key, value in result_dict.items():
                    if value is None:
                        continue
                    if key not in processed_multi_modal_data:
                        processed_multi_modal_data[key] = []
                    processed_multi_modal_data[key].extend(value)

        # Update the http_body with the processed multimodal data
        http_body["multi_modal_data"] = processed_multi_modal_data

    return http_body