From 699d208de54f14ebfced903898be6946f52eb4e1 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 28 Oct 2026 14:24:01 +0900 Subject: [PATCH 1/3] fix(distributed): harden cluster pre-auth admission and rejections Five follow-ups from the #2674 security review on the authenticated cluster transport. None changes what a party without the secret can read or inject; they close DoS and robustness gaps on the accept path and the router demux. - Pending-handshake per-source limit keys IPv6 sources per /64 (`pending_source_key`), so a host that owns a /64 cannot take one pending slot per address; nodes inside one /64 share the budget (documented). - The unread `active_conns` counter and its manual decrements are gone; every per-connection resource is RAII, guarded by `panicking_rpc_handler_releases_its_connection_slot`. - Handshake warnings move to `distributed::handshake_warn` with separate `Noise` and `Diagnostic` budgets (5 per 60 s each) and a rollover summary, so junk connections cannot push the secret-mismatch line down to debug. - The disaggregated router refuses an unauthenticated transport at build time and its demux delivers result frames only from configured `--prefill-peers` / `--decode-peers` identities; other holders of the secret can no longer inject tokens through the sequential request ids. - Admission runs inside the handshake after the client authenticated: `server_handshake_admitting` answers a full cap with the new 0x02 verdict, the client gets the retryable `ServerAtCapacity`, and `PROTOCOL_VERSION` goes 3 to 4 (no released version carries 3). Verified on GB10: fmt, clippy, the distributed:: and server::router_front lib modules, zero_config_cluster_init, the real-model remote pipeline suites, both disaggregated e2e suites, and a manual router + prefill + decode loopback (HTTP 200, no dropped result frame). Closes #2692 --- docs/CONTINUOUS_BATCHING.md | 6 +- docs/distributed.md | 65 +++++-- src/bin/mlx_server.rs | 10 +- src/distributed/cluster_handshake.rs | 109 +++++++++-- src/distributed/cluster_handshake_tests.rs | 68 +++++++ src/distributed/handshake_warn.rs | 198 ++++++++++++++++++++ src/distributed/handshake_warn_tests.rs | 145 ++++++++++++++ src/distributed/mod.rs | 7 +- src/distributed/pending_handshakes.rs | 66 +++++-- src/distributed/pending_handshakes_tests.rs | 55 ++++++ src/distributed/stage_assignment.rs | 15 +- src/distributed/tcp_transport.rs | 140 +++++++------- src/distributed/tcp_transport_tests.rs | 116 ++++++++++++ src/server/router_front.rs | 84 ++++++++- src/server/router_front_tests.rs | 153 +++++++++++++-- 15 files changed, 1086 insertions(+), 151 deletions(-) create mode 100644 src/distributed/handshake_warn.rs create mode 100644 src/distributed/handshake_warn_tests.rs diff --git a/docs/CONTINUOUS_BATCHING.md b/docs/CONTINUOUS_BATCHING.md index 6f5edc980..6fcff902e 100644 --- a/docs/CONTINUOUS_BATCHING.md +++ b/docs/CONTINUOUS_BATCHING.md @@ -857,8 +857,10 @@ Networking flags: | Flag | Role | Purpose | |------|------|---------| | `--serving-bind ` | prefill, decode, router | This node's own role-transport listener (`host:port`). | -| `--decode-peers ` | prefill, router | On a router, the decode pool it balances; on a prefill node, the static fallback used only when the router does not pick a decode target for a request. | -| `--prefill-peers ` | decode, router | On a router, the prefill pool it balances; on a decode node, the prefill node(s) it accepts handoffs from. | +| `--decode-peers ` | prefill, router | On a router, the decode pool it balances and the only decode addresses it accepts result frames from; on a prefill node, the static fallback used only when the router does not pick a decode target for a request. | +| `--prefill-peers ` | decode, router | On a router, the prefill pool it balances and the only prefill addresses it accepts result frames from; on a decode node, the prefill node(s) it accepts handoffs from. | + +On a router, each peer entry must be the address the node connects from: its `--serving-bind` address, with an unspecified bind IP (`0.0.0.0`, `::`) replaced by the IP the router sees the connection come from. A result frame from any other authenticated address is dropped with a `router: dropping result frame` warning and the request fails with the handoff timeout (see `docs/distributed.md`). ### Multi-node routing, load balancing, and failover diff --git a/docs/distributed.md b/docs/distributed.md index acd664a19..27c9754bb 100644 --- a/docs/distributed.md +++ b/docs/distributed.md @@ -335,17 +335,29 @@ What the secret protects: reads an offer, and the peer proves it before the coordinator sends one, so a coordinator with another secret cannot assign a stage or a model, and the peer keeps waiting for the real one. The peer's assignment listener - draws unauthenticated sockets from a pending pool of 8 (4 per source IP) - and closes the excess at accept. + draws unauthenticated sockets from a pending pool of 8 (4 per source: an + IPv4 address or an IPv6 /64) and closes the excess at accept. - **Disaggregated role connections.** Since issue #2655 the router, prefill and decode role transports run the same connection handshake: the router's request frames, the prefill node's KV handoff to a decode node, and every result frame travel only over authenticated connections, and a raw TCP client that skips the handshake never gets a frame to any role's receive - loop. Frames from authenticated members stay trusted (no per-frame source - check inside the disaggregated protocol); `MLXCEL_DECODE_ALLOWLIST` is - unchanged. + loop. The router refuses to start on a transport bound without the secret, + and its result demux delivers a frame only when the sender's + handshake-bound identity is one of its `--prefill-peers` or + `--decode-peers` addresses (canonical IP and port): request ids are + sequential, so any other holder of the secret, such as a pipeline stage + that shares it, could otherwise inject tokens into a client's response. A + frame from another member is dropped with a warning (`router: dropping + result frame from ...`) and the request fails with the handoff timeout. The + bound identity is the sender's `--serving-bind` address with an unspecified + bind IP replaced by the connection's source IP, so a peer must be listed + under the address it connects from (a multi-homed node, or one bound to + `0.0.0.0` that reaches the router over another interface, needs that + interface's address in the router's peer list). The decode role's handoff + intake and the prefill role's request intake keep trusting any + authenticated member; `MLXCEL_DECODE_ALLOWLIST` is unchanged. Fail-closed scope: `--pp-auto`, `--pp-peer`, any `--distributed-config` whose `pipeline_parallel_size` is 2 or more, and `--node-role prefill`, `decode` and @@ -375,19 +387,39 @@ in the log: without a rejection record. - *Closed before the handshake*: a connect-and-close liveness probe (the router health monitor, `probe_peers`) is logged at debug only. - -The handshake wire protocol is versioned (`PROTOCOL_VERSION` 3 since #2656, -which added the frame protection suite and the derived per-frame keys); a pair -of processes built from different versions fails with an explicit version -message rather than a confusing I/O error. Mixed versions in one cluster are -not supported. +- *At capacity*: a peer that authenticates while the accepting node already + holds its `max_connections` authenticated connections is answered with the + at-capacity verdict inside the handshake. The connecting side fails with + "server authenticated this node but refused the connection: it already + holds its maximum number of authenticated connections; retry later", which + the connection pool and `Transport::connect` treat as retryable (a warning, + not a fatal rejection); the accepting side logs "max connections (N) + reached". The verdict is sent only after the client's authenticator + verified, so a party without the secret never learns the capacity state. + +These warnings are rate limited in two classes so that junk connections (bad +magic, a malformed hello, a stall to the 5 second timeout, a pending-pool +refusal) cannot push the typed refusals (probable secret mismatch, identity +not owned, suite mismatch, a frame that failed authentication) down to debug: +each class may log 5 warnings per minute, the rest go to debug, and the first +warning of a new minute reports how many lines of its class were suppressed +in the previous one (`N cluster handshake warnings (diagnostic) were logged +at debug in the last 60s`). + +The handshake wire protocol is versioned (`PROTOCOL_VERSION` 4: 3 added the +frame protection suite and the derived per-frame keys in #2656, 4 added the +at-capacity verdict in #2692); a pair of processes built from different +versions fails with an explicit version message rather than a confusing I/O +error. Mixed versions in one cluster are not supported. Every version bump so +far landed after v0.8.0-beta.1, so no released version is affected. Residual risks: - Per-frame protection stops an on-path attacker from reading, altering, injecting, replaying or reordering frames on an authenticated connection; it does not stop one from dropping or delaying them. A dropped frame surfaces as the existing stage timeout or absolute deadline, exactly like a lost connection. - Frame lengths and kinds (`TensorData`, `Control`, RPC request or response) are visible on the wire, as is the timing of every decode step. An observer can tell how many tokens a cluster is producing even though it cannot read any of them. -- Frames from authenticated members stay trusted inside the pipeline and disaggregated protocols: a host that holds the secret is a full member. Rotate the secret when a member is decommissioned. -- Connections that have not finished the handshake do not count against the 512 authenticated-connection cap. They draw from a separate pending pool (64 in flight, at most 8 per source IP, each for at most the 5 second handshake timeout); a socket beyond either bound is closed at accept. A flood of unauthenticated connections from enough source IPs (8 or more at the defaults) can fill the pending pool and so delay new connections from any source, including a real peer reconnecting, by up to the 5 second timeout per slot while the flood lasts; it cannot take slots from established authenticated connections or inject frames. Failed inbound handshakes and pending-pool refusals are logged at warn for the first few per minute and at debug afterwards. +- Frames from authenticated members stay trusted inside the pipeline protocol and on the prefill and decode intakes of the disaggregated protocol: a host that holds the secret is a full member there. The router's result demux is the exception (it accepts frames only from its configured peers). Rotate the secret when a member is decommissioned. +- Connections that have not finished the handshake do not count against the 512 authenticated-connection cap. They draw from a separate pending pool (64 in flight, at most 8 per source, each for at most the 5 second handshake timeout); a socket beyond either bound is closed at accept. A source is an IPv4 address or an IPv6 /64 (a host that owns a /64 would otherwise get one pending slot per address), which also means cluster nodes inside one IPv6 /64 share the per-source budget and a large IPv6 cluster in a single /64 may need a higher `max_pending_handshakes_per_ip`. A flood of unauthenticated connections from enough sources (8 or more at the defaults) can fill the pending pool and so delay new connections from any source, including a real peer reconnecting, by up to the 5 second timeout per slot while the flood lasts; it cannot take slots from established authenticated connections or inject frames. Failed inbound handshakes and pending-pool refusals are logged at warn for the first few per minute (per class, see above) and at debug afterwards. +- A party without the secret can still produce the "probable secret mismatch" warning (send a well-formed hello naming its own IP, then close after the server's authenticator), so it can saturate the diagnostic warning budget too. The per-minute summary of suppressed warnings is what keeps the operator informed then. - The receiver refuses an oversize `len` from the header alone, but a frame whose `len` is within the cap (256 MiB by default) makes it allocate and read the full body before the tag is checked. Only a peer that completed the handshake can reach this path, so it is a resource bound for authenticated members, not an exposure to the open network. - `MLXCEL_CLUSTER_SECRET` is inherited by child processes of the server. Prefer `--cluster-secret-file` or `MLXCEL_CLUSTER_SECRET_FILE` so the value never sits in the environment. - The server identity is not bound to the socket IP on the connecting side, so a stage address that is reached through NAT or a port forward still authenticates by the secret alone. @@ -410,7 +442,7 @@ What each layer covers: | Layer | Covers | Does not cover | |-------|--------|----------------| -| Connect-time handshake (`PROTOCOL_VERSION` 3) | Only a holder of the secret can open or accept a connection; the accepted identity is bound to the TCP source IP; the frame suite is agreed inside the authenticated transcript. | Anything that happens on the TCP stream after the handshake. | +| Connect-time handshake (`PROTOCOL_VERSION` 4) | Only a holder of the secret can open or accept a connection; the accepted identity is bound to the TCP source IP; the frame suite is agreed inside the authenticated transcript; a full connection cap is reported to an authenticated peer as a typed, retryable verdict. | Anything that happens on the TCP stream after the handshake. | | Per-frame protection (#2656) | Every frame after the handshake: confidentiality of the payload, integrity of header and payload, replay and reorder rejection through an implicit sequence number, binding to the connection whose handshake derived the keys. Defeats the relay, which cannot derive keys without the secret. Covers pooled outbound connections, `send_stream`, RPC requests and responses, the Thunderbolt and RDMA wrappers (same TCP core), the disaggregated router, prefill and decode transports, and the `--pp-auto` stage assignment channel. | Dropping or delaying frames; frame lengths, kinds and timing. | | Discovery beacons (`--cluster-discovery=mdns`) | Each UDP beacon carries an HMAC, a timestamp and a nonce; the coordinator rejects an invalid MAC, a 60 s clock skew, a replay inside that window, another mlxcel version, or an advertised address that is not the datagram's source. | Hiding the fact that a cluster is forming; beacon payloads are readable. | @@ -427,8 +459,9 @@ sequence number (starting at 0, never sent, never wrapped: a sender that reaches the end closes the connection), and the clear-text `kind || len` header is the associated data, so the receiver refuses an oversize `len` before any decryption work. A frame that fails to decrypt closes the -connection; the event is logged through the same rate limiter as a failed -handshake, never with payload bytes. Keys are per connection, so a frame +connection; the event is logged through the diagnostic class of the same +rate limiter as a failed handshake, never with payload bytes. Keys are per +connection, so a frame recorded on one connection fails on every other. The suite is announced in the handshake hello and covered by both authenticators; a node configured for a different suite is refused with an explicit record. Nothing new is diff --git a/src/bin/mlx_server.rs b/src/bin/mlx_server.rs index 56a4ad20f..dfba807a1 100644 --- a/src/bin/mlx_server.rs +++ b/src/bin/mlx_server.rs @@ -1403,13 +1403,19 @@ struct ServerArgs { peers: Vec, /// Comma-separated prefill-node addresses. Decode nodes use this to identify - /// accepted handoff sources; routers use it to select a prefill target. + /// accepted handoff sources; routers use it to select a prefill target and + /// accept result frames only from these addresses, so each entry must be + /// the address the prefill node connects from (its `--serving-bind`, with + /// an unspecified bind IP replaced by its source IP). /// Consumed when `--node-role decode` or `--node-role router`. #[arg(long, value_delimiter = ',', value_name = "ADDR")] prefill_peers: Vec, /// Comma-separated decode-node addresses. Prefill nodes hand KV state to one - /// of these targets; routers use it to route decode continuations. + /// of these targets; routers use it to route decode continuations and + /// accept result frames only from these addresses, so each entry must be + /// the address the decode node connects from (its `--serving-bind`, with + /// an unspecified bind IP replaced by its source IP). /// Consumed when `--node-role prefill` or `--node-role router`. #[arg(long, value_delimiter = ',', value_name = "ADDR")] decode_peers: Vec, diff --git a/src/distributed/cluster_handshake.rs b/src/distributed/cluster_handshake.rs index a2c451b69..eca2e75db 100644 --- a/src/distributed/cluster_handshake.rs +++ b/src/distributed/cluster_handshake.rs @@ -26,7 +26,8 @@ //! server -> client : MAGIC(8) | version u16 | suite u8 | nonce_s [32] | id_len u16 | identity //! | tag_s = HMAC(K, "mlxcel-cluster-auth-v1/server", transcript) //! client -> server : tag_c = HMAC(K, "mlxcel-cluster-auth-v1/client", transcript || tag_s) -//! server -> client : verdict u8 (0x01 accept; anything else or EOF is a rejection) +//! server -> client : verdict u8 (0x01 accept; 0x02 at capacity, retry later; +//! anything else or EOF is a rejection) //! ``` //! //! `identity` is the sender's listener address (`host:port`). The accept side @@ -71,8 +72,17 @@ //! [`IdentityNotOwned`] (wrong bind or advertised address, or a NAT between //! the nodes), [`ClientAbortedAfterServerProof`] (the client most likely holds //! a different secret: it verified this server's authenticator, failed, and -//! closed), and [`HandshakeRejected`] (a client authenticator that does not -//! verify although the server's own did, i.e. the transcript was altered). +//! closed), [`HandshakeRejected`] (a client authenticator that does not +//! verify although the server's own did, i.e. the transcript was altered), +//! and [`ConnectionLimitReached`] (the client authenticated but the listener +//! is at its connection cap). +//! +//! Admission is part of the handshake (#2692): the accept side's `admit` +//! callback ([`server_handshake_admitting`]) runs only after the client's +//! authenticator verified, and a `None` from it is answered with the +//! `0x02` verdict, so the client gets a typed, retryable [`ServerAtCapacity`] +//! instead of an accepted handshake followed by a closed socket, and a party +//! without the secret never learns the capacity state. //! //! Used by: tcp_transport (accept side), connection_pool (connect side), //! stage_assignment (both sides). Both entry points take a bare `TcpStream`, @@ -93,9 +103,11 @@ const MAGIC: &[u8; 8] = b"MLXCAUTH"; /// Wire version of the handshake. 2 added the identity rejection record /// (#2655); 3 added the frame suite byte and the derived per-frame keys /// (#2656), so a build without frame protection fails with the explicit -/// version message (mixed versions are unsupported; discovery beacons already -/// pin `CARGO_PKG_VERSION`). -pub const PROTOCOL_VERSION: u16 = 3; +/// version message; 4 added the at-capacity verdict (#2692). Mixed versions +/// are unsupported (discovery beacons already pin `CARGO_PKG_VERSION`), and +/// every bump so far landed after v0.8.0-beta.1, so no released version is +/// affected. +pub const PROTOCOL_VERSION: u16 = 4; const NONCE_LEN: usize = 32; const TAG_LEN: usize = 32; const MAX_IDENTITY_LEN: usize = 255; @@ -106,6 +118,10 @@ const SERVER_LABEL: &[u8] = b"mlxcel-cluster-auth-v1/server"; const CLIENT_LABEL: &[u8] = b"mlxcel-cluster-auth-v1/client"; const VERDICT_ACCEPT: u8 = 0x01; const VERDICT_REJECT: u8 = 0x00; +/// The client authenticated but the server is at its connection cap. Sent +/// only after the client's authenticator verified, so it tells nothing to a +/// party without the secret. Retryable on the client side. +const VERDICT_AT_CAPACITY: u8 = 0x02; /// First 8 bytes of the server's reply when it refuses the client's advertised /// identity or frame suite instead of sending a server hello. Followed by one @@ -174,6 +190,30 @@ pub struct FrameSuiteMismatch { )] pub struct ClientAbortedAfterServerProof; +/// Connect side: the server verified this node's authenticator and then +/// answered with the at-capacity verdict. Unlike [`HandshakeRejected`] this is +/// transient: the connection pool and `Transport::connect` treat it as +/// retryable (warn and continue), not as a configuration error. +#[derive(Debug, thiserror::Error)] +#[error( + "server authenticated this node but refused the connection: it already holds its maximum \ + number of authenticated connections; retry later" +)] +pub struct ServerAtCapacity; + +/// Accept side: the client authenticated, but the caller's `admit` callback +/// ([`server_handshake_admitting`]) found no free connection slot. The client +/// was told so with the at-capacity verdict. +#[derive(Debug, thiserror::Error)] +#[error( + "authenticated peer {peer} refused: this node already holds its maximum number of \ + authenticated connections" +)] +pub struct ConnectionLimitReached { + /// The client's bound listener identity. + pub peer: String, +} + /// Bind an advertised listener identity to the TCP source IP it arrived from. /// /// Returns the identity the transport should report for the connection. A @@ -426,13 +466,17 @@ async fn client_handshake_inner( .read_exact(&mut verdict) .await .context("server closed the connection instead of accepting the cluster handshake")?; - if verdict[0] != VERDICT_ACCEPT { - return Err(HandshakeRejected( - "server rejected this node's authenticator although the server's own authenticator \ - verified; the handshake was altered in transit" - .to_string(), - ) - .into()); + match verdict[0] { + VERDICT_ACCEPT => {} + VERDICT_AT_CAPACITY => return Err(ServerAtCapacity.into()), + _ => { + return Err(HandshakeRejected( + "server rejected this node's authenticator although the server's own \ + authenticator verified; the handshake was altered in transit" + .to_string(), + ) + .into()); + } } Ok(SessionKeys::derive( auth, @@ -470,21 +514,44 @@ pub async fn server_handshake_on( suite: FrameSuite, channel: FrameChannel, ) -> Result<(String, SessionKeys)> { + server_handshake_admitting(stream, auth, local_identity, suite, channel, || Some(())) + .await + .map(|(identity, keys, ())| (identity, keys)) +} + +/// [`server_handshake_on`] with admission inside the handshake (#2692). +/// +/// `admit` is called exactly once, and only after the client's authenticator +/// verified, so a party without the secret never learns whether this node has +/// a free slot. `Some(permit)` accepts the connection and the permit is +/// returned with the identity and keys; `None` answers the client with the +/// at-capacity verdict (the client fails with [`ServerAtCapacity`]) and fails +/// here with [`ConnectionLimitReached`]. A listener that admits everything +/// passes `|| Some(())`. +pub(crate) async fn server_handshake_admitting

