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
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
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#[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#[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#[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 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
307pub 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 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 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 #[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
407pub 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
460pub 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 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 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 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 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 loop {
688 while let Some(result) = route_open_tasks.try_join_next() {
689 finish_route_open_task(result)?;
690 }
691
692 let read = read_frame(&mut *read_half);
695 tokio::pin!(read);
696 let frame = loop {
697 tokio::select! {
698 close = &mut close_receiver => {
699 return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
700 }
701 result = route_open_tasks.join_next(), if !route_open_tasks.is_empty() => {
702 finish_route_open_task(
703 result.expect("a non-empty route.open JoinSet has a next task")
704 )?;
705 }
706 result = &mut read => {
707 break match result.map_err(ConnectionError::FrameIo)? {
708 Some(frame) => frame,
709 None => return Ok(ConnectionLoopExit::PeerClosed),
710 };
711 }
712 }
713 };
714
715 if let Some(target_module_id) = router.route_open_target(&frame) {
716 while let Some(result) = route_open_tasks.try_join_next() {
720 finish_route_open_task(result)?;
721 }
722
723 if route_open_tasks.len() >= MAX_PENDING_ROUTE_OPENS_PER_CONNECTION {
724 let refusal = router
727 .route_open_capacity_refusal(
728 &ctx,
729 &frame,
730 &target_module_id,
731 route_open_tasks.len(),
732 MAX_PENDING_ROUTE_OPENS_PER_CONNECTION,
733 )
734 .map_err(ConnectionError::Router)?;
735 let send_result = tokio::select! {
736 close = &mut close_receiver => {
737 return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
738 }
739 result = ctx.egress.send(refusal) => result,
740 };
741 send_result.map_err(ConnectionError::Router)?;
742 continue;
743 }
744
745 let task_router = Arc::clone(&router);
746 let task_ctx = ctx.clone();
747 let dispatch_started_at = Instant::now();
750 route_open_tasks.spawn(async move {
751 route_open_tail(task_router, task_ctx, frame, dispatch_started_at).await
752 });
753 continue;
754 }
755
756 let route_result = tokio::select! {
759 close = &mut close_receiver => {
760 return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
761 }
762 result = router.route_for_connection(&ctx, frame) => result,
763 };
764
765 if let Err(err) = route_result {
766 if let Some(error_frame) = err.to_error_frame() {
767 warn!(
768 connection_id = ctx.connection_id.get(),
769 error = %err,
770 "routing failure recovered with ERROR frame"
771 );
772 let send_result = tokio::select! {
773 close = &mut close_receiver => {
774 return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
775 }
776 result = ctx.egress.send(error_frame) => result,
777 };
778 send_result.map_err(ConnectionError::Router)?;
779 } else {
780 debug!(
781 connection_id = ctx.connection_id.get(),
782 error = %err,
783 "fatal routing failure"
784 );
785 return Err(ConnectionError::Router(err));
786 }
787 }
788 }
789}
790
791async fn route_open_tail(
792 router: Arc<Router>,
793 ctx: RouteCtx,
794 frame: crate::Frame,
795 dispatch_started_at: Instant,
796) -> Result<(), RouterError> {
797 match router
798 .route_for_connection_started(&ctx, frame, Some(dispatch_started_at))
799 .await
800 {
801 Ok(()) => Ok(()),
802 Err(err) => {
803 let Some(error_frame) = err.to_error_frame() else {
804 return Err(err);
805 };
806 warn!(
807 connection_id = ctx.connection_id.get(),
808 error = %err,
809 "routing failure recovered with ERROR frame"
810 );
811 ctx.egress.send(error_frame).await
812 }
813 }
814}
815
816fn finish_route_open_task(
817 result: Result<Result<(), RouterError>, tokio::task::JoinError>,
818) -> Result<(), ConnectionError> {
819 match result {
820 Ok(Ok(())) => Ok(()),
821 Ok(Err(err)) => Err(ConnectionError::Router(err)),
822 Err(err) => Err(ConnectionError::Router(RouterError::backend(
823 0,
824 0,
825 format!("route.open task failed: {err}"),
826 ))),
827 }
828}
829
830fn close_reason(
831 result: Result<CloseReason, tokio::sync::oneshot::error::RecvError>,
832) -> CloseReason {
833 result.unwrap_or_else(|_| {
834 CloseReason::new(
835 "close_registry_dropped",
836 "connection close registration was dropped without a reason",
837 )
838 })
839}
840
841pub(crate) fn connection_egress() -> (FrameSink, mpsc::Receiver<crate::router::OutboundFrame>) {
845 let (tx, rx) = mpsc::channel::<crate::router::OutboundFrame>(CONNECTION_EGRESS_FRAME_CAP);
846 (
847 FrameSink::with_byte_budget(tx, CONNECTION_EGRESS_BYTE_BUDGET),
848 rx,
849 )
850}
851
852async fn drain_writer<W>(
853 write_half: W,
854 mut rx: mpsc::Receiver<crate::router::OutboundFrame>,
855) -> Result<(), FrameIoError>
856where
857 W: AsyncWrite + Unpin,
858{
859 let mut writer = BufWriter::new(write_half);
860 while let Some(outbound) = rx.recv().await {
861 write_outbound(&mut writer, outbound).await?;
862 while let Ok(outbound) = rx.try_recv() {
863 write_outbound(&mut writer, outbound).await?;
864 }
865 writer.flush().await.map_err(FrameIoError::Io)?;
866 }
867 writer.flush().await.map_err(FrameIoError::Io)?;
868 Ok(())
869}
870
871async fn write_outbound<W>(
878 writer: &mut BufWriter<W>,
879 outbound: crate::router::OutboundFrame,
880) -> Result<(), FrameIoError>
881where
882 W: AsyncWrite + Unpin,
883{
884 const SLOW_REPLY_QUEUE: Duration = Duration::from_millis(1000);
885 let queued = outbound.enqueued_at.elapsed();
886 let charge = outbound.charge;
889 if let Some(charge) = &charge {
890 charge.taken_by_writer();
891 }
892 let frame = outbound.frame;
893 if frame.header.channel == 0 && queued >= SLOW_REPLY_QUEUE {
894 let write_started = std::time::Instant::now();
895 let result = write_frame(writer, &frame).await;
896 tracing::warn!(
897 corr = frame.header.corr,
898 queued_ms = queued.as_millis() as u64,
899 write_ms = write_started.elapsed().as_millis() as u64,
900 "slow control reply write"
901 );
902 result?;
903 } else {
904 write_frame(writer, &frame).await?;
905 }
906 if let Some(flushed) = outbound.flushed {
907 writer.flush().await.map_err(FrameIoError::Io)?;
908 let _ = flushed.send(());
909 }
910 Ok(())
911}
912
913#[cfg(all(test, unix))]
914#[tokio::test]
915async fn shutdown_notice_ack_waits_for_socket_flush() {
916 let (socket, mut peer) = tokio::io::duplex(1);
917 let (tx, rx) = mpsc::channel(1);
918 let sink = FrameSink::new(tx);
919 let writer = tokio::spawn(drain_writer(socket, rx));
920 let frame = crate::Frame::build(
921 subc_protocol::FrameType::Push,
922 subc_protocol::Flags::new(false, subc_protocol::Priority::Interactive, false),
923 0,
924 0,
925 0,
926 b"notice".to_vec(),
927 )
928 .unwrap();
929 let send = sink.send_flushed(frame);
930 tokio::pin!(send);
931 assert!(
932 timeout(Duration::from_millis(20), &mut send).await.is_err(),
933 "queueing bytes is not a socket flush acknowledgement"
934 );
935 let (sent, received) = timeout(Duration::from_secs(1), async {
936 tokio::join!(&mut send, read_frame(&mut peer))
937 })
938 .await
939 .unwrap();
940 sent.unwrap();
941 assert_eq!(received.unwrap().unwrap().body, b"notice");
942 writer.abort();
943}
944
945#[derive(Debug)]
946pub enum ServerError {
947 NoListeners,
948 Accept {
949 local_addr: Option<SocketAddr>,
950 source: io::Error,
951 },
952}
953
954impl fmt::Display for ServerError {
955 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
956 match self {
957 Self::NoListeners => write!(f, "no TCP listeners were provided"),
958 Self::Accept { local_addr, source } => match local_addr {
959 Some(addr) => write!(f, "failed to accept TCP connection on {addr}: {source}"),
960 None => write!(f, "failed to accept TCP connection: {source}"),
961 },
962 }
963 }
964}
965
966impl Error for ServerError {
967 fn source(&self) -> Option<&(dyn Error + 'static)> {
968 match self {
969 Self::Accept { source, .. } => Some(source),
970 Self::NoListeners => None,
971 }
972 }
973}
974
975#[derive(Debug)]
976pub enum ConnectionError {
977 Auth(AuthError),
978 UnauthenticatedCapacity,
979 FrameIo(FrameIoError),
980 Router(RouterError),
981 WriterTask(tokio::task::JoinError),
982}
983
984impl ConnectionError {
985 fn is_quiet_reject(&self) -> bool {
986 matches!(self, Self::Auth(_) | Self::UnauthenticatedCapacity)
987 }
988
989 fn pre_auth_reject_stage(&self) -> Option<PreAuthRejectStage> {
990 match self {
991 Self::UnauthenticatedCapacity => Some(PreAuthRejectStage::WaitingForHandshakeSlot),
992 Self::Auth(AuthError::Timeout { .. }) => Some(PreAuthRejectStage::HandshakeDeadline),
993 Self::Auth(
994 AuthError::InvalidClientAuth
995 | AuthError::InvalidServerProof
996 | AuthError::DaemonIdMismatch,
997 ) => Some(PreAuthRejectStage::InvalidProof),
998 Self::Auth(AuthError::MessageTooLarge { .. } | AuthError::JsonDecode { .. }) => {
999 Some(PreAuthRejectStage::MalformedMessage)
1000 }
1001 Self::Auth(AuthError::UnexpectedEof { .. }) => Some(PreAuthRejectStage::PeerClosed),
1002 Self::Auth(AuthError::Io { source, .. })
1003 if matches!(
1004 source.kind(),
1005 io::ErrorKind::BrokenPipe
1006 | io::ErrorKind::ConnectionAborted
1007 | io::ErrorKind::ConnectionReset
1008 | io::ErrorKind::UnexpectedEof
1009 ) =>
1010 {
1011 Some(PreAuthRejectStage::PeerClosed)
1012 }
1013 Self::Auth(_) => Some(PreAuthRejectStage::OtherAuthenticationFailure),
1014 Self::FrameIo(_) | Self::Router(_) | Self::WriterTask(_) => None,
1015 }
1016 }
1017
1018 fn auth_stage(&self) -> Option<AuthStage> {
1019 match self {
1020 Self::Auth(
1021 AuthError::Io { stage, .. }
1022 | AuthError::Timeout { stage, .. }
1023 | AuthError::UnexpectedEof { stage, .. }
1024 | AuthError::MessageTooLarge { stage, .. }
1025 | AuthError::JsonEncode { stage, .. }
1026 | AuthError::JsonDecode { stage, .. },
1027 ) => Some(*stage),
1028 Self::Auth(
1029 AuthError::Random(_)
1030 | AuthError::KeyTooShort { .. }
1031 | AuthError::InvalidServerProof
1032 | AuthError::DaemonIdMismatch
1033 | AuthError::InvalidClientAuth,
1034 )
1035 | Self::UnauthenticatedCapacity
1036 | Self::FrameIo(_)
1037 | Self::Router(_)
1038 | Self::WriterTask(_) => None,
1039 }
1040 }
1041}
1042
1043impl fmt::Display for ConnectionError {
1044 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1045 match self {
1046 Self::Auth(err) => write!(f, "connection auth failed: {err}"),
1047 Self::UnauthenticatedCapacity => write!(
1048 f,
1049 "too many concurrent unauthenticated subc TCP connections"
1050 ),
1051 Self::FrameIo(err) => write!(f, "frame connection error: {err}"),
1052 Self::Router(err) => write!(f, "router connection error: {err}"),
1053 Self::WriterTask(err) => write!(f, "connection writer task failed: {err}"),
1054 }
1055 }
1056}
1057
1058impl Error for ConnectionError {
1059 fn source(&self) -> Option<&(dyn Error + 'static)> {
1060 match self {
1061 Self::Auth(err) => Some(err),
1062 Self::FrameIo(err) => Some(err),
1063 Self::Router(err) => Some(err),
1064 Self::WriterTask(err) => Some(err),
1065 Self::UnauthenticatedCapacity => None,
1066 }
1067 }
1068}
1069
1070#[cfg(test)]
1071mod tests {
1072 use std::{
1073 pin::Pin,
1074 sync::atomic::{AtomicUsize, Ordering},
1075 task::{Context, Poll},
1076 };
1077
1078 use super::*;
1079 use serde::{de::DeserializeOwned, Serialize};
1080 use subc_protocol::{
1081 DecodeError, ErrorBody, Flags, FrameType, Priority, HEADER_LEN, PROTOCOL_VERSION,
1082 };
1083 use tokio::io::{duplex, AsyncReadExt, AsyncWriteExt, ReadBuf};
1084
1085 use subc_transport::{
1086 authenticate_client, compute_proof, ClientAuth, ClientHello, ConnectionInfo, Endpoint,
1087 ServerProof, CLIENT_AUTH_DOMAIN, NONCE_LEN, SCHEMA_VERSION,
1088 };
1089
1090 use crate::{ControlHandler, EchoBackend, Frame, ReadStage, Registry};
1091
1092 const TEST_DEADLINE: Duration = Duration::from_secs(2);
1093 const TEST_DAEMON_VER: &str = "test-subc-server";
1094
1095 #[tokio::test]
1100 async fn stale_queued_control_reply_logs_slow_reply_write() {
1101 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::WARN);
1102 let (tx, rx) = mpsc::channel::<crate::router::OutboundFrame>(4);
1103 let reply = Frame::build_with_version(
1104 PROTOCOL_VERSION,
1105 FrameType::Error,
1106 Flags::new(false, Priority::Interactive, false),
1107 0,
1108 0,
1109 7,
1110 serde_json::to_vec(&ErrorBody {
1111 code: "test".into(),
1112 message: "reply".into(),
1113 detail: None,
1114 })
1115 .expect("body encodes"),
1116 )
1117 .expect("frame builds");
1118 tx.send(crate::router::OutboundFrame {
1119 frame: reply,
1120 enqueued_at: std::time::Instant::now() - Duration::from_millis(1500),
1121 flushed: None,
1122 charge: None,
1123 })
1124 .await
1125 .expect("queued");
1126 drop(tx);
1127 let (write_half, mut read_half) = duplex(64 * 1024);
1128 drain_writer(write_half, rx).await.expect("writer drains");
1129 let mut sink = Vec::new();
1130 read_half.read_to_end(&mut sink).await.expect("read");
1131 assert!(!sink.is_empty(), "frame reached the socket");
1132 let captured = crate::router::test_log::captured_logs(&logs);
1133 assert!(
1134 captured.contains("slow control reply write") && captured.contains("corr=7"),
1135 "expected slow reply WARN naming corr, got: {captured}"
1136 );
1137 let queued_ms: u64 = captured
1138 .split("queued_ms=")
1139 .nth(1)
1140 .and_then(|s| s.split_whitespace().next())
1141 .and_then(|s| s.parse().ok())
1142 .expect("queued_ms present");
1143 assert!(
1144 queued_ms >= 1500,
1145 "queued_ms reflects residency: {queued_ms}"
1146 );
1147 }
1148
1149 #[tokio::test]
1153 async fn fresh_control_and_stale_data_frames_log_nothing() {
1154 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::WARN);
1155 let (tx, rx) = mpsc::channel::<crate::router::OutboundFrame>(4);
1156 let control = Frame::build_with_version(
1157 PROTOCOL_VERSION,
1158 FrameType::Error,
1159 Flags::new(false, Priority::Interactive, false),
1160 0,
1161 0,
1162 8,
1163 serde_json::to_vec(&ErrorBody {
1164 code: "test".into(),
1165 message: "fresh".into(),
1166 detail: None,
1167 })
1168 .expect("body encodes"),
1169 )
1170 .expect("frame builds");
1171 tx.send(crate::router::OutboundFrame {
1172 frame: control,
1173 enqueued_at: std::time::Instant::now(),
1174 flushed: None,
1175 charge: None,
1176 })
1177 .await
1178 .expect("queued");
1179 let data = Frame::build_with_version(
1180 PROTOCOL_VERSION,
1181 FrameType::Error,
1182 Flags::new(false, Priority::Interactive, false),
1183 9,
1184 1,
1185 9,
1186 serde_json::to_vec(&ErrorBody {
1187 code: "test".into(),
1188 message: "data".into(),
1189 detail: None,
1190 })
1191 .expect("body encodes"),
1192 )
1193 .expect("frame builds");
1194 tx.send(crate::router::OutboundFrame {
1195 frame: data,
1196 enqueued_at: std::time::Instant::now() - Duration::from_millis(5000),
1197 flushed: None,
1198 charge: None,
1199 })
1200 .await
1201 .expect("queued");
1202 drop(tx);
1203 let (write_half, _read_half) = duplex(64 * 1024);
1204 drain_writer(write_half, rx).await.expect("writer drains");
1205 let captured = crate::router::test_log::captured_logs(&logs);
1206 assert!(
1207 !captured.contains("slow control reply write"),
1208 "no WARN for fresh control or stale data frames, got: {captured}"
1209 );
1210 }
1211
1212 struct CountingReader {
1213 bytes: Vec<u8>,
1214 offset: usize,
1215 first_read_end: Option<usize>,
1216 reads: Arc<AtomicUsize>,
1217 }
1218
1219 impl CountingReader {
1220 fn new(bytes: Vec<u8>, first_read_end: Option<usize>) -> (Self, Arc<AtomicUsize>) {
1221 let reads = Arc::new(AtomicUsize::new(0));
1222 (
1223 Self {
1224 bytes,
1225 offset: 0,
1226 first_read_end,
1227 reads: Arc::clone(&reads),
1228 },
1229 reads,
1230 )
1231 }
1232 }
1233
1234 impl AsyncRead for CountingReader {
1235 fn poll_read(
1236 mut self: Pin<&mut Self>,
1237 _cx: &mut Context<'_>,
1238 buf: &mut ReadBuf<'_>,
1239 ) -> Poll<io::Result<()>> {
1240 self.reads.fetch_add(1, Ordering::Relaxed);
1241 let available = self.bytes.len().saturating_sub(self.offset);
1242 let first_read_remaining = self
1243 .first_read_end
1244 .filter(|end| self.offset < *end)
1245 .map_or(available, |end| end - self.offset);
1246 let count = available.min(first_read_remaining).min(buf.remaining());
1247 let end = self.offset + count;
1248 buf.put_slice(&self.bytes[self.offset..end]);
1249 self.offset = end;
1250 Poll::Ready(Ok(()))
1251 }
1252 }
1253
1254 fn encode_frames(frames: &[Frame]) -> Vec<u8> {
1255 let mut bytes = Vec::new();
1256 for frame in frames {
1257 bytes.extend_from_slice(&frame.header.encode());
1258 bytes.extend_from_slice(&frame.body);
1259 }
1260 bytes
1261 }
1262
1263 async fn read_frames<R>(reader: &mut R, count: usize) -> Vec<Frame>
1264 where
1265 R: AsyncRead + Unpin,
1266 {
1267 let mut frames = Vec::with_capacity(count);
1268 for _ in 0..count {
1269 frames.push(read_frame(reader).await.unwrap().unwrap());
1270 }
1271 frames
1272 }
1273
1274 fn request(channel: u16, corr: u64, body: &[u8]) -> Frame {
1275 Frame::build(
1276 FrameType::Request,
1277 Flags::new(true, Priority::Interactive, false),
1278 channel,
1279 0,
1280 corr,
1281 body.to_vec(),
1282 )
1283 .unwrap()
1284 }
1285
1286 fn echo_router() -> Arc<Router> {
1287 let mut router = Router::with_default_self_handler();
1288 router.register_backend(7, EchoBackend).unwrap();
1289 router.register_backend(9, EchoBackend).unwrap();
1290 Arc::new(router)
1291 }
1292
1293 fn test_auth() -> (ServerAuth, ConnectionInfo) {
1294 test_auth_with_limit(4)
1295 }
1296
1297 fn test_auth_with_limit(max_unauthenticated: usize) -> (ServerAuth, ConnectionInfo) {
1298 test_auth_with_deadline_and_limit(TEST_DEADLINE, max_unauthenticated)
1299 }
1300
1301 fn test_auth_with_deadline_and_limit(
1302 deadline: Duration,
1303 max_unauthenticated: usize,
1304 ) -> (ServerAuth, ConnectionInfo) {
1305 let key = vec![0x42; 32];
1306 let daemon_id = [0x24; 16];
1307 let conn = ConnectionInfo {
1308 schema: SCHEMA_VERSION,
1309 wire_version: None,
1310 endpoints: vec![Endpoint {
1311 host: "127.0.0.1".to_owned(),
1312 port: 1,
1313 }],
1314 key: key.clone(),
1315 daemon_id,
1316 pid: std::process::id(),
1317 daemon_ver: TEST_DAEMON_VER.to_owned(),
1318 };
1319 (
1320 ServerAuth::with_limits(
1321 key,
1322 daemon_id,
1323 TEST_DAEMON_VER,
1324 deadline,
1325 max_unauthenticated,
1326 ),
1327 conn,
1328 )
1329 }
1330
1331 async fn authenticate<S>(stream: &mut S, conn: &ConnectionInfo)
1332 where
1333 S: AsyncRead + AsyncWrite + Unpin,
1334 {
1335 authenticate_client(stream, conn, TEST_DEADLINE)
1336 .await
1337 .expect("test client should authenticate")
1338 }
1339
1340 async fn write_auth_message<S, T>(stream: &mut S, value: &T)
1341 where
1342 S: AsyncWrite + Unpin,
1343 T: Serialize,
1344 {
1345 let body = serde_json::to_vec(value).expect("auth message encodes");
1346 stream
1347 .write_all(&(body.len() as u32).to_le_bytes())
1348 .await
1349 .expect("auth message length writes");
1350 stream
1351 .write_all(&body)
1352 .await
1353 .expect("auth message body writes");
1354 }
1355
1356 async fn read_auth_message<S, T>(stream: &mut S) -> T
1357 where
1358 S: AsyncRead + Unpin,
1359 T: DeserializeOwned,
1360 {
1361 let mut length = [0; 4];
1362 stream
1363 .read_exact(&mut length)
1364 .await
1365 .expect("auth message length reads");
1366 let mut body = vec![0; u32::from_le_bytes(length) as usize];
1367 stream
1368 .read_exact(&mut body)
1369 .await
1370 .expect("auth message body reads");
1371 serde_json::from_slice(&body).expect("auth message decodes")
1372 }
1373
1374 async fn send_wrong_key_client_auth<S>(stream: &mut S, conn: &ConnectionInfo)
1375 where
1376 S: AsyncRead + AsyncWrite + Unpin,
1377 {
1378 let hello = ClientHello {
1379 client_nonce: [0x57; NONCE_LEN],
1380 role: "client".to_owned(),
1381 };
1382 write_auth_message(stream, &hello).await;
1383 let proof: ServerProof = read_auth_message(stream).await;
1384 let mut wrong_key = conn.key.clone();
1385 wrong_key[0] ^= 0xff;
1386 let client_auth = ClientAuth {
1387 client_auth: compute_proof(
1388 &wrong_key,
1389 CLIENT_AUTH_DOMAIN,
1390 &hello.client_nonce,
1391 &proof.server_nonce,
1392 &conn.daemon_id,
1393 ),
1394 };
1395 write_auth_message(stream, &client_auth).await;
1396 }
1397
1398 fn assert_pre_auth_log(logs: &Arc<std::sync::Mutex<Vec<u8>>>, stage: &str, port: u16) {
1399 let captured = crate::router::test_log::captured_logs(logs);
1400 assert!(
1401 captured.contains(&format!("stage=\"{stage}\"")),
1402 "expected rejection stage {stage}, got: {captured}"
1403 );
1404 assert!(
1405 captured.contains(&format!("peer_addr=127.0.0.1:{port}")),
1406 "expected peer address, got: {captured}"
1407 );
1408 assert!(
1409 captured.contains("elapsed_ms=") && captured.contains("suppressed_since_last="),
1410 "expected elapsed time and suppression count, got: {captured}"
1411 );
1412 assert!(
1413 !captured.contains("WARN"),
1414 "pre-auth rejection must not be logged as a warning: {captured}"
1415 );
1416 }
1417
1418 #[tokio::test]
1419 async fn logs_handshake_deadline_pre_auth_rejection() {
1420 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1421 let deadline = Duration::from_millis(60);
1422 let (auth, _) = test_auth_with_deadline_and_limit(deadline, 4);
1423 let (client, server_stream) = duplex(512);
1424 let port = 41001;
1425 let server = handle_accepted_connection(
1426 server_stream,
1427 echo_router(),
1428 auth,
1429 SocketAddr::from(([127, 0, 0, 1], port)),
1430 Instant::now(),
1431 );
1432
1433 let error = tokio::time::timeout(deadline * 3, server)
1434 .await
1435 .expect("idle handshake should end at its deadline")
1436 .expect_err("idle handshake must be rejected");
1437 drop(client);
1438 assert!(matches!(
1439 error,
1440 ConnectionError::Auth(AuthError::Timeout {
1441 stage: AuthStage::ClientHello,
1442 ..
1443 })
1444 ));
1445 let captured = crate::router::test_log::captured_logs(&logs);
1446 assert!(
1447 captured.contains("reason=\"handshake_deadline_expired\""),
1448 "{captured}"
1449 );
1450 assert_pre_auth_log(&logs, "handshake_deadline", port);
1451 }
1452
1453 #[tokio::test]
1454 async fn logs_invalid_client_proof_pre_auth_rejection() {
1455 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1456 let (auth, conn) = test_auth_with_deadline_and_limit(Duration::from_millis(200), 4);
1457 let (mut client, server_stream) = duplex(512);
1458 let port = 41002;
1459 let server = handle_accepted_connection(
1460 server_stream,
1461 echo_router(),
1462 auth,
1463 SocketAddr::from(([127, 0, 0, 1], port)),
1464 Instant::now(),
1465 );
1466
1467 let (error, ()) = tokio::join!(server, send_wrong_key_client_auth(&mut client, &conn));
1468 assert!(matches!(
1469 error,
1470 Err(ConnectionError::Auth(AuthError::InvalidClientAuth))
1471 ));
1472 let captured = crate::router::test_log::captured_logs(&logs);
1473 assert!(
1474 captured.contains("reason=\"client_proof_mismatch\""),
1475 "{captured}"
1476 );
1477 assert_pre_auth_log(&logs, "client_authentication", port);
1478 }
1479
1480 #[tokio::test]
1481 async fn logs_malformed_pre_auth_message() {
1482 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1483 let (auth, _) = test_auth_with_deadline_and_limit(Duration::from_millis(200), 4);
1484 let (mut client, server_stream) = duplex(512);
1485 let port = 41003;
1486 let server = handle_accepted_connection(
1487 server_stream,
1488 echo_router(),
1489 auth,
1490 SocketAddr::from(([127, 0, 0, 1], port)),
1491 Instant::now(),
1492 );
1493 let send_malformed_message = async {
1494 let body = b"{";
1495 client
1496 .write_all(&(body.len() as u32).to_le_bytes())
1497 .await
1498 .expect("message length writes");
1499 client.write_all(body).await.expect("message body writes");
1500 };
1501
1502 let (error, ()) = tokio::join!(server, send_malformed_message);
1503 assert!(matches!(
1504 error,
1505 Err(ConnectionError::Auth(AuthError::JsonDecode {
1506 stage: AuthStage::ClientHello,
1507 ..
1508 }))
1509 ));
1510 let captured = crate::router::test_log::captured_logs(&logs);
1511 assert!(
1512 captured.contains("reason=\"malformed_or_oversized_message\""),
1513 "{captured}"
1514 );
1515 assert_pre_auth_log(&logs, "malformed_message", port);
1516 }
1517
1518 #[tokio::test]
1519 async fn logs_peer_close_during_handshake() {
1520 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1521 let (auth, _) = test_auth_with_deadline_and_limit(Duration::from_millis(200), 4);
1522 let (client, server_stream) = duplex(512);
1523 let port = 41004;
1524 drop(client);
1525
1526 let error = handle_accepted_connection(
1527 server_stream,
1528 echo_router(),
1529 auth,
1530 SocketAddr::from(([127, 0, 0, 1], port)),
1531 Instant::now(),
1532 )
1533 .await
1534 .expect_err("closed peer must be rejected");
1535 assert!(matches!(
1536 error,
1537 ConnectionError::Auth(AuthError::UnexpectedEof {
1538 stage: AuthStage::ClientHello,
1539 actual: 0,
1540 ..
1541 })
1542 ));
1543 let captured = crate::router::test_log::captured_logs(&logs);
1544 assert!(
1545 captured.contains("reason=\"peer_closed_before_authentication_completed\""),
1546 "{captured}"
1547 );
1548 assert_pre_auth_log(&logs, "peer_closed", port);
1549 }
1550
1551 #[tokio::test]
1555 async fn each_listener_rate_limits_its_own_rejection_logs() {
1556 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1557 for port in [41011, 41012] {
1558 let (auth, _) = test_auth_with_deadline_and_limit(Duration::from_millis(200), 4);
1559 let (client, server_stream) = duplex(512);
1560 drop(client);
1561 handle_accepted_connection(
1562 server_stream,
1563 echo_router(),
1564 auth,
1565 SocketAddr::from(([127, 0, 0, 1], port)),
1566 Instant::now(),
1567 )
1568 .await
1569 .expect_err("closed peer must be rejected");
1570 }
1571 let captured = crate::router::test_log::captured_logs(&logs);
1572 for port in [41011, 41012] {
1573 assert!(
1574 captured.contains(&format!("127.0.0.1:{port}")),
1575 "listener on port {port} logged nothing: {captured}"
1576 );
1577 }
1578 }
1579
1580 #[tokio::test]
1581 async fn logs_unauthenticated_capacity_rejection() {
1582 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1583 let deadline = Duration::from_millis(60);
1584 let (auth, _) = test_auth_with_deadline_and_limit(deadline, 1);
1585 let held_slot = auth
1586 .unauthenticated
1587 .clone()
1588 .try_acquire_owned()
1589 .expect("the only pre-auth slot is available");
1590 let (client, server_stream) = duplex(512);
1591 let port = 41005;
1592 let server = handle_accepted_connection(
1593 server_stream,
1594 echo_router(),
1595 auth,
1596 SocketAddr::from(([127, 0, 0, 1], port)),
1597 Instant::now(),
1598 );
1599
1600 let error = tokio::time::timeout(deadline * 3, server)
1601 .await
1602 .expect("queued connection should end at the capacity deadline")
1603 .expect_err("connection without a free slot must be rejected");
1604 drop(held_slot);
1605 drop(client);
1606 assert!(matches!(error, ConnectionError::UnauthenticatedCapacity));
1607 let captured = crate::router::test_log::captured_logs(&logs);
1608 assert!(
1609 captured.contains("reason=\"capacity_wait_timed_out\""),
1610 "{captured}"
1611 );
1612 assert_pre_auth_log(&logs, "waiting_for_handshake_slot", port);
1613 }
1614
1615 #[test]
1616 fn pre_auth_rejection_limiter_suppresses_bursts_and_reports_count() {
1617 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::INFO);
1618 let limiter = PreAuthRejectLimiter::new();
1619 let capacity = ConnectionError::UnauthenticatedCapacity;
1620 let handshake_timeout = ConnectionError::Auth(AuthError::Timeout {
1621 stage: AuthStage::ClientHello,
1622 deadline: Duration::from_millis(200),
1623 });
1624 let peer_addr = SocketAddr::from(([127, 0, 0, 1], 41006));
1625
1626 log_pre_auth_rejection(
1627 &capacity,
1628 peer_addr,
1629 Duration::from_millis(1),
1630 Duration::ZERO,
1631 &limiter,
1632 );
1633 for _ in 0..4 {
1634 log_pre_auth_rejection(
1635 &capacity,
1636 peer_addr,
1637 Duration::from_millis(2),
1638 Duration::from_millis(100),
1639 &limiter,
1640 );
1641 }
1642 log_pre_auth_rejection(
1643 &handshake_timeout,
1644 peer_addr,
1645 Duration::from_millis(3),
1646 Duration::from_millis(100),
1647 &limiter,
1648 );
1649 log_pre_auth_rejection(
1650 &capacity,
1651 peer_addr,
1652 Duration::from_millis(4),
1653 Duration::from_secs(1),
1654 &limiter,
1655 );
1656
1657 let captured = crate::router::test_log::captured_logs(&logs);
1658 assert_eq!(
1659 captured
1660 .matches("subc TCP connection rejected before routing")
1661 .count(),
1662 3,
1663 "one line per stage in a burst, then a later capacity line: {captured}"
1664 );
1665 assert!(
1666 captured.contains("suppressed_since_last=4"),
1667 "later capacity line should report the four hidden rejects: {captured}"
1668 );
1669 assert!(
1670 captured.contains("stage=\"handshake_deadline\""),
1671 "{captured}"
1672 );
1673 }
1674
1675 #[tokio::test]
1676 async fn buffered_frame_reader_coalesces_reads_and_preserves_short_reads() {
1677 let frames = vec![
1678 request(7, 1, b"first"),
1679 request(9, 2, b"second"),
1680 request(7, 3, b"third"),
1681 request(9, 4, b"fourth"),
1682 ];
1683 let bytes = encode_frames(&frames);
1684
1685 let (mut direct, direct_reads) = CountingReader::new(bytes.clone(), None);
1686 assert_eq!(read_frames(&mut direct, frames.len()).await, frames);
1687 assert_eq!(direct_reads.load(Ordering::Relaxed), frames.len() * 3);
1688
1689 let (buffered_source, buffered_reads) = CountingReader::new(bytes, None);
1690 let mut buffered = BufReader::new(buffered_source);
1691 assert_eq!(read_frames(&mut buffered, frames.len()).await, frames);
1692 assert_eq!(buffered_reads.load(Ordering::Relaxed), 1);
1693
1694 let split_frame = request(7, 5, b"split-body");
1695 let split_bytes = encode_frames(std::slice::from_ref(&split_frame));
1696 let (split_source, split_reads) = CountingReader::new(split_bytes, Some(10));
1697 let mut split_reader = BufReader::new(split_source);
1698 assert_eq!(
1699 read_frame(&mut split_reader).await.unwrap(),
1700 Some(split_frame)
1701 );
1702 assert_eq!(split_reads.load(Ordering::Relaxed), 2);
1703 }
1704
1705 #[tokio::test]
1706 async fn interleaved_channels_on_one_stream_demux_byte_identically_after_auth() {
1707 let (mut client, server_stream) = duplex(4096);
1708 let (auth, conn) = test_auth();
1709 let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
1710 authenticate(&mut client, &conn).await;
1711 let frames = [
1712 request(7, 1, b"chan7-first\0opaque"),
1713 request(9, 2, b"chan9-middle-{json?}"),
1714 request(7, 3, b"chan7-second\xffbytes"),
1715 ];
1716
1717 for frame in &frames {
1718 crate::write_frame(&mut client, frame).await.unwrap();
1719 }
1720
1721 for expected in &frames {
1722 let response = crate::read_frame(&mut client).await.unwrap().unwrap();
1723 assert_eq!(response.header.ty, FrameType::Response);
1724 assert_eq!(response.header.channel, expected.header.channel);
1725 assert_eq!(response.header.corr, expected.header.corr);
1726 assert_eq!(response.body, expected.body);
1727 }
1728
1729 drop(client);
1730 server.await.unwrap().unwrap();
1731 }
1732
1733 #[tokio::test]
1734 async fn channel_zero_goes_to_subc_self_handler_after_auth() {
1735 let (mut client, server_stream) = duplex(512);
1736 let (auth, conn) = test_auth();
1737 let server = tokio::spawn(handle_connection(
1738 server_stream,
1739 Arc::new(Router::with_default_self_handler()),
1740 auth,
1741 ));
1742 authenticate(&mut client, &conn).await;
1743 let ping = Frame::build(
1744 FrameType::Ping,
1745 Flags::new(false, Priority::Passive, false),
1746 0,
1747 0,
1748 55,
1749 Vec::new(),
1750 )
1751 .unwrap();
1752
1753 crate::write_frame(&mut client, &ping).await.unwrap();
1754 let response = crate::read_frame(&mut client).await.unwrap().unwrap();
1755
1756 assert_eq!(response.header.ty, FrameType::Pong);
1757 assert_eq!(response.header.channel, 0);
1758 assert_eq!(response.header.corr, 55);
1759 assert!(response.body.is_empty());
1760
1761 drop(client);
1762 server.await.unwrap().unwrap();
1763 }
1764
1765 #[tokio::test]
1766 async fn unauthenticated_connection_is_rejected_before_routing() {
1767 let (mut client, server_stream) = duplex(512);
1768 let (auth, _conn) = test_auth();
1769 let registry = Arc::new(Registry::default());
1770 let router = Arc::new(Router::with_control_handler(Arc::new(ControlHandler::new(
1771 Arc::clone(®istry),
1772 ))));
1773 let server = tokio::spawn(handle_connection(server_stream, router, auth));
1774 let ping = Frame::build(
1775 FrameType::Ping,
1776 Flags::new(false, Priority::Passive, false),
1777 0,
1778 0,
1779 66,
1780 Vec::new(),
1781 )
1782 .unwrap();
1783
1784 crate::write_frame(&mut client, &ping).await.unwrap();
1785 if let Ok(Ok(Some(frame))) =
1786 tokio::time::timeout(Duration::from_millis(200), crate::read_frame(&mut client)).await
1787 {
1788 panic!("unauthenticated frame reached router: {frame:?}");
1789 }
1790
1791 let err = server.await.unwrap().unwrap_err();
1792 assert!(matches!(err, ConnectionError::Auth(_)));
1793 assert_eq!(registry.active_registration_count().unwrap(), 0);
1794 }
1795
1796 #[tokio::test]
1797 async fn over_cap_peer_queues_for_a_slot_and_authenticates_when_one_frees() {
1798 let (mut first_client, first_server_stream) = duplex(2048);
1803 let (mut second_client, second_server_stream) = duplex(2048);
1804 let (auth, conn) = test_auth_with_limit(1);
1805 let registry = Arc::new(Registry::default());
1806 let router = Arc::new(Router::with_control_handler(Arc::new(ControlHandler::new(
1807 Arc::clone(®istry),
1808 ))));
1809
1810 let first_server = tokio::spawn(handle_connection(
1811 first_server_stream,
1812 Arc::clone(&router),
1813 auth.clone(),
1814 ));
1815
1816 let second_server = tokio::spawn(handle_connection(
1817 second_server_stream,
1818 Arc::clone(&router),
1819 auth.clone(),
1820 ));
1821
1822 tokio::time::sleep(Duration::from_millis(50)).await;
1824 assert!(
1825 !second_server.is_finished(),
1826 "queued peer must not be reset"
1827 );
1828
1829 authenticate(&mut first_client, &conn).await;
1832 authenticate(&mut second_client, &conn).await;
1833
1834 drop(first_client);
1835 drop(second_client);
1836 let _ = first_server.await;
1837 let _ = second_server.await;
1838 assert_eq!(registry.active_registration_count().unwrap(), 0);
1839 }
1840
1841 #[tokio::test]
1842 async fn over_cap_peer_is_rejected_when_no_slot_frees_within_deadline() {
1843 let (mut second_client, second_server_stream) = duplex(512);
1849 let (auth, conn) = test_auth_with_limit(1);
1850 let registry = Arc::new(Registry::default());
1851 let router = Arc::new(Router::with_control_handler(Arc::new(ControlHandler::new(
1852 Arc::clone(®istry),
1853 ))));
1854
1855 let held_slot = auth
1856 .unauthenticated
1857 .clone()
1858 .try_acquire_owned()
1859 .expect("sole pre-auth slot");
1860
1861 let second_server = tokio::spawn(handle_connection(
1862 second_server_stream,
1863 Arc::clone(&router),
1864 auth.clone(),
1865 ));
1866 let second_err = tokio::time::timeout(TEST_DEADLINE * 2, second_server)
1867 .await
1868 .expect("capacity reject should settle at the deadline")
1869 .expect("second connection task should not panic")
1870 .expect_err("queued peer must be rejected when no slot frees");
1871 drop(held_slot);
1872 assert!(matches!(
1873 second_err,
1874 ConnectionError::UnauthenticatedCapacity
1875 ));
1876 let mut closed = [0u8; 1];
1877 assert_eq!(
1878 second_client.read(&mut closed).await.unwrap(),
1879 0,
1880 "capacity-rejected peer should observe a closed stream"
1881 );
1882 assert_eq!(registry.active_registration_count().unwrap(), 0);
1883
1884 let (mut authed_client, authed_server_stream) = duplex(2048);
1885 let authed_server = tokio::spawn(handle_connection(
1886 authed_server_stream,
1887 Arc::clone(&router),
1888 auth,
1889 ));
1890 authenticate(&mut authed_client, &conn).await;
1891 let ping = Frame::build(
1892 FrameType::Ping,
1893 Flags::new(false, Priority::Passive, false),
1894 0,
1895 0,
1896 77,
1897 Vec::new(),
1898 )
1899 .unwrap();
1900 crate::write_frame(&mut authed_client, &ping).await.unwrap();
1901 let pong = crate::read_frame(&mut authed_client)
1902 .await
1903 .unwrap()
1904 .unwrap();
1905 assert_eq!(pong.header.ty, FrameType::Pong);
1906 assert_eq!(pong.header.channel, 0);
1907 assert_eq!(pong.header.corr, 77);
1908 assert_eq!(registry.active_registration_count().unwrap(), 0);
1909
1910 drop(authed_client);
1911 authed_server.await.unwrap().unwrap();
1912 }
1913
1914 #[tokio::test]
1915 async fn bind_ack_during_partial_frame_preserves_next_request() {
1916 use subc_control::ClientControlRequest;
1917 use subc_protocol::{
1918 manifest::{
1919 Concurrency, ExecutionMode, IdentityScope, ModuleManifest, ProviderRole, Tool,
1920 },
1921 session::{ModuleControlRequest, ModuleControlResponse},
1922 BindIdentity, ModuleHelloBody, RouteTarget,
1923 };
1924
1925 let mut configured_router = Router::with_default_self_handler();
1926 configured_router.register_backend(7, EchoBackend).unwrap();
1927 let router = Arc::new(configured_router);
1928 let (auth, conn) = test_auth();
1929 let (mut module, module_stream) = duplex(4096);
1930 let module_server = tokio::spawn(handle_connection(
1931 module_stream,
1932 Arc::clone(&router),
1933 auth.clone(),
1934 ));
1935 authenticate(&mut module, &conn).await;
1936 let manifest = ModuleManifest::builder("frame-test", "0.1.0")
1937 .protocol_ver(PROTOCOL_VERSION)
1938 .provides(vec![ProviderRole::ToolProvider {
1939 tools: vec![Tool {
1940 name: "read".into(),
1941 description: None,
1942 execution_mode: ExecutionMode::Pure,
1943 schema: serde_json::json!({"type": "object"}),
1944 }],
1945 identity_scope: vec![IdentityScope::Project, IdentityScope::Session],
1946 concurrency: Concurrency::ModuleManaged,
1947 emits_push: true,
1948 sub_supervises: true,
1949 }])
1950 .build();
1951 let hello = Frame::build(
1952 FrameType::Hello,
1953 Flags::new(false, Priority::Passive, false),
1954 0,
1955 0,
1956 1,
1957 serde_json::to_vec(&ModuleHelloBody {
1958 manifest,
1959 protocol_ver: PROTOCOL_VERSION,
1960 control_ops: None,
1961 launch_nonce: None,
1962 })
1963 .unwrap(),
1964 )
1965 .unwrap();
1966 crate::write_frame(&mut module, &hello).await.unwrap();
1967 assert_eq!(
1968 read_frame(&mut module).await.unwrap().unwrap().header.ty,
1969 FrameType::HelloAck
1970 );
1971
1972 let (mut client, client_stream) = duplex(4096);
1973 let client_server = tokio::spawn(handle_connection(client_stream, router, auth));
1974 authenticate(&mut client, &conn).await;
1975 let open = Frame::build(
1976 FrameType::Request,
1977 Flags::new(false, Priority::Passive, false),
1978 0,
1979 0,
1980 2,
1981 serde_json::to_vec(&ClientControlRequest::RouteOpen {
1982 target: RouteTarget::ToolProvider {
1983 module_id: "frame-test".into(),
1984 },
1985 identity: BindIdentity::new(std::env::current_dir().unwrap(), "unit", "session"),
1986 consumer_identity: None,
1987 consumer_capabilities: None,
1988 role_versions: None,
1989 admission_facts: None,
1990 scope: None,
1991 })
1992 .unwrap(),
1993 )
1994 .unwrap();
1995 crate::write_frame(&mut client, &open).await.unwrap();
1996 let bind = timeout(TEST_DEADLINE, read_frame(&mut module))
1997 .await
1998 .unwrap()
1999 .unwrap()
2000 .unwrap();
2001 assert!(matches!(
2002 serde_json::from_slice::<ModuleControlRequest>(&bind.body).unwrap(),
2003 ModuleControlRequest::RouteBind { .. }
2004 ));
2005
2006 let ping = request(7, 3, b"partial-frame-body");
2007 client.write_all(&ping.header.encode()).await.unwrap();
2008 tokio::time::sleep(Duration::from_millis(30)).await;
2010 let ack = Frame::build(
2011 FrameType::Response,
2012 Flags::new(false, Priority::Passive, false),
2013 0,
2014 0,
2015 bind.header.corr,
2016 serde_json::to_vec(&ModuleControlResponse::RouteBindAck {}).unwrap(),
2017 )
2018 .unwrap();
2019 crate::write_frame(&mut module, &ack).await.unwrap();
2020 let opened = timeout(TEST_DEADLINE, read_frame(&mut client))
2021 .await
2022 .unwrap()
2023 .unwrap();
2024 if opened.is_none() {
2025 panic!("client closed: {:?}", client_server.await);
2026 }
2027 let opened = opened.unwrap();
2028 assert_eq!(opened.header.corr, 2);
2029 client.write_all(&ping.body).await.unwrap();
2030 let pong = timeout(TEST_DEADLINE, read_frame(&mut client))
2031 .await
2032 .expect("partial frame must reach the router after the bind ack")
2033 .expect("frame must decode")
2034 .expect("connection must remain open");
2035 assert_eq!(pong.header.ty, FrameType::Response);
2036 assert_eq!(pong.header.corr, 3);
2037 assert_eq!(pong.body, ping.body);
2038 drop(client);
2039 drop(module);
2040 client_server.await.unwrap().unwrap();
2041 let _ = module_server.await.unwrap();
2042 }
2043
2044 #[tokio::test]
2045 async fn aborted_accept_does_not_end_listener() {
2046 let listener = Arc::new(TcpListener::bind("127.0.0.1:0").await.unwrap());
2047 let addr = listener.local_addr().unwrap();
2048 let (auth, conn) = test_auth();
2049 let mut attempts = 0;
2050 let server = tokio::spawn(serve_listener_with_accept(
2051 Some(addr),
2052 echo_router(),
2053 auth,
2054 move || {
2055 attempts += 1;
2056 let result = if attempts == 1 {
2057 Some(io::Error::from(io::ErrorKind::ConnectionAborted))
2058 } else {
2059 None
2060 };
2061 let listener = Arc::clone(&listener);
2062 async move {
2063 match result {
2064 Some(err) => Err(err),
2065 None => listener.accept().await,
2066 }
2067 }
2068 },
2069 ));
2070 let mut client = tokio::net::TcpStream::connect(addr).await.unwrap();
2071 authenticate(&mut client, &conn).await;
2072 let ping = Frame::build(
2073 FrameType::Ping,
2074 Flags::new(false, Priority::Passive, false),
2075 0,
2076 0,
2077 77,
2078 Vec::new(),
2079 )
2080 .unwrap();
2081 crate::write_frame(&mut client, &ping).await.unwrap();
2082 let pong = timeout(TEST_DEADLINE, read_frame(&mut client))
2083 .await
2084 .unwrap()
2085 .unwrap()
2086 .unwrap();
2087 assert_eq!(pong.header.ty, FrameType::Pong);
2088 assert!(
2089 !server.is_finished(),
2090 "a temporary accept failure must not stop the listener"
2091 );
2092 server.abort();
2093 }
2094
2095 #[tokio::test]
2096 async fn exhausted_accept_backs_off_then_serves_connection() {
2097 let listener = Arc::new(TcpListener::bind("127.0.0.1:0").await.unwrap());
2098 let addr = listener.local_addr().unwrap();
2099 let (auth, conn) = test_auth();
2100 let mut attempts = 0;
2101 let server = tokio::spawn(serve_listener_with_accept(
2102 Some(addr),
2103 echo_router(),
2104 auth,
2105 move || {
2106 attempts += 1;
2107 #[cfg(unix)]
2110 let emfile = rustix::io::Errno::MFILE.raw_os_error();
2111 #[cfg(windows)]
2112 let emfile = 10024;
2113 let error = (attempts == 1).then(|| io::Error::from_raw_os_error(emfile));
2114 let listener = Arc::clone(&listener);
2115 async move {
2116 match error {
2117 Some(err) => Err(err),
2118 None => listener.accept().await,
2119 }
2120 }
2121 },
2122 ));
2123 let started = Instant::now();
2124 let mut client = tokio::net::TcpStream::connect(addr).await.unwrap();
2125 authenticate(&mut client, &conn).await;
2126 assert!(
2127 started.elapsed() >= Duration::from_millis(50),
2128 "fd exhaustion must back off before retrying"
2129 );
2130 assert!(!server.is_finished());
2131 server.abort();
2132 }
2133
2134 #[tokio::test]
2135 async fn fatal_accept_error_still_stops_listener() {
2136 let (auth, _) = test_auth();
2137 let err = serve_listener_with_accept(None, echo_router(), auth, || async {
2138 Err(io::Error::from(io::ErrorKind::PermissionDenied))
2139 })
2140 .await
2141 .unwrap_err();
2142 assert!(
2143 matches!(err, ServerError::Accept { source, .. } if source.kind() == io::ErrorKind::PermissionDenied)
2144 );
2145 }
2146
2147 #[tokio::test]
2148 async fn serve_listeners_with_no_listeners_returns_typed_error() {
2149 let (auth, _conn) = test_auth();
2150 let err = serve_listeners(Vec::new(), echo_router(), auth)
2151 .await
2152 .expect_err("empty listener set must fail loudly");
2153 assert!(matches!(err, ServerError::NoListeners));
2154 }
2155
2156 #[tokio::test]
2157 async fn malformed_header_returns_typed_error_no_panic() {
2158 let (mut client, server_stream) = duplex(128);
2159 let (auth, conn) = test_auth();
2160 let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
2161 authenticate(&mut client, &conn).await;
2162 let mut header = [0u8; HEADER_LEN];
2163 header[4] = PROTOCOL_VERSION;
2164 header[5] = 250;
2165
2166 client.write_all(&header).await.unwrap();
2167 drop(client);
2168
2169 let err = server.await.unwrap().unwrap_err();
2170 assert!(matches!(
2171 err,
2172 ConnectionError::FrameIo(FrameIoError::DecodeHeader(DecodeError::UnknownFrameType {
2173 byte: 250
2174 }))
2175 ));
2176 }
2177
2178 #[tokio::test]
2179 async fn truncated_body_returns_typed_error_no_panic() {
2180 let (mut client, server_stream) = duplex(128);
2181 let (auth, conn) = test_auth();
2182 let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
2183 authenticate(&mut client, &conn).await;
2184 let frame = request(7, 8, b"abcd");
2185
2186 client.write_all(&frame.header.encode()).await.unwrap();
2187 client.write_all(b"ab").await.unwrap();
2188 drop(client);
2189
2190 let err = server.await.unwrap().unwrap_err();
2191 assert!(matches!(
2192 err,
2193 ConnectionError::FrameIo(FrameIoError::UnexpectedEof {
2194 stage: ReadStage::Body,
2195 expected: 4,
2196 actual: 2
2197 })
2198 ));
2199 }
2200
2201 #[tokio::test]
2202 async fn unknown_channel_is_returned_as_error_frame_and_connection_continues() {
2203 let (mut client, server_stream) = duplex(1024);
2204 let (auth, conn) = test_auth();
2205 let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
2206 authenticate(&mut client, &conn).await;
2207 let unknown = request(42, 10, b"lost");
2208 let known = request(7, 11, b"still-routes");
2209
2210 crate::write_frame(&mut client, &unknown).await.unwrap();
2211 crate::write_frame(&mut client, &known).await.unwrap();
2212
2213 let error = crate::read_frame(&mut client).await.unwrap().unwrap();
2214 assert_eq!(error.header.ty, FrameType::Error);
2215 assert_eq!(error.header.channel, 42);
2216 assert_eq!(error.header.corr, 10);
2217 let error_body: ErrorBody = serde_json::from_slice(&error.body).unwrap();
2218 assert_eq!(error_body.code, "unknown_channel");
2219
2220 let response = crate::read_frame(&mut client).await.unwrap().unwrap();
2221 assert_eq!(response.header.ty, FrameType::Response);
2222 assert_eq!(response.header.channel, 7);
2223 assert_eq!(response.header.corr, 11);
2224 assert_eq!(response.body, b"still-routes");
2225
2226 drop(client);
2227 server.await.unwrap().unwrap();
2228 }
2229}