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