( + stream: &mut TcpStream, + auth: &ClusterAuth, + local_identity: &str, + suite: FrameSuite, + channel: FrameChannel, + admit: impl FnOnce() -> Option

, +) -> Result<(String, SessionKeys, P)> { tokio::time::timeout( HANDSHAKE_TIMEOUT, - server_handshake_inner(stream, auth, local_identity, suite, channel), + server_handshake_inner(stream, auth, local_identity, suite, channel, admit), ) .await .map_err(|_| anyhow::anyhow!("cluster handshake timed out after {HANDSHAKE_TIMEOUT:?}"))? } -async fn server_handshake_inner( +async fn server_handshake_inner

( stream: &mut TcpStream, auth: &ClusterAuth, local_identity: &str, suite: FrameSuite, channel: FrameChannel, -) -> Result<(String, SessionKeys)> { + admit: impl FnOnce() -> Option

, +) -> Result<(String, SessionKeys, P)> { let client_hello = read_hello(stream, HelloFrom::Client).await?; let source_ip = stream.peer_addr().context("peer address")?.ip(); // Bind before replying so a client that lies about its address never @@ -538,6 +605,14 @@ async fn server_handshake_inner( ) .into()); } + // Admission comes after authentication: only a holder of the secret + // learns whether this node is at its connection cap, and it learns it + // as a typed verdict instead of an accept followed by a closed socket. + let Some(permit) = admit() else { + let _ = stream.write_all(&[VERDICT_AT_CAPACITY]).await; + let _ = stream.flush().await; + return Err(ConnectionLimitReached { peer: bound }.into()); + }; stream.write_all(&[VERDICT_ACCEPT]).await?; stream.flush().await?; let keys = SessionKeys::derive( @@ -548,7 +623,7 @@ async fn server_handshake_inner( &server_nonce, HandshakeRole::Server, ); - Ok((bound, keys)) + Ok((bound, keys, permit)) } /// Whether an I/O error (possibly wrapped in context) is the peer ending the diff --git a/src/distributed/cluster_handshake_tests.rs b/src/distributed/cluster_handshake_tests.rs index e8933a854..608abf5f2 100644 --- a/src/distributed/cluster_handshake_tests.rs +++ b/src/distributed/cluster_handshake_tests.rs @@ -479,3 +479,71 @@ async fn handshake_rejects_frame_suite_mismatch_with_typed_errors() { .unwrap_or_else(|| panic!("expected HandshakeRejected, got: {client_err:#}")); assert!(rejected.0.contains("frame protection suite"), "{rejected}"); } + +/// A full connection cap is reported inside the handshake, after the client +/// authenticated: the client gets the typed, retryable `ServerAtCapacity` +/// (not `HandshakeRejected`) and the server `ConnectionLimitReached` naming +/// the peer it turned away. +#[tokio::test] +async fn server_at_capacity_sends_capacity_verdict() { + let auth = ClusterAuth::new(&ClusterSecret::random()); + let (mut server, mut client) = pair().await; + let (server_res, client_res) = tokio::join!( + server_handshake_admitting::<()>( + &mut server, + &auth, + "127.0.0.1:1", + FrameSuite::default(), + FrameChannel::Transport, + || None, + ), + client_handshake(&mut client, &auth, "127.0.0.1:2", FrameSuite::default()), + ); + let client_err = client_res.unwrap_err(); + assert!( + client_err.downcast_ref::().is_some(), + "expected ServerAtCapacity, got: {client_err:#}" + ); + assert!(client_err.downcast_ref::().is_none()); + assert!( + client_err.to_string().contains("retry later"), + "{client_err}" + ); + let server_err = server_res.unwrap_err(); + let full = server_err + .downcast_ref::() + .unwrap_or_else(|| panic!("expected ConnectionLimitReached, got: {server_err:#}")); + assert_eq!(full.peer, "127.0.0.1:2"); +} + +/// `admit` runs only after the client's authenticator verified: a client with +/// another secret never triggers it, so it learns nothing about capacity. +#[tokio::test] +async fn admission_is_not_consulted_before_the_client_authenticates() { + let server_auth = ClusterAuth::new(&ClusterSecret::random()); + let client_auth = ClusterAuth::new(&ClusterSecret::random()); + let (mut server, mut client) = pair().await; + let admitted = std::sync::atomic::AtomicBool::new(false); + let (server_res, client_res) = tokio::join!( + server_handshake_admitting::<()>( + &mut server, + &server_auth, + "127.0.0.1:1", + FrameSuite::default(), + FrameChannel::Transport, + || { + admitted.store(true, std::sync::atomic::Ordering::SeqCst); + None + }, + ), + client_handshake( + &mut client, + &client_auth, + "127.0.0.1:2", + FrameSuite::default() + ), + ); + assert!(server_res.is_err()); + assert!(client_res.is_err()); + assert!(!admitted.load(std::sync::atomic::Ordering::SeqCst)); +} diff --git a/src/distributed/handshake_warn.rs b/src/distributed/handshake_warn.rs new file mode 100644 index 000000000..5d0728e38 --- /dev/null +++ b/src/distributed/handshake_warn.rs @@ -0,0 +1,198 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Rate limiter for the warnings an authenticated listener logs about +//! connections it refused (#2655, #2692). +//! +//! A party without the cluster secret can open sockets at will, so every +//! refusal it causes (a pending-pool refusal, a hello with bad magic, a stall +//! to the handshake timeout) is logged at warn only a few times per minute +//! and at debug afterwards. The budget is split in two classes so that a +//! trickle of such junk cannot push the one line an operator needs down to +//! debug: +//! +//! - [`HandshakeWarnClass::Noise`]: anything a stranger can produce without +//! the secret and without completing the protocol (pending-pool refusals, +//! malformed or absent hellos, timeouts, unparseable identities). +//! - [`HandshakeWarnClass::Diagnostic`]: the typed refusals that tell the +//! operator what to fix: a probable secret mismatch +//! ([`ClientAbortedAfterServerProof`], [`HandshakeRejected`]), an identity +//! or suite that does not match ([`IdentityNotOwned`], +//! [`FrameSuiteMismatch`]), and a frame that failed authentication after the +//! handshake ([`FrameRejected`]). +//! +//! Each class has its own window of [`HANDSHAKE_WARN_PER_WINDOW`] warnings +//! per [`HANDSHAKE_WARN_WINDOW`]. When a window rolls over and the previous +//! one suppressed lines, the first call of the new window also logs how many +//! were suppressed, so a sustained flood still leaves one summary line per +//! minute in the log. +//! +//! Residual risk: a party without the secret can still produce +//! `ClientAbortedAfterServerProof` (send a well-formed hello naming its own +//! IP, then close after the server proof), so it can saturate the +//! `Diagnostic` window too. The rollover summary is what keeps the operator +//! informed then. Warning on the first event per source instead would need a +//! per-source map that an IPv6 source defeats by rotating addresses. +//! +//! Used by: `tcp_transport::accept_loop` and `handle_connection`, +//! `stage_assignment` (its listener and rejected-frame paths), and the +//! disaggregated router's result demux (`server::router_front`). + +use std::fmt; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +use super::cluster_handshake::{ + ClientAbortedAfterServerProof, FrameSuiteMismatch, HandshakeRejected, IdentityNotOwned, +}; +use super::frame_protection::FrameRejected; + +/// Warnings each class may log at warn per window; the rest go to debug. +pub(crate) const HANDSHAKE_WARN_PER_WINDOW: u32 = 5; +/// Length of one warning window. +pub(crate) const HANDSHAKE_WARN_WINDOW: Duration = Duration::from_secs(60); + +/// Which budget a refusal draws from. See the module doc. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum HandshakeWarnClass { + /// Producible by anyone without the secret; bounded tightly. + Noise, + /// A typed refusal the operator needs to see. + Diagnostic, +} + +impl HandshakeWarnClass { + fn index(self) -> usize { + match self { + Self::Noise => 0, + Self::Diagnostic => 1, + } + } +} + +impl fmt::Display for HandshakeWarnClass { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Noise => "noise", + Self::Diagnostic => "diagnostic", + }) + } +} + +/// Classify a failed inbound handshake for the warning budget: `Diagnostic` +/// when the error chain holds one of the typed refusals, `Noise` otherwise +/// (bad magic, a malformed hello, a timeout, an I/O error). +pub(crate) fn handshake_warn_class(err: &anyhow::Error) -> HandshakeWarnClass { + let diagnostic = err.chain().any(|cause| { + cause + .downcast_ref::() + .is_some() + || cause.downcast_ref::().is_some() + || cause.downcast_ref::().is_some() + || cause.downcast_ref::().is_some() + || cause.downcast_ref::().is_some() + }); + if diagnostic { + HandshakeWarnClass::Diagnostic + } else { + HandshakeWarnClass::Noise + } +} + +/// What [`WarnBudget::allow_at`] decided for one event. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct WarnDecision { + /// Log this event at warn (`true`) or at debug (`false`). + pub warn: bool, + /// Events of the same class that were logged at debug in the window that + /// just ended. Non-zero only on the first event of a new window. + pub suppressed_in_previous_window: u32, +} + +/// One window of a class: when it started and how many events it has seen. +#[derive(Debug, Clone, Copy)] +struct Window { + start: Instant, + count: u32, +} + +/// Per-class warning windows. Thread-safe; one instance is process-wide +/// ([`handshake_warn_allowed`]), and tests build their own to inject time. +#[derive(Debug)] +pub(crate) struct WarnBudget { + windows: Mutex<[Option; 2]>, + per_window: u32, + window: Duration, +} + +impl WarnBudget { + /// A budget of `per_window` warnings per `window` for each class. + pub(crate) const fn new(per_window: u32, window: Duration) -> Self { + Self { + windows: Mutex::new([None, None]), + per_window, + window, + } + } + + /// Account one event of `class` happening at `now`. + pub(crate) fn allow_at(&self, class: HandshakeWarnClass, now: Instant) -> WarnDecision { + let mut windows = self.windows.lock().unwrap_or_else(|e| e.into_inner()); + let slot = &mut windows[class.index()]; + match slot { + Some(window) if now.saturating_duration_since(window.start) < self.window => { + window.count = window.count.saturating_add(1); + WarnDecision { + warn: window.count <= self.per_window, + suppressed_in_previous_window: 0, + } + } + _ => { + let suppressed = slot + .map(|window| window.count.saturating_sub(self.per_window)) + .unwrap_or(0); + *slot = Some(Window { + start: now, + count: 1, + }); + WarnDecision { + warn: true, + suppressed_in_previous_window: suppressed, + } + } + } + } +} + +static BUDGET: WarnBudget = WarnBudget::new(HANDSHAKE_WARN_PER_WINDOW, HANDSHAKE_WARN_WINDOW); + +/// Whether a refusal of `class` may be logged at warn right now (debug +/// otherwise), against the process-wide budget. When the class's window has +/// just rolled over and the previous window suppressed lines, one summary +/// warn with their count is logged here. +pub(crate) fn handshake_warn_allowed(class: HandshakeWarnClass) -> bool { + let decision = BUDGET.allow_at(class, Instant::now()); + if decision.suppressed_in_previous_window > 0 { + tracing::warn!( + "{} cluster handshake warnings ({class}) were logged at debug in the last {}s", + decision.suppressed_in_previous_window, + HANDSHAKE_WARN_WINDOW.as_secs() + ); + } + decision.warn +} + +#[cfg(test)] +#[path = "handshake_warn_tests.rs"] +mod tests; diff --git a/src/distributed/handshake_warn_tests.rs b/src/distributed/handshake_warn_tests.rs new file mode 100644 index 000000000..03e8aeddd --- /dev/null +++ b/src/distributed/handshake_warn_tests.rs @@ -0,0 +1,145 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::net::{IpAddr, Ipv4Addr}; + +use super::*; +use crate::distributed::cluster_handshake::HANDSHAKE_TIMEOUT; +use crate::distributed::frame_protection::FrameSuite; + +fn budget() -> WarnBudget { + WarnBudget::new(HANDSHAKE_WARN_PER_WINDOW, HANDSHAKE_WARN_WINDOW) +} + +/// Junk connections cannot push the typed refusals down to debug: each class +/// has its own window. +#[test] +fn noise_flood_does_not_suppress_diagnostic_warnings() { + let budget = budget(); + let now = Instant::now(); + let mut noise_warned = 0; + for i in 0..50u64 { + let at = now + Duration::from_millis(i * 10); + if budget.allow_at(HandshakeWarnClass::Noise, at).warn { + noise_warned += 1; + } + } + assert_eq!(noise_warned, HANDSHAKE_WARN_PER_WINDOW); + let decision = budget.allow_at(HandshakeWarnClass::Diagnostic, now + Duration::from_secs(1)); + assert_eq!( + decision, + WarnDecision { + warn: true, + suppressed_in_previous_window: 0, + } + ); + // And the diagnostic window has its own full budget. + let mut diagnostic_warned = 1; + for i in 0..10u64 { + let at = now + Duration::from_secs(2) + Duration::from_millis(i); + if budget.allow_at(HandshakeWarnClass::Diagnostic, at).warn { + diagnostic_warned += 1; + } + } + assert_eq!(diagnostic_warned, HANDSHAKE_WARN_PER_WINDOW); +} + +/// The first event after a window rolls over reports how many events of its +/// class the previous window logged at debug. +#[test] +fn rollover_reports_suppressed_count() { + let budget = budget(); + let now = Instant::now(); + for i in 0..8u64 { + let decision = budget.allow_at(HandshakeWarnClass::Noise, now + Duration::from_secs(i)); + assert_eq!(decision.warn, i < u64::from(HANDSHAKE_WARN_PER_WINDOW)); + assert_eq!(decision.suppressed_in_previous_window, 0); + } + let later = budget.allow_at(HandshakeWarnClass::Noise, now + Duration::from_secs(61)); + assert_eq!( + later, + WarnDecision { + warn: true, + suppressed_in_previous_window: 3, + } + ); + // The summary is reported once: the next event in the new window carries + // no count, and a window that suppressed nothing reports zero. + let next = budget.allow_at(HandshakeWarnClass::Noise, now + Duration::from_secs(62)); + assert_eq!(next.suppressed_in_previous_window, 0); + let quiet = budget.allow_at(HandshakeWarnClass::Noise, now + Duration::from_secs(200)); + assert_eq!( + quiet, + WarnDecision { + warn: true, + suppressed_in_previous_window: 0, + } + ); + // The other class's window is untouched by the rollover. + let other = budget.allow_at( + HandshakeWarnClass::Diagnostic, + now + Duration::from_secs(200), + ); + assert_eq!( + other, + WarnDecision { + warn: true, + suppressed_in_previous_window: 0, + } + ); +} + +/// The typed refusals are diagnostic, also when wrapped in context; anything +/// a stranger can produce without the secret is noise. +#[test] +fn classifies_handshake_errors() { + let diagnostic: Vec = vec![ + ClientAbortedAfterServerProof.into(), + HandshakeRejected("secret mismatch".to_string()).into(), + IdentityNotOwned { + advertised: "10.0.0.9:1".to_string(), + source_ip: IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), + } + .into(), + FrameSuiteMismatch { + peer_suite_id: 0, + local: FrameSuite::default(), + } + .into(), + FrameRejected.into(), + anyhow::Error::from(HandshakeRejected("wrapped".to_string())) + .context("cluster handshake with 10.0.0.9:1 failed"), + ]; + for err in &diagnostic { + assert_eq!( + handshake_warn_class(err), + HandshakeWarnClass::Diagnostic, + "{err:#}" + ); + } + let noise: Vec = vec![ + anyhow::anyhow!("cluster handshake timed out after {HANDSHAKE_TIMEOUT:?}"), + anyhow::anyhow!("bad magic"), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "eof").into(), + ]; + for err in &noise { + assert_eq!( + handshake_warn_class(err), + HandshakeWarnClass::Noise, + "{err:#}" + ); + } + assert_eq!(HandshakeWarnClass::Noise.to_string(), "noise"); + assert_eq!(HandshakeWarnClass::Diagnostic.to_string(), "diagnostic"); +} diff --git a/src/distributed/mod.rs b/src/distributed/mod.rs index c5a6ef878..e866914a8 100644 --- a/src/distributed/mod.rs +++ b/src/distributed/mod.rs @@ -57,6 +57,7 @@ pub mod discovery_beacon; pub mod failure_detector; pub mod frame_protection; pub mod handoff_queue; +pub(crate) mod handshake_warn; pub mod heartbeat; pub mod kv_cache_serde; pub mod kv_cache_transfer; @@ -101,9 +102,9 @@ pub use cluster_auth::{ MIN_CLUSTER_SECRET_BYTES, }; pub use cluster_handshake::{ - ClientAbortedAfterServerProof, FrameSuiteMismatch, HANDSHAKE_TIMEOUT, HandshakeRejected, - IdentityNotOwned, PeerClosedBeforeHandshake, client_handshake, client_handshake_on, - server_handshake, server_handshake_on, + ClientAbortedAfterServerProof, ConnectionLimitReached, FrameSuiteMismatch, HANDSHAKE_TIMEOUT, + HandshakeRejected, IdentityNotOwned, PeerClosedBeforeHandshake, ServerAtCapacity, + client_handshake, client_handshake_on, server_handshake, server_handshake_on, }; pub use cluster_init::{ ClusterDiscoveryMode, ClusterInitPlan, ClusterInitRequest, DEFAULT_CONTROL_BASE_PORT, diff --git a/src/distributed/pending_handshakes.rs b/src/distributed/pending_handshakes.rs index 8f1c19dc7..5f080bcbd 100644 --- a/src/distributed/pending_handshakes.rs +++ b/src/distributed/pending_handshakes.rs @@ -24,11 +24,16 @@ //! While the pending pool is full, a new connection from any source //! (including a real peer reconnecting) is closed at accept, so a flood from //! enough source IPs delays new connections until slots free up, at most -//! [`HANDSHAKE_TIMEOUT`] per slot. A per-IP sub-limit keeps one source from -//! taking the whole pending pool; it is not sufficient on its own (real -//! stages and an attacker can share a NAT IP, and an IPv6 source can rotate -//! addresses within its prefix), which is why the global pending pool is -//! separate from the main cap. +//! [`HANDSHAKE_TIMEOUT`] per slot. A per-source sub-limit keeps one source +//! from taking the whole pending pool. The source key is the canonical IPv4 +//! address, or the /64 prefix of an IPv6 address ([`pending_source_key`]): a +//! host that owns a /64 can otherwise open one pending slot per address and +//! bypass the per-source bound by rotating addresses (#2692). The flip side +//! is that cluster nodes inside one IPv6 /64 share one per-source budget, so +//! a large IPv6 cluster in a single /64 may need a higher +//! `max_pending_handshakes_per_ip`. The sub-limit is still not sufficient on +//! its own (real stages and an attacker can share a NAT IP or a /64), which +//! is why the global pending pool is separate from the main cap. //! //! Used by: `tcp_transport::accept_loop`. The type is `pub(crate)` with a //! small API so another authenticated listener inside `distributed` (for @@ -37,22 +42,40 @@ //! [`HANDSHAKE_TIMEOUT`]: super::cluster_handshake::HANDSHAKE_TIMEOUT use std::collections::HashMap; -use std::net::IpAddr; +use std::net::{IpAddr, Ipv6Addr}; use std::sync::{Arc, Mutex}; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; /// Default number of handshakes in flight across all sources. pub const DEFAULT_MAX_PENDING_HANDSHAKES: usize = 64; -/// Default number of handshakes in flight from one source IP. +/// Default number of handshakes in flight from one source (an IPv4 address +/// or an IPv6 /64, see [`pending_source_key`]). pub const DEFAULT_MAX_PENDING_HANDSHAKES_PER_IP: usize = 8; +/// The key the per-source pending count is kept under. +/// +/// The address is canonicalized first, so an IPv4-mapped IPv6 source keeps +/// counting as its IPv4 address (the same key the handshake binds identities +/// to). An IPv4 address is its own key. An IPv6 address is masked to its /64: +/// a single host or tenant normally holds a whole /64 and can pick any +/// address inside it, so counting per full address would let it take one +/// pending slot per address. +pub(crate) fn pending_source_key(ip: IpAddr) -> IpAddr { + match ip.to_canonical() { + v4 @ IpAddr::V4(_) => v4, + IpAddr::V6(v6) => IpAddr::V6(Ipv6Addr::from(u128::from(v6) & !((1u128 << 64) - 1))), + } +} + /// Why a socket was refused a pending-handshake slot. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum PendingSlotRefusal { /// The global pending pool is full. Global { max: usize }, - /// This source IP already holds its share of the pending pool. + /// This source already holds its share of the pending pool. `ip` is the + /// source key ([`pending_source_key`]): the IPv4 address, or the /64 + /// prefix of the IPv6 source. PerIp { ip: IpAddr, max: usize }, } @@ -60,9 +83,18 @@ impl std::fmt::Display for PendingSlotRefusal { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Global { max } => write!(f, "max pending handshakes ({max}) reached"), - Self::PerIp { ip, max } => { + Self::PerIp { + ip: ip @ IpAddr::V4(_), + max, + } => { write!(f, "max pending handshakes from {ip} ({max}) reached") } + Self::PerIp { + ip: ip @ IpAddr::V6(_), + max, + } => { + write!(f, "max pending handshakes from {ip}/64 ({max}) reached") + } } } } @@ -101,13 +133,14 @@ impl PendingHandshakeLimiter { /// Reserve a slot for a socket from `peer_ip`, or say why it cannot have /// one. Never waits: a full pool means the socket is dropped at accept. /// - /// IPv4-mapped IPv6 sources count against their IPv4 key - /// (`to_canonical`), the same key the handshake binds identities to. + /// The per-source count is keyed by [`pending_source_key`]: IPv4-mapped + /// IPv6 sources count against their IPv4 address and IPv6 sources against + /// their /64. pub(crate) fn try_acquire( &self, peer_ip: IpAddr, ) -> Result { - let ip = peer_ip.to_canonical(); + let ip = pending_source_key(peer_ip); let per_ip_guard = { let mut counts = self.per_ip.lock().unwrap_or_else(|e| e.into_inner()); let count = counts.entry(ip).or_insert(0); @@ -137,13 +170,14 @@ impl PendingHandshakeLimiter { }) } - /// Current number of pending handshakes from `peer_ip` (tests). + /// Current number of pending handshakes from `peer_ip`'s source key + /// (tests). #[cfg(test)] pub(crate) fn pending_from(&self, peer_ip: IpAddr) -> usize { self.per_ip .lock() .unwrap_or_else(|e| e.into_inner()) - .get(&peer_ip.to_canonical()) + .get(&pending_source_key(peer_ip)) .copied() .unwrap_or(0) } @@ -163,8 +197,8 @@ pub(crate) struct PendingHandshakeSlot { _per_ip: PerIpGuard, } -/// RAII decrement of one source IP's pending count; removes the entry when it -/// reaches zero so the map never grows with past sources. +/// RAII decrement of one source key's pending count; removes the entry when +/// it reaches zero so the map never grows with past sources. #[derive(Debug)] struct PerIpGuard { counts: Arc>>, diff --git a/src/distributed/pending_handshakes_tests.rs b/src/distributed/pending_handshakes_tests.rs index 39a134fd3..b216a2893 100644 --- a/src/distributed/pending_handshakes_tests.rs +++ b/src/distributed/pending_handshakes_tests.rs @@ -20,6 +20,10 @@ fn ip(last: u8) -> IpAddr { IpAddr::V4(Ipv4Addr::new(10, 0, 0, last)) } +fn ip6(s: &str) -> IpAddr { + s.parse().unwrap() +} + #[test] fn per_ip_limit_refuses_the_excess_and_releases_on_drop() { let limiter = PendingHandshakeLimiter::new(8, 2).unwrap(); @@ -66,6 +70,57 @@ fn ipv4_mapped_ipv6_counts_against_the_ipv4_key() { ); } +/// A host that owns a /64 cannot get one pending slot per address: every +/// address inside the /64 draws from the same per-source budget, and the +/// refusal names the /64 key. +#[test] +fn ipv6_sources_in_one_64_share_the_per_source_limit() { + let limiter = PendingHandshakeLimiter::new(64, 2).unwrap(); + let a = limiter.try_acquire(ip6("2001:db8:1:2::1")).unwrap(); + let _b = limiter.try_acquire(ip6("2001:db8:1:2::2")).unwrap(); + let refusal = limiter + .try_acquire(ip6("2001:db8:1:2:ffff::3")) + .unwrap_err(); + assert_eq!( + refusal, + PendingSlotRefusal::PerIp { + ip: ip6("2001:db8:1:2::"), + max: 2, + } + ); + assert_eq!( + refusal.to_string(), + "max pending handshakes from 2001:db8:1:2::/64 (2) reached" + ); + assert_eq!(limiter.pending_from(ip6("2001:db8:1:2::9")), 2); + // Dropping one slot frees the shared budget for any address in the /64. + drop(a); + let _c = limiter.try_acquire(ip6("2001:db8:1:2:ffff::3")).unwrap(); + assert_eq!(limiter.pending_from(ip6("2001:db8:1:2::")), 2); +} + +#[test] +fn ipv6_sources_in_different_64s_are_counted_separately() { + let limiter = PendingHandshakeLimiter::new(64, 1).unwrap(); + let _a = limiter.try_acquire(ip6("2001:db8:1:2::1")).unwrap(); + let _b = limiter.try_acquire(ip6("2001:db8:1:3::1")).unwrap(); + assert_eq!(limiter.pending_from(ip6("2001:db8:1:2::1")), 1); + assert_eq!(limiter.pending_from(ip6("2001:db8:1:3::1")), 1); + assert_eq!( + limiter.try_acquire(ip6("2001:db8:1:3::2")).unwrap_err(), + PendingSlotRefusal::PerIp { + ip: ip6("2001:db8:1:3::"), + max: 1, + } + ); + // IPv4 keys are unchanged and print without a prefix length. + assert_eq!(pending_source_key(ip(1)), ip(1)); + assert_eq!( + PendingSlotRefusal::PerIp { ip: ip(1), max: 1 }.to_string(), + "max pending handshakes from 10.0.0.1 (1) reached" + ); +} + #[test] fn zero_bounds_are_refused_at_construction() { for (total, per_ip) in [(0, 8), (64, 0), (0, 0)] { diff --git a/src/distributed/stage_assignment.rs b/src/distributed/stage_assignment.rs index 7b791955d..9e4755287 100644 --- a/src/distributed/stage_assignment.rs +++ b/src/distributed/stage_assignment.rs @@ -72,12 +72,12 @@ use super::cluster_init::{ }; use super::config::{ClusterConfig, NodeRole}; use super::frame_protection::{FrameChannel, FrameRejected, FrameSuite}; +use super::handshake_warn::{HandshakeWarnClass, handshake_warn_allowed, handshake_warn_class}; use super::pending_handshakes::PendingHandshakeLimiter; use super::pipeline::{ resolve_in_process_pipeline_num_layers, resolve_in_process_stage_assignments_for_model, }; use super::tcp_frame::ProtectedStream; -use super::tcp_transport::handshake_warn_allowed; use super::transport::MessageKind; /// Wire version of the assignment frames (independent of the handshake's). @@ -680,14 +680,15 @@ struct PeerShared { } /// Log a frame that failed authentication on the assignment channel through -/// the shared handshake rate limiter (warn, then debug once the window is -/// exhausted), without payload bytes. Returns whether `err` was such a -/// rejection; any other error is left to the caller's own log line. +/// the shared handshake rate limiter's diagnostic budget (warn, then debug +/// once the window is exhausted), without payload bytes. Returns whether +/// `err` was such a rejection; any other error is left to the caller's own +/// log line. fn log_rejected_frame(bound: &str, err: &anyhow::Error) -> bool { if err.downcast_ref::().is_none() { return false; } - if handshake_warn_allowed() { + if handshake_warn_allowed(HandshakeWarnClass::Diagnostic) { tracing::warn!("assignment listener: closing connection from {bound}: {err}"); } else { tracing::debug!("assignment listener: closing connection from {bound}: {err}"); @@ -717,7 +718,7 @@ async fn handle_coordinator_connection( Err(err) => { if err.downcast_ref::().is_some() { tracing::debug!("assignment listener: {err}"); - } else if handshake_warn_allowed() { + } else if handshake_warn_allowed(handshake_warn_class(&err)) { tracing::warn!( "assignment listener: rejecting connection from {}: {err:#}", source_ip.map(|ip| ip.to_string()).unwrap_or_default() @@ -974,7 +975,7 @@ pub async fn await_stage_assignment(config: PeerJoinConfig) -> Result { - if handshake_warn_allowed() { + if handshake_warn_allowed(HandshakeWarnClass::Noise) { tracing::warn!( "assignment listener: rejecting unauthenticated connection from {}: {refusal}", peer_addr.ip() diff --git a/src/distributed/tcp_transport.rs b/src/distributed/tcp_transport.rs index 2ead236ec..add0a4471 100644 --- a/src/distributed/tcp_transport.rs +++ b/src/distributed/tcp_transport.rs @@ -45,10 +45,16 @@ //! therefore delay new connections from any source while the pending pool is //! full (each slot frees within the handshake timeout), but cannot take slots //! from established authenticated connections or inject frames (#2655). A -//! zero bound is refused at `bind` when `auth` is set. A failed handshake is -//! logged with its class (closed before the handshake at debug; identity not -//! owned, probable secret mismatch, altered transcript at warn, rate -//! limited). +//! zero bound is refused at `bind` when `auth` is set. The per-source bound +//! counts IPv4 sources per address and IPv6 sources per /64 (#2692). A failed +//! handshake is logged with its class (closed before the handshake at debug; +//! identity not owned, probable secret mismatch, altered transcript at warn, +//! rate limited through [`super::handshake_warn`] with a budget per class so +//! junk connections cannot hide those typed refusals). The `max_connections` +//! check happens inside the handshake, after the peer authenticated: a peer +//! that finds this node full gets the typed at-capacity verdict +//! ([`super::cluster_handshake::ServerAtCapacity`], retryable) instead of an +//! accepted handshake followed by a closed socket. //! //! Without `auth` (the default, used by single-node serving) the transport //! performs no authentication, and the reported `from` is the raw socket @@ -74,7 +80,7 @@ use std::collections::HashMap; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, Ordering}; use anyhow::{Context, Result}; use bytes::Bytes; @@ -82,9 +88,13 @@ use tokio::net::TcpListener; use tokio::sync::{Mutex, Notify, Semaphore, mpsc}; use super::cluster_auth::ClusterAuth; -use super::cluster_handshake::{HandshakeRejected, PeerClosedBeforeHandshake, server_handshake}; +use super::cluster_handshake::{ + ConnectionLimitReached, HandshakeRejected, PeerClosedBeforeHandshake, + server_handshake_admitting, +}; use super::connection_pool::{ConnectionPool, PoolConfig}; -use super::frame_protection::{FrameRejected, FrameSuite}; +use super::frame_protection::{FrameChannel, FrameRejected, FrameSuite}; +use super::handshake_warn::{HandshakeWarnClass, handshake_warn_allowed, handshake_warn_class}; use super::pending_handshakes::PendingHandshakeLimiter; pub use super::pending_handshakes::{ DEFAULT_MAX_PENDING_HANDSHAKES, DEFAULT_MAX_PENDING_HANDSHAKES_PER_IP, @@ -107,7 +117,9 @@ pub struct TcpTransportConfig { /// Maximum allowed frame size in bytes (default 256 MiB). pub max_frame_size: usize, /// Maximum concurrent inbound connections (DoS protection). With `auth` - /// set, only authenticated connections count against it. + /// set, only authenticated connections count against it, and a peer that + /// authenticates while the cap is full is refused with the typed + /// at-capacity verdict inside the handshake. pub max_connections: usize, /// With `auth` set: maximum inbound connections that may be inside the /// cluster handshake at once, across all sources (default @@ -116,8 +128,12 @@ pub struct TcpTransportConfig { /// Ignored without `auth`. pub max_pending_handshakes: usize, /// With `auth` set: maximum inbound connections inside the handshake from - /// one source IP (default [`DEFAULT_MAX_PENDING_HANDSHAKES_PER_IP`]). - /// IPv4-mapped IPv6 sources count against their IPv4 address. Must be + /// one source (default [`DEFAULT_MAX_PENDING_HANDSHAKES_PER_IP`]). An + /// IPv4 source is counted per address (IPv4-mapped IPv6 sources count + /// against their IPv4 address); an IPv6 source is counted per /64, so a + /// host cannot bypass the bound by rotating addresses inside its prefix. + /// Cluster nodes inside one IPv6 /64 therefore share this budget, and a + /// large IPv6 cluster in a single /64 may need a higher value. Must be /// non-zero when `auth` is set (`bind` fails otherwise). Ignored without /// `auth`. pub max_pending_handshakes_per_ip: usize, @@ -265,6 +281,14 @@ impl TcpTransport { Ok(transport) } + + /// Whether this transport runs the cluster handshake on every connection + /// ([`TcpTransportConfig::auth`] set). A role that must only ever talk to + /// authenticated peers (the disaggregated router) checks this at build + /// time instead of trusting its caller. + pub fn is_authenticated(&self) -> bool { + self.config.auth.is_some() + } } impl Transport for TcpTransport { @@ -452,9 +476,13 @@ struct AcceptAuth { /// its own task before `handle_connection` reads a single frame; a failed /// handshake drops the socket and logs the peer IP (never secret material). /// While inside the handshake a connection holds a pending-handshake slot -/// instead of a `max_connections` permit; the permit is taken only once the -/// peer is authenticated. Without `accept_auth` the permit is taken at -/// accept, as before. +/// instead of a `max_connections` permit; the permit is taken inside the +/// handshake, once the peer is authenticated, and a full cap is answered +/// with the at-capacity verdict. Without `accept_auth` the permit is taken at +/// accept, as before. Every per-connection resource (the permit, the pending +/// slot, the socket) is released by drop, so a connection task that panics +/// (for example inside a registered RPC handler) leaks nothing; there is +/// deliberately no manual connection counter to keep in step. #[allow(clippy::too_many_arguments)] async fn accept_loop( listener: TcpListener, @@ -469,7 +497,6 @@ async fn accept_loop( shutdown: Arc, ) { let conn_semaphore = Arc::new(Semaphore::new(max_connections)); - let active_conns = Arc::new(AtomicUsize::new(0)); loop { tokio::select! { @@ -497,7 +524,7 @@ async fn accept_loop( Some(accept_auth) => match accept_auth.pending.try_acquire(peer_addr.ip()) { Ok(slot) => (None, Some(slot)), Err(refusal) => { - if handshake_warn_allowed() { + if handshake_warn_allowed(HandshakeWarnClass::Noise) { tracing::warn!( "Rejecting unauthenticated connection from {}: {refusal}", peer_addr.ip() @@ -520,31 +547,35 @@ async fn accept_loop( let rpc_handler = rpc_handler.clone(); let shutdown_flag = shutdown_flag.clone(); let shutdown = shutdown.clone(); - let active = active_conns.clone(); let accept_auth = accept_auth.clone(); let conn_semaphore = conn_semaphore.clone(); - active.fetch_add(1, Ordering::Relaxed); tokio::spawn(async move { // Held until the task exits: the accept-time - // permit without auth, the post-handshake permit + // permit without auth, the in-handshake permit // with auth. let (_permit, mut stream) = match accept_auth { None => (permit, ProtectedStream::plain(stream)), Some(accept_auth) => { - let outcome = server_handshake( + // The connection cap is checked inside the + // handshake, after the peer authenticated, + // so a full cap is a typed verdict for the + // peer and a stranger learns nothing. + let outcome = server_handshake_admitting( &mut stream, &accept_auth.auth, &accept_auth.local_identity, accept_auth.suite, + FrameChannel::Transport, + || conn_semaphore.try_acquire_owned().ok(), ) .await; // The pending slot covers the handshake only. drop(pending_slot); - let stream = match outcome { - Ok((identity, keys)) => { + match outcome { + Ok((identity, keys, acquired)) => { peer = identity; - ProtectedStream::protected(stream, keys) + (Some(acquired), ProtectedStream::protected(stream, keys)) } Err(err) => { if err.downcast_ref::().is_some() { @@ -553,7 +584,17 @@ async fn accept_loop( "Connection from {} closed before the cluster handshake", peer_addr.ip() ); - } else if handshake_warn_allowed() { + } else if let Some(full) = + err.downcast_ref::() + { + // An authenticated peer, so not + // rate limited: a stranger cannot + // reach this line. + tracing::warn!( + "Rejecting connection from {}: max connections ({max_connections}) reached", + full.peer + ); + } else if handshake_warn_allowed(handshake_warn_class(&err)) { tracing::warn!( "Rejecting connection from {}: cluster handshake failed: {err}", peer_addr.ip() @@ -564,19 +605,6 @@ async fn accept_loop( peer_addr.ip() ); } - active.fetch_sub(1, Ordering::Relaxed); - return; - } - }; - // Authenticated: only now does it count - // against the connection cap. - match conn_semaphore.try_acquire_owned() { - Ok(acquired) => (Some(acquired), stream), - Err(_) => { - tracing::warn!( - "Rejecting connection from {peer}: max connections ({max_connections}) reached" - ); - active.fetch_sub(1, Ordering::Relaxed); return; } } @@ -593,7 +621,6 @@ async fn accept_loop( &shutdown, ) .await; - active.fetch_sub(1, Ordering::Relaxed); }); } Err(e) => { @@ -608,10 +635,10 @@ async fn accept_loop( /// Handle a single inbound TCP connection, reading frames in a loop. /// /// The loop ends on the first read error. A frame that fails authentication -/// ([`FrameRejected`]) is logged at warn through the handshake rate limiter -/// (debug once the window is exhausted), without payload bytes; any other -/// error (peer closed, malformed or oversize header) ends the loop quietly as -/// before. +/// ([`FrameRejected`]) is logged at warn through the handshake rate limiter's +/// diagnostic budget (debug once the window is exhausted), without payload +/// bytes; any other error (peer closed, malformed or oversize header) ends the +/// loop quietly as before. async fn handle_connection( stream: &mut ProtectedStream, peer: &str, @@ -685,7 +712,7 @@ async fn handle_connection( } Err(err) => { if err.downcast_ref::().is_some() { - if handshake_warn_allowed() { + if handshake_warn_allowed(HandshakeWarnClass::Diagnostic) { tracing::warn!( "Closing connection from {peer}: {err}" ); @@ -703,35 +730,6 @@ async fn handle_connection( } } -/// Failed inbound handshakes logged at warn per one-minute window. Beyond -/// this the rejection is logged at debug so an unauthenticated peer cannot -/// flood the log. -const HANDSHAKE_WARN_PER_WINDOW: u32 = 5; -const HANDSHAKE_WARN_WINDOW_SECS: u64 = 60; - -/// Process-wide limiter for the failed-handshake and rejected-frame warnings. -/// Shared with the stage assignment listener (`stage_assignment`), which -/// accepts unauthenticated sockets the same way. -pub(crate) fn handshake_warn_allowed() -> bool { - use std::sync::Mutex; - use std::time::Instant; - static STATE: Mutex> = Mutex::new(None); - let now = Instant::now(); - let mut guard = STATE.lock().unwrap_or_else(|e| e.into_inner()); - match guard.as_mut() { - Some((start, count)) - if now.duration_since(*start).as_secs() < HANDSHAKE_WARN_WINDOW_SECS => - { - *count = count.saturating_add(1); - *count <= HANDSHAKE_WARN_PER_WINDOW - } - _ => { - *guard = Some((now, 1)); - true - } - } -} - #[cfg(test)] #[path = "tcp_transport_tests.rs"] mod tests; diff --git a/src/distributed/tcp_transport_tests.rs b/src/distributed/tcp_transport_tests.rs index 11b8e8e26..180072c73 100644 --- a/src/distributed/tcp_transport_tests.rs +++ b/src/distributed/tcp_transport_tests.rs @@ -188,6 +188,8 @@ mod cluster_auth { use tokio::io::{AsyncReadExt, AsyncWriteExt}; use crate::distributed::cluster_auth::{ClusterAuth, ClusterSecret}; + use crate::distributed::cluster_handshake::ServerAtCapacity; + use crate::distributed::connection_pool::PoolConfig; use crate::distributed::tcp_transport::{TcpTransport, TcpTransportConfig}; use crate::distributed::transport::{MessageKind, Transport, TransportMessage}; @@ -465,6 +467,120 @@ mod cluster_auth { plain.shutdown().await.unwrap(); } + /// Regression guard for the connection-slot accounting (#2692 item 2): + /// every per-connection resource is released by drop, so a connection + /// task that panics inside a registered RPC handler frees its + /// `max_connections` slot. With one slot, client A's request makes the + /// handler panic, and client B still gets through afterwards. + /// + /// Client A's `rpc_call` fails by timing out: `TcpTransport` reads RPC + /// responses only on accepted connections, never on the pooled outbound + /// one the request went out on, so no echo could reach A even without the + /// panic. B's observable is therefore a plain `send` reaching the server. + #[tokio::test] + async fn panicking_rpc_handler_releases_its_connection_slot() { + let shared = auth(); + let server = TcpTransport::bind(TcpTransportConfig { + bind_address: "127.0.0.1:0".to_string(), + auth: Some(Arc::clone(&shared)), + max_connections: 1, + ..Default::default() + }) + .await + .unwrap(); + server + .serve_rpc(Box::new(|request| { + if request == b"boom" { + panic!("rpc handler panic (expected by the test)"); + } + request.to_vec() + })) + .await + .unwrap(); + let server_addr = server.local_addr().unwrap(); + + let client_a = TcpTransport::bind(TcpTransportConfig { + bind_address: "127.0.0.1:0".to_string(), + auth: Some(Arc::clone(&shared)), + pool_config: PoolConfig { + connect_timeout: Duration::from_secs(1), + ..PoolConfig::default() + }, + ..Default::default() + }) + .await + .unwrap(); + let failed = tokio::time::timeout( + Duration::from_secs(5), + client_a.rpc_call(&server_addr, b"boom"), + ) + .await + .expect("rpc_call must return"); + assert!(failed.is_err(), "the panicking handler cannot answer"); + + let client_b = bind(Some(shared)).await; + client_b.send(&server_addr, ping()).await.unwrap(); + let (from, _) = tokio::time::timeout(Duration::from_secs(2), server.recv()) + .await + .expect("client B must get the slot the panicked connection task held") + .unwrap(); + assert_eq!(from, client_b.local_addr().unwrap()); + + client_b.shutdown().await.unwrap(); + client_a.shutdown().await.unwrap(); + server.shutdown().await.unwrap(); + } + + /// A full connection cap is answered inside the handshake: client B's + /// `send` fails with the typed, retryable `ServerAtCapacity` instead of + /// an accepted handshake followed by a closed socket. + #[tokio::test] + async fn full_connection_cap_reports_capacity_to_the_client() { + let shared = auth(); + let server = TcpTransport::bind(TcpTransportConfig { + bind_address: "127.0.0.1:0".to_string(), + auth: Some(Arc::clone(&shared)), + max_connections: 1, + ..Default::default() + }) + .await + .unwrap(); + let server_addr = server.local_addr().unwrap(); + + let client_a = bind(Some(Arc::clone(&shared))).await; + client_a.send(&server_addr, ping()).await.unwrap(); + let (from, _) = tokio::time::timeout(Duration::from_secs(5), server.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(from, client_a.local_addr().unwrap()); + // A keeps its pooled connection, so the single slot stays taken. + + let client_b = bind(Some(shared)).await; + let err = tokio::time::timeout(Duration::from_secs(1), client_b.send(&server_addr, ping())) + .await + .expect("send must fail within 1 s, not hang on a closed socket") + .expect_err("the server is at its connection cap"); + assert!( + err.chain() + .any(|cause| cause.downcast_ref::().is_some()), + "expected ServerAtCapacity in the chain, got: {err:#}" + ); + + // Once A releases its connection the slot is free for B. + client_a.shutdown().await.unwrap(); + tokio::time::sleep(Duration::from_millis(200)).await; + client_b.send(&server_addr, ping()).await.unwrap(); + let (from, _) = tokio::time::timeout(Duration::from_secs(5), server.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(from, client_b.local_addr().unwrap()); + + client_b.shutdown().await.unwrap(); + server.shutdown().await.unwrap(); + } + #[tokio::test] async fn default_config_stays_unauthenticated() { let server = bind(None).await; diff --git a/src/server/router_front.rs b/src/server/router_front.rs index 91b65f7ae..47840ed51 100644 --- a/src/server/router_front.rs +++ b/src/server/router_front.rs @@ -40,6 +40,17 @@ //! background demux task (`spawn_result_demux`) routes each incoming frame to //! the per-request channel keyed by `request_id`. //! +//! The transport must be bound with the cluster secret ([`RouterState::build`] +//! refuses an unauthenticated one), and the demux delivers a frame only when +//! its authenticated sender identity is one of the configured +//! `--prefill-peers` or `--decode-peers` addresses +//! ([`RouterState::accepts_result_from`], #2692): request ids are sequential, +//! so without the check any other holder of the secret (a pipeline stage that +//! shares it) could inject tokens into a client's response. The sender +//! identity is the peer's `--serving-bind` address as bound by the handshake +//! (an unspecified bind IP is replaced by the connection's source IP), so a +//! peer must be configured under the address it connects from. +//! //! # Output filtering (issue #198) //! //! Every decode text piece (and the prefill first token) passes through the @@ -72,7 +83,7 @@ //! control; when the prefill queue is full or no prefill node is available, the //! router returns HTTP 503 instead of dispatching. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::convert::Infallible; use std::net::SocketAddr; use std::path::PathBuf; @@ -100,6 +111,7 @@ use crate::distributed::disaggregated::{ BackpressureAction, DisaggRoutingStrategy, PrefillRequestFrame, RequestRouter, ResultFrame, RouterConfig, StreamBridge, TokenEvent, TokenSource, control_parts, sampling_to_serializable, }; +use crate::distributed::handshake_warn::{HandshakeWarnClass, handshake_warn_allowed}; use crate::distributed::metrics::ClusterMetrics; use crate::distributed::registry::{NodeRegistry, NodeStatus}; use crate::distributed::request_tracker::RequestId; @@ -136,13 +148,19 @@ const HEALTH_PROBE_TIMEOUT: Duration = Duration::from_secs(2); /// Injected into axum handlers via `axum::extract::State>`. pub struct RouterState { /// Transport bound to the router's own `--serving-bind` address; receives - /// [`ResultFrame`]s from the prefill and decode nodes. + /// [`ResultFrame`]s from the prefill and decode nodes. Always bound with + /// the cluster secret ([`Self::build`] refuses anything else). pub transport: Arc, /// Demux map: `request_id -> per-request result channel sender`. /// Keyed by the same `u64` that goes into [`PrefillRequestFrame::request_id`]. pub pending: Mutex>>, + /// The authenticated sender identities the demux accepts a + /// [`ResultFrame`] from: every `--prefill-peers` and `--decode-peers` + /// address with its IP canonicalized. See [`Self::accepts_result_from`]. + pub(crate) result_senders: HashSet, + /// Monotonically increasing counter used to assign unique request ids. pub next_id: AtomicU64, @@ -200,7 +218,12 @@ impl RouterState { /// /// Registers every address from `config.prefill_peers` and /// `config.decode_peers` in the [`NodeRegistry`] as `Online` so - /// `route_to_prefill` can find them. + /// `route_to_prefill` can find them, and as the only senders the result + /// demux accepts frames from. + /// + /// Fails when `transport` was bound without the cluster secret: every + /// frame the router acts on must come over an authenticated connection, + /// and the sender check below relies on the handshake-bound identity. pub fn build( config: Arc, transport: Arc, @@ -209,6 +232,17 @@ impl RouterState { tokenizer: Arc, model_path: PathBuf, ) -> Result { + if !transport.is_authenticated() { + anyhow::bail!( + "the disaggregated router requires a role transport bound with the cluster secret" + ); + } + let result_senders: HashSet = config + .prefill_peers + .iter() + .chain(config.decode_peers.iter()) + .map(|addr| canonical_socket_addr(*addr)) + .collect(); // Build a registry of ONLY the prefill and decode peers the router can // route to. The router itself is deliberately NOT registered: it is a // front-end, not an inference worker, so registering it (as Hybrid) @@ -287,6 +321,7 @@ impl RouterState { Ok(Self { transport, pending: Mutex::new(HashMap::new()), + result_senders, next_id: AtomicU64::new(1), reply_to, router: request_router, @@ -305,6 +340,24 @@ impl RouterState { }) } + /// Whether a [`ResultFrame`] from the authenticated sender identity `from` + /// (what the transport's `recv()` reports) may reach a request channel. + /// + /// `from` must parse as a socket address and, with its IP canonicalized, + /// be one of the configured prefill or decode peer addresses; anything + /// else (another cluster member that holds the secret, an unparseable + /// identity) is refused. The identity the handshake binds is the sender's + /// `--serving-bind` address with an unspecified IP replaced by the + /// connection's source IP, so a multi-homed peer, or one bound to + /// `0.0.0.0` that reaches the router over an interface other than the one + /// in `--prefill-peers` / `--decode-peers`, has to be configured under the + /// address it connects from. + pub(crate) fn accepts_result_from(&self, from: &str) -> bool { + from.parse::() + .map(|addr| self.result_senders.contains(&canonical_socket_addr(addr))) + .unwrap_or(false) + } + /// Clone one immutable settings snapshot without holding the lock across /// prompt rendering, network dispatch, or streaming. #[must_use] @@ -405,15 +458,40 @@ struct DecodeTarget { // ── Background demux task ──────────────────────────────────────────────── +/// `addr` with its IP canonicalized (an IPv4-mapped IPv6 address becomes the +/// IPv4 address), the form both `result_senders` and the sender check use. +fn canonical_socket_addr(addr: SocketAddr) -> SocketAddr { + SocketAddr::new(addr.ip().to_canonical(), addr.port()) +} + /// Spawn a background task that forwards incoming [`ResultFrame`]s to the /// per-request channels registered in `state.pending`. /// +/// A frame whose authenticated sender is not a configured prefill or decode +/// peer ([`RouterState::accepts_result_from`]) is dropped with a warn through +/// the handshake warning budget's diagnostic class (#2692); the request it +/// named then fails with the usual handoff timeout, and the warn line is the +/// diagnostic for a peer configured under the wrong address. +/// /// The task exits when the transport's receive loop returns an error (e.g. /// on shutdown). pub fn spawn_result_demux(state: Arc) { tokio::spawn(async move { while let Ok((from, msg)) = state.transport.recv().await { tracing::debug!(%from, "router demux: received a transport frame"); + if !state.accepts_result_from(&from) { + let line = format!( + "router: dropping result frame from {from:?}: not a configured prefill or \ + decode peer (check --prefill-peers, --decode-peers and the sender's \ + --serving-bind)" + ); + if handshake_warn_allowed(HandshakeWarnClass::Diagnostic) { + tracing::warn!("{line}"); + } else { + tracing::debug!("{line}"); + } + continue; + } if let Ok((op, payload)) = control_parts(msg) && op == ResultFrame::OPERATION && let Ok(frame) = ResultFrame::decode(&payload) diff --git a/src/server/router_front_tests.rs b/src/server/router_front_tests.rs index 7cb74ffe1..f2f4eb9a9 100644 --- a/src/server/router_front_tests.rs +++ b/src/server/router_front_tests.rs @@ -30,23 +30,36 @@ use axum::http::{Method, Request, StatusCode}; use serde_json::{Value, json}; use tower::ServiceExt; +use crate::distributed::cluster_auth::{ClusterAuth, ClusterSecret}; +use crate::distributed::disaggregated::ResultPhase; use crate::distributed::tcp_transport::TcpTransportConfig; const DEFAULT_ROUTER_TEST_TEMPLATE: &str = "{% for message in messages %}{{ message.content }}{% endfor %}"; -async fn router_test_state_with_config_and_template( - mut config: ServerConfig, - template: &str, -) -> Arc { - let transport = Arc::new( +fn test_auth() -> Arc { + Arc::new(ClusterAuth::new(&ClusterSecret::random())) +} + +/// Bind a loopback transport the way every disaggregated role does: with the +/// cluster secret (`RouterState::build` refuses anything else). +async fn bind_authenticated(auth: Arc) -> Arc { + Arc::new( TcpTransport::bind(TcpTransportConfig { bind_address: "127.0.0.1:0".to_string(), + auth: Some(auth), ..TcpTransportConfig::default() }) .await .expect("bind router test transport"), - ); + ) +} + +async fn router_test_state_with_config_and_template( + mut config: ServerConfig, + template: &str, +) -> Arc { + let transport = bind_authenticated(test_auth()).await; let reply_to = transport.local_addr().expect("router transport address"); if config.prefill_peers.is_empty() { config.prefill_peers = vec!["127.0.0.1:9".parse().expect("prefill address")]; @@ -131,6 +144,125 @@ async fn response_json(response: Response) -> Value { serde_json::from_slice(&bytes).expect("JSON response") } +// --------------------------------------------------------------------------- +// Result demux admission (#2692) +// --------------------------------------------------------------------------- + +/// Build a router state over `transport` with the given prefill peers and the +/// default template and stub tokenizer. +fn build_router_state( + transport: Arc, + prefill_peers: Vec, +) -> Result { + let reply_to = transport.local_addr().expect("router transport address"); + RouterState::build( + Arc::new(ServerConfig { + prefill_peers, + ..ServerConfig::default() + }), + transport, + reply_to, + Arc::new(ChatTemplateProcessor::with_template( + DEFAULT_ROUTER_TEST_TEMPLATE.to_string(), + )), + Arc::new(MlxcelTokenizer::stub()), + PathBuf::from("router-test-model"), + ) +} + +/// The router never trusts its caller to have bound the transport with the +/// secret: an unauthenticated transport is refused at build time. +#[tokio::test] +async fn build_refuses_an_unauthenticated_transport() { + let plain = Arc::new( + TcpTransport::bind(TcpTransportConfig { + bind_address: "127.0.0.1:0".to_string(), + ..TcpTransportConfig::default() + }) + .await + .expect("bind plain transport"), + ); + let err = build_router_state(plain, vec!["127.0.0.1:9".parse().unwrap()]) + .err() + .expect("an unauthenticated transport must be refused") + .to_string(); + assert!(err.contains("cluster secret"), "{err}"); +} + +/// Only the configured prefill and decode peer addresses may deliver result +/// frames; the comparison is on the canonical IP and the port. +#[tokio::test] +async fn accepts_result_from_only_configured_peers() { + let state = build_router_state( + bind_authenticated(test_auth()).await, + vec!["127.0.0.1:9".parse().unwrap()], + ) + .expect("build router state"); + assert!(state.accepts_result_from("127.0.0.1:9")); + assert!(state.accepts_result_from("[::ffff:127.0.0.1]:9")); + assert!(!state.accepts_result_from("127.0.0.1:10")); + assert!(!state.accepts_result_from("not-an-addr")); +} + +/// Two authenticated members hold the secret, but only the one whose listener +/// address is a configured peer reaches a request's channel; the other's frame +/// for the same request id is dropped. +#[tokio::test] +async fn result_demux_drops_frames_from_unconfigured_senders() { + let auth = test_auth(); + let configured = bind_authenticated(Arc::clone(&auth)).await; + let stranger = bind_authenticated(Arc::clone(&auth)).await; + let configured_addr: SocketAddr = configured + .local_addr() + .expect("configured peer address") + .parse() + .expect("socket address"); + let state = Arc::new( + build_router_state(bind_authenticated(auth).await, vec![configured_addr]) + .expect("build router state"), + ); + let router_addr = state.transport.local_addr().expect("router address"); + + let request_id = state.next_id.fetch_add(1, Ordering::Relaxed); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + state.pending.lock().unwrap().insert(request_id, tx); + spawn_result_demux(Arc::clone(&state)); + + let frame_from = |tokens: &str| ResultFrame { + request_id, + phase: ResultPhase::FirstToken, + tokens: vec![tokens.to_string()], + start_sequence: 0, + done: true, + error: None, + generated_tokens: Some(1), + }; + stranger + .send(&router_addr, frame_from("injected").encode().unwrap()) + .await + .expect("stranger send"); + // The stranger's frame is dropped: nothing reaches the channel. + let dropped = tokio::time::timeout(Duration::from_millis(500), rx.recv()).await; + assert!( + dropped.is_err(), + "unconfigured sender's frame was delivered: {dropped:?}" + ); + + configured + .send(&router_addr, frame_from("legit").encode().unwrap()) + .await + .expect("configured send"); + let delivered = tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .expect("configured peer's frame must be delivered") + .expect("channel open"); + assert_eq!(delivered.tokens, vec!["legit".to_string()]); + + stranger.shutdown().await.unwrap(); + configured.shutdown().await.unwrap(); + state.transport.shutdown().await.unwrap(); +} + #[tokio::test] async fn router_settings_route_is_absent_until_enabled() { let response = create_router_app(router_test_state().await) @@ -925,14 +1057,7 @@ fn kimi_k3_router_tokenizer(dir: &std::path::Path) -> MlxcelTokenizer { } async fn router_test_state_with_tokenizer(tokenizer: MlxcelTokenizer) -> Arc { - let transport = Arc::new( - TcpTransport::bind(TcpTransportConfig { - bind_address: "127.0.0.1:0".to_string(), - ..TcpTransportConfig::default() - }) - .await - .expect("bind router test transport"), - ); + let transport = bind_authenticated(test_auth()).await; let reply_to = transport.local_addr().expect("router transport address"); let config = ServerConfig { prefill_peers: vec!["127.0.0.1:9".parse().expect("prefill address")], From 80974c3ccf1d4c3aa1e70c4c267574e34762fce6 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 28 Oct 2026 14:46:19 +0900 Subject: [PATCH 2/3] docs: tighten admission hardening wording and demux drop logging Correct the residual-risk text: IdentityNotOwned and FrameSuiteMismatch are decided from the hello before the client authenticates, so a party without the secret can trigger them and saturate the Diagnostic window, not only ClientAbortedAfterServerProof; the Noise class definition no longer claims to be the only stranger-producible class. Describe ServerAtCapacity precisely: it is not fatal like HandshakeRejected, connect_with_retry returns it without backoff, the next acquire reconnects, and Transport::connect warns. The "failed handshake is terminal" comment in connect_with_retry now covers this case. State that the router sender check authenticates the source IP only; the port is the advertised listener port and is not verified. Note in startup.rs that RouterState::build also enforces the cluster secret. Merge a stray re-wrapped line in docs/distributed.md. spawn_result_demux now logs the drop through tracing fields in both branches instead of building a String with format!. No behavior change. Validation: cargo fmt --check, cargo clippy --profile test-fast --features cuda -p mlxcel --lib --tests -- -D warnings, and the server::router_front:: and distributed::handshake_warn:: lib tests (38 passed). Refs #2692 --- docs/distributed.md | 18 +++++++++++------- src/distributed/cluster_handshake.rs | 10 ++++++---- src/distributed/connection_pool.rs | 9 ++++++--- src/distributed/handshake_warn.rs | 18 +++++++++++------- src/distributed/tcp_transport.rs | 2 +- src/server/router_front.rs | 12 +++++------- src/server/startup.rs | 4 +++- 7 files changed, 43 insertions(+), 30 deletions(-) diff --git a/docs/distributed.md b/docs/distributed.md index 27c9754bb..162e19946 100644 --- a/docs/distributed.md +++ b/docs/distributed.md @@ -346,7 +346,11 @@ What the secret protects: loop. The router refuses to start on a transport bound without the secret, and its result demux delivers a frame only when the sender's handshake-bound identity is one of its `--prefill-peers` or - `--decode-peers` addresses (canonical IP and port): request ids are + `--decode-peers` addresses (canonical IP and port). Only the source IP is + authenticated (the handshake verifies it against the TCP connection); the + port is the sender's advertised listener port and is not verified, so the + check tells cluster members apart by source IP plus claimed port, not by + port-level proof. Request ids are sequential, so any other holder of the secret, such as a pipeline stage that shares it, could otherwise inject tokens into a client's response. A frame from another member is dropped with a warning (`router: dropping @@ -392,8 +396,9 @@ in the log: at-capacity verdict inside the handshake. The connecting side fails with "server authenticated this node but refused the connection: it already holds its maximum number of authenticated connections; retry later", which - the connection pool and `Transport::connect` treat as retryable (a warning, - not a fatal rejection); the accepting side logs "max connections (N) + is not treated as fatal like a handshake rejection: the connection pool + returns it without backoff and the next `acquire` reconnects, and + `Transport::connect` logs a warning and continues; the accepting side logs "max connections (N) reached". The verdict is sent only after the client's authenticator verified, so a party without the secret never learns the capacity state. @@ -419,7 +424,7 @@ Residual risks: - Frame lengths and kinds (`TensorData`, `Control`, RPC request or response) are visible on the wire, as is the timing of every decode step. An observer can tell how many tokens a cluster is producing even though it cannot read any of them. - Frames from authenticated members stay trusted inside the pipeline protocol and on the prefill and decode intakes of the disaggregated protocol: a host that holds the secret is a full member there. The router's result demux is the exception (it accepts frames only from its configured peers). Rotate the secret when a member is decommissioned. - Connections that have not finished the handshake do not count against the 512 authenticated-connection cap. They draw from a separate pending pool (64 in flight, at most 8 per source, each for at most the 5 second handshake timeout); a socket beyond either bound is closed at accept. A source is an IPv4 address or an IPv6 /64 (a host that owns a /64 would otherwise get one pending slot per address), which also means cluster nodes inside one IPv6 /64 share the per-source budget and a large IPv6 cluster in a single /64 may need a higher `max_pending_handshakes_per_ip`. A flood of unauthenticated connections from enough sources (8 or more at the defaults) can fill the pending pool and so delay new connections from any source, including a real peer reconnecting, by up to the 5 second timeout per slot while the flood lasts; it cannot take slots from established authenticated connections or inject frames. Failed inbound handshakes and pending-pool refusals are logged at warn for the first few per minute (per class, see above) and at debug afterwards. -- A party without the secret can still produce the "probable secret mismatch" warning (send a well-formed hello naming its own IP, then close after the server's authenticator), so it can saturate the diagnostic warning budget too. The per-minute summary of suppressed warnings is what keeps the operator informed then. +- A party without the secret can still produce the "probable secret mismatch" warning (send a well-formed hello naming its own IP, then close after the server's authenticator). The "identity not owned" and "suite mismatch" refusals are decided from the hello alone, before the client authenticates, so it can trigger those as well. Any of them can saturate the diagnostic warning budget. The per-minute summary of suppressed warnings is what keeps the operator informed then. - The receiver refuses an oversize `len` from the header alone, but a frame whose `len` is within the cap (256 MiB by default) makes it allocate and read the full body before the tag is checked. Only a peer that completed the handshake can reach this path, so it is a resource bound for authenticated members, not an exposure to the open network. - `MLXCEL_CLUSTER_SECRET` is inherited by child processes of the server. Prefer `--cluster-secret-file` or `MLXCEL_CLUSTER_SECRET_FILE` so the value never sits in the environment. - The server identity is not bound to the socket IP on the connecting side, so a stage address that is reached through NAT or a port forward still authenticates by the secret alone. @@ -442,7 +447,7 @@ What each layer covers: | Layer | Covers | Does not cover | |-------|--------|----------------| -| Connect-time handshake (`PROTOCOL_VERSION` 4) | Only a holder of the secret can open or accept a connection; the accepted identity is bound to the TCP source IP; the frame suite is agreed inside the authenticated transcript; a full connection cap is reported to an authenticated peer as a typed, retryable verdict. | Anything that happens on the TCP stream after the handshake. | +| Connect-time handshake (`PROTOCOL_VERSION` 4) | Only a holder of the secret can open or accept a connection; the accepted identity is bound to the TCP source IP; the frame suite is agreed inside the authenticated transcript; a full connection cap is reported to an authenticated peer as a typed verdict that is not fatal. | Anything that happens on the TCP stream after the handshake. | | Per-frame protection (#2656) | Every frame after the handshake: confidentiality of the payload, integrity of header and payload, replay and reorder rejection through an implicit sequence number, binding to the connection whose handshake derived the keys. Defeats the relay, which cannot derive keys without the secret. Covers pooled outbound connections, `send_stream`, RPC requests and responses, the Thunderbolt and RDMA wrappers (same TCP core), the disaggregated router, prefill and decode transports, and the `--pp-auto` stage assignment channel. | Dropping or delaying frames; frame lengths, kinds and timing. | | Discovery beacons (`--cluster-discovery=mdns`) | Each UDP beacon carries an HMAC, a timestamp and a nonce; the coordinator rejects an invalid MAC, a 60 s clock skew, a replay inside that window, another mlxcel version, or an advertised address that is not the datagram's source. | Hiding the fact that a cluster is forming; beacon payloads are readable. | @@ -461,8 +466,7 @@ header is the associated data, so the receiver refuses an oversize `len` before any decryption work. A frame that fails to decrypt closes the connection; the event is logged through the diagnostic class of the same rate limiter as a failed handshake, never with payload bytes. Keys are per -connection, so a frame -recorded on one connection fails on every other. The suite is announced in +connection, so a frame recorded on one connection fails on every other. The suite is announced in the handshake hello and covered by both authenticators; a node configured for a different suite is refused with an explicit record. Nothing new is configured: the secret sources above are the only input, and there is no diff --git a/src/distributed/cluster_handshake.rs b/src/distributed/cluster_handshake.rs index eca2e75db..f05189322 100644 --- a/src/distributed/cluster_handshake.rs +++ b/src/distributed/cluster_handshake.rs @@ -80,7 +80,7 @@ //! Admission is part of the handshake (#2692): the accept side's `admit` //! callback ([`server_handshake_admitting`]) runs only after the client's //! authenticator verified, and a `None` from it is answered with the -//! `0x02` verdict, so the client gets a typed, retryable [`ServerAtCapacity`] +//! `0x02` verdict, so the client gets a typed [`ServerAtCapacity`] (not fatal, see its docs) //! instead of an accepted handshake followed by a closed socket, and a party //! without the secret never learns the capacity state. //! @@ -191,9 +191,11 @@ pub struct FrameSuiteMismatch { pub struct ClientAbortedAfterServerProof; /// Connect side: the server verified this node's authenticator and then -/// answered with the at-capacity verdict. Unlike [`HandshakeRejected`] this is -/// transient: the connection pool and `Transport::connect` treat it as -/// retryable (warn and continue), not as a configuration error. +/// answered with the at-capacity verdict. Unlike [`HandshakeRejected`] it is +/// not treated as fatal: `connect_with_retry` returns it to the caller +/// immediately, without backoff, and the next `acquire` opens a fresh +/// connection; `Transport::connect` logs a warning and continues instead of +/// failing as it does for a rejection. #[derive(Debug, thiserror::Error)] #[error( "server authenticated this node but refused the connection: it already holds its maximum \ diff --git a/src/distributed/connection_pool.rs b/src/distributed/connection_pool.rs index 61e0aa1aa..f5808fcf2 100644 --- a/src/distributed/connection_pool.rs +++ b/src/distributed/connection_pool.rs @@ -205,9 +205,12 @@ impl ConnectionPool { let _ = stream.set_nodelay(true); } let stream = match &self.auth { - // A failed handshake is terminal, not a transient - // network fault: a wrong secret never becomes right - // on retry, so do not burn the backoff budget on it. + // A failed handshake is not retried here, so the + // backoff budget is not burned on it. A rejection + // (wrong secret, identity or suite) is terminal: it + // never becomes right on retry. `ServerAtCapacity` is + // transient but also returned without backoff: the + // peer is full now, and the next `acquire` reconnects. Some(pool_auth) => match client_handshake( &mut stream, &pool_auth.auth, diff --git a/src/distributed/handshake_warn.rs b/src/distributed/handshake_warn.rs index 5d0728e38..61eb4cc32 100644 --- a/src/distributed/handshake_warn.rs +++ b/src/distributed/handshake_warn.rs @@ -22,9 +22,10 @@ //! trickle of such junk cannot push the one line an operator needs down to //! debug: //! -//! - [`HandshakeWarnClass::Noise`]: anything a stranger can produce without -//! the secret and without completing the protocol (pending-pool refusals, -//! malformed or absent hellos, timeouts, unparseable identities). +//! - [`HandshakeWarnClass::Noise`]: refusals a stranger can produce with no +//! protocol knowledge at all and that carry no operator action +//! (pending-pool refusals, malformed or absent hellos, timeouts, +//! unparseable identities). //! - [`HandshakeWarnClass::Diagnostic`]: the typed refusals that tell the //! operator what to fix: a probable secret mismatch //! ([`ClientAbortedAfterServerProof`], [`HandshakeRejected`]), an identity @@ -38,11 +39,14 @@ //! were suppressed, so a sustained flood still leaves one summary line per //! minute in the log. //! -//! Residual risk: a party without the secret can still produce +//! Residual risk: the `Diagnostic` class is not reserved for authenticated +//! parties. A party without the secret can still produce //! `ClientAbortedAfterServerProof` (send a well-formed hello naming its own -//! IP, then close after the server proof), so it can saturate the -//! `Diagnostic` window too. The rollover summary is what keeps the operator -//! informed then. Warning on the first event per source instead would need a +//! IP, then close after the server proof), and the `IdentityNotOwned` and +//! `FrameSuiteMismatch` refusals are decided from the hello alone, before +//! the client authenticates, so it can trigger those too. Any of them can +//! saturate the `Diagnostic` window. The rollover summary is what keeps the +//! operator informed then. Warning on the first event per source instead would need a //! per-source map that an IPv6 source defeats by rotating addresses. //! //! Used by: `tcp_transport::accept_loop` and `handle_connection`, diff --git a/src/distributed/tcp_transport.rs b/src/distributed/tcp_transport.rs index add0a4471..f838bc0fd 100644 --- a/src/distributed/tcp_transport.rs +++ b/src/distributed/tcp_transport.rs @@ -53,7 +53,7 @@ //! junk connections cannot hide those typed refusals). The `max_connections` //! check happens inside the handshake, after the peer authenticated: a peer //! that finds this node full gets the typed at-capacity verdict -//! ([`super::cluster_handshake::ServerAtCapacity`], retryable) instead of an +//! ([`super::cluster_handshake::ServerAtCapacity`], not fatal) instead of an //! accepted handshake followed by a closed socket. //! //! Without `auth` (the default, used by single-node serving) the transport diff --git a/src/server/router_front.rs b/src/server/router_front.rs index 47840ed51..ee661f8e2 100644 --- a/src/server/router_front.rs +++ b/src/server/router_front.rs @@ -480,15 +480,13 @@ pub fn spawn_result_demux(state: Arc) { while let Ok((from, msg)) = state.transport.recv().await { tracing::debug!(%from, "router demux: received a transport frame"); if !state.accepts_result_from(&from) { - let line = format!( - "router: dropping result frame from {from:?}: not a configured prefill or \ - decode peer (check --prefill-peers, --decode-peers and the sender's \ - --serving-bind)" - ); + const HINT: &str = "not a configured prefill or decode peer (check \ + --prefill-peers, --decode-peers and the sender's \ + --serving-bind)"; if handshake_warn_allowed(HandshakeWarnClass::Diagnostic) { - tracing::warn!("{line}"); + tracing::warn!(?from, "router: dropping result frame: {HINT}"); } else { - tracing::debug!("{line}"); + tracing::debug!(?from, "router: dropping result frame: {HINT}"); } continue; } diff --git a/src/server/startup.rs b/src/server/startup.rs index 54d2553ec..d228daec2 100644 --- a/src/server/startup.rs +++ b/src/server/startup.rs @@ -3770,7 +3770,9 @@ pub async fn start_server(mut startup: ServerStartupConfig) -> Result<()> { })?; // The router's role transport authenticates every connection with // the cluster secret (#2655); `cluster_auth` is `Some` here because - // the resolver above fails closed for the router role. + // the resolver above fails closed for the router role, and + // `RouterState::build` enforces it again by refusing a transport bound + // without the secret. let transport = std::sync::Arc::new( crate::distributed::tcp_transport::TcpTransport::bind( crate::distributed::tcp_transport::TcpTransportConfig { From ef1040d3cbd66d481fd10c446c644d9a7b8a506f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 28 Oct 2026 14:55:56 +0900 Subject: [PATCH 3/3] docs: add technical report for PR #2705 --- ...-cluster-pre-auth-admission-20261028.en.md | 143 ++++++++++++++++++ ...-cluster-pre-auth-admission-20261028.ko.md | 143 ++++++++++++++++++ 2 files changed, 286 insertions(+) create mode 100644 TECHNICAL_REPORTS/2705-harden-cluster-pre-auth-admission-20261028.en.md create mode 100644 TECHNICAL_REPORTS/2705-harden-cluster-pre-auth-admission-20261028.ko.md diff --git a/TECHNICAL_REPORTS/2705-harden-cluster-pre-auth-admission-20261028.en.md b/TECHNICAL_REPORTS/2705-harden-cluster-pre-auth-admission-20261028.en.md new file mode 100644 index 000000000..f4775b101 --- /dev/null +++ b/TECHNICAL_REPORTS/2705-harden-cluster-pre-auth-admission-20261028.en.md @@ -0,0 +1,143 @@ +# PR #2705: Harden cluster pre-auth admission and rejections + +**Date**: 2026-10-28 (written before merge; the PR is open as a draft) +**Status**: Implementation complete; implementation and security reviews found no CRITICAL or HIGH issues; local tests passed on GB10 (spark-102); hosted CI in progress at the time of writing +**Languages**: Rust, Markdown +**Risk Level**: Medium (handshake wire version bump 3 to 4, so nodes built before and after this change cannot pair; a router whose peer list does not match the senders' bound identities now drops their result frames) + +## Executive Summary + +PR #2705 (closes #2692) closes five residuals that the security review of #2674 left on the authenticated cluster transport in `src/distributed/` and on the disaggregated router's result demux. None of them lets a party without the cluster secret read or inject frames. They are DoS, log-visibility and member-isolation hardening on a surface that already requires the secret, grouped into one PR because they share one subsystem and one verification pass. + +1. The per-source pending-handshake limit keys IPv6 sources by /64 instead of by full address, so a host that owns a /64 can no longer bypass it by rotating addresses. +2. The unread `active_conns` counter in `accept_loop`, which leaked when a connection task panicked, is deleted. Every per-connection resource is now released by drop. +3. The single process-wide "5 warnings per minute" limiter is replaced by a new `handshake_warn` module with two independent budgets, `Noise` and `Diagnostic`, so junk connections cannot push a probable-secret-mismatch line down to debug. +4. The router refuses to build on a transport bound without the secret, and its result demux drops frames whose authenticated sender is not one of `--prefill-peers` or `--decode-peers`. +5. The `max_connections` check moves inside the handshake, after the client authenticates. A full listener now answers with a new `VERDICT_AT_CAPACITY` (0x02) byte and the client gets a typed `ServerAtCapacity` error instead of an accepted handshake followed by a closed socket. `PROTOCOL_VERSION` goes from 3 to 4. + +The PR is 17 files, +1104 / -156, in two authored commits (`699d208d` implementation, `80974c3c` review wording) plus a merge from `main` (`9900bb89`). + +## 1. Problem Statement + +All five items were confirmed on `main` @ 54d63b0a when #2692 was filed. + +### 1.1 IPv6 rotation bypassed the per-source pending limit + +`PendingHandshakeLimiter::try_acquire` (from #2655) counted in-flight handshakes per canonical source address, capped by `max_pending_handshakes_per_ip` (default 8). An IPv6 host normally controls a whole /64 and can source each connection from a different address, so it could take every slot in the global pending pool (default 64) by itself. The module doc already admitted the gap. + +### 1.2 A counter that leaked and was never read + +`accept_loop` kept `active_conns: Arc`, incremented on accept and decremented by hand on three exit paths. A panic inside a connection task, for example inside a registered RPC handler, skipped the decrement. The workspace builds with `panic = "unwind"`, so tokio drops the panicked task's future; the semaphore permit and the pending slot are RAII and were already released correctly. The only casualty was `active_conns`, which nothing read. It was a latent trap: the first metric or limit built on it would have drifted upward on every panic. + +### 1.3 Junk connections hid the line an operator needs + +`handshake_warn_allowed()` in `tcp_transport.rs` was one budget of 5 warnings per 60 s shared by every rejection: pending-pool refusals, all handshake failures, and frames that failed authentication. Anyone who can open sockets can spend that budget with bad magic bytes, malformed hellos or stalls to the 5 s timeout. A genuine `ClientAbortedAfterServerProof` ("the peer most likely holds a different cluster secret") arriving after a few junk sockets was logged only at debug. + +### 1.4 The router demux trusted any secret holder + +`spawn_result_demux` delivered a `ResultFrame` to whichever pending request matched its `request_id`, regardless of `from`. Request ids come from `next_id.fetch_add`, so they are sequential and guessable. Any holder of the secret (for example a pipeline stage sharing the same secret file) could inject tokens into another client's streamed response. Separately, `RouterState::build` accepted any `TcpTransport`; the only thing keeping the router's transport authenticated was the resolver upstream and a comment in `startup.rs`, and both router test fixtures bound without auth. + +### 1.5 A full connection cap looked like a successful handshake + +`server_handshake_inner` wrote `VERDICT_ACCEPT`, and only after it returned did `accept_loop` try `conn_semaphore.try_acquire_owned()`. On failure it dropped the socket. The client saw a successful handshake, then a broken connection on its first send, with nothing in the error to say the peer was full. + +## 2. Change Summary + +| Area | Change | +|------|--------| +| `src/distributed/pending_handshakes.rs` | New `pending_source_key(ip)`: canonicalize (IPv4-mapped IPv6 stays IPv4), return IPv4 as is, mask IPv6 to its /64 with `u128::from(v6) & !((1u128 << 64) - 1)`. Used by `try_acquire` and the test helper `pending_from`. `PendingSlotRefusal::PerIp` carries the key and prints IPv6 keys as `2001:db8:1:2::/64`. | +| `src/distributed/tcp_transport.rs` | `active_conns` and its three decrements removed with the `AtomicUsize` import. `accept_loop` calls `server_handshake_admitting(..., FrameChannel::Transport, \|\| conn_semaphore.try_acquire_owned().ok())` and logs `ConnectionLimitReached` unconditionally (only an authenticated peer can reach it). New `TcpTransport::is_authenticated()`. Old limiter deleted; call sites pass a class. | +| `src/distributed/handshake_warn.rs` (new, 202 lines) | `HandshakeWarnClass { Noise, Diagnostic }`, `handshake_warn_class(&anyhow::Error)` that walks the error chain for the five typed refusals, `WarnBudget` with a two-slot per-class window and an injectable `allow_at(class, now) -> WarnDecision`, and the process-wide `handshake_warn_allowed(class)` that also logs `N cluster handshake warnings () were logged at debug in the last 60s` on rollover. | +| `src/distributed/cluster_handshake.rs` | `VERDICT_AT_CAPACITY = 0x02`, `PROTOCOL_VERSION = 4`. New `pub(crate) server_handshake_admitting

(..., admit: impl FnOnce() -> Option

) -> Result<(String, SessionKeys, P)>`; `server_handshake_on` wraps it with `\|\| Some(())`. Client maps verdict 0x02 to `ServerAtCapacity`; server returns `ConnectionLimitReached { peer }`. Wire diagram and module doc updated. | +| `src/distributed/stage_assignment.rs` | Imports the new limiter; pending refusals draw `Noise`, handshake failures `handshake_warn_class(&err)`, rejected frames `Diagnostic`. Handshake admission unchanged (wrapper admits everything). | +| `src/distributed/connection_pool.rs`, `mod.rs` | Comment on the no-retry handshake path covers the transient `ServerAtCapacity` case; `ServerAtCapacity` and `ConnectionLimitReached` re-exported; `handshake_warn` is `pub(crate)`. | +| `src/server/router_front.rs` | `RouterState::build` bails when `!transport.is_authenticated()`. New `result_senders: HashSet` (prefill and decode peers with canonical IPs) and `accepts_result_from(&str)`. The demux drops other senders with a `router: dropping result frame` warn through the `Diagnostic` budget, using tracing fields rather than a formatted string. | +| `src/bin/mlx_server.rs`, `src/server/startup.rs` | `--prefill-peers` / `--decode-peers` help says each entry must be the address the peer connects from. Startup comment credits `RouterState::build` with enforcing auth. | +| `docs/distributed.md`, `docs/CONTINUOUS_BATCHING.md` | /64 counting and its shared-subnet consequence, the router sender check and what it does and does not authenticate, the at-capacity outcome, the two warning classes, version 4, and the residual risk that a stranger can still reach the `Diagnostic` class. | +| Tests (5 `_tests.rs` files, +523) | 13 new tests across `pending_handshakes`, `handshake_warn`, `cluster_handshake`, `tcp_transport::cluster_auth` and `server::router_front`; both router fixtures now bind with a random `ClusterSecret`. | + +## 3. Technical Decisions + +### 3.1 Key by /64, not a second per-prefix tier + +The issue rejected adding a separate per-prefix limit with its own config field. One knob, keyed by a coarser source, covers the case. The cost is explicit in the docs: cluster nodes inside one IPv6 /64 now share the per-source budget of 8 in-flight handshakes, so a large IPv6 cluster in one /64 may need a higher `max_pending_handshakes_per_ip`. Since the pending pool only covers the few seconds of a handshake, this matters only during a mass reconnect. Canonicalizing before masking matters: without it an IPv4-mapped IPv6 address (`::ffff:a.b.c.d`) would be masked to `::/64` (its upper 64 bits are zero) and every IPv4 client on a dual-stack socket would share one key. + +### 3.2 Delete the counter, do not fix it + +Making the decrement panic-safe (a drop guard) would have kept a value nobody reads. The issue instead records that a future active-connection metric should derive from `max_connections - conn_semaphore.available_permits()`, which cannot drift because the permit is already RAII. The regression guard `panicking_rpc_handler_releases_its_connection_slot` passes before and after the change by design; it pins the property (a panicking handler leaks no slot) rather than the deletion. + +The guard's observable had to change from the issue's text. `TcpTransport` parses `RpcResponse` frames only in `handle_connection`, which runs on accepted connections; the pooled outbound connection an `rpc_call` writes on is never read. So no TCP client can receive an RPC echo, and the test uses client B's `send` reaching `server.recv()` within 2 s instead. This is a pre-existing limitation of RPC over TCP (the only in-repo caller, `bench.rs`, ignores the result), recorded in the PR and the issue but not fixed here. + +### 3.3 Two warning classes instead of per-source first-event logging + +The alternative, always warning on the first event from each source, needs a per-source map that is either unbounded or evicting, and IPv6 rotation defeats both. Two fixed windows (a `Mutex<[Option; 2]>`) cost no per-source state. `WarnBudget::allow_at` takes `now` so the tests can drive 50 noise events and a 61 s rollover without sleeping; the process-wide wrapper is the only caller of `Instant::now()`. + +Classification walks the `anyhow` chain with `downcast_ref` for `ClientAbortedAfterServerProof`, `HandshakeRejected`, `IdentityNotOwned`, `FrameSuiteMismatch` and `FrameRejected`. Everything untyped (bad magic, malformed hello, timeout, I/O) is `Noise`. Pending refusals are always `Noise`; post-handshake frame rejections and router drops are always `Diagnostic`. + +The split is not a security boundary. `IdentityNotOwned` and `FrameSuiteMismatch` are decided from the client hello alone, and `ClientAbortedAfterServerProof` only needs a well-formed hello naming the attacker's own IP followed by a close. A stranger can therefore still saturate `Diagnostic`. The review commit corrected the docs to say so; the rollover summary line is the remaining signal in that case. + +### 3.4 Admission as a callback inside the handshake + +`server_handshake_admitting` takes `admit: impl FnOnce() -> Option

` and returns the permit with the keys. Three properties follow from that shape: + +- `admit()` is called after the client authenticator verifies (a mismatch returns `VERDICT_REJECT` first), so only a secret holder can learn the capacity state. `admission_is_not_consulted_before_the_client_authenticates` pins this. +- The permit is a local that travels out with the result. If the `VERDICT_ACCEPT` write fails or the 5 s timeout fires after `admit()`, the permit is dropped with the future; nothing is counted by hand. +- Listeners that do not cap connections (the stage assignment channel) keep calling `server_handshake_on`, which passes `|| Some(())`, so their behavior is unchanged. + +A generic permit type keeps `cluster_handshake.rs` free of any dependency on the semaphore in `tcp_transport.rs`. + +### 3.5 `ServerAtCapacity` is not fatal, and not retried with backoff + +`Transport::connect` treats `HandshakeRejected` as fatal because no retry fixes a secret mismatch. `ServerAtCapacity` is a separate type so `connect` logs a warning and continues. In `connect_with_retry` it takes the same no-backoff exit as every handshake failure: the error goes straight to the caller and the next `acquire` opens a fresh connection. The review commit adjusted the docs and the pool comment, which had first called this "retryable" without saying who retries. + +### 3.6 Bump the wire version rather than reuse a reject byte + +Before this PR, any verdict other than 0x01 was a rejection. An old client receiving 0x02 would report `HandshakeRejected` ("handshake was altered in transit"), which is misleading and fatal at `connect`. Bumping `PROTOCOL_VERSION` to 4 makes a mixed pair fail at the hello with the explicit version message. #2650, #2674, #2684 and this change all landed after v0.8.0-beta.1, so no released build speaks version 3. Builds from `main` between those PRs and this one cannot pair with builds after it, which matches the existing "mixed versions are unsupported" rule. + +### 3.7 Router sender check on canonical IP and port + +`result_senders` stores each configured peer as `SocketAddr::new(ip.to_canonical(), port)`, and `accepts_result_from` parses `from` the same way. That strips IPv4-mapped forms, scope ids and flowinfo on both sides; an unparseable `from` is refused, and a configured `0.0.0.0` entry matches nothing because a bound identity is never unspecified. Only the source IP is authenticated (the handshake binds it to the TCP connection); the port is the listener port the sender advertised. Two secret holders behind one IP can therefore claim each other's port. The docs state this, and it is accepted under the threat model in which secret holders are members. + +Rejected: binding each request to the exact prefill or decode node it was routed to. That needs per-request routing state in the demux, and the protocol does not fix which role sends which phase. The check also covers only the router's result intake; the prefill request intake and the decode handoff intake keep trusting any authenticated member (out of scope in #2692). + +The behavior change for operators: a peer whose bound identity differs from its configured address (a multi-homed node, or one bound to `0.0.0.0` that reaches the router over another interface) now has its frames dropped, and the request fails with the existing handoff timeout. The `router: dropping result frame` warn is the diagnostic. + +## 4. Verification + +From the PR body (GB10 spark-102, `--profile test-fast`, builds outside `gpu-lock`, runs under `gpu-lock run --tag issue-2692`): + +- `cargo fmt --check` and `cargo clippy --profile test-fast --features cuda -p mlxcel --lib --tests -- -D warnings`: clean, rerun after the review commit. +- The 13 new or touched tests by name: 13 passed. `--lib distributed::`: 1473 passed, 0 failed, 4 ignored. `--lib server::router_front`: 35 passed. `--test zero_config_cluster_init`: 4 passed, 3 ignored. `tests/disaggregated_serving_e2e.rs`: 2 passed. `tests/disaggregated_router_e2e.rs`: 7 passed. +- Real-model parity (`MLXCEL_REQUIRE_MODELS=1 MLXCEL_SDPA_DETERMINISTIC=1`, `--ignored`): `pipeline_remote_real_models` 10 passed and 1 failed, `pipeline_server_remote_real_models` 7 passed and 1 failed, `pipeline_ci_multi_stage_real_models` 2 passed. Every llama, gemma3, gemma4, qwen3, qwen3.5 and glm-4.7 parity test passed. +- The two failures are not a regression. They are the `glm_moe_dsa` tests panicking in `repo_model_dir` because the `glm5-4bit` checkpoint is absent: `mlx-community/GLM-5-4bit` is 418.6 GB and cannot be loaded on this 128 GB GB10 box, and the tests predate #2674. +- Manual three-process loopback (router, prefill and decode sharing one `--cluster-secret-file`, `qwen3-0.6b-4bit`, 127.0.0.1): `POST /v1/chat/completions` returned HTTP 200 with `"2 + 2 = 4."`, and the router log had no `dropping result frame` line. This shows the sender check does not drop legitimate frames in the loopback setup. + +Failing-first evidence: the fix hunks were reverted temporarily on the branch and the tests run by name. `ipv6_sources_in_one_64_share_the_per_source_limit`, `ipv6_sources_in_different_64s_are_counted_separately`, `noise_flood_does_not_suppress_diagnostic_warnings`, `build_refuses_an_unauthenticated_transport` and `result_demux_drops_frames_from_unconfigured_senders` all failed. The admission tests (`server_at_capacity_sends_capacity_verdict`, `admission_is_not_consulted_before_the_client_authenticates`, `full_connection_cap_reports_capacity_to_the_client`) and the classifier tests do not compile without the new API. `rollover_reports_suppressed_count` and the panic guard pass on both sides by design. + +Pre-merge gate on the merged head `9900bb89` (orchestrator log, partial at the time of writing): `verify-fmt`, `verify-versions`, `verify-env-var-docs`, `verify-kernel-dtype-keys`, `verify-kernel-port-dispatch`, `verify-binary-assets` and `verify-windows-paths` all exited 0. The contract test build was still compiling, so no test results from that run are cited here. + +## 5. Remaining Gaps and Risks + +- **Version 3 and 4 nodes cannot pair.** Rolling upgrades of a cluster built from `main` need every node updated together. No released version is affected. +- **The verdict byte is unauthenticated**, as it was before. An on-path attacker can swap 0x01 and 0x02, but it could already drop the connection, so this adds no capability. +- **The `Diagnostic` budget is reachable without the secret** (section 3.3). The two classes protect against cheap junk, not against a stranger who sends a well-formed hello. +- **Router drops share the process-wide `Diagnostic` budget** with handshake diagnostics. A misconfigured peer sends a stream of result frames per request, each one dropped and counted, so it can use up the 5 warnings per minute and push a concurrent secret-mismatch line to debug. The rollover summary still counts them, but it does not say which kind of event was suppressed. +- **The rollover summary is emitted lazily.** It is logged on the first event of the next window for that class. If a flood stops, its suppressed count is not reported until the next event of that class, possibly much later, and the message's "in the last 60s" wording is then inaccurate. +- **`ServerAtCapacity` has no backoff.** A caller that retries `send` in a tight loop against a full peer will run a full handshake per attempt, and `connect_with_retry` logs `Cluster handshake with peer ... failed` at warn on each one without rate limiting. Only secret holders are involved, and no in-repo caller is known to loop this way, but nothing prevents it. +- **Port spoofing between co-located members** (section 3.7) is accepted by design. +- **RPC over TCP cannot return responses** (section 3.2). Pre-existing, documented in the PR, not tracked by a dedicated issue as far as this report found. +- **Real-model parity for `glm_moe_dsa`** remains unverified on this hardware because the checkpoint does not fit. + +## 6. Learning Points + +- **Put admission inside the authenticated exchange.** Checking capacity after the handshake leaks nothing but gives a bad client experience; checking it before leaks capacity state to strangers. A callback invoked between "client verified" and "verdict written" gets both right, and returning the permit with the session keys keeps the resource lifetime in RAII. +- **Rate limiters that share a budget across trust levels invert their purpose.** The original limiter protected the log from strangers by silencing the operator. Splitting budgets by what the line is worth, not by who sent it, is cheap when the classes are few and fixed. +- **Canonicalize before masking.** IPv4-mapped IPv6 addresses must become IPv4 before any prefix logic, or every IPv4 client on a dual-stack listener collapses into one /64. +- **Delete unread state instead of hardening it.** A counter nobody reads still has to be kept correct by every future edit; derive the value from the RAII resource when it is needed. +- **Inject time into windowed logic.** `allow_at(class, now)` let the rollover and flood tests run deterministically with no sleeps. + +## 7. References + +- PR #2705, issue #2692 +- Origin: #2674 (#2655 pre-auth pending limit and typed rejections), #2684 (#2656 per-frame AEAD), #2650 (cluster secret for distributed roles) +- Key files: `src/distributed/handshake_warn.rs`, `src/distributed/cluster_handshake.rs`, `src/distributed/pending_handshakes.rs`, `src/distributed/tcp_transport.rs`, `src/server/router_front.rs`, `docs/distributed.md` diff --git a/TECHNICAL_REPORTS/2705-harden-cluster-pre-auth-admission-20261028.ko.md b/TECHNICAL_REPORTS/2705-harden-cluster-pre-auth-admission-20261028.ko.md new file mode 100644 index 000000000..af36594db --- /dev/null +++ b/TECHNICAL_REPORTS/2705-harden-cluster-pre-auth-admission-20261028.ko.md @@ -0,0 +1,143 @@ +# PR #2705: 클러스터 인증 전 수락 단계와 거부 처리 강화 + +**날짜**: 2026-10-28 (머지 전 작성, PR은 draft 상태로 열려 있음) +**상태**: 구현 완료, 구현 리뷰와 보안 리뷰에서 CRITICAL/HIGH 이슈 없음, GB10 (spark-102)에서 로컬 테스트 통과, 작성 시점에 hosted CI 진행 중 +**언어**: Rust, Markdown +**위험도**: 중간 (핸드셰이크 wire 버전이 3에서 4로 올라가므로 이 변경 전후로 빌드된 노드는 서로 연결되지 않음, 피어 목록이 송신자의 바인딩된 identity와 맞지 않는 router는 이제 그 송신자의 result frame을 버림) + +## 요약 + +PR #2705 (#2692 해결)는 #2674 보안 리뷰가 `src/distributed/`의 인증된 클러스터 transport와 disaggregated router의 result demux에 남긴 다섯 가지 잔여 문제를 해결합니다. 어느 것도 클러스터 secret이 없는 쪽이 frame을 읽거나 주입할 수 있게 하는 문제는 아닙니다. 이미 secret을 요구하는 표면에 대한 DoS, 로그 가시성, 멤버 간 격리 강화이며, 같은 서브시스템과 같은 검증 절차를 공유하므로 하나의 PR로 묶었습니다. + +1. 소스별 pending handshake 한도가 IPv6 소스를 전체 주소가 아닌 /64 단위로 셉니다. /64를 가진 호스트가 주소를 바꿔 가며 한도를 우회할 수 없습니다. +2. connection task가 panic하면 누수되던, 아무도 읽지 않는 `accept_loop`의 `active_conns` 카운터를 삭제했습니다. 연결별 자원은 모두 drop으로 해제됩니다. +3. 프로세스 전체에서 하나였던 "분당 경고 5개" limiter를 새 `handshake_warn` 모듈의 독립된 두 예산(`Noise`, `Diagnostic`)으로 바꿨습니다. 쓰레기 연결이 secret 불일치 의심 경고를 debug로 밀어내지 못합니다. +4. router는 secret 없이 바인딩된 transport로는 build되지 않으며, result demux는 인증된 송신자가 `--prefill-peers` 또는 `--decode-peers`에 없으면 frame을 버립니다. +5. `max_connections` 검사가 핸드셰이크 안, 클라이언트 인증 이후로 옮겨졌습니다. 가득 찬 listener는 새 `VERDICT_AT_CAPACITY` (0x02) 바이트로 응답하고, 클라이언트는 "핸드셰이크 성공 후 소켓 종료" 대신 타입이 있는 `ServerAtCapacity` 에러를 받습니다. `PROTOCOL_VERSION`이 3에서 4로 올라갑니다. + +PR은 17개 파일, +1104 / -156이며, 작성 커밋 두 개(`699d208d` 구현, `80974c3c` 리뷰 문구 정리)와 `main` 머지 커밋(`9900bb89`)으로 구성됩니다. + +## 1. 문제 정의 + +#2692 작성 시점에 다섯 항목 모두 `main` @ 54d63b0a에서 확인되었습니다. + +### 1.1 IPv6 주소 순환으로 소스별 pending 한도 우회 + +#2655에서 도입된 `PendingHandshakeLimiter::try_acquire`는 진행 중인 핸드셰이크를 정규화된 소스 주소별로 세고 `max_pending_handshakes_per_ip`(기본 8)로 제한했습니다. IPv6 호스트는 보통 /64 전체를 쓰며 연결마다 다른 주소를 소스로 쓸 수 있으므로, 혼자서 전역 pending pool(기본 64) 전체를 차지할 수 있었습니다. 모듈 문서에도 이 공백이 이미 적혀 있었습니다. + +### 1.2 누수되면서 아무도 읽지 않던 카운터 + +`accept_loop`는 `active_conns: Arc`를 accept 시 증가시키고 세 종료 경로에서 수동으로 감소시켰습니다. connection task 안에서, 예를 들어 등록된 RPC handler 안에서 panic이 나면 감소가 건너뛰어졌습니다. 워크스페이스는 `panic = "unwind"`로 빌드되므로 tokio가 panic한 task의 future를 drop하고, semaphore permit과 pending slot은 RAII라 이미 정상 해제되었습니다. 피해를 보는 것은 아무도 읽지 않는 `active_conns`뿐이었습니다. 다만 이 값 위에 metric이나 한도를 처음 만드는 순간 panic마다 값이 위로 어긋나는 잠재적 함정이었습니다. + +### 1.3 쓰레기 연결이 운영자에게 필요한 한 줄을 가림 + +`tcp_transport.rs`의 `handshake_warn_allowed()`는 60초당 경고 5개라는 하나의 예산을 모든 거부가 나눠 썼습니다: pending pool 거부, 모든 핸드셰이크 실패, 인증에 실패한 frame. 소켓을 열 수 있는 누구나 잘못된 magic 바이트, 깨진 hello, 5초 timeout까지 버티기로 이 예산을 소진할 수 있습니다. 쓰레기 소켓 몇 개 뒤에 도착한 진짜 `ClientAbortedAfterServerProof`("피어가 다른 클러스터 secret을 가지고 있을 가능성이 높음")는 debug로만 기록되었습니다. + +### 1.4 router demux가 secret을 가진 누구든 신뢰 + +`spawn_result_demux`는 `from`과 상관없이 `request_id`가 맞는 대기 요청에 `ResultFrame`을 전달했습니다. request id는 `next_id.fetch_add`로 만들어지므로 순차적이고 추측 가능합니다. secret을 가진 누구든(예를 들어 같은 secret 파일을 공유하는 pipeline stage) 다른 클라이언트의 스트리밍 응답에 토큰을 주입할 수 있었습니다. 별개로 `RouterState::build`는 아무 `TcpTransport`나 받았습니다. router transport가 인증된 상태임을 보장하는 것은 상위의 resolver와 `startup.rs`의 주석뿐이었고, router 테스트 fixture 두 개 모두 auth 없이 바인딩했습니다. + +### 1.5 가득 찬 연결 한도가 성공한 핸드셰이크처럼 보임 + +`server_handshake_inner`가 `VERDICT_ACCEPT`를 쓰고 반환한 뒤에야 `accept_loop`가 `conn_semaphore.try_acquire_owned()`를 시도했습니다. 실패하면 소켓을 그냥 닫았습니다. 클라이언트는 핸드셰이크 성공을 보고, 첫 send에서 끊긴 연결을 만났으며, 에러에는 피어가 가득 찼다는 정보가 없었습니다. + +## 2. 변경 요약 + +| 영역 | 변경 | +|------|------| +| `src/distributed/pending_handshakes.rs` | 새 `pending_source_key(ip)`: 정규화(IPv4-mapped IPv6는 IPv4로), IPv4는 그대로, IPv6는 `u128::from(v6) & !((1u128 << 64) - 1)`로 /64 마스킹. `try_acquire`와 테스트 헬퍼 `pending_from`이 사용. `PendingSlotRefusal::PerIp`는 이 키를 담고 IPv6 키를 `2001:db8:1:2::/64`처럼 출력. | +| `src/distributed/tcp_transport.rs` | `active_conns`, 수동 감소 세 곳, `AtomicUsize` import 제거. `accept_loop`가 `server_handshake_admitting(..., FrameChannel::Transport, \|\| conn_semaphore.try_acquire_owned().ok())`를 호출하고 `ConnectionLimitReached`는 rate limit 없이 기록(인증된 피어만 도달 가능). 새 `TcpTransport::is_authenticated()`. 기존 limiter 삭제, 호출부는 class를 넘김. | +| `src/distributed/handshake_warn.rs` (신규, 202줄) | `HandshakeWarnClass { Noise, Diagnostic }`, 에러 체인에서 타입이 있는 거부 다섯 가지를 찾는 `handshake_warn_class(&anyhow::Error)`, class별 window 두 칸과 시간 주입이 가능한 `allow_at(class, now) -> WarnDecision`을 가진 `WarnBudget`, rollover 시 `N cluster handshake warnings () were logged at debug in the last 60s`를 기록하는 프로세스 전역 `handshake_warn_allowed(class)`. | +| `src/distributed/cluster_handshake.rs` | `VERDICT_AT_CAPACITY = 0x02`, `PROTOCOL_VERSION = 4`. 새 `pub(crate) server_handshake_admitting

(..., admit: impl FnOnce() -> Option

) -> Result<(String, SessionKeys, P)>`, `server_handshake_on`은 `\|\| Some(())`로 감싼 wrapper. 클라이언트는 verdict 0x02를 `ServerAtCapacity`로, 서버는 `ConnectionLimitReached { peer }`를 반환. wire 다이어그램과 모듈 문서 갱신. | +| `src/distributed/stage_assignment.rs` | 새 limiter를 import, pending 거부는 `Noise`, 핸드셰이크 실패는 `handshake_warn_class(&err)`, 거부된 frame은 `Diagnostic`. 핸드셰이크 수락 로직은 그대로(wrapper가 모두 수락). | +| `src/distributed/connection_pool.rs`, `mod.rs` | 재시도하지 않는 핸드셰이크 경로의 주석이 일시적인 `ServerAtCapacity`까지 설명. `ServerAtCapacity`, `ConnectionLimitReached` re-export, `handshake_warn`은 `pub(crate)`. | +| `src/server/router_front.rs` | `RouterState::build`가 `!transport.is_authenticated()`이면 실패. 새 `result_senders: HashSet`(IP를 정규화한 prefill, decode 피어)와 `accepts_result_from(&str)`. demux는 다른 송신자의 frame을 `Diagnostic` 예산을 통해 `router: dropping result frame` 경고와 함께 버리며, 문자열을 조립하지 않고 tracing field로 기록. | +| `src/bin/mlx_server.rs`, `src/server/startup.rs` | `--prefill-peers` / `--decode-peers` 도움말에 각 항목은 피어가 접속해 오는 주소여야 한다고 명시. startup 주석은 `RouterState::build`가 auth를 강제한다고 기술. | +| `docs/distributed.md`, `docs/CONTINUOUS_BATCHING.md` | /64 단위 집계와 같은 서브넷 공유의 결과, router 송신자 검사가 인증하는 것과 하지 않는 것, at-capacity 결과, 두 경고 class, 버전 4, secret 없는 쪽도 `Diagnostic` class에 도달할 수 있다는 잔여 위험. | +| 테스트 (`_tests.rs` 5개, +523) | `pending_handshakes`, `handshake_warn`, `cluster_handshake`, `tcp_transport::cluster_auth`, `server::router_front`에 새 테스트 13개. router fixture 두 개는 이제 임의의 `ClusterSecret`으로 바인딩. | + +## 3. 기술적 선택과 그 이유 + +### 3.1 별도 prefix 단계 대신 /64를 키로 사용 + +이슈는 자체 설정 필드를 가진 별도의 prefix별 한도를 추가하는 안을 기각했습니다. 하나의 설정값을 더 거친 소스 단위로 세는 것으로 충분합니다. 대가는 문서에 명시되어 있습니다: 같은 IPv6 /64 안의 클러스터 노드들은 이제 진행 중 핸드셰이크 8개라는 소스별 예산을 공유하므로, 하나의 /64에 있는 대규모 IPv6 클러스터는 `max_pending_handshakes_per_ip`를 높여야 할 수 있습니다. pending pool은 핸드셰이크의 몇 초 동안만 쓰이므로 대량 재연결 시에만 의미가 있습니다. 마스킹 전에 정규화하는 것이 중요합니다. 그렇지 않으면 IPv4-mapped IPv6 주소(`::ffff:a.b.c.d`)는 상위 64비트가 0이라 `::/64`로 마스킹되고, dual-stack 소켓의 모든 IPv4 클라이언트가 하나의 키를 공유하게 됩니다. + +### 3.2 카운터를 고치지 않고 삭제 + +감소를 panic에 안전하게 만드는 것(drop guard)은 아무도 읽지 않는 값을 유지할 뿐입니다. 이슈는 대신 향후 활성 연결 metric은 `max_connections - conn_semaphore.available_permits()`로 계산하라고 기록합니다. permit이 이미 RAII이므로 어긋날 수 없습니다. 회귀 방지 테스트 `panicking_rpc_handler_releases_its_connection_slot`은 설계상 변경 전후 모두 통과하며, 삭제 자체가 아니라 "panic한 handler가 slot을 누수하지 않는다"는 성질을 고정합니다. + +이 테스트의 관측 방법은 이슈 문구에서 바뀌었습니다. `TcpTransport`는 `RpcResponse` frame을 accept된 연결에서 도는 `handle_connection`에서만 파싱하고, `rpc_call`이 쓰는 pool의 outbound 연결은 읽지 않습니다. 따라서 어떤 TCP 클라이언트도 RPC 응답을 받을 수 없고, 테스트는 대신 클라이언트 B의 `send`가 2초 안에 `server.recv()`에 도달하는지를 봅니다. 이것은 TCP 위 RPC의 기존 한계이며(저장소 안의 유일한 호출자 `bench.rs`는 결과를 무시함), PR과 이슈에 기록되었지만 여기서 고치지는 않았습니다. + +### 3.3 소스별 첫 이벤트 경고 대신 두 경고 class + +대안인 "소스마다 첫 이벤트는 항상 경고"는 소스별 map이 필요하고, 그 map은 무한히 커지거나 축출을 해야 하며 IPv6 주소 순환은 둘 다 무력화합니다. 고정된 window 두 개(`Mutex<[Option; 2]>`)는 소스별 상태가 필요 없습니다. `WarnBudget::allow_at`이 `now`를 받으므로 테스트는 sleep 없이 noise 이벤트 50개와 61초 rollover를 재현할 수 있고, `Instant::now()`를 호출하는 곳은 프로세스 전역 wrapper뿐입니다. + +분류는 `anyhow` 체인을 `downcast_ref`로 훑어 `ClientAbortedAfterServerProof`, `HandshakeRejected`, `IdentityNotOwned`, `FrameSuiteMismatch`, `FrameRejected`를 찾습니다. 타입이 없는 나머지(잘못된 magic, 깨진 hello, timeout, I/O)는 `Noise`입니다. pending 거부는 항상 `Noise`, 핸드셰이크 이후 frame 거부와 router drop은 항상 `Diagnostic`입니다. + +이 구분은 보안 경계가 아닙니다. `IdentityNotOwned`와 `FrameSuiteMismatch`는 client hello만으로 결정되고, `ClientAbortedAfterServerProof`는 공격자 자신의 IP를 담은 정상 hello를 보낸 뒤 연결을 닫기만 하면 됩니다. 따라서 secret이 없는 쪽도 `Diagnostic`을 포화시킬 수 있습니다. 리뷰 커밋에서 문서를 이에 맞게 고쳤고, 그 경우 남는 신호는 rollover 요약 한 줄입니다. + +### 3.4 핸드셰이크 안의 callback으로 수락 결정 + +`server_handshake_admitting`은 `admit: impl FnOnce() -> Option

`를 받고 permit을 key와 함께 반환합니다. 이 형태에서 세 가지 성질이 나옵니다. + +- `admit()`은 클라이언트 authenticator가 검증된 뒤에 호출됩니다(불일치면 먼저 `VERDICT_REJECT`를 반환). 따라서 secret을 가진 쪽만 용량 상태를 알 수 있습니다. `admission_is_not_consulted_before_the_client_authenticates`가 이를 고정합니다. +- permit은 결과와 함께 밖으로 나가는 지역 변수입니다. `admit()` 이후 `VERDICT_ACCEPT` 쓰기가 실패하거나 5초 timeout이 걸리면 permit은 future와 함께 drop됩니다. 수동으로 세는 것은 없습니다. +- 연결 수를 제한하지 않는 listener(stage assignment 채널)는 계속 `|| Some(())`을 넘기는 `server_handshake_on`을 호출하므로 동작이 바뀌지 않습니다. + +permit 타입을 generic으로 두어 `cluster_handshake.rs`가 `tcp_transport.rs`의 semaphore에 의존하지 않습니다. + +### 3.5 `ServerAtCapacity`는 치명적이지 않고, backoff 재시도도 하지 않음 + +`Transport::connect`는 재시도로 secret 불일치를 고칠 수 없으므로 `HandshakeRejected`를 치명적으로 처리합니다. `ServerAtCapacity`는 별도 타입이므로 `connect`는 경고를 기록하고 계속합니다. `connect_with_retry`에서는 다른 핸드셰이크 실패와 같이 backoff 없이 빠져나갑니다: 에러는 바로 호출자에게 가고, 다음 `acquire`가 새 연결을 엽니다. 처음에는 문서와 pool 주석이 이를 누가 재시도하는지 말하지 않고 "retryable"이라고만 했는데, 리뷰 커밋에서 바로잡았습니다. + +### 3.6 reject 바이트 재사용 대신 wire 버전 상향 + +이 PR 전에는 0x01이 아닌 verdict는 모두 거부였습니다. 이전 클라이언트가 0x02를 받으면 `HandshakeRejected`("handshake was altered in transit")로 보고하는데, 이는 오해를 부르고 `connect`에서 치명적입니다. `PROTOCOL_VERSION`을 4로 올리면 버전이 섞인 쌍은 hello 단계에서 명시적인 버전 메시지로 실패합니다. #2650, #2674, #2684와 이 변경은 모두 v0.8.0-beta.1 이후에 들어갔으므로 버전 3을 쓰는 릴리스 빌드는 없습니다. 그 PR들과 이 PR 사이의 `main` 빌드는 이후 빌드와 연결되지 않으며, 이는 기존의 "버전 혼용 미지원" 규칙과 같습니다. + +### 3.7 정규화된 IP와 포트로 router 송신자 검사 + +`result_senders`는 설정된 피어를 `SocketAddr::new(ip.to_canonical(), port)`로 저장하고, `accepts_result_from`도 `from`을 같은 방식으로 파싱합니다. 양쪽에서 IPv4-mapped 형태, scope id, flowinfo가 제거됩니다. 파싱되지 않는 `from`은 거부되고, 바인딩된 identity는 unspecified일 수 없으므로 설정된 `0.0.0.0` 항목은 아무것과도 맞지 않습니다. 인증되는 것은 소스 IP뿐입니다(핸드셰이크가 TCP 연결에 바인딩). 포트는 송신자가 광고한 listener 포트입니다. 따라서 같은 IP 뒤의 secret 보유자 둘은 서로의 포트를 사칭할 수 있습니다. 문서에 이 점이 명시되어 있으며, secret 보유자를 멤버로 보는 위협 모델 아래에서 수용했습니다. + +기각한 안: 각 요청을 라우팅된 정확한 prefill 또는 decode 노드에 묶는 것. demux에 요청별 라우팅 상태가 필요하고, 프로토콜은 어느 role이 어느 단계를 보내는지 고정하지 않습니다. 이 검사는 router의 result 수신만 다루며, prefill의 요청 수신과 decode의 handoff 수신은 계속 인증된 멤버 누구든 신뢰합니다(#2692의 범위 밖). + +운영자가 체감하는 변화: 바인딩된 identity가 설정 주소와 다른 피어(multi-homed 노드, 또는 `0.0.0.0`에 바인딩되어 다른 인터페이스로 router에 접속하는 노드)의 frame은 이제 버려지고, 요청은 기존 handoff timeout으로 실패합니다. `router: dropping result frame` 경고가 진단 수단입니다. + +## 4. 검증 + +PR 본문 기준(GB10 spark-102, `--profile test-fast`, 빌드는 `gpu-lock` 밖에서, 실행은 `gpu-lock run --tag issue-2692` 아래에서): + +- `cargo fmt --check`와 `cargo clippy --profile test-fast --features cuda -p mlxcel --lib --tests -- -D warnings`: 깨끗함, 리뷰 커밋 후 재실행. +- 새로 추가되거나 수정된 테스트 13개 이름 지정 실행: 13개 통과. `--lib distributed::`: 1473 통과, 0 실패, 4 ignored. `--lib server::router_front`: 35 통과. `--test zero_config_cluster_init`: 4 통과, 3 ignored. `tests/disaggregated_serving_e2e.rs`: 2 통과. `tests/disaggregated_router_e2e.rs`: 7 통과. +- 실제 모델 parity(`MLXCEL_REQUIRE_MODELS=1 MLXCEL_SDPA_DETERMINISTIC=1`, `--ignored`): `pipeline_remote_real_models` 10 통과 1 실패, `pipeline_server_remote_real_models` 7 통과 1 실패, `pipeline_ci_multi_stage_real_models` 2 통과. llama, gemma3, gemma4, qwen3, qwen3.5, glm-4.7 parity 테스트는 모두 통과. +- 두 실패는 회귀가 아닙니다. `glm5-4bit` 체크포인트가 없어서 `glm_moe_dsa` 테스트가 `repo_model_dir`에서 panic한 것입니다. `mlx-community/GLM-5-4bit`는 418.6 GB라 128 GB GB10 장비에서 로드할 수 없으며, 이 테스트들은 #2674보다 먼저 있었습니다. +- 수동 3-프로세스 loopback(router, prefill, decode가 하나의 `--cluster-secret-file`을 공유, `qwen3-0.6b-4bit`, 127.0.0.1): `POST /v1/chat/completions`가 HTTP 200과 `"2 + 2 = 4."`를 반환했고, router 로그에 `dropping result frame` 줄이 없었습니다. loopback 구성에서 송신자 검사가 정상 frame을 버리지 않음을 보여 줍니다. + +실패 우선 증거: 브랜치에서 수정 hunk를 임시로 되돌리고 테스트를 이름으로 실행했습니다. `ipv6_sources_in_one_64_share_the_per_source_limit`, `ipv6_sources_in_different_64s_are_counted_separately`, `noise_flood_does_not_suppress_diagnostic_warnings`, `build_refuses_an_unauthenticated_transport`, `result_demux_drops_frames_from_unconfigured_senders`가 모두 실패했습니다. 수락 관련 테스트(`server_at_capacity_sends_capacity_verdict`, `admission_is_not_consulted_before_the_client_authenticates`, `full_connection_cap_reports_capacity_to_the_client`)와 분류 테스트는 새 API 없이는 컴파일되지 않습니다. `rollover_reports_suppressed_count`와 panic 회귀 테스트는 설계상 양쪽에서 통과합니다. + +머지된 head `9900bb89`의 머지 전 게이트(오케스트레이터 로그, 작성 시점에 진행 중): `verify-fmt`, `verify-versions`, `verify-env-var-docs`, `verify-kernel-dtype-keys`, `verify-kernel-port-dispatch`, `verify-binary-assets`, `verify-windows-paths`가 모두 exit 0이었습니다. contract 테스트 빌드는 아직 컴파일 중이었으므로 그 실행의 테스트 결과는 인용하지 않습니다. + +## 5. 남은 공백과 위험 + +- **버전 3과 4 노드는 연결되지 않습니다.** `main`에서 빌드한 클러스터의 롤링 업그레이드는 모든 노드를 함께 갱신해야 합니다. 릴리스된 버전은 영향이 없습니다. +- **verdict 바이트는 이전처럼 인증되지 않습니다.** 경로상의 공격자는 0x01과 0x02를 바꿀 수 있지만, 이미 연결을 끊을 수 있었으므로 새 능력은 아닙니다. +- **`Diagnostic` 예산은 secret 없이도 도달 가능합니다**(3.3절). 두 class는 값싼 쓰레기 연결을 막을 뿐, 정상 hello를 보내는 외부인은 막지 못합니다. +- **router drop은 프로세스 전역 `Diagnostic` 예산을 핸드셰이크 진단과 공유합니다.** 잘못 설정된 피어는 요청마다 result frame을 연달아 보내고 각각이 버려지며 집계되므로, 분당 5개 경고를 소진해 동시에 발생한 secret 불일치 경고를 debug로 밀어낼 수 있습니다. rollover 요약은 개수를 세지만 어떤 종류의 이벤트가 억제되었는지는 말하지 않습니다. +- **rollover 요약은 지연되어 기록됩니다.** 해당 class의 다음 window 첫 이벤트에서 기록됩니다. flood가 멈추면 억제된 개수는 그 class의 다음 이벤트가 올 때까지(훨씬 나중일 수 있음) 보고되지 않으며, 그때는 메시지의 "in the last 60s" 문구가 부정확합니다. +- **`ServerAtCapacity`에는 backoff가 없습니다.** 가득 찬 피어에 대해 `send`를 빠르게 반복 재시도하는 호출자는 시도마다 전체 핸드셰이크를 수행하고, `connect_with_retry`는 매번 rate limit 없이 `Cluster handshake with peer ... failed`를 warn으로 기록합니다. secret 보유자만 관련되고 저장소 안에 이렇게 반복하는 호출자는 알려져 있지 않지만, 막는 장치도 없습니다. +- **같은 위치 멤버 간 포트 사칭**(3.7절)은 설계상 수용했습니다. +- **TCP 위 RPC는 응답을 돌려받을 수 없습니다**(3.2절). 기존 한계이며 PR에 기록되었지만, 이 보고서가 확인한 범위에서는 전용 이슈로 추적되지 않습니다. +- **`glm_moe_dsa`의 실제 모델 parity**는 체크포인트가 맞지 않아 이 하드웨어에서 검증되지 않은 채로 남습니다. + +## 6. 학습 포인트 + +- **수락 판단은 인증된 교환 안에 둡니다.** 핸드셰이크 이후에 용량을 검사하면 정보는 새지 않지만 클라이언트 경험이 나쁘고, 이전에 검사하면 외부인에게 용량 상태가 샙니다. "클라이언트 검증 완료"와 "verdict 쓰기" 사이에 호출되는 callback은 둘 다 해결하며, permit을 session key와 함께 반환하면 자원 수명이 RAII 안에 머뭅니다. +- **신뢰 수준이 다른 이벤트가 예산을 공유하는 rate limiter는 목적이 뒤집힙니다.** 원래 limiter는 외부인으로부터 로그를 지키려다 운영자를 침묵시켰습니다. 누가 보냈는지가 아니라 그 줄의 가치로 예산을 나누는 것은 class가 적고 고정되어 있으면 비용이 거의 없습니다. +- **마스킹 전에 정규화합니다.** IPv4-mapped IPv6 주소는 prefix 처리 전에 IPv4로 바꿔야 합니다. 그렇지 않으면 dual-stack listener의 모든 IPv4 클라이언트가 하나의 /64로 합쳐집니다. +- **읽지 않는 상태는 강화하지 말고 삭제합니다.** 아무도 읽지 않는 카운터도 이후의 모든 수정에서 정확하게 유지되어야 합니다. 필요할 때 RAII 자원에서 값을 유도합니다. +- **window 로직에는 시간을 주입합니다.** `allow_at(class, now)` 덕분에 rollover와 flood 테스트가 sleep 없이 결정적으로 동작합니다. + +## 7. 참고 + +- PR #2705, 이슈 #2692 +- 출발점: #2674 (#2655 인증 전 pending 한도와 타입이 있는 거부), #2684 (#2656 frame별 AEAD), #2650 (분산 role의 클러스터 secret) +- 주요 파일: `src/distributed/handshake_warn.rs`, `src/distributed/cluster_handshake.rs`, `src/distributed/pending_handshakes.rs`, `src/distributed/tcp_transport.rs`, `src/server/router_front.rs`, `docs/distributed.md`