fix(security): stop pickling the multinode ip handshake (GH-1413)

lightllm/utils/multinode_utils.py

@@ -1,10 +1,35 @@
+import ipaddress
 import zmq
 import socket
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.shm_port_args import get_shm_port_args
 
 logger = init_logger(__name__)
 
+# 子节点上报的载荷只是一个 ip 字符串，正常不会超过 45 字节（ipv6 最长形式），
+# 这里留一些余量，避免异常端或恶意端发送超大报文。
+_MAX_CHILD_IP_BYTES = 64
+
+
+def _decode_child_ip(raw: bytes) -> str:
+    """把子节点上报的原始字节解析为 ip 字符串。
+
+    这里刻意不使用 recv_pyobj：该端口必须绑定在所有网卡上（子节点要跨机连接过来），
+    而 recv_pyobj 等价于对网络数据直接做 pickle.loads，任何能访问该端口的主机都可以
+    构造恶意 pickle 在主节点上执行任意代码。载荷本身只是一个 ip，用 utf-8 解码加
+    ipaddress 校验就足够，且能彻底消除反序列化执行路径。
+
+    Raises:
+        ValueError: 载荷过大、不是合法 utf-8 或不是合法 ip 时抛出，让启动直接失败，
+            而不是把非法值写进 args.child_ips 后在后续建连时才报错。
+    """
+    if len(raw) > _MAX_CHILD_IP_BYTES:
+        raise ValueError(f"child ip payload too large: {len(raw)} bytes > {_MAX_CHILD_IP_BYTES}")
+
+    ip_str = raw.decode("utf-8").strip()
+    ipaddress.ip_address(ip_str)  # 非法 ip 会抛 ValueError
+    return ip_str
+
 
 def send_and_receive_node_ip(args):
     # 在多节点tp的部署形式中，0 号节点作为主节点，其他节点作为
@@ -21,14 +46,16 @@ def send_and_receive_node_ip(args):
                 comm_socket = context.socket(zmq.PULL)
                 comm_socket.bind(f"tcp://*:{base_port + i + 100}")
                 logger.info(f"binding port {base_port + i + 100}")
-                args.child_ips.append(comm_socket.recv_pyobj())
-                comm_socket.close()
+                try:
+                    args.child_ips.append(_decode_child_ip(comm_socket.recv()))
+                finally:
+                    comm_socket.close()
             logger.info(f"Received child IPs: {args.child_ips}")
         else:
             local_ip = socket.gethostbyname(socket.gethostname())
             context = zmq.Context(2)
             comm_socket = context.socket(zmq.PUSH)
             comm_socket.connect(f"tcp://{args.nccl_host}:{base_port + args.node_rank + 100}")
             logger.info(f"connecting to {args.nccl_host}:{base_port + args.node_rank + 100}")
-            comm_socket.send_pyobj(local_ip)
+            comm_socket.send(local_ip.encode("utf-8"))
             comm_socket.close()

unit_tests/utils/test_multinode_utils.py

@@ -0,0 +1,98 @@
+import os
+import pickle
+import pytest
+
+from lightllm.utils.multinode_utils import _decode_child_ip, _MAX_CHILD_IP_BYTES
+
+
+class _RceProbe:
+    """Pickle payload that would run code if the receiver called pickle.loads()."""
+
+    MARKER = "/tmp/lightllm_multinode_rce_probe"
+
+    def __reduce__(self):
+        return (os.system, (f"touch {self.MARKER}",))
+
+
+@pytest.mark.parametrize("ip", ["127.0.0.1", "10.0.0.7", "192.168.1.255", "::1", "fe80::1"])
+def test_decode_child_ip_accepts_valid_addresses(ip):
+    assert _decode_child_ip(ip.encode("utf-8")) == ip
+
+
+def test_decode_child_ip_strips_surrounding_whitespace():
+    assert _decode_child_ip(b"  10.0.0.7\n") == "10.0.0.7"
+
+
+@pytest.mark.parametrize(
+    "raw",
+    [
+        b"",
+        b"not-an-ip",
+        b"10.0.0.256",
+        b"127.0.0.1; rm -rf /",
+        b"\xff\xfe\xfd",  # invalid utf-8
+    ],
+)
+def test_decode_child_ip_rejects_malformed_payloads(raw):
+    with pytest.raises((ValueError, UnicodeDecodeError)):
+        _decode_child_ip(raw)
+
+
+def test_decode_child_ip_rejects_oversized_payload():
+    with pytest.raises(ValueError, match="too large"):
+        _decode_child_ip(b"1" * (_MAX_CHILD_IP_BYTES + 1))
+
+
+@pytest.mark.parametrize("protocol", [0, 1, 2, pickle.HIGHEST_PROTOCOL])
+def test_malicious_pickle_payload_is_rejected_without_executing(protocol):
+    """The regression this module exists for.
+
+    A pickled payload must be rejected as a malformed IP. Critically, it must be
+    rejected *without* being deserialized -- the marker file proves no code ran.
+    """
+    payload = pickle.dumps(_RceProbe(), protocol=protocol)
+    if os.path.exists(_RceProbe.MARKER):
+        os.remove(_RceProbe.MARKER)
+
+    try:
+        with pytest.raises((ValueError, UnicodeDecodeError)):
+            _decode_child_ip(payload)
+
+        assert not os.path.exists(_RceProbe.MARKER), "pickle payload was executed -- deserialization still reachable"
+    finally:
+        if os.path.exists(_RceProbe.MARKER):
+            os.remove(_RceProbe.MARKER)
+
+
+class _CompactRceProbe:
+    """Worst case for the validator: a pickle small and plain enough to reach ip_address().
+
+    Protocol 0 serialises to pure ASCII, and with a short command the payload stays
+    under the size cap -- so neither the length check nor the utf-8 decode rejects it.
+    That leaves ipaddress.ip_address() as the only thing standing between this payload
+    and execution, which is exactly the property worth pinning down.
+    """
+
+    MARKER = "/tmp/lm_rce"
+
+    def __reduce__(self):
+        return (os.system, (f"touch {self.MARKER}",))
+
+
+def test_compact_ascii_pickle_reaches_and_is_stopped_by_ip_validation():
+    payload = pickle.dumps(_CompactRceProbe(), protocol=0)
+
+    # Preconditions: this payload really does slip past the two cheaper checks.
+    assert len(payload) <= _MAX_CHILD_IP_BYTES, f"payload {len(payload)}B no longer exercises the ip_address() path"
+    assert payload.isascii(), "payload must be valid utf-8 to exercise the ip_address() path"
+    payload.decode("utf-8")
+
+    if os.path.exists(_CompactRceProbe.MARKER):
+        os.remove(_CompactRceProbe.MARKER)
+    try:
+        with pytest.raises(ValueError):
+            _decode_child_ip(payload)
+        assert not os.path.exists(_CompactRceProbe.MARKER), "pickle payload was executed"
+    finally:
+        if os.path.exists(_CompactRceProbe.MARKER):
+            os.remove(_CompactRceProbe.MARKER)