[inureyes]

Summary

Five hardening items from the #2674 security review on the authenticated cluster transport (src/distributed/) and the disaggregated router demux, in one PR with one verification pass. None changes what a party without the secret can read or inject.

1. Per-source pending limit keyed per IPv6 /64. pending_source_key(ip) canonicalizes, keeps IPv4 as is and masks IPv6 to its /64; try_acquire and pending_from use it, PendingSlotRefusal::PerIp carries the key and prints .../64 for IPv6. Module doc, the max_pending_handshakes_per_ip field doc and docs/distributed.md say that nodes inside one /64 share the per-source budget.
2. active_conns deleted with its three manual decrements and the AtomicUsize import; no replacement counter. cluster_auth::panicking_rpc_handler_releases_its_connection_slot is the regression guard (see the note below on its observable).
3. src/distributed/handshake_warn.rs (HandshakeWarnClass {Noise, Diagnostic}, handshake_warn_class(&anyhow::Error), WarnBudget::allow_at(class, now) -> WarnDecision { warn, suppressed_in_previous_window }, process-wide handshake_warn_allowed(class) with the rollover summary N cluster handshake warnings (<class>) were logged at debug in the last 60s). Pending refusals draw Noise, handshake failures handshake_warn_class(&err), rejected frames Diagnostic; tcp_transport.rs and stage_assignment.rs call sites updated, the old limiter removed. Residual risk noted in docs/distributed.md.
4. Router sender check. TcpTransport::is_authenticated(); RouterState::build fails with the disaggregated router requires a role transport bound with the cluster secret; result_senders: HashSet<SocketAddr> (prefill ∪ decode peers, canonical IPs) and accepts_result_from(&str); the demux drops other senders with the router: dropping result frame from ... warn through the Diagnostic budget. Both router test fixtures bind with auth. docs/distributed.md, docs/CONTINUOUS_BATCHING.md and the --prefill-peers / --decode-peers help now say a peer must be listed under the address it connects from (its --serving-bind, with an unspecified IP replaced by the source IP). docs/distributed.md states that the check authenticates the source IP (handshake-verified) while the port is the sender's advertised listener port, so it tells members apart by source IP plus claimed port, not by port-level proof. The demux logs drops through tracing fields without building a message string.
5. Admission inside the handshake. VERDICT_AT_CAPACITY = 0x02, PROTOCOL_VERSION 3 to 4 (cluster_handshake.rs, constant unchanged in name and place), pub(crate) async fn server_handshake_admitting<P>(stream, auth, local_identity, suite, channel, admit: impl FnOnce() -> Option<P>) -> Result<(String, SessionKeys, P)> calling admit() only after the client authenticator verified; server_handshake / server_handshake_on wrap it with || Some(()). New errors ConnectionLimitReached { peer } (server) and ServerAtCapacity (client; not treated as fatal like HandshakeRejected: connect_with_retry returns it without backoff, the next pool acquire reconnects, and Transport::connect warns), the latter re-exported from distributed/mod.rs. accept_loop passes || conn_semaphore.try_acquire_owned().ok() and keeps the max connections (N) reached warn. Wire diagram and docs updated.

