Security hotfix: remediate authentication, code execution, and SSRF findings

3rdparty/FastChat/fastchat/serve/controller.py

@@ -14,14 +14,12 @@
 import threading
 
 from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse
+from fastapi.responses import JSONResponse
 import numpy as np
-import requests
 import uvicorn
 
 from fastchat.constants import (
     CONTROLLER_HEART_BEAT_EXPIRATION,
-    WORKER_API_TIMEOUT,
     ErrorCode,
     SERVER_ERROR_MSG,
 )
@@ -74,39 +72,8 @@ def __init__(self, dispatch_method: str):
     def register_worker(
         self, worker_name: str, check_heart_beat: bool, worker_status: dict
     ):
-        if worker_name not in self.worker_info:
-            logger.info(f"Register a new worker: {worker_name}")
-        else:
-            logger.info(f"Register an existing worker: {worker_name}")
-
-        if not worker_status:
-            worker_status = self.get_worker_status(worker_name)
-        if not worker_status:
-            return False
-
-        self.worker_info[worker_name] = WorkerInfo(
-            worker_status["model_names"],
-            worker_status["speed"],
-            worker_status["queue_length"],
-            check_heart_beat,
-            time.time(),
-        )
-
-        logger.info(f"Register done: {worker_name}, {worker_status}")
-        return True
-
-    def get_worker_status(self, worker_name: str):
-        try:
-            r = requests.post(worker_name + "/worker_get_status", timeout=5)
-        except requests.exceptions.RequestException as e:
-            logger.error(f"Get status fails: {worker_name}, {e}")
-            return None
-
-        if r.status_code != 200:
-            logger.error(f"Get status fails: {worker_name}, {r}")
-            return None
-
-        return r.json()
+        logger.warning("Dynamic worker registration is disabled.")
+        return False
 
     def remove_worker(self, worker_name: str):
         del self.worker_info[worker_name]
@@ -145,22 +112,6 @@ def get_worker_address(self, model_name: str):
                 worker_name = worker_names[pt]
                 return worker_name
 
-            # Check status before returning
-            while True:
-                pt = np.random.choice(np.arange(len(worker_names)), p=worker_speeds)
-                worker_name = worker_names[pt]
-
-                if self.get_worker_status(worker_name):
-                    break
-                else:
-                    self.remove_worker(worker_name)
-                    worker_speeds[pt] = 0
-                    norm = np.sum(worker_speeds)
-                    if norm < 1e-4:
-                        return ""
-                    worker_speeds = worker_speeds / norm
-                    continue
-            return worker_name
         elif self.dispatch_method == DispatchMethod.SHORTEST_QUEUE:
             worker_names = []
             worker_qlen = []
@@ -223,12 +174,10 @@ def worker_api_get_status(self):
         speed = 0
         queue_length = 0
 
-        for w_name in self.worker_info:
-            worker_status = self.get_worker_status(w_name)
-            if worker_status is not None:
-                model_names.update(worker_status["model_names"])
-                speed += worker_status["speed"]
-                queue_length += worker_status["queue_length"]
+        for worker_status in self.worker_info.values():
+            model_names.update(worker_status.model_names)
+            speed += worker_status.speed
+            queue_length += worker_status.queue_length
 
         model_names = sorted(list(model_names))
         return {
@@ -237,33 +186,16 @@ def worker_api_get_status(self):
             "queue_length": queue_length,
         }
 
-    def worker_api_generate_stream(self, params):
-        worker_addr = self.get_worker_address(params["model"])
-        if not worker_addr:
-            yield self.handle_no_worker(params)
-
-        try:
-            response = requests.post(
-                worker_addr + "/worker_generate_stream",
-                json=params,
-                stream=True,
-                timeout=WORKER_API_TIMEOUT,
-            )
-            for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
-                if chunk:
-                    yield chunk + b"\0"
-        except requests.exceptions.RequestException as e:
-            yield self.handle_worker_timeout(worker_addr)
-
-
 app = FastAPI()
 
 
 @app.post("/register_worker")
 async def register_worker(request: Request):
