Skip to main content

subc_daemon/
server.rs

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