CVE-2024-2912: Initialization of a Resource with an Insecure Default (CWE-1188)
value: An insecure deserialization vulnerability exists in the BentoML framework, allowing remote code execution (RCE) by sending a specially crafted POST request. By exploiting this vulnerability, attackers can execute arbitrary commands on the server hosting the BentoML application. The vulnerability is triggered when a serialized object, crafted to execute OS commands upon deserialization, is sent to any valid BentoML endpoint. This issue poses a significant security risk, enabling attackers to compromise the server and potentially gain unauthorized access or control.
CWE description: The product initializes or sets a resource with a default that is intended to be changed by the administrator, but the default is not secure.
Published: 2024-04-16T00:15Z
Repository: https://github.com/bentoml/bentoml
Language: Python
Files changed: src/_bentoml_impl/client/__init__.py, src/_bentoml_impl/client/http.py, src/_bentoml_impl/loader.py, src/_bentoml_impl/server/app.py, src/_bentoml_impl/server/serving.py, src/_bentoml_impl/worker/service.py
Commit date: 2024-02-27 00:41:13 +0800
Commit message: feat(client): disallow media type other than json for HTTP clients (#4520)

* feat(client): disallow media type other than json for HTTP clients

Signed-off-by: Frost Ming <me@frostming.com>

* fix: check media type on server side

Signed-off-by: Frost Ming <me@frostming.com>

---------

Signed-off-by: Frost Ming <me@frostming.com>
Diff: diff --git a/src/_bentoml_impl/client/__init__.py b/src/_bentoml_impl/client/__init__.py
index daea131d..37ad07ca 100644
--- a/src/_bentoml_impl/client/__init__.py
+++ b/src/_bentoml_impl/client/__init__.py
@@ -1,13 +1,66 @@
 from .base import AbstractClient
-from .http import AsyncHTTPClient
+from .http import AsyncHTTPClient as _AsyncHTTPClient
 from .http import HTTPClient
-from .http import SyncHTTPClient
+from .http import SyncHTTPClient as _SyncHTTPClient
 from .proxy import RemoteProxy
 
 __all__ = [
     "AsyncHTTPClient",
     "SyncHTTPClient",
     "HTTPClient",
     "AbstractClient",
     "RemoteProxy",
 ]
+
+
+class SyncHTTPClient(_SyncHTTPClient):
+    """A synchronous client for BentoML service.
+
+    Args:
+        url (str): URL of the BentoML service.
+        token (str, optional): Authentication token. Defaults to None.
+        timeout (float, optional): Timeout for the client. Defaults to 30.
+
+    Example::
+
+        with SyncHTTPClient("http://localhost:3000") as client:
+            resp = client.call("classify", input_series=[[1,2,3,4]])
+            assert resp == [0]
+            # Or using named method directly
+            resp = client.classify(input_series=[[1,2,3,4]])
+            assert resp == [0]
+    """
+
+    def __init__(
+        self, url: str, *, token: str | None = None, timeout: float = 30
+    ) -> None:
+        super().__init__(url, token=token, timeout=timeout)
+
+
+class AsyncHTTPClient(_AsyncHTTPClient):
+    """An asynchronous client for BentoML service.
+
+    Args:
+        url (str): URL of the BentoML service.
+        token (str, optional): Authentication token. Defaults to None.
+        timeout (float, optional): Timeout for the client. Defaults to 30.
+
+    Example::
+
+        async with AsyncHTTPClient("http://localhost:3000") as client:
+            resp = await client.call("classify", input_series=[[1,2,3,4]])
+            assert resp == [0]
+            # Or using named method directly
+            resp = await client.classify(input_series=[[1,2,3,4]])
+            assert resp == [0]
+
+            # Streaming
+            resp = client.stream(prompt="hello")
+            async for data in resp:
+                print(data)
+    """
+
+    def __init__(
+        self, url: str, *, token: str | None = None, timeout: float = 30
+    ) -> None:
+        super().__init__(url, token=token, timeout=timeout)
diff --git a/src/_bentoml_impl/client/http.py b/src/_bentoml_impl/client/http.py
index 37e4b3f6..729a0f75 100644
--- a/src/_bentoml_impl/client/http.py
+++ b/src/_bentoml_impl/client/http.py
@@ -338,18 +338,11 @@ class HTTPClient(AbstractClient, t.Generic[C]):
 
 
 class SyncHTTPClient(HTTPClient[httpx.Client]):
     """A synchronous client for BentoML service.
 
-    Example:
-
-        with SyncHTTPClient("http://localhost:3000") as client:
-            resp = client.call("classify", input_series=[[1,2,3,4]])
-            assert resp == [0]
-            # Or using named method directly
-            resp = client.classify(input_series=[[1,2,3,4]])
-            assert resp == [0]
+    .. note:: Inner usage ONLY
     """
 
     client_cls = httpx.Client
 
     def __enter__(self: T) -> T:
@@ -447,28 +440,11 @@ class SyncHTTPClient(HTTPClient[httpx.Client]):
 
 
 class AsyncHTTPClient(HTTPClient[httpx.AsyncClient]):
     """An asynchronous client for BentoML service.
 
-    Example:
-
-        async with AsyncHTTPClient("http://localhost:3000") as client:
-            resp = await client.call("classify", input_series=[[1,2,3,4]])
-            assert resp == [0]
-            # Or using named method directly
-            resp = await client.classify(input_series=[[1,2,3,4]])
-            assert resp == [0]
-
-    .. note::
-
-        If the endpoint returns an async generator, it should be awaited before iterating.
-
-        Example:
-
-            resp = await client.stream(prompt="hello")
-            async for data in resp:
-                print(data)
+    .. note:: Inner usage ONLY
     """
 
     client_cls = httpx.AsyncClient
 
     async def is_ready(self, timeout: int | None = None) -> bool:
diff --git a/src/_bentoml_impl/loader.py b/src/_bentoml_impl/loader.py
index ae6f01ec..12d88575 100644
--- a/src/_bentoml_impl/loader.py
+++ b/src/_bentoml_impl/loader.py
@@ -104,16 +104,17 @@ def import_service(
     `normalize_identifier` function.
     """
     from _bentoml_sdk import Service
 
     if bento_path is None:
-        bento_path = pathlib.Path(".").absolute()
+        bento_path = pathlib.Path(".")
+    bento_path = bento_path.absolute()
 
     # patch python path if needed
-    if bento_path != pathlib.Path("."):
+    if bento_path != pathlib.Path(".").absolute():
         # a project
-        extra_python_path = str(bento_path.absolute())
+        extra_python_path = str(bento_path)
         sys.path.insert(0, extra_python_path)
     else:
         # a project under current directory
         extra_python_path = None
 
diff --git a/src/_bentoml_impl/server/app.py b/src/_bentoml_impl/server/app.py
index 1120aa45..883c70a9 100644
--- a/src/_bentoml_impl/server/app.py
+++ b/src/_bentoml_impl/server/app.py
@@ -3,10 +3,11 @@ from __future__ import annotations
 import asyncio
 import functools
 import inspect
 import sys
 import typing as t
+from http import HTTPStatus
 from pathlib import Path
 
 import anyio
 import anyio.to_thread
 import pydantic
@@ -59,10 +60,11 @@ class ContextMiddleware:
 class ServiceAppFactory(BaseAppFactory):
     @inject
     def __init__(
         self,
         service: Service[t.Any],
+        is_main: bool = False,
         enable_metrics: bool = Provide[
             BentoMLContainer.api_server_config.metrics.enabled
         ],
         traffic: dict[str, t.Any] = Provide[BentoMLContainer.api_server_config.traffic],
         enable_access_control: bool = Provide[BentoMLContainer.http.cors.enabled],
@@ -72,10 +74,11 @@ class ServiceAppFactory(BaseAppFactory):
     ) -> None:
         from bentoml._internal.runner.container import AutoContainer
 
         self.service = service
         self.enable_metrics = enable_metrics
+        self.is_main = is_main
         timeout = traffic.get("timeout")
         max_concurrency = traffic.get("max_concurrency")
         self.enable_access_control = enable_access_control
         self.access_control_options = access_control_options
         super().__init__(timeout=timeout, max_concurrency=max_concurrency)
@@ -173,20 +176,20 @@ class ServiceAppFactory(BaseAppFactory):
                 status_code=status,
             )
         else:
             return JSONResponse("", status_code=status)
 
-    def __call__(self, is_main: bool = False) -> Starlette:
+    def __call__(self) -> Starlette:
         app = super().__call__()
 
         app.add_exception_handler(
             pydantic.ValidationError, self.handle_validation_error
         )
         app.add_exception_handler(BentoMLException, self.handle_bentoml_exception)
         app.add_exception_handler(Exception, self.handle_uncaught_exception)
         app.add_route("/schema.json", self.schema_view, name="schema")
-        if is_main:
+        if self.is_main:
             if BentoMLContainer.new_index:
                 assets = Path(__file__).parent / "assets"
                 app.mount("/assets", StaticFiles(directory=assets), name="assets")
             else:
                 from bentoml._internal import server
@@ -405,10 +408,16 @@ class ServiceAppFactory(BaseAppFactory):
 
         from ..serde import ALL_SERDE
 
         media_type = request.headers.get("Content-Type", "application/json")
         media_type = media_type.split(";")[0].strip()
+        if self.is_main and media_type == "application/vnd.bentoml+pickle":
+            # Disallow pickle media type for main service for security reasons
+            raise BentoMLException(
+                "Pickle media type is not allowed for main service",
+                error_code=HTTPStatus.UNSUPPORTED_MEDIA_TYPE,
+            )
 
         method = self.service.apis[name]
         func = getattr(self._service_instance, name)
         ctx = self.service.context
         serde = ALL_SERDE[media_type]()
diff --git a/src/_bentoml_impl/server/serving.py b/src/_bentoml_impl/server/serving.py
index 46850504..d504afc3 100644
--- a/src/_bentoml_impl/server/serving.py
+++ b/src/_bentoml_impl/server/serving.py
@@ -258,11 +258,10 @@ def serve_http(
             str(backlog),
             "--worker-id",
             "$(CIRCUS.WID)",
             "--prometheus-dir",
             prometheus_dir,
-            "--main",
             *ssl_args,
             *timeout_args,
         ]
         if worker_envs:
             server_args.extend(["--worker-env", json.dumps(worker_envs)])
diff --git a/src/_bentoml_impl/worker/service.py b/src/_bentoml_impl/worker/service.py
index ee50effe..70966061 100644
--- a/src/_bentoml_impl/worker/service.py
+++ b/src/_bentoml_impl/worker/service.py
@@ -93,11 +93,10 @@ import click
 @click.option(
     "--timeout",
     type=click.INT,
     help="Specify the timeout for API server",
 )
-@click.option("--main", "is_main", type=click.BOOL, default=False, is_flag=True)
 def main(
     bento_identifier: str,
     service_name: str,
     fd: int,
     runner_map: str | None,
@@ -112,11 +111,10 @@ def main(
     ssl_cert_reqs: int | None,
     ssl_ca_certs: str | None,
     ssl_ciphers: str | None,
     development_mode: bool,
     timeout: int,
-    is_main: bool = False,
 ):
     """
     Start a HTTP server worker for given service.
     """
     import psutil
@@ -163,12 +161,14 @@ def main(
 
     if prometheus_dir is not None:
         BentoMLContainer.prometheus_multiproc_dir.set(prometheus_dir)
     component_context.component_name = service.name
 
-    app_factory = ServiceAppFactory(service)
-    asgi_app = app_factory(is_main=is_main)
+    asgi_app = ServiceAppFactory(
+        service, is_main=component_context.component_type == "entry_service"
+    )()
+
     uvicorn_extra_options: dict[str, t.Any] = {}
     if ssl_version is not None:
         uvicorn_extra_options["ssl_version"] = ssl_version
     if ssl_cert_reqs is not None:
         uvicorn_extra_options["ssl_cert_reqs"] = ssl_cert_reqs

Vulnerable code: from .http import AsyncHTTPClient
from .http import SyncHTTPClient
    Example:

        with SyncHTTPClient("http://localhost:3000") as client:
            resp = client.call("classify", input_series=[[1,2,3,4]])
            assert resp == [0]
            # Or using named method directly
            resp = client.classify(input_series=[[1,2,3,4]])
            assert resp == [0]
    Example:

        async with AsyncHTTPClient("http://localhost:3000") as client:
            resp = await client.call("classify", input_series=[[1,2,3,4]])
            assert resp == [0]
            # Or using named method directly
            resp = await client.classify(input_series=[[1,2,3,4]])
            assert resp == [0]

    .. note::

        If the endpoint returns an async generator, it should be awaited before iterating.

        Example:

            resp = await client.stream(prompt="hello")
            async for data in resp:
                print(data)
        bento_path = pathlib.Path(".").absolute()
    if bento_path != pathlib.Path("."):
        extra_python_path = str(bento_path.absolute())
    def __call__(self, is_main: bool = False) -> Starlette:
        if is_main:
            "--main",
@click.option("--main", "is_main", type=click.BOOL, default=False, is_flag=True)
    is_main: bool = False,
    app_factory = ServiceAppFactory(service)
    asgi_app = app_factory(is_main=is_main)
Fixed code: from .http import AsyncHTTPClient as _AsyncHTTPClient
from .http import SyncHTTPClient as _SyncHTTPClient


class SyncHTTPClient(_SyncHTTPClient):
    """A synchronous client for BentoML service.

    Args:
        url (str): URL of the BentoML service.
        token (str, optional): Authentication token. Defaults to None.
        timeout (float, optional): Timeout for the client. Defaults to 30.

    Example::

        with SyncHTTPClient("http://localhost:3000") as client:
            resp = client.call("classify", input_series=[[1,2,3,4]])
            assert resp == [0]
            # Or using named method directly
            resp = client.classify(input_series=[[1,2,3,4]])
            assert resp == [0]
    """

    def __init__(
        self, url: str, *, token: str | None = None, timeout: float = 30
    ) -> None:
        super().__init__(url, token=token, timeout=timeout)


class AsyncHTTPClient(_AsyncHTTPClient):
    """An asynchronous client for BentoML service.

    Args:
        url (str): URL of the BentoML service.
        token (str, optional): Authentication token. Defaults to None.
        timeout (float, optional): Timeout for the client. Defaults to 30.

    Example::

        async with AsyncHTTPClient("http://localhost:3000") as client:
            resp = await client.call("classify", input_series=[[1,2,3,4]])
            assert resp == [0]
            # Or using named method directly
            resp = await client.classify(input_series=[[1,2,3,4]])
            assert resp == [0]

            # Streaming
            resp = client.stream(prompt="hello")
            async for data in resp:
                print(data)
    """

    def __init__(
        self, url: str, *, token: str | None = None, timeout: float = 30
    ) -> None:
        super().__init__(url, token=token, timeout=timeout)
    .. note:: Inner usage ONLY
    .. note:: Inner usage ONLY
        bento_path = pathlib.Path(".")
    bento_path = bento_path.absolute()
    if bento_path != pathlib.Path(".").absolute():
        extra_python_path = str(bento_path)
from http import HTTPStatus
        is_main: bool = False,
        self.is_main = is_main
    def __call__(self) -> Starlette:
        if self.is_main:
        if self.is_main and media_type == "application/vnd.bentoml+pickle":
            # Disallow pickle media type for main service for security reasons
            raise BentoMLException(
                "Pickle media type is not allowed for main service",
                error_code=HTTPStatus.UNSUPPORTED_MEDIA_TYPE,
            )
    asgi_app = ServiceAppFactory(
        service, is_main=component_context.component_type == "entry_service"
    )()