-    data = await request.json()
-    controller.register_worker(
-        data["worker_name"], data["check_heart_beat"], data.get("worker_status", None)
+    return JSONResponse(
+        status_code=403,
+        content={
+            "error": "Dynamic worker registration is disabled by the Parrot security hotfix."
+        },
     )
 
 
@@ -294,9 +226,12 @@ async def receive_heart_beat(request: Request):
 
 @app.post("/worker_generate_stream")
 async def worker_api_generate_stream(request: Request):
-    params = await request.json()
-    generator = controller.worker_api_generate_stream(params)
-    return StreamingResponse(generator)
+    return JSONResponse(
+        status_code=403,
+        content={
+            "error": "Worker proxying is disabled by the Parrot security hotfix."
+        },
+    )
 
 
 @app.post("/worker_get_status")

README.md

@@ -4,6 +4,38 @@
 
 > This repo is currently a research prototype and is not actively maintained. Please open issue or contact the authors when you need help.
 
+## Security hotfix notice
+
+The current branch disables several prototype features that were unsafe when
+exposed to untrusted clients:
+
+- Python native calls no longer accept or execute client-supplied code.
+- Dynamic Parrot engine registration and heartbeat endpoints are disabled.
+- Dynamic registration in the vendored FastChat controller is disabled.
+- ServeCore requires a bearer token for every HTTP endpoint. Set
+  `PARROT_API_KEY` to a strong random value in both ServeCore and trusted client
+  processes before starting the service.
+
+These changes intentionally break the affected APIs and have received only
+targeted security regression testing. Each session now receives a
+cryptographically random credential that is checked on every session, call, and
+semantic-variable operation. This research prototype should still not be
+exposed directly to an untrusted network.
+
+The pre-hotfix source remains available at commit
+[`2e1825e`](https://github.com/microsoft/ParrotServe/commit/2e1825ee2bc38cb783bab9d8ec3e5ae99a93ba46)
+for isolated historical reproduction only. That revision contains known
+arbitrary-code-execution, unauthenticated-access, and
+server-side-request-forgery vulnerabilities. It is unsupported and must not be
+deployed on an untrusted network.
+
+To clone the preserved legacy tag into an isolated security environment:
+
+```bash
+git clone --branch legacy-unsafe-pre-security-hotfix --single-branch \
+  https://github.com/microsoft/ParrotServe.git ParrotServe-legacy-unsafe
+```
+
 Parrot is a distributed, multi-tenant serving system for **LLM-based Applications**. With the Semantic Variable abstraction, Parrot can easily grasp the **app-level information** like LLM computation graph (DAG) or the prompt structure, which enables many interesting features like:
 - Automatically parallelize and batch LLM requests in complex LLM applications. Asynchronous communication between dependent requests.
 - Performance objective deduction and DAG-aware scheduling.

docs/user_docs/parrot_apis.md

@@ -99,54 +99,14 @@ Response body:
 
 ## Submit Native Function Call
 
-> NOTE: This API is expiermental
-
-We have some built-in native functions. We don’t allow user to submit their customized code because it may introduce safety problems.
+> **Disabled by the security hotfix.**
 
 Endpoint: `/{api_version}/py_native_call`
 
-- Submit a python native function call [POST].
-
-PS: The `"func_code"` must be a string dumped from a Python binary code, encoded by `base64`.
-- We recommend using `marshal` to dump a Python code (`func.__code__`) to bytes. See `parrot/utils/serialize_utils.py`, `serialize_func_code` function.
-- For encoding a bytes using `base64` (For safe transport via HTTP), see `parrot/utils/serialize_utils.py`, `bytes_to_encoded_b64str` function.
-
-Request body:
-
-```json
-{
-    "session_id": "xxx",
-    "session_auth": "yyy",
-    "func_name": "xxx", // Function name.
-    "func_code": "some code bytes", // Bytecode of the function. If the function is cached, you can omit this field.
-    "parameters": [
-        {
-            "name": "a",
-            "is_output": false / true,
-            "var_id": "bbb", // Optional if it is output
-        },
-        ...
-    ],
-}
-```
-
-Response body:
-```json
-{
-	"request_id": "xxx",
-	"session_id": "yyy",
-	"param_info": [
-        {
-            "placeholder_name": "fff",
-            "is_output": true / false,
-            "var_name": "ddd",
-            "var_id": "ccc",
-            "var_desc": "The first output of request xxx",
-            "var_scope": "eeee",
-        }
-	]
-}
-```
+- The endpoint returns `410 Gone`.
+- Client-supplied Python bytecode is not accepted, deserialized, or executed.
+- The former `func_code` request format is available only in the known-vulnerable
+  legacy revision documented in the repository README.
 
 ## Semantic Variable
 

parrot/frontend/pfunc/function.py

@@ -2,18 +2,13 @@
 # Licensed under the MIT license.
 
 
-import marshal
 from abc import ABC
-import types
 from typing import Tuple, Callable, List, Dict, Type, Optional, Any, Set, Union
 import regex as re
 from dataclasses import dataclass, asdict
 
 from parrot.utils import (
     get_logger,
-    serialize_func_code,
-    deserialize_func_code,
-    bytes_to_encoded_b64str,
 )
 
 from parrot.serve.graph.call_request import (
@@ -474,7 +469,7 @@ def __init__(
     ):
         super().__init__(name, params)
 
-        self.pyfunc_code_dumped = serialize_func_code(pyfunc.__code__)
+        self.pyfunc = pyfunc
         metadata_dict = NativeFuncMetadata.get_default_dict()
         metadata_dict.update(**metadata_kwargs)
         self.metadata = NativeFuncMetadata(**metadata_dict)
@@ -579,11 +574,7 @@ def display_signature(self) -> str:
         return f"{self.name}({', '.join([f'{param.name}: {param.typ}' for param in self.params])})"
 
     def get_pyfunc(self) -> Callable:
-        code_deserialized = deserialize_func_code(self.pyfunc_code_dumped)
-
-        # Here for the scope Dict, we pass {} because we don't want to pollute the scope.
-        # Hence the pyfunc we get is just a temporary one.
-        return types.FunctionType(code_deserialized, {}, self.name)
+        return self.pyfunc
 
 
 class PyNativeCall(BasicCall):
@@ -626,7 +617,4 @@ def to_request_payload(self, with_code: bool = True) -> Dict:
         payload["parameters"] = parameters
         payload["func_name"] = self.func.name
 
-        if with_code:
-            payload["func_code"] = bytes_to_encoded_b64str(self.func.pyfunc_code_dumped)
-
         return payload

parrot/frontend/pfunc/vm.py

@@ -8,6 +8,7 @@
 import traceback
 import importlib
 import inspect
+import os
 from typing import Callable, Optional, Literal, Dict, List, Any, Generator
 
 from parrot.constants import NONE_SESSION_ID
@@ -281,7 +282,12 @@ def session_registered(self) -> bool:
     def register_session(self) -> None:
         """Register a session to the ServeCore."""
 
-        resp = register_session(http_addr=self.core_http_addr, api_key="1")
+        api_key = os.environ.get("PARROT_API_KEY")
+        if not api_key:
+            raise RuntimeError(
+                "PARROT_API_KEY must be set before registering a session."
+            )
+        resp = register_session(http_addr=self.core_http_addr, api_key=api_key)
         self.session_id = resp.session_id
         self._session_auth = resp.session_auth