Pre-existing limitation found while writing the item 2 test: TcpTransport reads RpcResponse frames only on accepted connections, never on the pooled outbound connection an rpc_call went out on, so no RPC echo can reach a TCP client (the only in-repo RPC-over-TCP caller, bench.rs, ignores the result). The guard therefore keeps the panicking RPC handler as the trigger (client A's rpc_call(b"boom") fails) and uses client B's send reaching server.recv() within 2 s as the observable.

Changes during review

The implementation review and the security review found no CRITICAL or HIGH issues. The follow-up commit tightens wording only: the residual-risk notes now name IdentityNotOwned and FrameSuiteMismatch (decided from the hello alone) next to ClientAbortedAfterServerProof as refusals a party without the secret can push into the Diagnostic window; the ServerAtCapacity docs and the connection_pool.rs "failed handshake is terminal" comment describe the no-backoff reconnect precisely; the startup.rs router comment credits RouterState::build for enforcing cluster_auth; and the sender-check docs say only the source IP is verified.

Failing-first evidence (main 84e20ce)

Checked by temporarily reverting the fix hunks on this branch (key on the full address, one shared warn window, no is_authenticated check, no sender check), rebuilding the lib tests and running the tests by name, then restoring the files byte for byte (sha256 verified):

• ipv6_sources_in_one_64_share_the_per_source_limit: FAILED (third acquire from 2001:db8:1:2:ffff::3 succeeded, unwrap_err panicked at pending_handshakes_tests.rs:83).
• ipv6_sources_in_different_64s_are_counted_separately: FAILED (second address in the /64 acquired a slot, pending_handshakes_tests.rs:110).
• noise_flood_does_not_suppress_diagnostic_warnings: FAILED (left: WarnDecision { warn: false, .. }).
• build_refuses_an_unauthenticated_transport: FAILED (an unauthenticated transport must be refused).
• result_demux_drops_frames_from_unconfigured_senders: FAILED (unconfigured sender's frame was delivered: ... tokens: ["injected"]).
• rollover_reports_suppressed_count passes with a shared window too (only one class is exercised); classifies_handshake_errors and accepts_result_from_only_configured_peers call functions that do not exist on main, so they do not compile there.
• server_at_capacity_sends_capacity_verdict, admission_is_not_consulted_before_the_client_authenticates and cluster_auth::full_connection_cap_reports_capacity_to_the_client do not compile on main (server_handshake_admitting, ServerAtCapacity, ConnectionLimitReached are new).
• cluster_auth::panicking_rpc_handler_releases_its_connection_slot passes with and without the counter, by design.

Verification (GB10 spark-102, --profile test-fast, builds outside gpu-lock, runs under gpu-lock run --tag issue-2692)

• cargo fmt --check: clean.
• cargo clippy --profile test-fast --features cuda -p mlxcel --lib --tests -- -D warnings: clean (rerun after the review commit: clean; server::router_front:: + distributed::handshake_warn:: 38 passed, and the touched modules cluster_handshake, pending_handshakes, tcp_transport 53 passed at 699d208).
• 13 new or touched tests by name (--lib -- --test-threads=1): 13 passed.
• --lib distributed:: -- --test-threads=1: 1473 passed, 0 failed, 4 ignored.
• --lib server::router_front -- --test-threads=1: 35 passed.
• --test zero_config_cluster_init: 4 passed, 3 ignored.
• tests/disaggregated_serving_e2e.rs (--features cuda,test-utils -- --include-ignored --test-threads=1): 2 passed.
• tests/disaggregated_router_e2e.rs (same flags): 7 passed (49.7 s).
• MLXCEL_REQUIRE_MODELS=1 MLXCEL_SDPA_DETERMINISTIC=1, --features cuda,test-utils -- --ignored --test-threads=1: pipeline_remote_real_models 10 passed, 1 failed; pipeline_server_remote_real_models 7 passed, 1 failed; pipeline_ci_multi_stage_real_models 2 passed. The two failures are the glm_moe_dsa tests panicking in repo_model_dir because glm5-4bit is absent: mlx-community/GLM-5-4bit is 418.6 GB and cannot be loaded on this 128 GB box, and the tests predate fix(distributed): cap pre-auth slots, type rejections, auth DI roles #2674 (added 2026-04-13), so this is a missing checkpoint, not a regression. Every llama, gemma3, gemma4, qwen3, qwen3.5 and glm-4.7 parity test passed.
• Manual three-process loopback (router + prefill + decode, one shared --cluster-secret-file mode 0600, qwen3-0.6b-4bit, 127.0.0.1 ports 19801-19806, target/test-fast/mlxcel-server built from this branch): POST /v1/chat/completions returned HTTP 200 with "2 + 2 = 4.", and the router log contains no dropping result frame line. Processes killed by PID afterwards.

Closes #2692

[inureyes]

Implementation Review Summary

Intent

Implement the five #2692 hardening items on the authenticated cluster transport and router demux (IPv6 /64 pending key, drop active_conns, per-class handshake warn budget, router sender check, at-capacity verdict with PROTOCOL_VERSION 4).

Findings Addressed

None needed fixing at CRITICAL or HIGH. Checked specifically: admit() runs only after the client authenticator verifies; a permit taken by admit() is dropped on a later write failure or timeout (RAII, no leak); ServerAtCapacity is not HandshakeRejected, so Transport::connect warns instead of failing; the router sender check canonicalizes IP and port on both sides exactly as bind_advertised_identity produces from. The item 2 deviation is correct: RpcResponse frames are parsed only in handle_connection (accepted connections), and the pooled outbound stream an rpc_call writes on is never read, so no TCP client can receive an echo.

Remaining Items

• handshake_warn.rs module doc defines Noise as "anything a stranger can produce without the secret", and the residual-risk notes name only ClientAbortedAfterServerProof; IdentityNotOwned and FrameSuiteMismatch are also pre-auth refusals a stranger can trigger, so they can saturate the Diagnostic window too (MEDIUM, doc accuracy only; the classification itself follows the issue)
• ServerAtCapacity docs say the connection pool treats it as retryable; connect_with_retry returns it at once without backoff (the next acquire reconnects), and its "a failed handshake is terminal" comment now also covers this transient case (LOW, wording)
• docs/distributed.md per-frame paragraph re-wrap leaves a short "connection, so a frame" line; the startup.rs router comment still credits only the resolver for cluster_auth being set (LOW, cosmetic)

Verification

• All stated requirements implemented
• No placeholder/mock code remaining
• Integrated into project code flow
• Project conventions followed
• Existing modules reused where applicable
• No unintended structural changes
• Tests pass: cargo fmt --check and clippy --lib --tests -D warnings clean; under gpu-lock distributed::cluster_handshake 20, handshake_warn 3, pending_handshakes 6, tcp_transport 27, server::router_front 35, all passed

[inureyes]

Security and performance review of 699d208: no CRITICAL or HIGH findings, so no fixes were pushed. Points already raised in the implementation review are not repeated.

Checked and sound:

• Capacity state is not exposed before authentication. admit() runs only after the client authenticator verifies (a mismatch returns VERDICT_REJECT first), so a party without the secret cannot see or trigger VERDICT_AT_CAPACITY. The verdict byte is unauthenticated as before, so an on-path attacker can still swap 0x01/0x02, but it could already drop the connection.
• No permit or slot leaks. The permit is a local that is dropped if the VERDICT_ACCEPT write fails or the handshake times out, the pending slot is dropped right after the handshake, and nothing else is counted by hand.
• Router sender check: canonical_socket_addr rebuilds the address from ip().to_canonical() and the port. That removes IPv4-mapped forms, scope ids and flowinfo on both sides. The bound identity is never unspecified, so a configured 0.0.0.0 entry matches nothing, and an unparseable from is refused.
• The /64 key masks the upper 64 bits of the big-endian u128, and the shift cannot overflow. The per-source map is still capped by the global pending pool, and the warn budget is a fixed two-slot array with no per-source state.
• No lock is held across an await on the accept path. Both new mutexes are std and only cover synchronous updates.

LOW, not fixed:

• The router check matches IP and port, but only the IP is authenticated: the port in the bound identity is whatever the sender advertised. A secret holder on the same host as a configured peer, or behind the same NAT IP, can claim that peer's port and pass the check. That is acceptable under the stated threat model (secret holders are members, and per-request binding was rejected), but docs/distributed.md says "canonical IP and port" and could add that the check tells members apart by source IP only.
• spawn_result_demux builds the drop message with format! even when the budget sends it to debug. This costs one allocation per dropped frame, and only a secret holder can trigger it. Negligible.

Tests run (test-fast, cuda, --test-threads=1, under gpu-lock): distributed::cluster_handshake:: 20/20, distributed::handshake_warn:: 3/3, distributed::pending_handshakes:: 6/6, distributed::tcp_transport:: 27/27, server::router_front:: 35/35.