fix: replace pickle deserialization with RestrictedUnpickler in PD WebSocket endpoints (CVE-2026-26220)

lightllm/server/api_http.py

@@ -25,8 +25,9 @@
 import base64
 import os
 from io import BytesIO
-import pickle
+import pickle  # kept for non-network uses; WebSocket paths use safe_pickle
 import setproctitle
+from lightllm.utils.safe_pickle import safe_loads as _safe_pickle_loads
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 import ujson as json
@@ -426,7 +427,7 @@ async def register_and_keep_alive(websocket: WebSocket):
         while True:
             # 等待接收消息，设置超时为10秒
             data = await websocket.receive_bytes()
-            obj = pickle.loads(data)
+            obj = _safe_pickle_loads(data)  # CVE-2026-26220: restricted unpickling
             await g_objs.httpserver_manager.put_to_handle_queue(obj)
 
     except (WebSocketDisconnect, Exception, RuntimeError) as e:
@@ -447,7 +448,7 @@ async def kv_move_status(websocket: WebSocket):
         while True:
             # 等待接收消息，设置超时为10秒
             data = await websocket.receive_bytes()
-            upkv_status = pickle.loads(data)
+            upkv_status = _safe_pickle_loads(data)  # CVE-2026-26220: restricted unpickling
             logger.info(f"received upkv_status {upkv_status} from {(client_ip, client_port)}")
             await g_objs.httpserver_manager.update_req_status(upkv_status)
     except (WebSocketDisconnect, Exception, RuntimeError) as e:

lightllm/server/httpserver/pd_loop.py

@@ -1,6 +1,7 @@
 import asyncio
-import pickle
+import pickle  # kept for outbound pickle.dumps; inbound paths use safe_pickle
 import websockets
+from lightllm.utils.safe_pickle import safe_loads as _safe_pickle_loads
 import ujson as json
 import socket
 import httpx
@@ -103,7 +104,7 @@ async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_O
                 # 接收 pd master 发来的请求，并推理后，将生成的token转发回pd master。
                 while True:
                     recv_bytes = await websocket.recv()
-                    obj = pickle.loads(recv_bytes)
+                    obj = _safe_pickle_loads(recv_bytes)  # CVE-2026-26220: restricted unpickling
                     if obj[0] == ObjType.REQ:
                         prompt, sampling_params, multimodal_params = obj[1]
                         group_req_id = sampling_params.group_request_id
@@ -183,7 +184,7 @@ async def _get_pd_master_objs(args: StartArgs) -> Optional[Dict[int, PD_Master_O
             response = await client.get(uri)
             if response.status_code == 200:
                 base64data = response.json()["data"]
-                id_to_pd_master_obj = pickle.loads(base64.b64decode(base64data))
+                id_to_pd_master_obj = _safe_pickle_loads(base64.b64decode(base64data))  # CVE-2026-26220: restricted unpickling
                 return id_to_pd_master_obj
             else:
                 logger.error(f"get pd_master_objs error {response.status_code}")

lightllm/utils/safe_pickle.py

@@ -0,0 +1,78 @@
+"""
+safe_pickle.py — Restricted unpickling for LightLLM PD WebSocket endpoints.
+
+CVE-2026-26220: The PD disaggregation WebSocket endpoints and the config-server
+HTTP response handler called pickle.loads() on untrusted network data, enabling
+unauthenticated remote code execution.  This module replaces those bare
+pickle.loads() calls with a RestrictedUnpickler that whitelists only the
+internal LightLLM dataclass modules that legitimately flow through those
+channels.
+
+Usage:
+    from lightllm.utils.safe_pickle import safe_loads
+
+    obj = safe_loads(data)   # raises UnpicklingError for non-whitelisted types
+"""
+
+import io
+import pickle
+
+# ---------------------------------------------------------------------------
+# Allowlist: (module, name) pairs permitted to be deserialized.
+# Keep this list minimal — add entries only when a new type is deliberately
+# added to a PD WebSocket protocol message.
+# ---------------------------------------------------------------------------
+_ALLOWED_MODULES: dict[str, set[str]] = {
+    # Built-in safe types used as containers / primitives
+    "builtins": {"dict", "list", "tuple", "set", "int", "float", "str", "bool", "bytes", "NoneType"},
+    # LightLLM PD protocol structures
+    "lightllm.server.pd_io_struct": {
+        "ObjType",
+        "NodeRole",
+        "PD_Master_Obj",
+        "PD_Client_Obj",
+        "_PD_Client_RunStatus",
+        "UpKVStatus",
+        "NixlUpKVStatus",
+        "DecodeNodeInfo",
+        "NIXLDecodeNodeInfo",
+        "KVMoveTask",
+        "KVMoveTaskGroup",
+        "PDTransJoinInfo",
+        "PDTransLeaveInfo",
+        "NIXLChunckedTransTask",
+        "NIXLChunckedTransTaskRet",
+        "NIXLChunckedTransTaskGroup",
+    },
+    # SamplingParams / StartArgs sent inside REQ messages
+    "lightllm.server.core.objs.py_sampling_params": {"SamplingParams"},
+    "lightllm.server.core.objs.sampling_params": {"SamplingParams"},
+    "lightllm.server.core.objs.start_args_type": {"StartArgs"},
+    # Multimodal params can accompany REQ messages
+    "lightllm.server.multimodal_params": {"MultimodalParams"},
+    # enum base class (needed for ObjType / NodeRole reconstruction)
+    "enum": {"Enum"},
+}
+
+
+class _RestrictedUnpickler(pickle.Unpickler):
+    """Unpickler that raises UnpicklingError for any non-whitelisted class."""
+
+    def find_class(self, module: str, name: str):
+        allowed_names = _ALLOWED_MODULES.get(module)
+        if allowed_names is not None and name in allowed_names:
+            return super().find_class(module, name)
+        raise pickle.UnpicklingError(
+            f"Refusing to deserialize {module}.{name}: not in safe_pickle allowlist. "
+            f"This may indicate an attempted pickle-based RCE exploit (CVE-2026-26220)."
+        )
+
+
+def safe_loads(data: bytes) -> object:
+    """
+    Drop-in replacement for pickle.loads() that only permits whitelisted types.
+
+    Raises pickle.UnpicklingError if the payload attempts to instantiate a
+    class outside the allowlist.
+    """
+    return _RestrictedUnpickler(io.BytesIO(data)).load()