From a25468b3f37dbe608a6dba312713435fa9e9e81c Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Fri, 16 Oct 2026 11:56:21 +0800 Subject: [PATCH 1/2] feat(security): extract ZMQ HMAC auth into shared module and extend to NCCL engine - Create zmq_auth.py with generate_auth_key, send_auth_pyobj, recv_auth_pyobj - Refactor HCCLCheckpointEngine to use the shared helpers - Add HMAC authentication to NCCLCheckpointEngine's PUB/SUB channel - Broadcast auth_key via MasterMetadata so workers share the secret - Export new auth utilities from checkpoint_engine package --- src/twinkle/checkpoint_engine/__init__.py | 4 + .../hccl_checkpoint_engine.py | 46 +++------ .../nccl_checkpoint_engine.py | 27 +++++- src/twinkle/checkpoint_engine/zmq_auth.py | 95 +++++++++++++++++++ 4 files changed, 139 insertions(+), 33 deletions(-) create mode 100644 src/twinkle/checkpoint_engine/zmq_auth.py diff --git a/src/twinkle/checkpoint_engine/__init__.py b/src/twinkle/checkpoint_engine/__init__.py index 85febef20..42dcae37a 100644 --- a/src/twinkle/checkpoint_engine/__init__.py +++ b/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', ] diff --git a/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py b/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py index 148b0a594..0d51fec3e 100644 --- a/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py +++ b/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) diff --git a/src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py b/src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py index d6209c58b..7d29bf8a6 100644 --- a/src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py +++ b/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,6 +66,7 @@ def __init__( metadata: dict[str, TensorMeta], socket: zmq.Socket, topic: str, + auth_key: str = '', ) -> None: self.rank = rank self.pg = pg @@ -70,6 +74,7 @@ def __init__( 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) diff --git a/src/twinkle/checkpoint_engine/zmq_auth.py b/src/twinkle/checkpoint_engine/zmq_auth.py new file mode 100644 index 000000000..190c60e88 --- /dev/null +++ b/src/twinkle/checkpoint_engine/zmq_auth.py @@ -0,0 +1,95 @@ +# 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) From 99a25eb6a81a284b84e89ae42a23b730f3019e1e Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Fri, 16 Oct 2026 13:16:35 +0800 Subject: [PATCH 2/2] fix: reformat ZMQ auth error message string concatenation --- src/twinkle/checkpoint_engine/zmq_auth.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/twinkle/checkpoint_engine/zmq_auth.py b/src/twinkle/checkpoint_engine/zmq_auth.py index 190c60e88..ec61e56b2 100644 --- a/src/twinkle/checkpoint_engine/zmq_auth.py +++ b/src/twinkle/checkpoint_engine/zmq_auth.py @@ -80,16 +80,14 @@ def recv_auth_pyobj( 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}).') + 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.') + raise RuntimeError(f'HMAC verification failed for ZMQ message at {where}; ' + 'refusing to deserialize untrusted payload.') return pickle.loads(payload_bytes)