fix(security): bind balance_serve scheduler ZMQ socket to loopback (#2087)

[Anai-Guo]

Summary

Fixes #2087 (GHSA-83vp-v6wg-x93x): unauthenticated pickle-deserialization RCE on the balance_serve scheduler.

The scheduler binds its ZeroMQ ROUTER socket with tcp://* (i.e. 0.0.0.0, all interfaces) and calls pickle.loads() on every received message with no authentication. Any host that can reach the auto-assigned sched_port can send a crafted pickle whose __reduce__ runs a shell command, achieving remote code execution as the server process.

Fix

Bind the scheduler socket to loopback (tcp://127.0.0.1) instead of tcp://*.

This is safe because the scheduler RPC is purely inter-process on the same host:

• SchedulerClient.__init__ already connects to tcp://localhost:{sched_port} (sched_rpc.py), and
• the scheduler is launched as a local subprocess by the balance_serve backend (subprocess.Popen([... sched_rpc.py ...])).

So loopback binding is functionally identical for the legitimate local client, while removing the socket from every non-loopback interface and closing the network attack surface. This implements the primary remediation from the issue ("Bind the scheduler socket to loopback").

Scope

• Applied to both copies of the file that carry the identical bind:
  • archive/ktransformers/server/balance_serve/sched_rpc.py
  • archive/kt-sft/ktransformers/server/balance_serve/sched_rpc.py
• Each change is a single line (tcp://* → tcp://127.0.0.1).

Left out of scope (larger change): the pickle.loads() on the inter-process channel remains. With the socket now loopback-only the network RCE vector is closed, but replacing pickle with a schema-validated format (JSON/msgpack) or adding ZeroMQ CURVE auth is a broader design change best handled separately — happy to follow up if maintainers want it.

🤖 Generated with Claude Code

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[yyj6666667]

Good security hardening. Binding the scheduler ZMQ socket to 127.0.0.1 prevents exposing the internal control channel on external interfaces while preserving same-host communication.