feat(security): Add a communication auth_key for NCCL and HCCL.

src/twinkle/checkpoint_engine/__init__.py

@@ -19,6 +19,7 @@
 from .mixin import CheckpointEngineMixin
 # Import backend implementations to register them
 from .nccl_checkpoint_engine import NCCLCheckpointEngine
+from .zmq_auth import generate_auth_key, recv_auth_pyobj, send_auth_pyobj
 
 __all__ = [
     'CheckpointEngine',
@@ -27,4 +28,7 @@
     'NCCLCheckpointEngine',
     'HCCLCheckpointEngine',
     'TensorMeta',
+    'generate_auth_key',
+    'send_auth_pyobj',
+    'recv_auth_pyobj',
 ]

src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py

@@ -8,11 +8,7 @@
 
 from __future__ import annotations
 
-import hashlib
-import hmac
 import os
-import pickle
-import secrets
 import time
 import torch
 import zmq
@@ -23,6 +19,7 @@
 from twinkle.utils import find_free_port, find_node_ip, is_valid_ipv6_address, stateless_init_process_group
 from twinkle.utils.zmq_utils import configure_zmq_socket
 from .base import CheckpointEngine
+from .zmq_auth import generate_auth_key, recv_auth_pyobj, send_auth_pyobj
 
 logger = get_logger()
 
