Skip to main content

subc_daemon/
server.rs

1use std::{
2    any::Any,
3    error::Error,
4    fmt, io,
5    net::SocketAddr,
6    sync::{
7        atomic::{AtomicU64, Ordering as AtomicOrdering},
8        Arc,
9    },
10    time::{Duration, Instant},
11};
12
13use subc_transport::{
14    authenticate_server, AuthError, AuthStage, DAEMON_ID_LEN, WATCHDOG_CLIENT_ROLE,
15};
16use tokio::{
17    io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader, BufWriter},
18    net::TcpListener,
19    sync::{mpsc, Semaphore},
20    task::{JoinHandle, JoinSet},
21    time::timeout,
22};
23use tracing::{debug, info, warn};
24
25use crate::{
26    forwarding::{CloseReason, ConnectionCloseReceiver},
27    observability::ConnectedClients,
28    read_frame,
29    router::{FrameSink, RouteCtx, Router},
30    write_frame, FrameIoError, RouterError,
31};
32
33/// Bytes (envelope header plus body) one connection's egress queue may hold
34/// before a non-waiting enqueue is refused, which closes that connection.
35///
36/// Sized for a client that is slow but still reading: one connection can
37/// multiplex many token streams, and at roughly 200-byte frames (21-byte header
38/// plus body) 4 MiB is about 19,000 queued frames, i.e. thousands of frames on
39/// each of several routes, enough to ride out a pause of seconds. It is also the
40/// most a stuck client can make the daemon hold for it: 4 MiB per connection,
41/// so even 128 stuck connections stay near 512 MiB rather than growing without
42/// bound. A single frame larger than the whole budget is still admitted into an
43/// empty queue so it can be sent at all.
44pub const CONNECTION_EGRESS_BYTE_BUDGET: usize = 4 * 1024 * 1024;
45/// Frame-count backstop for the same queue (the tokio channel's capacity).
46/// The byte budget is the real bound; this one only matters for floods of tiny
47/// or empty frames, which cost bookkeeping per frame rather than bytes. 32,768
48/// frames times 128 bytes is exactly the 4 MiB budget, so for any mean frame of
49/// 128 bytes or more the byte budget binds first; below that (for example
50/// 21-byte empty frames, of which 4 MiB would be almost 200,000) the count
51/// does. The channel allocates slots lazily, so a large capacity costs nothing
52/// until it is used.
53pub const CONNECTION_EGRESS_FRAME_CAP: usize = 32 * 1024;
54/// A pending `route.open` holds one reserved slot in the connection's egress
55/// queue until its module answers, and its response is sent outside the byte
56/// budget. Eight concurrent opens per connection keep a reconnect burst
57/// parallel while bounding both the reserved slots and the bytes that bypass
58/// the budget.
59pub const MAX_PENDING_ROUTE_OPENS_PER_CONNECTION: usize = 8;
60/// A reconnect herd may spread one target's opens over many client connections.
61/// Two safe per-connection bursts retain useful parallelism without restoring
62/// the hundreds-of-binds fanout that serial dispatch used to suppress.
63pub(crate) const MAX_PENDING_ROUTE_BINDS_PER_TARGET: usize =
64    MAX_PENDING_ROUTE_OPENS_PER_CONNECTION * 2;
65pub const DEFAULT_AUTH_DEADLINE: Duration = Duration::from_secs(2);
66// Sized for the restart-herd shape: after a daemon bounce, every live client
67// connection plus all supervised children re-dial within the same second
68// (~120+ observed on the 2026-07-14 fleet). Handshakes are cheap loopback
69// HMAC exchanges; the deadline, not the permit count, is the DoS bound.
70pub const DEFAULT_MAX_UNAUTHENTICATED_CONNECTIONS: usize = 256;
71const CLOSE_DRAIN_GRACE: Duration = Duration::from_secs(2);
72const PRE_AUTH_REJECT_LOG_INTERVAL: Duration = Duration::from_secs(1);
73
74#[derive(Clone, Copy, Debug)]
75enum PreAuthRejectStage {
76    WaitingForHandshakeSlot,
77    HandshakeDeadline,
78    InvalidProof,
79    MalformedMessage,
80    PeerClosed,
81    OtherAuthenticationFailure,
82}
83
84impl PreAuthRejectStage {
85    const fn index(self) -> usize {
86        match self {
87            Self::WaitingForHandshakeSlot => 0,
88            Self::HandshakeDeadline => 1,
89            Self::InvalidProof => 2,
90            Self::MalformedMessage => 3,
91            Self::PeerClosed => 4,
92            Self::OtherAuthenticationFailure => 5,
93        }
94    }
95
96    const fn stage(self) -> &'static str {
97        match self {
98            Self::WaitingForHandshakeSlot => "waiting_for_handshake_slot",
99            Self::HandshakeDeadline => "handshake_deadline",
100            Self::InvalidProof => "client_authentication",
101            Self::MalformedMessage => "malformed_message",
102            Self::PeerClosed => "peer_closed",
103            Self::OtherAuthenticationFailure => "authentication",
104        }
105    }
106
107    const fn reason(self) -> &'static str {
108        match self {
109            Self::WaitingForHandshakeSlot => "capacity_wait_timed_out",
110            Self::HandshakeDeadline => "handshake_deadline_expired",
111            Self::InvalidProof => "client_proof_mismatch",
112            Self::MalformedMessage => "malformed_or_oversized_message",
113            Self::PeerClosed => "peer_closed_before_authentication_completed",
114            Self::OtherAuthenticationFailure => "authentication_failed_before_routing",
115        }
116    }
117}
118
119#[derive(Debug)]
120struct PreAuthRejectRateLimit {
121    next_allowed_ns: AtomicU64,
122    suppressed: AtomicU64,
123}
124
125impl PreAuthRejectRateLimit {
126    const fn new() -> Self {
127        Self {
128            next_allowed_ns: AtomicU64::new(0),
129            suppressed: AtomicU64::new(0),
130        }
131    }
132}
133
134/// Rate-limit state for pre-authentication rejection logs, owned by one
135/// `ServerAuth` so each listener (and each test) counts only its own rejects.
136#[derive(Debug)]
137struct PreAuthRejectLog {
138    epoch: Instant,
139    limiter: PreAuthRejectLimiter,
140}
141
142impl PreAuthRejectLog {
143    fn new() -> Self {
144        Self {
145            epoch: Instant::now(),
146            limiter: PreAuthRejectLimiter::new(),
147        }
148    }
149}
150
151// Report accumulated rejects on the next allowed line instead of running a timer task.
152#[derive(Debug)]
153struct PreAuthRejectLimiter {
154    stages: [PreAuthRejectRateLimit; 6],
155}
156
157impl PreAuthRejectLimiter {
158    fn new() -> Self {
159        Self {
160            stages: std::array::from_fn(|_| PreAuthRejectRateLimit::new()),
161        }
162    }
163
164    fn allow_at(&self, stage: PreAuthRejectStage, now: Duration) -> Option<u64> {
165        let slot = &self.stages[stage.index()];
166        let now_ns = now.as_nanos().min(u64::MAX as u128) as u64;
167        let interval_ns = PRE_AUTH_REJECT_LOG_INTERVAL.as_nanos() as u64;
168        let mut next_allowed_ns = slot.next_allowed_ns.load(AtomicOrdering::Relaxed);
169
170        loop {
171            if now_ns < next_allowed_ns {
172                slot.suppressed.fetch_add(1, AtomicOrdering::Relaxed);
173                return None;
174            }
175
176            match slot.next_allowed_ns.compare_exchange_weak(
177                next_allowed_ns,
178                now_ns.saturating_add(interval_ns),
179                AtomicOrdering::Relaxed,
180                AtomicOrdering::Relaxed,
181            ) {
182                Ok(_) => return Some(slot.suppressed.swap(0, AtomicOrdering::Relaxed)),
183                Err(actual) => next_allowed_ns = actual,
184            }
185        }
186    }
187}
188
189fn log_pre_auth_rejection(
190    error: &ConnectionError,
191    peer_addr: SocketAddr,
192    elapsed: Duration,
193    now: Duration,
194    limiter: &PreAuthRejectLimiter,
195) {
196    let Some(stage) = error.pre_auth_reject_stage() else {
197        return;
198    };
199    let Some(suppressed_since_last) = limiter.allow_at(stage, now) else {
200        return;
201    };
202    let elapsed_ms = elapsed.as_millis().min(u64::MAX as u128) as u64;
203
204    info!(
205        peer_addr = %peer_addr,
206        elapsed_ms,
207        stage = stage.stage(),
208        reason = stage.reason(),
209        auth_stage = ?error.auth_stage(),
210        suppressed_since_last,
211        error = %error,
212        "subc TCP connection rejected before routing"
213    );
214}
215
216async fn handle_accepted_connection<S>(
217    stream: S,
218    router: Arc<Router>,
219    auth: ServerAuth,
220    peer_addr: SocketAddr,
221    accepted_at: Instant,
222) -> Result<(), ConnectionError>
223where
224    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
225{
226    let reject_log = Arc::clone(&auth.pre_auth_reject_log);
227    let result = handle_connection(stream, router, auth).await;
228    if let Err(error) = &result {
229        if error.is_quiet_reject() {
230            log_pre_auth_rejection(
231                error,
232                peer_addr,
233                accepted_at.elapsed(),
234                reject_log.epoch.elapsed(),
235                &reject_log.limiter,
236            );
237        } else {
238            warn!(?peer_addr, error = %error, "subc connection ended with error");
239        }
240    }
241    result
242}
243
244/// Authentication material and DoS bounds applied before a TCP connection may
245/// reach the frame router.
246#[derive(Clone)]
247pub struct ServerAuth {
248    key: Arc<[u8]>,
249    daemon_id: [u8; DAEMON_ID_LEN],
250    daemon_ver: Arc<str>,
251    deadline: Duration,
252    unauthenticated: Arc<Semaphore>,
253    connected_clients: ConnectedClients,
254    pre_auth_reject_log: Arc<PreAuthRejectLog>,
255}
256
257impl ServerAuth {
258    pub fn new(
259        key: Vec<u8>,
260        daemon_id: [u8; DAEMON_ID_LEN],
261        daemon_ver: impl Into<String>,
262    ) -> Self {
263        Self::with_limits(
264            key,
265            daemon_id,
266            daemon_ver,
267            DEFAULT_AUTH_DEADLINE,
268            DEFAULT_MAX_UNAUTHENTICATED_CONNECTIONS,
269        )
270    }
271
272    // Production limits are deliberately not config-routed: loosening pre-auth DoS posture changes attack surface.
273    pub fn with_limits(
274        key: Vec<u8>,
275        daemon_id: [u8; DAEMON_ID_LEN],
276        daemon_ver: impl Into<String>,
277        deadline: Duration,
278        max_unauthenticated: usize,
279    ) -> Self {
280        Self {
281            key: Arc::from(key),
282            daemon_id,
283            daemon_ver: Arc::from(daemon_ver.into()),
284            deadline,
285            unauthenticated: Arc::new(Semaphore::new(max_unauthenticated.max(1))),
286            connected_clients: ConnectedClients::new(),
287            pre_auth_reject_log: Arc::new(PreAuthRejectLog::new()),
288        }
289    }
290
291    pub fn with_connected_clients(mut self, connected_clients: ConnectedClients) -> Self {
292        self.connected_clients = connected_clients;
293        self
294    }
295}
296
297impl fmt::Debug for ServerAuth {
298    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
299        f.debug_struct("ServerAuth")
300            .field("key", &"<redacted>")
301            .field("daemon_id", &self.daemon_id)
302            .field("daemon_ver", &self.daemon_ver)
303            .field("deadline", &self.deadline)
304            .finish_non_exhaustive()
305    }
306}
307
308/// Serve an already-bound TCP listener. Each accepted connection gets its own
309/// async task so concurrent clients do not block the accept loop.
310pub async fn serve_listener(
311    listener: TcpListener,
312    router: Arc<Router>,
313    auth: ServerAuth,
314) -> Result<(), ServerError> {
315    serve_listener_with_accept(listener.local_addr().ok(), router, auth, || {
316        listener.accept()
317    })
318    .await
319}
320
321async fn serve_listener_with_accept<A, F>(
322    local_addr: Option<SocketAddr>,
323    router: Arc<Router>,
324    auth: ServerAuth,
325    mut accept: A,
326) -> Result<(), ServerError>
327where
328    A: FnMut() -> F,
329    F: std::future::Future<Output = io::Result<(tokio::net::TcpStream, SocketAddr)>>,
330{
331    loop {
332        let (stream, peer_addr) = match accept().await {
333            Ok(accepted) => accepted,
334            Err(source) => {
335                let kind = source.kind();
336                // Aborted/reset connections and interrupted syscalls affect one accept only.
337                // File-descriptor or socket-buffer exhaustion needs a short backoff to avoid spinning;
338                // other errors may mean the listener itself is unusable.
339                let exhausted = accept_resource_exhausted(&source);
340                if matches!(
341                    kind,
342                    io::ErrorKind::ConnectionAborted
343                        | io::ErrorKind::ConnectionReset
344                        | io::ErrorKind::Interrupted
345                ) || exhausted
346                {
347                    warn!(?local_addr, error = %source, "temporary TCP accept failure");
348                    if exhausted {
349                        tokio::time::sleep(Duration::from_millis(75)).await;
350                    }
351                    continue;
352                }
353                return Err(ServerError::Accept { local_addr, source });
354            }
355        };
356        let accepted_at = Instant::now();
357        // Every route frame is a discrete message whose reply the peer is waiting
358        // for, so there is never a later write for Nagle to coalesce with -- it can
359        // only hold a frame back until an ACK arrives.
360        //
361        // MEASURED: no effect on this transport. A 50-sample-per-arm sweep from
362        // 1 to 32 KiB over the full client->daemon->module->client path showed a
363        // flat 0.28-0.38ms p50 with no step at any buffer boundary, because a
364        // loopback ACK returns in microseconds and never reaches the delayed-ACK
365        // timer that makes Nagle expensive on a real network. Kept anyway: it is
366        // one syscall at accept, it removes the mechanism rather than relying on
367        // loopback staying fast, and Windows loopback was not part of that sweep.
368        // Do not cite it as a latency fix -- the measurement says it is not one.
369        //
370        // Failure is not fatal: the connection works, and refusing to serve a
371        // client over a socket option would be worse than anything it saves.
372        if let Err(source) = stream.set_nodelay(true) {
373            warn!(?peer_addr, error = %source, "could not disable Nagle on accepted connection");
374        }
375        debug!(?peer_addr, ?local_addr, "accepted subc TCP connection");
376        let router = Arc::clone(&router);
377        let auth = auth.clone();
378        tokio::spawn(async move {
379            let _ = handle_accepted_connection(stream, router, auth, peer_addr, accepted_at).await;
380        });
381    }
382}
383
384fn accept_resource_exhausted(error: &io::Error) -> bool {
385    // ENFILE/EMFILE/ENOBUFS by name rather than number: the numbers differ by
386    // platform (ENOBUFS is 55 on macOS and 105 on Linux, and each number means
387    // something unrelated on the other), so a literal list misclassifies.
388    // Winsock: WSAEMFILE and WSAENOBUFS.
389    #[cfg(unix)]
390    {
391        use rustix::io::Errno;
392        let exhausted = [Errno::NFILE, Errno::MFILE, Errno::NOBUFS];
393        error
394            .raw_os_error()
395            .is_some_and(|code| exhausted.iter().any(|e| e.raw_os_error() == code))
396    }
397    #[cfg(windows)]
398    {
399        matches!(error.raw_os_error(), Some(10024 | 10055))
400    }
401    #[cfg(not(any(unix, windows)))]
402    {
403        let _ = error;
404        false
405    }
406}
407
408/// Serve all already-bound loopback TCP listeners until one accept loop fails.
409pub async fn serve_listeners(
410    listeners: Vec<TcpListener>,
411    router: Arc<Router>,
412    auth: ServerAuth,
413) -> Result<(), ServerError> {
414    if listeners.is_empty() {
415        return Err(ServerError::NoListeners);
416    }
417
418    let (tx, mut rx) = mpsc::channel(listeners.len());
419    let mut accept_tasks = AbortTasksOnDrop::default();
420    for listener in listeners {
421        let router = Arc::clone(&router);
422        let auth = auth.clone();
423        let tx = tx.clone();
424        accept_tasks.push(tokio::spawn(async move {
425            let result = serve_listener(listener, router, auth).await;
426            let _ = tx.send(result).await;
427        }));
428    }
429    drop(tx);
430
431    rx.recv().await.unwrap_or(Ok(()))
432}
433
434#[derive(Default)]
435struct AbortTasksOnDrop {
436    handles: Vec<JoinHandle<()>>,
437}
438
439impl AbortTasksOnDrop {
440    fn push(&mut self, handle: JoinHandle<()>) {
441        self.handles.push(handle);
442    }
443}
444
445impl Drop for AbortTasksOnDrop {
446    fn drop(&mut self) {
447        for handle in &self.handles {
448            if !handle.is_finished() {
449                handle.abort();
450            }
451        }
452    }
453}
454
455#[derive(Debug)]
456enum ConnectionLoopExit {
457    PeerClosed,
458    CloseRequested(CloseReason),
459}
460
461/// Run the authenticated frame read -> route loop for one connection.
462///
463/// Every accepted TCP connection must complete the key-auth prelude before any
464/// envelope bytes are read by the router. Outbound frames flow through a bounded
465/// [`FrameSink`] drained by one writer task. Inbound dispatch remains serial for
466/// every frame except `route.open`, activity operations and privileged requests;
467/// their bind, file and operator waits run in connection-owned tasks so
468/// unrelated frames can keep flowing.
469pub async fn handle_connection<S>(
470    mut stream: S,
471    router: Arc<Router>,
472    auth: ServerAuth,
473) -> Result<(), ConnectionError>
474where
475    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
476{
477    // Capture only socket facts here. Process-table lookup happens lazily on a
478    // blocking thread when privileged dispatch needs a caller description.
479    let endpoints = (&stream as &dyn Any)
480        .downcast_ref::<tokio::net::TcpStream>()
481        .and_then(|tcp| tcp.local_addr().ok().zip(tcp.peer_addr().ok()));
482    // Over-cap connections WAIT for a handshake slot (bounded by the auth deadline)
483    // instead of being reset. A restart herd — every client and supervised child
484    // re-dialing the fresh daemon at once — otherwise loses supervised modules to
485    // the permit lottery: each reset burns a module restart-budget slot, and a
486    // module that treats auth failure as fatal can exhaust its budget into
487    // state=failed within the boot window (2026-07-14 aft outage). On loopback
488    // with the pre-auth HMAC deadline, a bounded queue is strictly safer than a
489    // reset.
490    //
491    // Pre-auth time is governed by TWO deliberately separate budgets: the queue
492    // wait below is bounded by `auth.deadline`, and `authenticate_server` then
493    // starts a FRESH `auth.deadline` for the handshake itself. Total pre-auth
494    // occupancy per connection is therefore up to 2x the configured deadline.
495    // This is intentional, not an accounting slip: charging queue time against
496    // the handshake budget would hand a herd-queued supervised module a
497    // near-zero handshake window under CPU saturation, recreating the fatal
498    // auth-failure -> restart-budget-burn path this queue exists to prevent.
499    // On a loopback-only, key-authenticated listener the doubled bound is a
500    // per-connection occupancy cost, not a meaningful DoS surface.
501    let permit =
502        match tokio::time::timeout(auth.deadline, auth.unauthenticated.clone().acquire_owned())
503            .await
504        {
505            Ok(Ok(permit)) => permit,
506            Ok(Err(_)) | Err(_) => {
507                let _ = stream.shutdown().await;
508                return Err(ConnectionError::UnauthenticatedCapacity);
509            }
510        };
511
512    let authenticated = authenticate_server(
513        &mut stream,
514        auth.key.as_ref(),
515        &auth.daemon_id,
516        auth.daemon_ver.as_ref(),
517        auth.deadline,
518    )
519    .await
520    .map_err(ConnectionError::Auth)?;
521    drop(permit);
522
523    let mut connection = router.begin_connection();
524    let connection_id = connection.id();
525    router.record_caller_endpoints(connection_id, endpoints);
526    // `authenticated.role` is CLIENT-SUPPLIED and unverified -- the handshake
527    // proves possession of the connection key, nothing about who is calling. So
528    // this exclusion is safe ONLY while the role decides a reporting question and
529    // never an authorization one: the watchdog's own loopback probe would
530    // otherwise inflate the very count it exists to sanity-check.
531    //
532    // The consequence of it being unverified, stated so nobody has to rediscover
533    // it: ANY client holding the key can claim this role and omit itself from
534    // `connected_clients`. That is a gauge a caller can lie to, and it is
535    // acceptable because the gauge informs an operator rather than gating
536    // anything. IF A ROLE IS EVER USED TO DECIDE ADMISSION, CAPACITY, OR PRIVILEGE,
537    // this stops being safe and the role must be attested rather than declared --
538    // the daemon already has the mechanism for that in the spawn-nonce path used
539    // for module identity.
540    let _connected_client = (authenticated.role != WATCHDOG_CLIENT_ROLE)
541        .then(|| auth.connected_clients.open(connection_id));
542    let close_receiver = connection.take_close_receiver();
543    debug!(
544        connection_id = connection_id.get(),
545        "subc authenticated connection opened"
546    );
547
548    let (read_half, write_half) = tokio::io::split(stream);
549    // Authentication is complete before this buffer can read ahead. The connection loop
550    // retains a partial frame read across completed route.open tasks; dropping the read
551    // happens only when a frame finishes or the connection ends.
552    let mut read_half = BufReader::new(read_half);
553    let (egress, rx) = connection_egress();
554    let mut writer = tokio::spawn(drain_writer(write_half, rx));
555
556    let ctx = RouteCtx {
557        connection_id,
558        egress: egress.clone(),
559    };
560
561    let mut route_open_tasks = JoinSet::new();
562    let mut privileged_tasks = JoinSet::new();
563    let loop_result = connection_loop(
564        &mut read_half,
565        Arc::clone(&router),
566        ctx.clone(),
567        close_receiver,
568        &mut route_open_tasks,
569        &mut privileged_tasks,
570    )
571    .await;
572
573    // Every in-flight route.open must have FINISHED before `connection` drops,
574    // because that drop runs the forwarding cleanup for this id, and cleanup is
575    // also where the id's closing mark is lifted. Dropping a JoinSet only
576    // requests abort: a task already running on another worker keeps going
577    // until its next await, and could commit a route for this connection after
578    // cleanup has cleared it, with nothing left to refuse it and nothing to
579    // clean it up. shutdown() aborts and then waits.
580    route_open_tasks.shutdown().await;
581    // Abort held approvals and wait for their withdrawal guards before lifting
582    // this connection's cleanup state or removing its caller endpoints.
583    privileged_tasks.shutdown().await;
584
585    drop(ctx);
586    drop(egress);
587    drop(connection);
588
589    let close_reason = match &loop_result {
590        Ok(ConnectionLoopExit::CloseRequested(reason)) => Some(reason.to_string()),
591        Ok(ConnectionLoopExit::PeerClosed) | Err(_) => None,
592    };
593    let writer_result = if close_reason.is_some() {
594        match timeout(CLOSE_DRAIN_GRACE, &mut writer).await {
595            Ok(result) => Some(result.map_err(ConnectionError::WriterTask)),
596            Err(_) => {
597                warn!(
598                    connection_id = connection_id.get(),
599                    grace = ?CLOSE_DRAIN_GRACE,
600                    "connection writer did not drain after close request; aborting writer task"
601                );
602                writer.abort();
603                let _ = writer.await;
604                None
605            }
606        }
607    } else {
608        Some(writer.await.map_err(ConnectionError::WriterTask))
609    };
610
611    let result = if let Some(reason) = close_reason.as_deref() {
612        match writer_result {
613            Some(Ok(Ok(()))) | None => Ok(()),
614            Some(Ok(Err(writer_err))) => {
615                debug!(
616                    connection_id = connection_id.get(),
617                    close_reason = reason,
618                    writer_error = %writer_err,
619                    "writer failed after requested connection close"
620                );
621                Ok(())
622            }
623            Some(Err(join_err)) => {
624                warn!(
625                    connection_id = connection_id.get(),
626                    close_reason = reason,
627                    join_error = %join_err,
628                    "writer task join failed after requested connection close"
629                );
630                Ok(())
631            }
632        }
633    } else {
634        let writer_result =
635            writer_result.expect("writer result is present without a close request");
636        match (loop_result, writer_result) {
637            (Err(loop_err), Ok(Ok(()))) => Err(loop_err),
638            (Err(loop_err), Ok(Err(writer_err))) => {
639                warn!(
640                    connection_id = connection_id.get(),
641                    writer_error = %writer_err,
642                    "writer failed while closing after connection error"
643                );
644                Err(loop_err)
645            }
646            (Err(loop_err), Err(join_err)) => {
647                warn!(
648                    connection_id = connection_id.get(),
649                    join_error = %join_err,
650                    "writer task join failed while closing after connection error"
651                );
652                Err(loop_err)
653            }
654            (Ok(ConnectionLoopExit::PeerClosed), Ok(Ok(()))) => Ok(()),
655            (Ok(ConnectionLoopExit::PeerClosed), Ok(Err(writer_err))) => {
656                Err(ConnectionError::FrameIo(writer_err))
657            }
658            (Ok(ConnectionLoopExit::PeerClosed), Err(join_err)) => Err(join_err),
659            (Ok(ConnectionLoopExit::CloseRequested(_)), _) => {
660                unreachable!("close requests are handled before normal writer result matching")
661            }
662        }
663    };
664
665    match &result {
666        Ok(()) => {
667            if let Some(reason) = close_reason.as_deref() {
668                debug!(
669                    connection_id = connection_id.get(),
670                    close_reason = reason,
671                    "subc connection closed by request"
672                );
673            } else {
674                debug!(
675                    connection_id = connection_id.get(),
676                    "subc connection closed"
677                );
678            }
679        }
680        Err(err) => debug!(
681            connection_id = connection_id.get(),
682            error = %err,
683            "subc connection exited with error"
684        ),
685    }
686
687    result
688}
689
690async fn connection_loop<R>(
691    read_half: &mut R,
692    router: Arc<Router>,
693    ctx: RouteCtx,
694    mut close_receiver: ConnectionCloseReceiver,
695    route_open_tasks: &mut JoinSet<Result<(), RouterError>>,
696    privileged_tasks: &mut JoinSet<Result<(), RouterError>>,
697) -> Result<ConnectionLoopExit, ConnectionError>
698where
699    R: AsyncRead + Unpin,
700{
701    let mut activity_tasks = JoinSet::new();
702    loop {
703        while let Some(result) = activity_tasks.try_join_next() {
704            finish_route_open_task(result)?;
705        }
706        while let Some(result) = privileged_tasks.try_join_next() {
707            finish_route_open_task(result)?;
708        }
709        while let Some(result) = route_open_tasks.try_join_next() {
710            finish_route_open_task(result)?;
711        }
712
713        // Keep the same read future when a route.open completes: read_frame owns
714        // partial header/body buffers that would be lost if that future were dropped.
715        let read = read_frame(&mut *read_half);
716        tokio::pin!(read);
717        let frame = loop {
718            tokio::select! {
719                close = &mut close_receiver => {
720                    return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
721                }
722                result = route_open_tasks.join_next(), if !route_open_tasks.is_empty() => {
723                    finish_route_open_task(
724                        result.expect("a non-empty route.open JoinSet has a next task")
725                    )?;
726                }
727                result = activity_tasks.join_next(), if !activity_tasks.is_empty() => {
728                    finish_route_open_task(result.expect("non-empty activity task set"))?;
729                }
730                result = privileged_tasks.join_next(), if !privileged_tasks.is_empty() => {
731                    finish_route_open_task(result.expect("non-empty privileged task set"))?;
732                }
733                result = &mut read => {
734                    break match result.map_err(ConnectionError::FrameIo)? {
735                        Some(frame) => frame,
736                        None => return Ok(ConnectionLoopExit::PeerClosed),
737                    };
738                }
739            }
740        };
741
742        if frame.header.channel == 0
743            && frame.header.ty == subc_protocol::FrameType::Request
744            && serde_json::from_slice::<serde_json::Value>(&frame.body)
745                .ok()
746                .is_some_and(|body| {
747                    matches!(
748                        body["op"].as_str(),
749                        Some(
750                            subc_control::ops::ACTIVITY_RECORD
751                                | subc_control::ops::ACTIVITY_QUERY
752                                | subc_control::ops::ACTIVITY_SUBSCRIBE
753                        )
754                    )
755                })
756        {
757            // Bound outstanding activity tasks as well as the disk queue,
758            // so the reader never waits for either capacity when the disk stalls.
759            if activity_tasks.len() >= 128 {
760                let response = crate::control::control_error_frame(
761                    &frame,
762                    subc_control::ACTIVITY_UNAVAILABLE,
763                    "activity queue full",
764                )
765                .map_err(ConnectionError::Router)?;
766                ctx.egress
767                    .send(response)
768                    .await
769                    .map_err(ConnectionError::Router)?;
770            } else {
771                let task_router = Arc::clone(&router);
772                let task_ctx = ctx.clone();
773                activity_tasks.spawn(async move {
774                    route_open_tail(task_router, task_ctx, frame, Instant::now()).await
775                });
776            }
777            continue;
778        }
779
780        if let Some(target_module_id) = router.route_open_target(&frame) {
781            // A task can finish while the reader is waiting for the next frame.
782            // Reap it before checking admission so completed work never occupies
783            // one of the deliberately scarce connection slots.
784            while let Some(result) = route_open_tasks.try_join_next() {
785                finish_route_open_task(result)?;
786            }
787
788            if route_open_tasks.len() >= MAX_PENDING_ROUTE_OPENS_PER_CONNECTION {
789                // Never wait for capacity here: doing so would recreate the same
790                // reader head-of-line stall with a smaller constant.
791                let refusal = router
792                    .route_open_capacity_refusal(
793                        &ctx,
794                        &frame,
795                        &target_module_id,
796                        route_open_tasks.len(),
797                        MAX_PENDING_ROUTE_OPENS_PER_CONNECTION,
798                    )
799                    .map_err(ConnectionError::Router)?;
800                let send_result = tokio::select! {
801                    close = &mut close_receiver => {
802                        return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
803                    }
804                    result = ctx.egress.send(refusal) => result,
805                };
806                send_result.map_err(ConnectionError::Router)?;
807                continue;
808            }
809
810            let task_router = Arc::clone(&router);
811            let task_ctx = ctx.clone();
812            // Start slow-dispatch timing before spawn so it covers the task's
813            // full scheduler and handler lifetime, as inline timing did.
814            let dispatch_started_at = Instant::now();
815            route_open_tasks.spawn(async move {
816                route_open_tail(task_router, task_ctx, frame, dispatch_started_at).await
817            });
818            continue;
819        }
820
821        if router.defers_privileged_control(&ctx, &frame) {
822            let task_router = Arc::clone(&router);
823            let task_ctx = ctx.clone();
824            let started = Instant::now();
825            privileged_tasks
826                .spawn(async move { route_open_tail(task_router, task_ctx, frame, started).await });
827            continue;
828        }
829
830        // Other frames still route serially: routing can backpressure reads,
831        // and a connection close cancels the in-progress route.
832        let route_result = tokio::select! {
833            close = &mut close_receiver => {
834                return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
835            }
836            result = router.route_for_connection(&ctx, frame) => result,
837        };
838
839        if let Err(err) = route_result {
840            if let Some(error_frame) = err.to_error_frame() {
841                warn!(
842                    connection_id = ctx.connection_id.get(),
843                    error = %err,
844                    "routing failure recovered with ERROR frame"
845                );
846                let send_result = tokio::select! {
847                    close = &mut close_receiver => {
848                        return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
849                    }
850                    result = ctx.egress.send(error_frame) => result,
851                };
852                send_result.map_err(ConnectionError::Router)?;
853            } else {
854                debug!(
855                    connection_id = ctx.connection_id.get(),
856                    error = %err,
857                    "fatal routing failure"
858                );
859                return Err(ConnectionError::Router(err));
860            }
861        }
862    }
863}
864
865async fn route_open_tail(
866    router: Arc<Router>,
867    ctx: RouteCtx,
868    frame: crate::Frame,
869    dispatch_started_at: Instant,
870) -> Result<(), RouterError> {
871    match router
872        .route_for_connection_started(&ctx, frame, Some(dispatch_started_at))
873        .await
874    {
875        Ok(()) => Ok(()),
876        Err(err) => {
877            let Some(error_frame) = err.to_error_frame() else {
878                return Err(err);
879            };
880            warn!(
881                connection_id = ctx.connection_id.get(),
882                error = %err,
883                "routing failure recovered with ERROR frame"
884            );
885            ctx.egress.send(error_frame).await
886        }
887    }
888}
889
890fn finish_route_open_task(
891    result: Result<Result<(), RouterError>, tokio::task::JoinError>,
892) -> Result<(), ConnectionError> {
893    match result {
894        Ok(Ok(())) => Ok(()),
895        Ok(Err(err)) => Err(ConnectionError::Router(err)),
896        Err(err) => Err(ConnectionError::Router(RouterError::backend(
897            0,
898            0,
899            format!("route.open task failed: {err}"),
900        ))),
901    }
902}
903
904fn close_reason(
905    result: Result<CloseReason, tokio::sync::oneshot::error::RecvError>,
906) -> CloseReason {
907    result.unwrap_or_else(|_| {
908        CloseReason::new(
909            "close_registry_dropped",
910            "connection close registration was dropped without a reason",
911        )
912    })
913}
914
915/// The outbound queue for one daemon connection: a sink bounded by
916/// [`CONNECTION_EGRESS_BYTE_BUDGET`] and [`CONNECTION_EGRESS_FRAME_CAP`], and
917/// the receiver its writer drains.
918pub(crate) fn connection_egress() -> (FrameSink, mpsc::Receiver<crate::router::OutboundFrame>) {
919    let (tx, rx) = mpsc::channel::<crate::router::OutboundFrame>(CONNECTION_EGRESS_FRAME_CAP);
920    (
921        FrameSink::with_byte_budget(tx, CONNECTION_EGRESS_BYTE_BUDGET),
922        rx,
923    )
924}
925
926async fn drain_writer<W>(
927    write_half: W,
928    mut rx: mpsc::Receiver<crate::router::OutboundFrame>,
929) -> Result<(), FrameIoError>
930where
931    W: AsyncWrite + Unpin,
932{
933    let mut writer = BufWriter::new(write_half);
934    while let Some(outbound) = rx.recv().await {
935        write_outbound(&mut writer, outbound).await?;
936        while let Ok(outbound) = rx.try_recv() {
937            write_outbound(&mut writer, outbound).await?;
938        }
939        writer.flush().await.map_err(FrameIoError::Io)?;
940    }
941    writer.flush().await.map_err(FrameIoError::Io)?;
942    Ok(())
943}
944
945/// Reply-path half of slow-control diagnosis: a channel-0 reply that sat in
946/// the writer queue past the threshold is reported with its queue residency
947/// and its own write duration separated, because "writer task not scheduled"
948/// and "socket write blocked" are different defects and the sum hides which.
949/// Data-plane frames are exempt: their latency is the client's own flow
950/// control, and logging them would drown the control signal in bulk traffic.
951async fn write_outbound<W>(
952    writer: &mut BufWriter<W>,
953    outbound: crate::router::OutboundFrame,
954) -> Result<(), FrameIoError>
955where
956    W: AsyncWrite + Unpin,
957{
958    const SLOW_REPLY_QUEUE: Duration = Duration::from_millis(1000);
959    let queued = outbound.enqueued_at.elapsed();
960    // Held until this frame has been written, then dropped at the end of this
961    // function, which gives its bytes back to the connection's egress budget.
962    let charge = outbound.charge;
963    if let Some(charge) = &charge {
964        charge.taken_by_writer();
965    }
966    let frame = outbound.frame;
967    if frame.header.channel == 0 && queued >= SLOW_REPLY_QUEUE {
968        let write_started = std::time::Instant::now();
969        let result = write_frame(writer, &frame).await;
970        tracing::warn!(
971            corr = frame.header.corr,
972            queued_ms = queued.as_millis() as u64,
973            write_ms = write_started.elapsed().as_millis() as u64,
974            "slow control reply write"
975        );
976        result?;
977    } else {
978        write_frame(writer, &frame).await?;
979    }
980    if let Some(flushed) = outbound.flushed {
981        writer.flush().await.map_err(FrameIoError::Io)?;
982        let _ = flushed.send(());
983    }
984    Ok(())
985}
986
987#[cfg(all(test, unix))]
988#[tokio::test]
989async fn shutdown_notice_ack_waits_for_socket_flush() {
990    let (socket, mut peer) = tokio::io::duplex(1);
991    let (tx, rx) = mpsc::channel(1);
992    let sink = FrameSink::new(tx);
993    let writer = tokio::spawn(drain_writer(socket, rx));
994    let frame = crate::Frame::build(
995        subc_protocol::FrameType::Push,
996        subc_protocol::Flags::new(false, subc_protocol::Priority::Interactive, false),
997        0,
998        0,
999        0,
1000        b"notice".to_vec(),
1001    )
1002    .unwrap();
1003    let send = sink.send_flushed(frame);
1004    tokio::pin!(send);
1005    assert!(
1006        timeout(Duration::from_millis(20), &mut send).await.is_err(),
1007        "queueing bytes is not a socket flush acknowledgement"
1008    );
1009    let (sent, received) = timeout(Duration::from_secs(1), async {
1010        tokio::join!(&mut send, read_frame(&mut peer))
1011    })
1012    .await
1013    .unwrap();
1014    sent.unwrap();
1015    assert_eq!(received.unwrap().unwrap().body, b"notice");
1016    writer.abort();
1017}
1018
1019#[derive(Debug)]
1020pub enum ServerError {
1021    NoListeners,
1022    Accept {
1023        local_addr: Option<SocketAddr>,
1024        source: io::Error,
1025    },
1026}
1027
1028impl fmt::Display for ServerError {
1029    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1030        match self {
1031            Self::NoListeners => write!(f, "no TCP listeners were provided"),
1032            Self::Accept { local_addr, source } => match local_addr {
1033                Some(addr) => write!(f, "failed to accept TCP connection on {addr}: {source}"),
1034                None => write!(f, "failed to accept TCP connection: {source}"),
1035            },
1036        }
1037    }
1038}
1039
1040impl Error for ServerError {
1041    fn source(&self) -> Option<&(dyn Error + 'static)> {
1042        match self {
1043            Self::Accept { source, .. } => Some(source),
1044            Self::NoListeners => None,
1045        }
1046    }
1047}
1048
1049#[derive(Debug)]
1050pub enum ConnectionError {
1051    Auth(AuthError),
1052    UnauthenticatedCapacity,
1053    FrameIo(FrameIoError),
1054    Router(RouterError),
1055    WriterTask(tokio::task::JoinError),
1056}
1057
1058impl ConnectionError {
1059    fn is_quiet_reject(&self) -> bool {
1060        matches!(self, Self::Auth(_) | Self::UnauthenticatedCapacity)
1061    }
1062
1063    fn pre_auth_reject_stage(&self) -> Option<PreAuthRejectStage> {
1064        match self {
1065            Self::UnauthenticatedCapacity => Some(PreAuthRejectStage::WaitingForHandshakeSlot),
1066            Self::Auth(AuthError::Timeout { .. }) => Some(PreAuthRejectStage::HandshakeDeadline),
1067            Self::Auth(
1068                AuthError::InvalidClientAuth
1069                | AuthError::InvalidServerProof
1070                | AuthError::DaemonIdMismatch,
1071            ) => Some(PreAuthRejectStage::InvalidProof),
1072            Self::Auth(AuthError::MessageTooLarge { .. } | AuthError::JsonDecode { .. }) => {
1073                Some(PreAuthRejectStage::MalformedMessage)
1074            }
1075            Self::Auth(AuthError::UnexpectedEof { .. }) => Some(PreAuthRejectStage::PeerClosed),
1076            Self::Auth(AuthError::Io { source, .. })
1077                if matches!(
1078                    source.kind(),
1079                    io::ErrorKind::BrokenPipe
1080                        | io::ErrorKind::ConnectionAborted
1081                        | io::ErrorKind::ConnectionReset
1082                        | io::ErrorKind::UnexpectedEof
1083                ) =>
1084            {
1085                Some(PreAuthRejectStage::PeerClosed)
1086            }
1087            Self::Auth(_) => Some(PreAuthRejectStage::OtherAuthenticationFailure),
1088            Self::FrameIo(_) | Self::Router(_) | Self::WriterTask(_) => None,
1089        }
1090    }
1091
1092    fn auth_stage(&self) -> Option<AuthStage> {
1093        match self {
1094            Self::Auth(
1095                AuthError::Io { stage, .. }
1096                | AuthError::Timeout { stage, .. }
1097                | AuthError::UnexpectedEof { stage, .. }
1098                | AuthError::MessageTooLarge { stage, .. }
1099                | AuthError::JsonEncode { stage, .. }
1100                | AuthError::JsonDecode { stage, .. },
1101            ) => Some(*stage),
1102            Self::Auth(
1103                AuthError::Random(_)
1104                | AuthError::KeyTooShort { .. }
1105                | AuthError::InvalidServerProof
1106                | AuthError::DaemonIdMismatch
1107                | AuthError::InvalidClientAuth,
1108            )
1109            | Self::UnauthenticatedCapacity
1110            | Self::FrameIo(_)
1111            | Self::Router(_)
1112            | Self::WriterTask(_) => None,
1113        }
1114    }
1115}
1116
1117impl fmt::Display for ConnectionError {
1118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1119        match self {
1120            Self::Auth(err) => write!(f, "connection auth failed: {err}"),
1121            Self::UnauthenticatedCapacity => write!(
1122                f,
1123                "too many concurrent unauthenticated subc TCP connections"
1124            ),
1125            Self::FrameIo(err) => write!(f, "frame connection error: {err}"),
1126            Self::Router(err) => write!(f, "router connection error: {err}"),
1127            Self::WriterTask(err) => write!(f, "connection writer task failed: {err}"),
1128        }
1129    }
1130}
1131
1132impl Error for ConnectionError {
1133    fn source(&self) -> Option<&(dyn Error + 'static)> {
1134        match self {
1135            Self::Auth(err) => Some(err),
1136            Self::FrameIo(err) => Some(err),
1137            Self::Router(err) => Some(err),
1138            Self::WriterTask(err) => Some(err),
1139            Self::UnauthenticatedCapacity => None,
1140        }
1141    }
1142}
1143
1144#[cfg(test)]
1145mod tests {
1146    use std::{
1147        pin::Pin,
1148        sync::atomic::{AtomicUsize, Ordering},
1149        task::{Context, Poll},
1150    };
1151
1152    use super::*;
1153    use serde::{de::DeserializeOwned, Serialize};
1154    use subc_protocol::{
1155        DecodeError, ErrorBody, Flags, FrameType, Priority, HEADER_LEN, PROTOCOL_VERSION,
1156    };
1157    use tokio::io::{duplex, AsyncReadExt, AsyncWriteExt, ReadBuf};
1158
1159    use subc_transport::{
1160        authenticate_client, compute_proof, ClientAuth, ClientHello, ConnectionInfo, Endpoint,
1161        ServerProof, CLIENT_AUTH_DOMAIN, NONCE_LEN, SCHEMA_VERSION,
1162    };
1163
1164    use crate::{ControlHandler, EchoBackend, Frame, ReadStage, Registry};
1165
1166    const TEST_DEADLINE: Duration = Duration::from_secs(2);
1167    const TEST_DAEMON_VER: &str = "test-subc-server";
1168
1169    /// Reply-path stamp, slow polarity: a channel-0 reply whose queue residency
1170    /// exceeds the threshold must produce the slow-control-reply-write WARN with
1171    /// queued_ms reflecting the residency. Uses a backdated stamp rather than a
1172    /// real stall so the test is fast and deterministic.
1173    #[tokio::test]
1174    async fn stale_queued_control_reply_logs_slow_reply_write() {
1175        let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::WARN);
1176        let (tx, rx) = mpsc::channel::<crate::router::OutboundFrame>(4);
1177        let reply = Frame::build_with_version(
1178            PROTOCOL_VERSION,
1179            FrameType::Error,
1180            Flags::new(false, Priority::Interactive, false),
1181            0,
1182            0,
1183            7,
1184            serde_json::to_vec(&ErrorBody {
1185                code: "test".into(),
1186                message: "reply".into(),
1187                detail: None,
1188            })
1189            .expect("body encodes"),
1190        )
1191        .expect("frame builds");
1192        tx.send(crate::router::OutboundFrame {
1193            frame: reply,
1194            enqueued_at: std::time::Instant::now() - Duration::from_millis(1500),
1195            flushed: None,
1196            charge: None,
1197        })
1198        .await
1199        .expect("queued");
1200        drop(tx);
1201        let (write_half, mut read_half) = duplex(64 * 1024);
1202        drain_writer(write_half, rx).await.expect("writer drains");
1203        let mut sink = Vec::new();
1204        read_half.read_to_end(&mut sink).await.expect("read");
1205        assert!(!sink.is_empty(), "frame reached the socket");
1206        let captured = crate::router::test_log::captured_logs(&logs);
1207        assert!(
1208            captured.contains("slow control reply write") && captured.contains("corr=7"),
1209            "expected slow reply WARN naming corr, got: {captured}"
1210        );
1211        let queued_ms: u64 = captured
1212            .split("queued_ms=")
1213            .nth(1)
1214            .and_then(|s| s.split_whitespace().next())
1215            .and_then(|s| s.parse().ok())
1216            .expect("queued_ms present");
1217        assert!(
1218            queued_ms >= 1500,
1219            "queued_ms reflects residency: {queued_ms}"
1220        );
1221    }
1222
1223    /// Fast polarity: a promptly-drained control reply and a stale DATA-PLANE
1224    /// frame must both stay silent — the WARN is channel-0-only by design, and
1225    /// a healthy queue must add zero log volume.
1226    #[tokio::test]
1227    async fn fresh_control_and_stale_data_frames_log_nothing() {
1228        let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::WARN);
1229        let (tx, rx) = mpsc::channel::<crate::router::OutboundFrame>(4);
1230        let control = Frame::build_with_version(
1231            PROTOCOL_VERSION,
1232            FrameType::Error,
1233            Flags::new(false, Priority::Interactive, false),
1234            0,
1235            0,
1236            8,
1237            serde_json::to_vec(&ErrorBody {
1238                code: "test".into(),
1239                message: "fresh".into(),
1240                detail: None,
1241            })
1242            .expect("body encodes"),
1243        )
1244        .expect("frame builds");
1245        tx.send(crate::router::OutboundFrame {
1246            frame: control,
1247            enqueued_at: std::time::Instant::now(),
1248            flushed: None,
1249            charge: None,
1250        })
1251        .await
1252        .expect("queued");
1253        let data = Frame::build_with_version(
1254            PROTOCOL_VERSION,
1255            FrameType::Error,
1256            Flags::new(false, Priority::Interactive, false),
1257            9,
1258            1,
1259            9,
1260            serde_json::to_vec(&ErrorBody {
1261                code: "test".into(),
1262                message: "data".into(),
1263                detail: None,
1264            })
1265            .expect("body encodes"),
1266        )
1267        .expect("frame builds");
1268        tx.send(crate::router::OutboundFrame {
1269            frame: data,
1270            enqueued_at: std::time::Instant::now() - Duration::from_millis(5000),
1271            flushed: None,
1272            charge: None,
1273        })
1274        .await
1275        .expect("queued");
1276        drop(tx);
1277        let (write_half, _read_half) = duplex(64 * 1024);
1278        drain_writer(write_half, rx).await.expect("writer drains");
1279        let captured = crate::router::test_log::captured_logs(&logs);
1280        assert!(
1281            !captured.contains("slow control reply write"),
1282            "no WARN for fresh control or stale data frames, got: {captured}"
1283        );
1284    }
1285
1286    struct CountingReader {
1287        bytes: Vec<u8>,
1288        offset: usize,
1289        first_read_end: Option<usize>,
1290        reads: Arc<AtomicUsize>,
1291    }
1292
1293    impl CountingReader {
1294        fn new(bytes: Vec<u8>, first_read_end: Option<usize>) -> (Self, Arc<AtomicUsize>) {
1295            let reads = Arc::new(AtomicUsize::new(0));
1296            (
1297                Self {
1298                    bytes,
1299                    offset: 0,
1300                    first_read_end,
1301                    reads: Arc::clone(&reads),
1302                },
1303                reads,
1304            )
1305        }
1306    }
1307
1308    impl AsyncRead for CountingReader {
1309        fn poll_read(
1310            mut self: Pin<&mut Self>,
1311            _cx: &mut Context<'_>,
1312            buf: &mut ReadBuf<'_>,
1313        ) -> Poll<io::Result<()>> {
1314            self.reads.fetch_add(1, Ordering::Relaxed);
1315            let available = self.bytes.len().saturating_sub(self.offset);
1316            let first_read_remaining = self
1317                .first_read_end
1318                .filter(|end| self.offset < *end)
1319                .map_or(available, |end| end - self.offset);
1320            let count = available.min(first_read_remaining).min(buf.remaining());
1321            let end = self.offset + count;
1322            buf.put_slice(&self.bytes[self.offset..end]);
1323            self.offset = end;
1324            Poll::Ready(Ok(()))
1325        }
1326    }
1327
1328    fn encode_frames(frames: &[Frame]) -> Vec<u8> {
1329        let mut bytes = Vec::new();
1330        for frame in frames {
1331            bytes.extend_from_slice(&frame.header.encode());
1332            bytes.extend_from_slice(&frame.body);
1333        }
1334        bytes
1335    }
1336
1337    async fn read_frames<R>(reader: &mut R, count: usize) -> Vec<Frame>
1338    where
1339        R: AsyncRead + Unpin,
1340    {
1341        let mut frames = Vec::with_capacity(count);
1342        for _ in 0..count {
1343            frames.push(read_frame(reader).await.unwrap().unwrap());
1344        }
1345        frames
1346    }
1347
1348    fn request(channel: u16, corr: u64, body: &[u8]) -> Frame {
1349        Frame::build(
1350            FrameType::Request,
1351            Flags::new(true, Priority::Interactive, false),
1352            channel,
1353            0,
1354            corr,
1355            body.to_vec(),
1356        )
1357        .unwrap()
1358    }
1359
1360    fn echo_router() -> Arc<Router> {
1361        let mut router = Router::with_default_self_handler();
1362        router.register_backend(7, EchoBackend).unwrap();
1363        router.register_backend(9, EchoBackend).unwrap();
1364        Arc::new(router)
1365    }
1366
1367    fn test_auth() -> (ServerAuth, ConnectionInfo) {
1368        test_auth_with_limit(4)
1369    }
1370
1371    fn test_auth_with_limit(max_unauthenticated: usize) -> (ServerAuth, ConnectionInfo) {
1372        test_auth_with_deadline_and_limit(TEST_DEADLINE, max_unauthenticated)
1373    }
1374
1375    fn test_auth_with_deadline_and_limit(
1376        deadline: Duration,
1377        max_unauthenticated: usize,
1378    ) -> (ServerAuth, ConnectionInfo) {
1379        let key = vec![0x42; 32];
1380        let daemon_id = [0x24; 16];
1381        let conn = ConnectionInfo {
1382            schema: SCHEMA_VERSION,
1383            wire_version: None,
1384            endpoints: vec![Endpoint {
1385                host: "127.0.0.1".to_owned(),
1386                port: 1,
1387            }],
1388            key: key.clone(),
1389            daemon_id,
1390            pid: std::process::id(),
1391            daemon_ver: TEST_DAEMON_VER.to_owned(),
1392        };
1393        (
1394            ServerAuth::with_limits(
1395                key,
1396                daemon_id,
1397                TEST_DAEMON_VER,
1398                deadline,
1399                max_unauthenticated,
1400            ),
1401            conn,
1402        )
1403    }
1404
1405    async fn authenticate<S>(stream: &mut S, conn: &ConnectionInfo)
1406    where
1407        S: AsyncRead + AsyncWrite + Unpin,
1408    {
1409        authenticate_client(stream, conn, TEST_DEADLINE)
1410            .await
1411            .expect("test client should authenticate")
1412    }
1413
1414    async fn write_auth_message<S, T>(stream: &mut S, value: &T)
1415    where
1416        S: AsyncWrite + Unpin,
1417        T: Serialize,
1418    {
1419        let body = serde_json::to_vec(value).expect("auth message encodes");
1420        stream
1421            .write_all(&(body.len() as u32).to_le_bytes())
1422            .await
1423            .expect("auth message length writes");
1424        stream
1425            .write_all(&body)
1426            .await
1427            .expect("auth message body writes");
1428    }
1429
1430    async fn read_auth_message<S, T>(stream: &mut S) -> T
1431    where
1432        S: AsyncRead + Unpin,
1433        T: DeserializeOwned,
1434    {
1435        let mut length = [0; 4];
1436        stream
1437            .read_exact(&mut length)
1438            .await
1439            .expect("auth message length reads");
1440        let mut body = vec![0; u32::from_le_bytes(length) as usize];
1441        stream
1442            .read_exact(&mut body)
1443            .await
1444            .expect("auth message body reads");
1445        serde_json::from_slice(&body).expect("auth message decodes")
1446    }
1447
1448    async fn send_wrong_key_client_auth<S>(stream: &mut S, conn: &ConnectionInfo)
1449    where
1450        S: AsyncRead + AsyncWrite + Unpin,
1451    {
1452        let hello = ClientHello {
1453            client_nonce: [0x57; NONCE_LEN],
1454            role: "client".to_owned(),
1455        };
1456        write_auth_message(stream, &hello).await;
1457        let proof: ServerProof = read_auth_message(stream).await;
1458        let mut wrong_key = conn.key.clone();
1459        wrong_key[0] ^= 0xff;
1460        let client_auth = ClientAuth {
1461            client_auth: compute_proof(
1462                &wrong_key,
1463                CLIENT_AUTH_DOMAIN,
1464                &hello.client_nonce,
1465                &proof.server_nonce,
1466                &conn.daemon_id,
1467            ),
1468        };
1469        write_auth_message(stream, &client_auth).await;
1470    }
1471
1472    fn assert_pre_auth_log(logs: &Arc<std::sync::Mutex<Vec<u8>>>, stage: &str, port: u16) {
1473        let captured = crate::router::test_log::captured_logs(logs);
1474        assert!(
1475            captured.contains(&format!("stage=\"{stage}\"")),
1476            "expected rejection stage {stage}, got: {captured}"
1477        );
1478        assert!(
1479            captured.contains(&format!("peer_addr=127.0.0.1:{port}")),
1480            "expected peer address, got: {captured}"
1481        );
1482        assert!(
1483            captured.contains("elapsed_ms=") && captured.contains("suppressed_since_last="),
1484            "expected elapsed time and suppression count, got: {captured}"
1485        );
1486        assert!(
1487            !captured.contains("WARN"),
1488            "pre-auth rejection must not be logged as a warning: {captured}"
1489        );
1490    }
1491
1492    #[tokio::test]
1493    async fn logs_handshake_deadline_pre_auth_rejection() {
1494        let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1495        let deadline = Duration::from_millis(60);
1496        let (auth, _) = test_auth_with_deadline_and_limit(deadline, 4);
1497        let (client, server_stream) = duplex(512);
1498        let port = 41001;
1499        let server = handle_accepted_connection(
1500            server_stream,
1501            echo_router(),
1502            auth,
1503            SocketAddr::from(([127, 0, 0, 1], port)),
1504            Instant::now(),
1505        );
1506
1507        let error = tokio::time::timeout(deadline * 3, server)
1508            .await
1509            .expect("idle handshake should end at its deadline")
1510            .expect_err("idle handshake must be rejected");
1511        drop(client);
1512        assert!(matches!(
1513            error,
1514            ConnectionError::Auth(AuthError::Timeout {
1515                stage: AuthStage::ClientHello,
1516                ..
1517            })
1518        ));
1519        let captured = crate::router::test_log::captured_logs(&logs);
1520        assert!(
1521            captured.contains("reason=\"handshake_deadline_expired\""),
1522            "{captured}"
1523        );
1524        assert_pre_auth_log(&logs, "handshake_deadline", port);
1525    }
1526
1527    #[tokio::test]
1528    async fn logs_invalid_client_proof_pre_auth_rejection() {
1529        let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1530        let (auth, conn) = test_auth_with_deadline_and_limit(Duration::from_millis(200), 4);
1531        let (mut client, server_stream) = duplex(512);
1532        let port = 41002;
1533        let server = handle_accepted_connection(
1534            server_stream,
1535            echo_router(),
1536            auth,
1537            SocketAddr::from(([127, 0, 0, 1], port)),
1538            Instant::now(),
1539        );
1540
1541        let (error, ()) = tokio::join!(server, send_wrong_key_client_auth(&mut client, &conn));
1542        assert!(matches!(
1543            error,
1544            Err(ConnectionError::Auth(AuthError::InvalidClientAuth))
1545        ));
1546        let captured = crate::router::test_log::captured_logs(&logs);
1547        assert!(
1548            captured.contains("reason=\"client_proof_mismatch\""),
1549            "{captured}"
1550        );
1551        assert_pre_auth_log(&logs, "client_authentication", port);
1552    }
1553
1554    #[tokio::test]
1555    async fn logs_malformed_pre_auth_message() {
1556        let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1557        let (auth, _) = test_auth_with_deadline_and_limit(Duration::from_millis(200), 4);
1558        let (mut client, server_stream) = duplex(512);
1559        let port = 41003;
1560        let server = handle_accepted_connection(
1561            server_stream,
1562            echo_router(),
1563            auth,
1564            SocketAddr::from(([127, 0, 0, 1], port)),
1565            Instant::now(),
1566        );
1567        let send_malformed_message = async {
1568            let body = b"{";
1569            client
1570                .write_all(&(body.len() as u32).to_le_bytes())
1571                .await
1572                .expect("message length writes");
1573            client.write_all(body).await.expect("message body writes");
1574        };
1575
1576        let (error, ()) = tokio::join!(server, send_malformed_message);
1577        assert!(matches!(
1578            error,
1579            Err(ConnectionError::Auth(AuthError::JsonDecode {
1580                stage: AuthStage::ClientHello,
1581                ..
1582            }))
1583        ));
1584        let captured = crate::router::test_log::captured_logs(&logs);
1585        assert!(
1586            captured.contains("reason=\"malformed_or_oversized_message\""),
1587            "{captured}"
1588        );
1589        assert_pre_auth_log(&logs, "malformed_message", port);
1590    }
1591
1592    #[tokio::test]
1593    async fn logs_peer_close_during_handshake() {
1594        let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1595        let (auth, _) = test_auth_with_deadline_and_limit(Duration::from_millis(200), 4);
1596        let (client, server_stream) = duplex(512);
1597        let port = 41004;
1598        drop(client);
1599
1600        let error = handle_accepted_connection(
1601            server_stream,
1602            echo_router(),
1603            auth,
1604            SocketAddr::from(([127, 0, 0, 1], port)),
1605            Instant::now(),
1606        )
1607        .await
1608        .expect_err("closed peer must be rejected");
1609        assert!(matches!(
1610            error,
1611            ConnectionError::Auth(AuthError::UnexpectedEof {
1612                stage: AuthStage::ClientHello,
1613                actual: 0,
1614                ..
1615            })
1616        ));
1617        let captured = crate::router::test_log::captured_logs(&logs);
1618        assert!(
1619            captured.contains("reason=\"peer_closed_before_authentication_completed\""),
1620            "{captured}"
1621        );
1622        assert_pre_auth_log(&logs, "peer_closed", port);
1623    }
1624
1625    // Rejection-log rate limits belong to one listener. A process-wide limiter
1626    // let one listener's rejects silence another's for a second, which hid the
1627    // second listener's first rejection entirely.
1628    #[tokio::test]
1629    async fn each_listener_rate_limits_its_own_rejection_logs() {
1630        let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1631        for port in [41011, 41012] {
1632            let (auth, _) = test_auth_with_deadline_and_limit(Duration::from_millis(200), 4);
1633            let (client, server_stream) = duplex(512);
1634            drop(client);
1635            handle_accepted_connection(
1636                server_stream,
1637                echo_router(),
1638                auth,
1639                SocketAddr::from(([127, 0, 0, 1], port)),
1640                Instant::now(),
1641            )
1642            .await
1643            .expect_err("closed peer must be rejected");
1644        }
1645        let captured = crate::router::test_log::captured_logs(&logs);
1646        for port in [41011, 41012] {
1647            assert!(
1648                captured.contains(&format!("127.0.0.1:{port}")),
1649                "listener on port {port} logged nothing: {captured}"
1650            );
1651        }
1652    }
1653
1654    #[tokio::test]
1655    async fn logs_unauthenticated_capacity_rejection() {
1656        let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1657        let deadline = Duration::from_millis(60);
1658        let (auth, _) = test_auth_with_deadline_and_limit(deadline, 1);
1659        let held_slot = auth
1660            .unauthenticated
1661            .clone()
1662            .try_acquire_owned()
1663            .expect("the only pre-auth slot is available");
1664        let (client, server_stream) = duplex(512);
1665        let port = 41005;
1666        let server = handle_accepted_connection(
1667            server_stream,
1668            echo_router(),
1669            auth,
1670            SocketAddr::from(([127, 0, 0, 1], port)),
1671            Instant::now(),
1672        );
1673
1674        let error = tokio::time::timeout(deadline * 3, server)
1675            .await
1676            .expect("queued connection should end at the capacity deadline")
1677            .expect_err("connection without a free slot must be rejected");
1678        drop(held_slot);
1679        drop(client);
1680        assert!(matches!(error, ConnectionError::UnauthenticatedCapacity));
1681        let captured = crate::router::test_log::captured_logs(&logs);
1682        assert!(
1683            captured.contains("reason=\"capacity_wait_timed_out\""),
1684            "{captured}"
1685        );
1686        assert_pre_auth_log(&logs, "waiting_for_handshake_slot", port);
1687    }
1688
1689    #[test]
1690    fn pre_auth_rejection_limiter_suppresses_bursts_and_reports_count() {
1691        let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1692        let limiter = PreAuthRejectLimiter::new();
1693        let capacity = ConnectionError::UnauthenticatedCapacity;
1694        let handshake_timeout = ConnectionError::Auth(AuthError::Timeout {
1695            stage: AuthStage::ClientHello,
1696            deadline: Duration::from_millis(200),
1697        });
1698        let peer_addr = SocketAddr::from(([127, 0, 0, 1], 41006));
1699
1700        log_pre_auth_rejection(
1701            &capacity,
1702            peer_addr,
1703            Duration::from_millis(1),
1704            Duration::ZERO,
1705            &limiter,
1706        );
1707        for _ in 0..4 {
1708            log_pre_auth_rejection(
1709                &capacity,
1710                peer_addr,
1711                Duration::from_millis(2),
1712                Duration::from_millis(100),
1713                &limiter,
1714            );
1715        }
1716        log_pre_auth_rejection(
1717            &handshake_timeout,
1718            peer_addr,
1719            Duration::from_millis(3),
1720            Duration::from_millis(100),
1721            &limiter,
1722        );
1723        log_pre_auth_rejection(
1724            &capacity,
1725            peer_addr,
1726            Duration::from_millis(4),
1727            Duration::from_secs(1),
1728            &limiter,
1729        );
1730
1731        let captured = crate::router::test_log::captured_logs(&logs);
1732        assert_eq!(
1733            captured
1734                .matches("subc TCP connection rejected before routing")
1735                .count(),
1736            3,
1737            "one line per stage in a burst, then a later capacity line: {captured}"
1738        );
1739        assert!(
1740            captured.contains("suppressed_since_last=4"),
1741            "later capacity line should report the four hidden rejects: {captured}"
1742        );
1743        assert!(
1744            captured.contains("stage=\"handshake_deadline\""),
1745            "{captured}"
1746        );
1747    }
1748
1749    #[tokio::test]
1750    async fn buffered_frame_reader_coalesces_reads_and_preserves_short_reads() {
1751        let frames = vec![
1752            request(7, 1, b"first"),
1753            request(9, 2, b"second"),
1754            request(7, 3, b"third"),
1755            request(9, 4, b"fourth"),
1756        ];
1757        let bytes = encode_frames(&frames);
1758
1759        let (mut direct, direct_reads) = CountingReader::new(bytes.clone(), None);
1760        assert_eq!(read_frames(&mut direct, frames.len()).await, frames);
1761        assert_eq!(direct_reads.load(Ordering::Relaxed), frames.len() * 3);
1762
1763        let (buffered_source, buffered_reads) = CountingReader::new(bytes, None);
1764        let mut buffered = BufReader::new(buffered_source);
1765        assert_eq!(read_frames(&mut buffered, frames.len()).await, frames);
1766        assert_eq!(buffered_reads.load(Ordering::Relaxed), 1);
1767
1768        let split_frame = request(7, 5, b"split-body");
1769        let split_bytes = encode_frames(std::slice::from_ref(&split_frame));
1770        let (split_source, split_reads) = CountingReader::new(split_bytes, Some(10));
1771        let mut split_reader = BufReader::new(split_source);
1772        assert_eq!(
1773            read_frame(&mut split_reader).await.unwrap(),
1774            Some(split_frame)
1775        );
1776        assert_eq!(split_reads.load(Ordering::Relaxed), 2);
1777    }
1778
1779    #[tokio::test]
1780    async fn interleaved_channels_on_one_stream_demux_byte_identically_after_auth() {
1781        let (mut client, server_stream) = duplex(4096);
1782        let (auth, conn) = test_auth();
1783        let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
1784        authenticate(&mut client, &conn).await;
1785        let frames = [
1786            request(7, 1, b"chan7-first\0opaque"),
1787            request(9, 2, b"chan9-middle-{json?}"),
1788            request(7, 3, b"chan7-second\xffbytes"),
1789        ];
1790
1791        for frame in &frames {
1792            crate::write_frame(&mut client, frame).await.unwrap();
1793        }
1794
1795        for expected in &frames {
1796            let response = crate::read_frame(&mut client).await.unwrap().unwrap();
1797            assert_eq!(response.header.ty, FrameType::Response);
1798            assert_eq!(response.header.channel, expected.header.channel);
1799            assert_eq!(response.header.corr, expected.header.corr);
1800            assert_eq!(response.body, expected.body);
1801        }
1802
1803        drop(client);
1804        server.await.unwrap().unwrap();
1805    }
1806
1807    #[tokio::test]
1808    async fn channel_zero_goes_to_subc_self_handler_after_auth() {
1809        let (mut client, server_stream) = duplex(512);
1810        let (auth, conn) = test_auth();
1811        let server = tokio::spawn(handle_connection(
1812            server_stream,
1813            Arc::new(Router::with_default_self_handler()),
1814            auth,
1815        ));
1816        authenticate(&mut client, &conn).await;
1817        let ping = Frame::build(
1818            FrameType::Ping,
1819            Flags::new(false, Priority::Passive, false),
1820            0,
1821            0,
1822            55,
1823            Vec::new(),
1824        )
1825        .unwrap();
1826
1827        crate::write_frame(&mut client, &ping).await.unwrap();
1828        let response = crate::read_frame(&mut client).await.unwrap().unwrap();
1829
1830        assert_eq!(response.header.ty, FrameType::Pong);
1831        assert_eq!(response.header.channel, 0);
1832        assert_eq!(response.header.corr, 55);
1833        assert!(response.body.is_empty());
1834
1835        drop(client);
1836        server.await.unwrap().unwrap();
1837    }
1838
1839    #[tokio::test]
1840    async fn unauthenticated_connection_is_rejected_before_routing() {
1841        let (mut client, server_stream) = duplex(512);
1842        let (auth, _conn) = test_auth();
1843        let registry = Arc::new(Registry::default());
1844        let router = Arc::new(Router::with_control_handler(Arc::new(ControlHandler::new(
1845            Arc::clone(&registry),
1846        ))));
1847        let server = tokio::spawn(handle_connection(server_stream, router, auth));
1848        let ping = Frame::build(
1849            FrameType::Ping,
1850            Flags::new(false, Priority::Passive, false),
1851            0,
1852            0,
1853            66,
1854            Vec::new(),
1855        )
1856        .unwrap();
1857
1858        crate::write_frame(&mut client, &ping).await.unwrap();
1859        if let Ok(Ok(Some(frame))) =
1860            tokio::time::timeout(Duration::from_millis(200), crate::read_frame(&mut client)).await
1861        {
1862            panic!("unauthenticated frame reached router: {frame:?}");
1863        }
1864
1865        let err = server.await.unwrap().unwrap_err();
1866        assert!(matches!(err, ConnectionError::Auth(_)));
1867        assert_eq!(registry.active_registration_count().unwrap(), 0);
1868    }
1869
1870    #[tokio::test]
1871    async fn over_cap_peer_queues_for_a_slot_and_authenticates_when_one_frees() {
1872        // Restart-herd contract: an over-cap pre-auth connection WAITS (bounded by
1873        // the auth deadline) instead of being reset. When the slot-holder finishes,
1874        // the queued peer must complete a normal handshake — the 2026-07-14 aft
1875        // outage was exactly a supervised child being reset out of this lottery.
1876        let (mut first_client, first_server_stream) = duplex(2048);
1877        let (mut second_client, second_server_stream) = duplex(2048);
1878        let (auth, conn) = test_auth_with_limit(1);
1879        let registry = Arc::new(Registry::default());
1880        let router = Arc::new(Router::with_control_handler(Arc::new(ControlHandler::new(
1881            Arc::clone(&registry),
1882        ))));
1883
1884        let first_server = tokio::spawn(handle_connection(
1885            first_server_stream,
1886            Arc::clone(&router),
1887            auth.clone(),
1888        ));
1889
1890        let second_server = tokio::spawn(handle_connection(
1891            second_server_stream,
1892            Arc::clone(&router),
1893            auth.clone(),
1894        ));
1895
1896        // The second peer must still be pending (not reset) while the slot is held.
1897        tokio::time::sleep(Duration::from_millis(50)).await;
1898        assert!(
1899            !second_server.is_finished(),
1900            "queued peer must not be reset"
1901        );
1902
1903        // First peer completes its handshake, freeing the slot; the queued second
1904        // peer then authenticates normally.
1905        authenticate(&mut first_client, &conn).await;
1906        authenticate(&mut second_client, &conn).await;
1907
1908        drop(first_client);
1909        drop(second_client);
1910        let _ = first_server.await;
1911        let _ = second_server.await;
1912        assert_eq!(registry.active_registration_count().unwrap(), 0);
1913    }
1914
1915    #[tokio::test]
1916    async fn over_cap_peer_is_rejected_when_no_slot_frees_within_deadline() {
1917        // The deadline stays the DoS bound: if no slot frees, the queued peer is
1918        // rejected at the deadline with a closed stream. The slot is held DIRECTLY
1919        // (not by another connection) so nothing can free it mid-test — a squatting
1920        // connection's own auth deadline would release the slot at exactly the
1921        // waiter's timeout, making the outcome racy.
1922        let (mut second_client, second_server_stream) = duplex(512);
1923        let (auth, conn) = test_auth_with_limit(1);
1924        let registry = Arc::new(Registry::default());
1925        let router = Arc::new(Router::with_control_handler(Arc::new(ControlHandler::new(
1926            Arc::clone(&registry),
1927        ))));
1928
1929        let held_slot = auth
1930            .unauthenticated
1931            .clone()
1932            .try_acquire_owned()
1933            .expect("sole pre-auth slot");
1934
1935        let second_server = tokio::spawn(handle_connection(
1936            second_server_stream,
1937            Arc::clone(&router),
1938            auth.clone(),
1939        ));
1940        let second_err = tokio::time::timeout(TEST_DEADLINE * 2, second_server)
1941            .await
1942            .expect("capacity reject should settle at the deadline")
1943            .expect("second connection task should not panic")
1944            .expect_err("queued peer must be rejected when no slot frees");
1945        drop(held_slot);
1946        assert!(matches!(
1947            second_err,
1948            ConnectionError::UnauthenticatedCapacity
1949        ));
1950        let mut closed = [0u8; 1];
1951        assert_eq!(
1952            second_client.read(&mut closed).await.unwrap(),
1953            0,
1954            "capacity-rejected peer should observe a closed stream"
1955        );
1956        assert_eq!(registry.active_registration_count().unwrap(), 0);
1957
1958        let (mut authed_client, authed_server_stream) = duplex(2048);
1959        let authed_server = tokio::spawn(handle_connection(
1960            authed_server_stream,
1961            Arc::clone(&router),
1962            auth,
1963        ));
1964        authenticate(&mut authed_client, &conn).await;
1965        let ping = Frame::build(
1966            FrameType::Ping,
1967            Flags::new(false, Priority::Passive, false),
1968            0,
1969            0,
1970            77,
1971            Vec::new(),
1972        )
1973        .unwrap();
1974        crate::write_frame(&mut authed_client, &ping).await.unwrap();
1975        let pong = crate::read_frame(&mut authed_client)
1976            .await
1977            .unwrap()
1978            .unwrap();
1979        assert_eq!(pong.header.ty, FrameType::Pong);
1980        assert_eq!(pong.header.channel, 0);
1981        assert_eq!(pong.header.corr, 77);
1982        assert_eq!(registry.active_registration_count().unwrap(), 0);
1983
1984        drop(authed_client);
1985        authed_server.await.unwrap().unwrap();
1986    }
1987
1988    #[tokio::test]
1989    async fn bind_ack_during_partial_frame_preserves_next_request() {
1990        use subc_control::ClientControlRequest;
1991        use subc_protocol::{
1992            manifest::{
1993                Concurrency, ExecutionMode, IdentityScope, ModuleManifest, ProviderRole, Tool,
1994            },
1995            session::{ModuleControlRequest, ModuleControlResponse},
1996            BindIdentity, ModuleHelloBody, RouteTarget,
1997        };
1998
1999        let mut configured_router = Router::with_default_self_handler();
2000        configured_router.register_backend(7, EchoBackend).unwrap();
2001        let router = Arc::new(configured_router);
2002        let (auth, conn) = test_auth();
2003        let (mut module, module_stream) = duplex(4096);
2004        let module_server = tokio::spawn(handle_connection(
2005            module_stream,
2006            Arc::clone(&router),
2007            auth.clone(),
2008        ));
2009        authenticate(&mut module, &conn).await;
2010        let manifest = ModuleManifest::builder("frame-test", "0.1.0")
2011            .protocol_ver(PROTOCOL_VERSION)
2012            .provides(vec![ProviderRole::ToolProvider {
2013                tools: vec![Tool {
2014                    name: "read".into(),
2015                    description: None,
2016                    execution_mode: ExecutionMode::Pure,
2017                    schema: serde_json::json!({"type": "object"}),
2018                }],
2019                identity_scope: vec![IdentityScope::Project, IdentityScope::Session],
2020                concurrency: Concurrency::ModuleManaged,
2021                emits_push: true,
2022                sub_supervises: true,
2023            }])
2024            .build();
2025        let hello = Frame::build(
2026            FrameType::Hello,
2027            Flags::new(false, Priority::Passive, false),
2028            0,
2029            0,
2030            1,
2031            serde_json::to_vec(&ModuleHelloBody {
2032                manifest,
2033                protocol_ver: PROTOCOL_VERSION,
2034                control_ops: None,
2035                launch_nonce: None,
2036            })
2037            .unwrap(),
2038        )
2039        .unwrap();
2040        crate::write_frame(&mut module, &hello).await.unwrap();
2041        assert_eq!(
2042            read_frame(&mut module).await.unwrap().unwrap().header.ty,
2043            FrameType::HelloAck
2044        );
2045
2046        let (mut client, client_stream) = duplex(4096);
2047        let client_server = tokio::spawn(handle_connection(client_stream, router, auth));
2048        authenticate(&mut client, &conn).await;
2049        let open = Frame::build(
2050            FrameType::Request,
2051            Flags::new(false, Priority::Passive, false),
2052            0,
2053            0,
2054            2,
2055            serde_json::to_vec(&ClientControlRequest::RouteOpen {
2056                target: RouteTarget::ToolProvider {
2057                    module_id: "frame-test".into(),
2058                },
2059                identity: BindIdentity::new(std::env::current_dir().unwrap(), "unit", "session"),
2060                consumer_identity: None,
2061                consumer_capabilities: None,
2062                role_versions: None,
2063                admission_facts: None,
2064                scope: None,
2065            })
2066            .unwrap(),
2067        )
2068        .unwrap();
2069        crate::write_frame(&mut client, &open).await.unwrap();
2070        let bind = timeout(TEST_DEADLINE, read_frame(&mut module))
2071            .await
2072            .unwrap()
2073            .unwrap()
2074            .unwrap();
2075        assert!(matches!(
2076            serde_json::from_slice::<ModuleControlRequest>(&bind.body).unwrap(),
2077            ModuleControlRequest::RouteBind { .. }
2078        ));
2079
2080        let ping = request(7, 3, b"partial-frame-body");
2081        client.write_all(&ping.header.encode()).await.unwrap();
2082        // Give the connection reader time to consume the header while the bind is pending.
2083        tokio::time::sleep(Duration::from_millis(30)).await;
2084        let ack = Frame::build(
2085            FrameType::Response,
2086            Flags::new(false, Priority::Passive, false),
2087            0,
2088            0,
2089            bind.header.corr,
2090            serde_json::to_vec(&ModuleControlResponse::RouteBindAck {}).unwrap(),
2091        )
2092        .unwrap();
2093        crate::write_frame(&mut module, &ack).await.unwrap();
2094        let opened = timeout(TEST_DEADLINE, read_frame(&mut client))
2095            .await
2096            .unwrap()
2097            .unwrap();
2098        if opened.is_none() {
2099            panic!("client closed: {:?}", client_server.await);
2100        }
2101        let opened = opened.unwrap();
2102        assert_eq!(opened.header.corr, 2);
2103        client.write_all(&ping.body).await.unwrap();
2104        let pong = timeout(TEST_DEADLINE, read_frame(&mut client))
2105            .await
2106            .expect("partial frame must reach the router after the bind ack")
2107            .expect("frame must decode")
2108            .expect("connection must remain open");
2109        assert_eq!(pong.header.ty, FrameType::Response);
2110        assert_eq!(pong.header.corr, 3);
2111        assert_eq!(pong.body, ping.body);
2112        drop(client);
2113        drop(module);
2114        client_server.await.unwrap().unwrap();
2115        let _ = module_server.await.unwrap();
2116    }
2117
2118    #[tokio::test]
2119    async fn aborted_accept_does_not_end_listener() {
2120        let listener = Arc::new(TcpListener::bind("127.0.0.1:0").await.unwrap());
2121        let addr = listener.local_addr().unwrap();
2122        let (auth, conn) = test_auth();
2123        let mut attempts = 0;
2124        let server = tokio::spawn(serve_listener_with_accept(
2125            Some(addr),
2126            echo_router(),
2127            auth,
2128            move || {
2129                attempts += 1;
2130                let result = if attempts == 1 {
2131                    Some(io::Error::from(io::ErrorKind::ConnectionAborted))
2132                } else {
2133                    None
2134                };
2135                let listener = Arc::clone(&listener);
2136                async move {
2137                    match result {
2138                        Some(err) => Err(err),
2139                        None => listener.accept().await,
2140                    }
2141                }
2142            },
2143        ));
2144        let mut client = tokio::net::TcpStream::connect(addr).await.unwrap();
2145        authenticate(&mut client, &conn).await;
2146        let ping = Frame::build(
2147            FrameType::Ping,
2148            Flags::new(false, Priority::Passive, false),
2149            0,
2150            0,
2151            77,
2152            Vec::new(),
2153        )
2154        .unwrap();
2155        crate::write_frame(&mut client, &ping).await.unwrap();
2156        let pong = timeout(TEST_DEADLINE, read_frame(&mut client))
2157            .await
2158            .unwrap()
2159            .unwrap()
2160            .unwrap();
2161        assert_eq!(pong.header.ty, FrameType::Pong);
2162        assert!(
2163            !server.is_finished(),
2164            "a temporary accept failure must not stop the listener"
2165        );
2166        server.abort();
2167    }
2168
2169    #[tokio::test]
2170    async fn exhausted_accept_backs_off_then_serves_connection() {
2171        let listener = Arc::new(TcpListener::bind("127.0.0.1:0").await.unwrap());
2172        let addr = listener.local_addr().unwrap();
2173        let (auth, conn) = test_auth();
2174        let mut attempts = 0;
2175        let server = tokio::spawn(serve_listener_with_accept(
2176            Some(addr),
2177            echo_router(),
2178            auth,
2179            move || {
2180                attempts += 1;
2181                // The platform's "too many open files" code: EMFILE on Unix,
2182                // WSAEMFILE on Windows. A bare 24 is ERROR_BAD_LENGTH on Windows.
2183                #[cfg(unix)]
2184                let emfile = rustix::io::Errno::MFILE.raw_os_error();
2185                #[cfg(windows)]
2186                let emfile = 10024;
2187                let error = (attempts == 1).then(|| io::Error::from_raw_os_error(emfile));
2188                let listener = Arc::clone(&listener);
2189                async move {
2190                    match error {
2191                        Some(err) => Err(err),
2192                        None => listener.accept().await,
2193                    }
2194                }
2195            },
2196        ));
2197        let started = Instant::now();
2198        let mut client = tokio::net::TcpStream::connect(addr).await.unwrap();
2199        authenticate(&mut client, &conn).await;
2200        assert!(
2201            started.elapsed() >= Duration::from_millis(50),
2202            "fd exhaustion must back off before retrying"
2203        );
2204        assert!(!server.is_finished());
2205        server.abort();
2206    }
2207
2208    #[tokio::test]
2209    async fn fatal_accept_error_still_stops_listener() {
2210        let (auth, _) = test_auth();
2211        let err = serve_listener_with_accept(None, echo_router(), auth, || async {
2212            Err(io::Error::from(io::ErrorKind::PermissionDenied))
2213        })
2214        .await
2215        .unwrap_err();
2216        assert!(
2217            matches!(err, ServerError::Accept { source, .. } if source.kind() == io::ErrorKind::PermissionDenied)
2218        );
2219    }
2220
2221    #[tokio::test]
2222    async fn serve_listeners_with_no_listeners_returns_typed_error() {
2223        let (auth, _conn) = test_auth();
2224        let err = serve_listeners(Vec::new(), echo_router(), auth)
2225            .await
2226            .expect_err("empty listener set must fail loudly");
2227        assert!(matches!(err, ServerError::NoListeners));
2228    }
2229
2230    #[tokio::test]
2231    async fn malformed_header_returns_typed_error_no_panic() {
2232        let (mut client, server_stream) = duplex(128);
2233        let (auth, conn) = test_auth();
2234        let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
2235        authenticate(&mut client, &conn).await;
2236        let mut header = [0u8; HEADER_LEN];
2237        header[4] = PROTOCOL_VERSION;
2238        header[5] = 250;
2239
2240        client.write_all(&header).await.unwrap();
2241        drop(client);
2242
2243        let err = server.await.unwrap().unwrap_err();
2244        assert!(matches!(
2245            err,
2246            ConnectionError::FrameIo(FrameIoError::DecodeHeader(DecodeError::UnknownFrameType {
2247                byte: 250
2248            }))
2249        ));
2250    }
2251
2252    #[tokio::test]
2253    async fn truncated_body_returns_typed_error_no_panic() {
2254        let (mut client, server_stream) = duplex(128);
2255        let (auth, conn) = test_auth();
2256        let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
2257        authenticate(&mut client, &conn).await;
2258        let frame = request(7, 8, b"abcd");
2259
2260        client.write_all(&frame.header.encode()).await.unwrap();
2261        client.write_all(b"ab").await.unwrap();
2262        drop(client);
2263
2264        let err = server.await.unwrap().unwrap_err();
2265        assert!(matches!(
2266            err,
2267            ConnectionError::FrameIo(FrameIoError::UnexpectedEof {
2268                stage: ReadStage::Body,
2269                expected: 4,
2270                actual: 2
2271            })
2272        ));
2273    }
2274
2275    #[tokio::test]
2276    async fn unknown_channel_is_returned_as_error_frame_and_connection_continues() {
2277        let (mut client, server_stream) = duplex(1024);
2278        let (auth, conn) = test_auth();
2279        let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
2280        authenticate(&mut client, &conn).await;
2281        let unknown = request(42, 10, b"lost");
2282        let known = request(7, 11, b"still-routes");
2283
2284        crate::write_frame(&mut client, &unknown).await.unwrap();
2285        crate::write_frame(&mut client, &known).await.unwrap();
2286
2287        let error = crate::read_frame(&mut client).await.unwrap().unwrap();
2288        assert_eq!(error.header.ty, FrameType::Error);
2289        assert_eq!(error.header.channel, 42);
2290        assert_eq!(error.header.corr, 10);
2291        let error_body: ErrorBody = serde_json::from_slice(&error.body).unwrap();
2292        assert_eq!(error_body.code, "unknown_channel");
2293
2294        let response = crate::read_frame(&mut client).await.unwrap().unwrap();
2295        assert_eq!(response.header.ty, FrameType::Response);
2296        assert_eq!(response.header.channel, 7);
2297        assert_eq!(response.header.corr, 11);
2298        assert_eq!(response.body, b"still-routes");
2299
2300        drop(client);
2301        server.await.unwrap().unwrap();
2302    }
2303}