diff --git a/3rdparty/FastChat/fastchat/serve/controller.py b/3rdparty/FastChat/fastchat/serve/controller.py index 8c35e41..e585afb 100644 --- a/3rdparty/FastChat/fastchat/serve/controller.py +++ b/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") diff --git a/README.md b/README.md index 7caa174..e3914e0 100644 --- a/README.md +++ b/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. diff --git a/docs/user_docs/parrot_apis.md b/docs/user_docs/parrot_apis.md index ffa1ab2..1b4d4db 100644 --- a/docs/user_docs/parrot_apis.md +++ b/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 diff --git a/parrot/frontend/pfunc/function.py b/parrot/frontend/pfunc/function.py index 567ba1c..44f8160 100644 --- a/parrot/frontend/pfunc/function.py +++ b/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 diff --git a/parrot/frontend/pfunc/vm.py b/parrot/frontend/pfunc/vm.py index 7076324..25206a2 100644 --- a/parrot/frontend/pfunc/vm.py +++ b/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 diff --git a/parrot/protocol/http_utils.py b/parrot/protocol/http_utils.py index 137f6fd..ad1b8b6 100644 --- a/parrot/protocol/http_utils.py +++ b/parrot/protocol/http_utils.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. +import os from typing import Type, Optional, Literal import requests import aiohttp @@ -13,6 +14,13 @@ logger = get_logger("API") +AUTH_TOKEN_ENV = "PARROT_API_KEY" + + +def _authorization_headers(): + token = os.environ.get(AUTH_TOKEN_ENV) + return {"Authorization": f"Bearer {token}"} if token else {} + def send_http_request( response_cls: Type[BaseResponse], @@ -29,11 +37,26 @@ def send_http_request( for _ in range(retry_times): try: if method == "GET": - resp = requests.get(url, json=kwargs, timeout=timeout) + resp = requests.get( + url, + json=kwargs, + headers=_authorization_headers(), + timeout=timeout, + ) elif method == "POST": - resp = requests.post(url, json=kwargs, timeout=timeout) + resp = requests.post( + url, + json=kwargs, + headers=_authorization_headers(), + timeout=timeout, + ) elif method == "DELETE": - resp = requests.delete(url, json=kwargs, timeout=timeout) + resp = requests.delete( + url, + json=kwargs, + headers=_authorization_headers(), + timeout=timeout, + ) else: raise ValueError(f"Invalid http method: {method}") @@ -66,15 +89,30 @@ async def async_send_http_request( ) -> BaseResponse: url = http_addr + api_url if method == "GET": - async with client_session.get(url, json=kwargs, timeout=timeout) as resp: + async with client_session.get( + url, + json=kwargs, + headers=_authorization_headers(), + timeout=timeout, + ) as resp: assert resp.ok, f"Send http request error: {resp.reason}" return await async_make_response(response_cls, resp) elif method == "POST": - async with client_session.post(url, json=kwargs, timeout=timeout) as resp: + async with client_session.post( + url, + json=kwargs, + headers=_authorization_headers(), + timeout=timeout, + ) as resp: assert resp.ok, f"Send http request error: {resp.reason}" return await async_make_response(response_cls, resp) elif method == "DELETE": - async with client_session.delete(url, json=kwargs, timeout=timeout) as resp: + async with client_session.delete( + url, + json=kwargs, + headers=_authorization_headers(), + timeout=timeout, + ) as resp: assert resp.ok, f"Send http request error: {resp.reason}" return await async_make_response(response_cls, resp) else: @@ -90,7 +128,9 @@ async def async_send_http_request_streaming( url = http_addr + api_url # NOTE(chaofan): Only POST now - async with client_session.post(url, json=kwargs) as reader: + async with client_session.post( + url, json=kwargs, headers=_authorization_headers() + ) as reader: # assert resp.ok, "Send http request error." async for chunk in reader.content.iter_chunked(4): yield int().from_bytes(chunk, "big") diff --git a/parrot/protocol/public/apis.py b/parrot/protocol/public/apis.py index a78d142..149ddc4 100644 --- a/parrot/protocol/public/apis.py +++ b/parrot/protocol/public/apis.py @@ -132,6 +132,7 @@ def submit_semantic_call( f"/{API_VERSION}/semantic_call", retry_times=1, session_id=session_id, + session_auth=session_auth, **payload, ) except BaseException as e: @@ -153,6 +154,7 @@ async def asubmit_semantic_call( f"/{API_VERSION}/semantic_call", retry_times=1, session_id=session_id, + session_auth=session_auth, **payload, ) except BaseException as e: @@ -172,6 +174,7 @@ def submit_py_native_call( f"/{API_VERSION}/py_native_call", retry_times=1, session_id=session_id, + session_auth=session_auth, **payload, ) except BaseException as e: @@ -193,6 +196,7 @@ async def asubmit_py_native_call( f"/{API_VERSION}/submit_py_native_call", retry_times=1, session_id=session_id, + session_auth=session_auth, **payload, ) except BaseException as e: diff --git a/parrot/serve/core.py b/parrot/serve/core.py index 556bd6a..30d2581 100644 --- a/parrot/serve/core.py +++ b/parrot/serve/core.py @@ -3,14 +3,14 @@ import json +import os +import secrets from typing import Dict import asyncio from parrot.utils import get_logger from parrot.constants import CORE_LOOP_INTERVAL -from parrot.protocol.internal.runtime_info import EngineRuntimeInfo -from parrot.engine.config import EngineConfig -from parrot.exceptions import ParrotCoreInternalError +from parrot.exceptions import ParrotCoreInternalError, ParrotCoreUserError from parrot.serve.graph import ( SVProducer, @@ -94,43 +94,6 @@ def __init__(self, config: Dict): ) ) - # ---------- APIs to Engine Layer ---------- - - def register_engine(self, payload: Dict) -> Dict: - """Register a new engine in the OS. - - Args: - config: EngineConfig. The engine config. - - Returns: - Dict. The response. - """ - - logger.debug(f"Register engine received.") - engine_config = EngineConfig(**payload["engine_config"]) - engine_id = self.engine_mgr.register_engine(engine_config) - return {"engine_id": engine_id} - - def engine_heartbeat(self, payload: Dict) -> Dict: - """Update the last seen time of an engine and other engine info. - - Args: - engine_id: int. The engine ID. - engine_runtime_info: EngineRuntimeInfo. The engine runtime info. - - Returns: - Dict. The response. - """ - - engine_id = payload["engine_id"] - engine_name = payload["engine_name"] - logger.debug(f"Engine {engine_name} (id={engine_id}) heartbeat received.") - engine_info = EngineRuntimeInfo(**payload["runtime_info"]) - - self.engine_mgr.engine_heartbeat(engine_id, engine_info) - - return {} - # ---------- Public Serving APIs ---------- # ---------- Session Management ---------- @@ -145,8 +108,22 @@ def register_session(self, payload: Dict) -> Dict: Dict. The response. """ - session_id = self.session_mgr.register_session() - return {"session_id": session_id, "session_auth": "1"} + expected_api_key = os.environ.get("PARROT_API_KEY") + supplied_api_key = payload.get("api_key") + if ( + not expected_api_key + or not isinstance(supplied_api_key, str) + or not secrets.compare_digest( + supplied_api_key.encode("utf-8"), + expected_api_key.encode("utf-8"), + ) + ): + raise ParrotCoreUserError( + RuntimeError("A valid API key is required to create a session.") + ) + + session_id, session_auth = self.session_mgr.register_session() + return {"session_id": session_id, "session_auth": session_auth} def remove_session(self, session_id: int, payload: Dict) -> Dict: """Remove a session in Serve Core. @@ -159,7 +136,7 @@ def remove_session(self, session_id: int, payload: Dict) -> Dict: Dict. The response. """ - self.session_mgr.check_session_status(session_id) + self.session_mgr.check_session_status(session_id, payload.get("session_auth")) self.session_mgr.remove_session(session_id) return {} @@ -175,7 +152,8 @@ def get_session_info(self, session_id: int, payload: Dict) -> Dict: Dict. The response. """ - return {} + self.session_mgr.check_session_status(session_id, payload.get("session_auth")) + return {"session_id": session_id} # ---------- Function Call ---------- @@ -212,7 +190,9 @@ def submit_semantic_call(self, payload: Dict) -> Dict: # This is for get the partial DAG and do optimized scheduling. # Update session last access time - self.session_mgr.check_session_status(session_id) + self.session_mgr.check_session_status( + session_id, payload.get("session_auth") + ) self.session_mgr.session_access_update(session_id) # Add the request to the session. @@ -224,31 +204,6 @@ def submit_semantic_call(self, payload: Dict) -> Dict: "param_info": param_info, } - def submit_py_native_call(self, payload: Dict) -> Dict: - """Submit a Python native call in a session to the ServeCore. - - Args: - payload: Dict. The request payload. - - Returns: - Dict. The response. - """ - - session_id = payload["session_id"] - - # Update session last access time - self.session_mgr.check_session_status(session_id) - self.session_mgr.session_access_update(session_id) - - # Add the request to the session. - session = self.session_mgr.get_session(session_id) - request_id, param_info = session.add_request(payload, is_native=True) - - return { - "request_id": request_id, - "param_info": param_info, - } - # ---------- Semantic Variable ---------- def register_semantic_variable(self, payload: Dict) -> Dict: @@ -264,7 +219,9 @@ def register_semantic_variable(self, payload: Dict) -> Dict: session_id = payload["session_id"] name = payload["var_name"] - self.session_mgr.check_session_status(session_id) + self.session_mgr.check_session_status( + session_id, payload.get("session_auth") + ) self.session_mgr.session_access_update(session_id) var = self.var_mgr.create_var(session_id, name) @@ -288,7 +245,9 @@ def set_semantic_variable(self, var_id: str, payload: Dict) -> Dict: session_id = payload["session_id"] content = payload["content"] - self.session_mgr.check_session_status(session_id) + self.session_mgr.check_session_status( + session_id, payload.get("session_auth") + ) self.session_mgr.session_access_update(session_id) var = self.var_mgr.get_var(session_id, var_id) @@ -315,7 +274,9 @@ async def get_semantic_variable(self, var_id: str, payload: Dict) -> Dict: session_id = payload["session_id"] criteria = payload["criteria"] - self.session_mgr.check_session_status(session_id) + self.session_mgr.check_session_status( + session_id, payload.get("session_auth") + ) self.session_mgr.session_access_update(session_id) var = self.var_mgr.get_var(session_id, var_id) diff --git a/parrot/serve/graph/call_request.py b/parrot/serve/graph/call_request.py index c110f5c..ebce0be 100644 --- a/parrot/serve/graph/call_request.py +++ b/parrot/serve/graph/call_request.py @@ -3,13 +3,11 @@ from dataclasses import dataclass from typing import Union, Dict, Optional, List, Type, Any -from types import CodeType, FunctionType import re from parrot.exceptions import parrot_assert, ParrotCoreUserError from parrot.sampling_config import SamplingConfig -from parrot.utils import deserialize_func_code, encoded_b64str_to_bytes from .semantic_variable import SemanticVariable from .perf_criteria import PerformanceCriteria @@ -387,19 +385,13 @@ def __init__( request_id: int, session_id: int, func_name: str, - func_code: Optional[CodeType], metadata: NativeCallMetadata = NativeCallMetadata.get_default(), ) -> None: self.request_id = request_id self.session_id = session_id self.func_name = func_name - # Construct the function. - # Here for the scope Dict, we pass {} because we don't want to pollute the scope. - # Hence the executable_func we get is just a temporary one. - self.executable_func: Optional[FunctionType] = None - if func_code is not None: - self.executable_func = FunctionType(func_code, {}, func_name) + self.executable_func = None # Metadata: additional information of the request. self.metadata = metadata @@ -414,6 +406,10 @@ def _preprocess(payload: Dict) -> Dict: # Check format. parrot_assert("func_name" in payload, "Missing field 'func_name' in request.") parrot_assert("parameters" in payload, "Missing field 'parameters' in request.") + parrot_assert( + "func_code" not in payload, + "Client-supplied Python code is not supported.", + ) processed_payload = payload.copy() @@ -446,18 +442,13 @@ def parse_from_payload( metadata_dict[key] = payload[key] metadata = NativeCallMetadata(**metadata_dict) - # Step 2. Extract the name and the code. + # Step 2. Extract the name. func_name = payload["func_name"] - func_code_serialized = payload.get("func_code", None) - if func_code_serialized is not None: - func_code_serialized = encoded_b64str_to_bytes(func_code_serialized) - func_code = deserialize_func_code(func_code_serialized) pynative_request = cls( request_id=request_id, session_id=session_id, func_name=func_name, - func_code=func_code, metadata=metadata, ) diff --git a/parrot/serve/graph/node.py b/parrot/serve/graph/node.py index 0a53a19..b6afbbd 100644 --- a/parrot/serve/graph/node.py +++ b/parrot/serve/graph/node.py @@ -355,7 +355,7 @@ def from_variables( """ pseudo_native_func = PyNativeCallRequest( - request_id=0, session_id=0, func_name=func_name, func_code=None + request_id=0, session_id=0, func_name=func_name ) node = cls(pseudo_native_func) node.input_vars = input_vars diff --git a/parrot/serve/http_server.py b/parrot/serve/http_server.py index 774cae9..75a56fa 100644 --- a/parrot/serve/http_server.py +++ b/parrot/serve/http_server.py @@ -4,38 +4,63 @@ import argparse import asyncio +import logging +import secrets import traceback -from typing import Optional +from typing import Optional, TYPE_CHECKING from fastapi import FastAPI, Request from fastapi.responses import JSONResponse from uvicorn import Config, Server import os -from parrot.serve.core import ParrotServeCore, create_serve_core -from parrot.protocol.internal.runtime_info import EngineRuntimeInfo from parrot.protocol.public.api_version import API_VERSION -from parrot.engine.config import EngineConfig -from parrot.utils import ( - get_logger, - create_task_in_loop, - set_log_output_file, - redirect_stdout_stderr_to_file, -) from parrot.exceptions import ParrotCoreUserError, ParrotCoreInternalError -from parrot.testing.latency_simulator import get_latency +if TYPE_CHECKING: + from parrot.serve.core import ParrotServeCore -logger = get_logger("Parrot ServeCore Server") + +logger = logging.getLogger("Parrot ServeCore Server") # FastAPI app app = FastAPI() # Core -pcore: Optional[ParrotServeCore] = None +pcore: Optional["ParrotServeCore"] = None # Mode release_mode = False +AUTH_TOKEN_ENV = "PARROT_API_KEY" + + +def _error_response(status_code: int, message: str) -> JSONResponse: + return JSONResponse( + status_code=status_code, + content={"error": message, "traceback": ""}, + ) + + +@app.middleware("http") +async def require_api_key(request: Request, call_next): + expected_token = os.environ.get(AUTH_TOKEN_ENV) + if not expected_token: + return _error_response( + 503, + f"ServeCore is locked because {AUTH_TOKEN_ENV} is not configured.", + ) + + authorization = request.headers.get("Authorization", "") + scheme, _, supplied_token = authorization.partition(" ") + if scheme.lower() != "bearer" or not secrets.compare_digest( + supplied_token.encode("utf-8"), expected_token.encode("utf-8") + ): + response = _error_response(401, "Missing or invalid bearer token.") + response.headers["WWW-Authenticate"] = "Bearer" + return response + + return await call_next(request) + @app.exception_handler(ParrotCoreUserError) async def parrot_core_internal_error_handler( @@ -79,7 +104,8 @@ async def remove_session(session_id: int, request: Request): @app.get(f"/{API_VERSION}" + "/session/{session_id}") async def get_session_info(session_id: int, request: Request): - raise NotImplementedError("Not implemented yet.") + payload = await request.json() + return pcore.get_session_info(session_id, payload) @app.post(f"/{API_VERSION}/semantic_call") @@ -97,6 +123,8 @@ async def submit_semantic_call(request: Request): # RTT if latency_open == 1: + from parrot.testing.latency_simulator import get_latency + latency = get_latency() await asyncio.sleep(latency / 2) @@ -111,9 +139,10 @@ async def submit_semantic_call(request: Request): @app.post(f"/{API_VERSION}/py_native_call") async def submit_py_native_call(request: Request): - payload = await request.json() - response = pcore.submit_py_native_call(payload) - return response + return _error_response( + 410, + "Python native calls are disabled because they execute client-supplied code.", + ) @app.post(f"/{API_VERSION}/semantic_var") @@ -149,16 +178,18 @@ async def get_semantic_variable_list(request: Request): @app.post("/engine_heartbeat") async def engine_heartbeat(request: Request): - payload = await request.json() - response = pcore.engine_heartbeat(payload) - return response + return _error_response( + 403, + "Dynamic engine registration and heartbeats are disabled by this hotfix.", + ) @app.post("/register_engine") async def register_engine(request: Request): - payload = await request.json() - response = pcore.register_engine(payload) - return response + return _error_response( + 403, + "Dynamic engine registration and heartbeats are disabled by this hotfix.", + ) def start_server( @@ -169,6 +200,14 @@ def start_server( global pcore global app + if not os.environ.get(AUTH_TOKEN_ENV): + raise RuntimeError( + f"{AUTH_TOKEN_ENV} must be set before starting ServeCore." + ) + + from parrot.serve.core import create_serve_core + from parrot.utils.async_utils import create_task_in_loop + # Create ServeCore pcore = create_serve_core( core_config_path=core_config_path, @@ -239,14 +278,15 @@ def start_server( if release_mode: # Disable logging - import logging - # We don't disable the error log logging.disable(logging.DEBUG) logging.disable(logging.INFO) # Set log output file if args.log_dir is not None: + from parrot.utils.logging import set_log_output_file + from parrot.utils.misc import redirect_stdout_stderr_to_file + set_log_output_file( log_file_dir_path=args.log_dir, log_file_name=args.log_filename, diff --git a/parrot/serve/session_auth.py b/parrot/serve/session_auth.py new file mode 100644 index 0000000..64c0e65 --- /dev/null +++ b/parrot/serve/session_auth.py @@ -0,0 +1,14 @@ +import secrets +from typing import Optional + + +def create_session_auth() -> str: + return secrets.token_urlsafe(32) + + +def verify_session_auth(supplied: Optional[str], expected: str) -> bool: + if not isinstance(supplied, str): + return False + return secrets.compare_digest( + supplied.encode("utf-8"), expected.encode("utf-8") + ) diff --git a/parrot/serve/session_manager.py b/parrot/serve/session_manager.py index 9354684..53ea472 100644 --- a/parrot/serve/session_manager.py +++ b/parrot/serve/session_manager.py @@ -8,6 +8,7 @@ from parrot.utils import RecyclePool, get_logger, time_counter_in_nanoseconds from .session.session import Session, SessionStatus +from .session_auth import create_session_auth, verify_session_auth from .scheduler.global_scheduler import GlobalScheduler from .variable_manager import SemanticVariableManager @@ -32,6 +33,7 @@ def __init__(self, **session_create_kwargs) -> None: # session_id -> last_access_time (nanoseconds) self._session_last_access_time: Dict[int, int] = {} + self._session_auth: Dict[int, str] = {} # ---------- Arguments for Creating Session ---------- self._session_create_kwargs = session_create_kwargs @@ -39,6 +41,7 @@ def __init__(self, **session_create_kwargs) -> None: def _remove_session(self, session_id: int) -> None: session = self.sessions.pop(session_id) self._session_last_access_time.pop(session_id) + self._session_auth.pop(session_id) session.free_session_resources() self._session_id_pool.free(session_id) @@ -46,11 +49,11 @@ def _remove_session(self, session_id: int) -> None: # ---------- Methods for Core ---------- - def register_session(self) -> int: + def register_session(self) -> Tuple[int, str]: """Create a new session. Returns: - int: The session ID. + Tuple[int, str]: The session ID and its authentication credential. """ # Create session object @@ -60,9 +63,11 @@ def register_session(self) -> int: # Maintain session info self.sessions[session_id] = session self._session_last_access_time[session_id] = time_counter_in_nanoseconds() + session_auth = create_session_auth() + self._session_auth[session_id] = session_auth logger.debug(f"Session (session_id={session_id}) registered.") - return session_id + return session_id, session_auth def remove_session(self, session_id: int) -> None: """Remove a session. @@ -106,7 +111,9 @@ def session_access_update(self, session_id: int) -> None: ) self._session_last_access_time[session_id] = time_counter_in_nanoseconds() - def check_session_status(self, session_id: int) -> None: + def check_session_status( + self, session_id: int, session_auth: Optional[str] + ) -> None: """Check the status of the session. Args: @@ -118,6 +125,14 @@ def check_session_status(self, session_id: int) -> None: RuntimeError(f"Session (session_id={session_id}) not found.") ) + expected_auth = self._session_auth[session_id] + if not verify_session_auth(session_auth, expected_auth): + raise ParrotCoreUserError( + RuntimeError( + f"Session credential for session_id={session_id} is invalid." + ) + ) + session = self.sessions[session_id] if session.status != SessionStatus.RUNNING: raise ParrotCoreUserError( diff --git a/parrot/utils/__init__.py b/parrot/utils/__init__.py index b056508..ea5cb69 100644 --- a/parrot/utils/__init__.py +++ b/parrot/utils/__init__.py @@ -11,13 +11,6 @@ from .profile import cprofile, torch_profile -from .serialize_utils import ( - bytes_to_encoded_b64str, - encoded_b64str_to_bytes, - serialize_func_code, - deserialize_func_code, -) - from .misc import ( set_random_seed, redirect_stdout_stderr_to_file, diff --git a/parrot/utils/serialize_utils.py b/parrot/utils/serialize_utils.py deleted file mode 100644 index 82525c6..0000000 --- a/parrot/utils/serialize_utils.py +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright (c) 2023 by Microsoft Corporation. -# Licensed under the MIT license. - -import base64 -import marshal - -from types import CodeType - - -def bytes_to_encoded_b64str(serialized: bytes) -> str: - """Convert bytes to string using base64 ASCII encoding. - - Args: - input_bytes (bytes): input bytes - - Returns: - str: base64 string - """ - - # Encode the serialized code to base64 - encoded_code = base64.b64encode(serialized).decode("utf-8") - return encoded_code - - -def encoded_b64str_to_bytes(b64str: str) -> bytes: - """Convert base64 string (encoded by method "bytes_to_str") to bytes.""" - - b64bytes = base64.b64decode(b64str) - return b64bytes - - -def serialize_func_code(pyfunc_code: CodeType) -> bytes: - """Serialize a Python function code to bytes.""" - - # Serialize the function - serialized_code = marshal.dumps(pyfunc_code) - return serialized_code - - -def deserialize_func_code(serialized_code: bytes) -> CodeType: - """Deserialize a Python function code from bytes.""" - - # Deserialize the function - code = marshal.loads(serialized_code) - return code diff --git a/tests/serve/test_http_security.py b/tests/serve/test_http_security.py new file mode 100644 index 0000000..7b4c62b --- /dev/null +++ b/tests/serve/test_http_security.py @@ -0,0 +1,82 @@ +import asyncio + +from starlette.requests import Request + +from parrot.serve.http_server import ( + AUTH_TOKEN_ENV, + engine_heartbeat, + register_engine, + require_api_key, + submit_py_native_call, +) +from parrot.serve.session_auth import create_session_auth, verify_session_auth + + +def _request(authorization=None): + headers = [] + if authorization is not None: + headers.append((b"authorization", authorization.encode())) + return Request({"type": "http", "headers": headers}) + + +def test_requests_are_rejected_without_server_token(monkeypatch): + monkeypatch.delenv(AUTH_TOKEN_ENV, raising=False) + + async def call_next(request): + raise AssertionError("Locked requests must not reach an endpoint.") + + response = asyncio.run(require_api_key(_request(), call_next)) + + assert response.status_code == 503 + + +def test_requests_require_valid_bearer_token(monkeypatch): + monkeypatch.setenv(AUTH_TOKEN_ENV, "test-token") + + async def call_next(request): + raise AssertionError("Unauthorized requests must not reach an endpoint.") + + response = asyncio.run(require_api_key(_request(), call_next)) + + assert response.status_code == 401 + assert response.headers["WWW-Authenticate"] == "Bearer" + + +def test_valid_bearer_token_reaches_endpoint(monkeypatch): + monkeypatch.setenv(AUTH_TOKEN_ENV, "test-token") + expected_response = object() + + async def call_next(request): + return expected_response + + response = asyncio.run( + require_api_key(_request("Bearer" + " " + "test-token"), call_next) + ) + + assert response is expected_response + + +def test_python_native_calls_are_disabled(): + response = asyncio.run(submit_py_native_call(_request())) + + assert response.status_code == 410 + assert b"disabled" in response.body + + +def test_dynamic_engine_registration_is_disabled(): + register_response = asyncio.run(register_engine(_request())) + heartbeat_response = asyncio.run(engine_heartbeat(_request())) + + assert register_response.status_code == 403 + assert heartbeat_response.status_code == 403 + + +def test_session_credentials_are_unique_and_unpredictable(): + first = create_session_auth() + second = create_session_auth() + + assert first != second + assert len(first) >= 32 + assert verify_session_auth(first, first) + assert not verify_session_auth(second, first) + assert not verify_session_auth(None, first) diff --git a/tests/serve/test_session.py b/tests/serve/test_session.py index 27269e9..9ed46c2 100644 --- a/tests/serve/test_session.py +++ b/tests/serve/test_session.py @@ -59,17 +59,20 @@ def test_session_manager(): ) # Test session registration - session_id = session_mgr.register_session() + session_id, session_auth = session_mgr.register_session() session = session_mgr.get_session(session_id) assert session.session_id == session_id + with pytest.raises(ParrotCoreUserError): + session_mgr.check_session_status(session_id, "wrong-session-auth") + session_mgr.check_session_status(session_id, session_auth) # Test session expiration time.sleep(11) session_mgr.check_running_sessions() with pytest.raises(ParrotCoreUserError): - session_mgr.check_session_status(session_id) + session_mgr.check_session_status(session_id, session_auth) def test_graph_executor():