@@ -92,42 +89,29 @@ def _recv_pyobj(self, where: str) -> Any:
         ``pickle.loads`` runs, so a crafted pickle from an untrusted
         peer can never reach the deserializer."""
         assert self.socket is not None
-        try:
-            msg = self.socket.recv(copy=True)
-        except zmq.error.Again as e:
-            raise RuntimeError(f'HCCL metadata timeout ({self.meta_timeout_s}s) waiting at {where}.') from e
-        if not self.auth_key:
-            # No key configured — fall back to original recv_pyobj behaviour.
-            return pickle.loads(msg) if isinstance(msg, (bytes, bytearray)) else self.socket.recv_pyobj()
-        if len(msg) < 32:
-            raise RuntimeError(f'Rejecting short ZMQ message at {where}: {len(msg)} bytes (< HMAC size 32).')
-        received_hmac = msg[:32]
-        payload = msg[32:]
-        expected_hmac = hmac.new(self.auth_key.encode(), payload, hashlib.sha256).digest()
-        if not hmac.compare_digest(received_hmac, expected_hmac):
-            raise RuntimeError(f'HMAC verification failed for ZMQ message at {where}; '
-                               'refusing to deserialize untrusted payload.')
-        return pickle.loads(payload)
+        return recv_auth_pyobj(
+            self.socket,
+            self.auth_key,
+            where=where,
+            timeout_s=self.meta_timeout_s,
+        )
 
     def _send_pyobj(self, payload: Any, where: str) -> None:
         """Send a pickled object with an HMAC signature prepended."""
         assert self.socket is not None
-        data = pickle.dumps(payload)
-        if self.auth_key:
-            sig = hmac.new(self.auth_key.encode(), data, hashlib.sha256).digest()
-            msg = sig + data
-        else:
-            msg = data
-        try:
-            self.socket.send(msg)
-        except zmq.error.Again as e:
-            raise RuntimeError(f'HCCL metadata timeout ({self.meta_timeout_s}s) sending at {where}.') from e
+        send_auth_pyobj(
+            self.socket,
+            payload,
+            self.auth_key,
+            where=where,
+            timeout_s=self.meta_timeout_s,
+        )
 
     def _start_zmq_server(self):
         self.ip = find_node_ip()
         self.zmq_port = find_free_port()
         self.dist_port = find_free_port()
-        self.auth_key = secrets.token_hex(32)  # 256-bit shared secret
+        self.auth_key = generate_auth_key()  # 256-bit shared secret
 
         self._zmq_ctx = zmq.Context()
         self.socket = self._new_socket(zmq.REP)

src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py

@@ -12,6 +12,7 @@
 from twinkle import get_logger
 from twinkle.utils.network import find_free_port, is_valid_ipv6_address
 from .base import CheckpointEngine, TensorMeta
+from .zmq_auth import generate_auth_key, recv_auth_pyobj, send_auth_pyobj
 
 logger = get_logger()
 
@@ -23,6 +24,7 @@ class MasterMetadata:
     # TCPStore address for the checkpoint NCCL process group
     nccl_store_host: str = ''
     nccl_store_port: int = 0
+    auth_key: str = ''  # HMAC shared secret for authenticating ZMQ control messages
 
 
 def _pg_broadcast(pg: dist.ProcessGroup, tensor: torch.Tensor, src: int = 0):
@@ -53,6 +55,7 @@ class BroadcastOperation:
         metadata: The metadata of tensors in the bucket.
         socket: The ZMQ socket for metadata communication.
         topic: The ZMQ topic for pub/sub.
+        auth_key: HMAC shared secret for authenticating ZMQ messages.
     """
 
     def __init__(
@@ -63,13 +66,15 @@ def __init__(
         metadata: dict[str, TensorMeta],
         socket: zmq.Socket,
         topic: str,
+        auth_key: str = '',
     ) -> None:
         self.rank = rank
         self.pg = pg
         self.bucket = bucket
         self.metadata = metadata
         self.socket = socket
         self.topic = topic
+        self.auth_key = auth_key
 
         loop = asyncio.get_running_loop()
         self._task = loop.run_in_executor(None, self._run)
@@ -78,10 +83,19 @@ def _run(self):
         # Broadcast tensor metadata via ZMQ PUB/SUB
         if self.rank == 0:
             self.socket.send_string(self.topic, flags=zmq.SNDMORE)
-            self.socket.send_pyobj(self.metadata)
+            send_auth_pyobj(
+                self.socket,
+                self.metadata,
+                self.auth_key,
+                where='BroadcastOperation.send',
+            )
         else:
             self.socket.recv_string()
-            self.metadata = self.socket.recv_pyobj()
+            self.metadata = recv_auth_pyobj(
+                self.socket,
+                self.auth_key,
+                where='BroadcastOperation.recv',
+            )
 
         # Broadcast tensor data via NCCL
         _pg_broadcast(self.pg, self.bucket, src=0)
@@ -121,6 +135,7 @@ def __init__(
         self.send_buf = None
         self.recv_buf = None
         self.socket = None
+        self.auth_key: str = ''  # HMAC shared secret, generated by the master
 
         # torch.distributed process group for checkpoint NCCL ops
         self._pg: dist.ProcessGroup | None = None
@@ -137,6 +152,7 @@ def _start_zmq_server(self):
         import ray
         self.ip = ray.util.get_node_ip_address().strip('[]')
         self.listen_port = find_free_port()
+        self.auth_key = generate_auth_key()  # 256-bit shared secret
 
         context = zmq.Context()
         self.socket = context.socket(zmq.PUB)
@@ -150,6 +166,7 @@ def _start_zmq_server(self):
 
     def _connect_zmq_client(self, metadata: MasterMetadata):
         """Connect to the ZMQ PUB server as a subscriber (receiver only)."""
+        self.auth_key = metadata.auth_key  # inherit the master's HMAC secret
         context = zmq.Context()
         self.socket = context.socket(zmq.SUB)
         if is_valid_ipv6_address(metadata.zmq_ip):
@@ -181,6 +198,7 @@ def prepare(self) -> MasterMetadata | None:
                     zmq_port=self.listen_port,
                     nccl_store_host=self._nccl_store_host,
                     nccl_store_port=self._nccl_store_port,
+                    auth_key=self.auth_key,
                 )
             return None
 
@@ -200,6 +218,7 @@ def prepare(self) -> MasterMetadata | None:
                 zmq_port=self.listen_port,
                 nccl_store_host=self._nccl_store_host,
                 nccl_store_port=self._nccl_store_port,
+                auth_key=self.auth_key,
             )
         else:
             self.send_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device='cuda')
@@ -395,6 +414,7 @@ async def send_weights(
                     },
                     socket=self.socket,
                     topic=self.topic,
+                    auth_key=self.auth_key,
                 )
 
                 # Swap buffers
@@ -433,6 +453,7 @@ async def send_weights(
             },
             socket=self.socket,
             topic=self.topic,
+            auth_key=self.auth_key,
         )
         await broadcast_op.wait_for_complete()
 
