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
32pub const CONNECTION_EGRESS_BYTE_BUDGET: usize = 4 * 1024 * 1024;
44pub const CONNECTION_EGRESS_FRAME_CAP: usize = 32 * 1024;
53pub const MAX_PENDING_ROUTE_OPENS_PER_CONNECTION: usize = 8;
59pub(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);
65pub 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
135struct 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#[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 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
284pub 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 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 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 #[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
384pub 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
437pub 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 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 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 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 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 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 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 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 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 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
818pub(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
848async 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 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 #[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 #[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(®istry),
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 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(®istry),
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 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 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 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(®istry),
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 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 #[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}