@@ -465,6 +486,7 @@ async def receive_weights(self, ) -> AsyncGenerator[tuple[str, torch.Tensor], No
             metadata=None,
             socket=self.socket,
             topic=self.topic,
+            auth_key=self.auth_key,
         )
         metadata = await broadcast_op.wait_for_complete()
         total_bytes += self.bucket_size
@@ -482,6 +504,7 @@ async def receive_weights(self, ) -> AsyncGenerator[tuple[str, torch.Tensor], No
                 metadata=None,
                 socket=self.socket,
                 topic=self.topic,
+                auth_key=self.auth_key,
             )
 
             # 2. Yield tensors from current buffer (send_buf)

src/twinkle/checkpoint_engine/zmq_auth.py

@@ -0,0 +1,93 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""HMAC authentication for ZMQ control-channel messages.
+
+Provides ``send_auth_pyobj`` / ``recv_auth_pyobj`` helpers that prepend an
+HMAC-SHA256 signature to pickled payloads, so that a crafted pickle from an
+untrusted peer can never reach the deserializer.
+
+When ``auth_key`` is empty the helpers fall back to plain ``pickle.dumps`` /
+``pickle.loads`` — i.e. the original, unauthenticated behaviour — so callers
+that have not opted into authentication are unaffected.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import pickle
+import secrets
+import zmq
+from typing import Any
+
+_HMAC_SIZE = 32  # SHA-256 digest size in bytes
+
+
+def generate_auth_key() -> str:
+    """Generate a 256-bit shared secret as a hex string."""
+    return secrets.token_hex(32)
+
+
+def send_auth_pyobj(
+    socket: zmq.Socket,
+    payload: Any,
+    auth_key: str = '',
+    *,
+    where: str = '',
+    timeout_s: float | None = None,
+) -> None:
+    """Send a pickled object, optionally with an HMAC signature prepended.
+
+    When ``auth_key`` is set, the wire format is ``HMAC-SHA256(key, data) || data``.
+    When empty, plain ``pickle.dumps(payload)`` is sent (backward compatible).
+    """
+    data = pickle.dumps(payload)
+    if auth_key:
+        sig = hmac.new(auth_key.encode(), data, hashlib.sha256).digest()
+        msg = sig + data
+    else:
+        msg = data
+    try:
+        socket.send(msg)
+    except zmq.error.Again as e:
+        ts = f'{timeout_s}s' if timeout_s is not None else 'unknown'
+        raise RuntimeError(f'ZMQ metadata timeout ({ts}) sending at {where}.') from e
+
+
+def recv_auth_pyobj(
+    socket: zmq.Socket,
+    auth_key: str = '',
+    *,
+    where: str = '',
+    timeout_s: float | None = None,
+) -> Any:
+    """Receive a pickled object after verifying its HMAC signature (if configured).
+
+    The raw bytes are received first; the HMAC is checked *before*
+    ``pickle.loads`` runs, so a crafted pickle from an untrusted peer can
+    never reach the deserializer.
+
+    When ``auth_key`` is empty, falls back to plain ``pickle.loads`` (backward
+    compatible with the original ``recv_pyobj`` behaviour).
+    """
+    try:
+        msg = socket.recv(copy=True)
+    except zmq.error.Again as e:
+        ts = f'{timeout_s}s' if timeout_s is not None else 'unknown'
+        raise RuntimeError(f'ZMQ metadata timeout ({ts}) waiting at {where}.') from e
+
+    if not auth_key:
+        # No key configured — fall back to original recv_pyobj behaviour.
+        return pickle.loads(msg) if isinstance(msg, (bytes, bytearray)) else socket.recv_pyobj()
+
+    if len(msg) < _HMAC_SIZE:
+        raise RuntimeError(f'Rejecting short ZMQ message at {where}: {len(msg)} bytes '
+                           f'(< HMAC size {_HMAC_SIZE}).')
+
+    received_hmac = msg[:_HMAC_SIZE]
+    payload_bytes = msg[_HMAC_SIZE:]
+    expected_hmac = hmac.new(auth_key.encode(), payload_bytes, hashlib.sha256).digest()
+    if not hmac.compare_digest(received_hmac, expected_hmac):
+        raise RuntimeError(f'HMAC verification failed for ZMQ message at {where}; '
+                           'refusing to deserialize untrusted payload.')
+
+    return pickle.loads(payload_bytes)