1use std::{
2 error::Error,
3 fmt, io,
4 net::SocketAddr,
5 sync::Arc,
6 time::{Duration, Instant},
7};
8
9use subc_transport::{authenticate_server, AuthError, DAEMON_ID_LEN, WATCHDOG_CLIENT_ROLE};
10use tokio::{
11 io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader, BufWriter},
12 net::TcpListener,
13 sync::{mpsc, Semaphore},
14 task::{JoinHandle, JoinSet},
15 time::timeout,
16};
17use tracing::{debug, warn};
18
19use crate::{
20 forwarding::{CloseReason, ConnectionCloseReceiver},
21 observability::ConnectedClients,
22 read_frame,
23 router::{FrameSink, RouteCtx, Router},
24 write_frame, FrameIoError, RouterError,
25};
26
27pub const CONNECTION_EGRESS_BYTE_BUDGET: usize = 4 * 1024 * 1024;
39pub const CONNECTION_EGRESS_FRAME_CAP: usize = 32 * 1024;
48pub const MAX_PENDING_ROUTE_OPENS_PER_CONNECTION: usize = 8;
54pub(crate) const MAX_PENDING_ROUTE_BINDS_PER_TARGET: usize =
58 MAX_PENDING_ROUTE_OPENS_PER_CONNECTION * 2;
59pub const DEFAULT_AUTH_DEADLINE: Duration = Duration::from_secs(2);
60pub const DEFAULT_MAX_UNAUTHENTICATED_CONNECTIONS: usize = 256;
65const CLOSE_DRAIN_GRACE: Duration = Duration::from_secs(2);
66
67#[derive(Clone)]
70pub struct ServerAuth {
71 key: Arc<[u8]>,
72 daemon_id: [u8; DAEMON_ID_LEN],
73 daemon_ver: Arc<str>,
74 deadline: Duration,
75 unauthenticated: Arc<Semaphore>,
76 connected_clients: ConnectedClients,
77}
78
79impl ServerAuth {
80 pub fn new(
81 key: Vec<u8>,
82 daemon_id: [u8; DAEMON_ID_LEN],
83 daemon_ver: impl Into<String>,
84 ) -> Self {
85 Self::with_limits(
86 key,
87 daemon_id,
88 daemon_ver,
89 DEFAULT_AUTH_DEADLINE,
90 DEFAULT_MAX_UNAUTHENTICATED_CONNECTIONS,
91 )
92 }
93
94 pub fn with_limits(
96 key: Vec<u8>,
97 daemon_id: [u8; DAEMON_ID_LEN],
98 daemon_ver: impl Into<String>,
99 deadline: Duration,
100 max_unauthenticated: usize,
101 ) -> Self {
102 Self {
103 key: Arc::from(key),
104 daemon_id,
105 daemon_ver: Arc::from(daemon_ver.into()),
106 deadline,
107 unauthenticated: Arc::new(Semaphore::new(max_unauthenticated.max(1))),
108 connected_clients: ConnectedClients::new(),
109 }
110 }
111
112 pub fn with_connected_clients(mut self, connected_clients: ConnectedClients) -> Self {
113 self.connected_clients = connected_clients;
114 self
115 }
116}
117
118impl fmt::Debug for ServerAuth {
119 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
120 f.debug_struct("ServerAuth")
121 .field("key", &"<redacted>")
122 .field("daemon_id", &self.daemon_id)
123 .field("daemon_ver", &self.daemon_ver)
124 .field("deadline", &self.deadline)
125 .finish_non_exhaustive()
126 }
127}
128
129pub async fn serve_listener(
132 listener: TcpListener,
133 router: Arc<Router>,
134 auth: ServerAuth,
135) -> Result<(), ServerError> {
136 serve_listener_with_accept(listener.local_addr().ok(), router, auth, || {
137 listener.accept()
138 })
139 .await
140}
141
142async fn serve_listener_with_accept<A, F>(
143 local_addr: Option<SocketAddr>,
144 router: Arc<Router>,
145 auth: ServerAuth,
146 mut accept: A,
147) -> Result<(), ServerError>
148where
149 A: FnMut() -> F,
150 F: std::future::Future<Output = io::Result<(tokio::net::TcpStream, SocketAddr)>>,
151{
152 loop {
153 let (stream, peer_addr) = match accept().await {
154 Ok(accepted) => accepted,
155 Err(source) => {
156 let kind = source.kind();
157 let exhausted = accept_resource_exhausted(&source);
161 if matches!(
162 kind,
163 io::ErrorKind::ConnectionAborted
164 | io::ErrorKind::ConnectionReset
165 | io::ErrorKind::Interrupted
166 ) || exhausted
167 {
168 warn!(?local_addr, error = %source, "temporary TCP accept failure");
169 if exhausted {
170 tokio::time::sleep(Duration::from_millis(75)).await;
171 }
172 continue;
173 }
174 return Err(ServerError::Accept { local_addr, source });
175 }
176 };
177 if let Err(source) = stream.set_nodelay(true) {
193 warn!(?peer_addr, error = %source, "could not disable Nagle on accepted connection");
194 }
195 debug!(?peer_addr, ?local_addr, "accepted subc TCP connection");
196 let router = Arc::clone(&router);
197 let auth = auth.clone();
198 tokio::spawn(async move {
199 if let Err(err) = handle_connection(stream, router, auth).await {
200 if err.is_quiet_reject() {
201 debug!(?peer_addr, error = %err, "subc TCP connection rejected before routing");
202 } else {
203 warn!(?peer_addr, error = %err, "subc connection ended with error");
204 }
205 }
206 });
207 }
208}
209
210fn accept_resource_exhausted(error: &io::Error) -> bool {
211 #[cfg(unix)]
216 {
217 use rustix::io::Errno;
218 let exhausted = [Errno::NFILE, Errno::MFILE, Errno::NOBUFS];
219 error
220 .raw_os_error()
221 .is_some_and(|code| exhausted.iter().any(|e| e.raw_os_error() == code))
222 }
223 #[cfg(windows)]
224 {
225 matches!(error.raw_os_error(), Some(10024 | 10055))
226 }
227 #[cfg(not(any(unix, windows)))]
228 {
229 let _ = error;
230 false
231 }
232}
233
234pub async fn serve_listeners(
236 listeners: Vec<TcpListener>,
237 router: Arc<Router>,
238 auth: ServerAuth,
239) -> Result<(), ServerError> {
240 if listeners.is_empty() {
241 return Err(ServerError::NoListeners);
242 }
243
244 let (tx, mut rx) = mpsc::channel(listeners.len());
245 let mut accept_tasks = AbortTasksOnDrop::default();
246 for listener in listeners {
247 let router = Arc::clone(&router);
248 let auth = auth.clone();
249 let tx = tx.clone();
250 accept_tasks.push(tokio::spawn(async move {
251 let result = serve_listener(listener, router, auth).await;
252 let _ = tx.send(result).await;
253 }));
254 }
255 drop(tx);
256
257 rx.recv().await.unwrap_or(Ok(()))
258}
259
260#[derive(Default)]
261struct AbortTasksOnDrop {
262 handles: Vec<JoinHandle<()>>,
263}
264
265impl AbortTasksOnDrop {
266 fn push(&mut self, handle: JoinHandle<()>) {
267 self.handles.push(handle);
268 }
269}
270
271impl Drop for AbortTasksOnDrop {
272 fn drop(&mut self) {
273 for handle in &self.handles {
274 if !handle.is_finished() {
275 handle.abort();
276 }
277 }
278 }
279}
280
281#[derive(Debug)]
282enum ConnectionLoopExit {
283 PeerClosed,
284 CloseRequested(CloseReason),
285}
286
287pub async fn handle_connection<S>(
295 mut stream: S,
296 router: Arc<Router>,
297 auth: ServerAuth,
298) -> Result<(), ConnectionError>
299where
300 S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
301{
302 let permit =
322 match tokio::time::timeout(auth.deadline, auth.unauthenticated.clone().acquire_owned())
323 .await
324 {
325 Ok(Ok(permit)) => permit,
326 Ok(Err(_)) | Err(_) => {
327 let _ = stream.shutdown().await;
328 return Err(ConnectionError::UnauthenticatedCapacity);
329 }
330 };
331
332 let authenticated = authenticate_server(
333 &mut stream,
334 auth.key.as_ref(),
335 &auth.daemon_id,
336 auth.daemon_ver.as_ref(),
337 auth.deadline,
338 )
339 .await
340 .map_err(ConnectionError::Auth)?;
341 drop(permit);
342
343 let mut connection = router.begin_connection();
344 let connection_id = connection.id();
345 let _connected_client = (authenticated.role != WATCHDOG_CLIENT_ROLE)
360 .then(|| auth.connected_clients.open(connection_id));
361 let close_receiver = connection.take_close_receiver();
362 debug!(
363 connection_id = connection_id.get(),
364 "subc authenticated connection opened"
365 );
366
367 let (read_half, write_half) = tokio::io::split(stream);
368 let mut read_half = BufReader::new(read_half);
372 let (egress, rx) = connection_egress();
373 let mut writer = tokio::spawn(drain_writer(write_half, rx));
374
375 let ctx = RouteCtx {
376 connection_id,
377 egress: egress.clone(),
378 };
379
380 let mut route_open_tasks = JoinSet::new();
381 let loop_result = connection_loop(
382 &mut read_half,
383 Arc::clone(&router),
384 ctx.clone(),
385 close_receiver,
386 &mut route_open_tasks,
387 )
388 .await;
389
390 route_open_tasks.shutdown().await;
398
399 drop(ctx);
400 drop(egress);
401 drop(connection);
402
403 let close_reason = match &loop_result {
404 Ok(ConnectionLoopExit::CloseRequested(reason)) => Some(reason.to_string()),
405 Ok(ConnectionLoopExit::PeerClosed) | Err(_) => None,
406 };
407 let writer_result = if close_reason.is_some() {
408 match timeout(CLOSE_DRAIN_GRACE, &mut writer).await {
409 Ok(result) => Some(result.map_err(ConnectionError::WriterTask)),
410 Err(_) => {
411 warn!(
412 connection_id = connection_id.get(),
413 grace = ?CLOSE_DRAIN_GRACE,
414 "connection writer did not drain after close request; aborting writer task"
415 );
416 writer.abort();
417 let _ = writer.await;
418 None
419 }
420 }
421 } else {
422 Some(writer.await.map_err(ConnectionError::WriterTask))
423 };
424
425 let result = if let Some(reason) = close_reason.as_deref() {
426 match writer_result {
427 Some(Ok(Ok(()))) | None => Ok(()),
428 Some(Ok(Err(writer_err))) => {
429 debug!(
430 connection_id = connection_id.get(),
431 close_reason = reason,
432 writer_error = %writer_err,
433 "writer failed after requested connection close"
434 );
435 Ok(())
436 }
437 Some(Err(join_err)) => {
438 warn!(
439 connection_id = connection_id.get(),
440 close_reason = reason,
441 join_error = %join_err,
442 "writer task join failed after requested connection close"
443 );
444 Ok(())
445 }
446 }
447 } else {
448 let writer_result =
449 writer_result.expect("writer result is present without a close request");
450 match (loop_result, writer_result) {
451 (Err(loop_err), Ok(Ok(()))) => Err(loop_err),
452 (Err(loop_err), Ok(Err(writer_err))) => {
453 warn!(
454 connection_id = connection_id.get(),
455 writer_error = %writer_err,
456 "writer failed while closing after connection error"
457 );
458 Err(loop_err)
459 }
460 (Err(loop_err), Err(join_err)) => {
461 warn!(
462 connection_id = connection_id.get(),
463 join_error = %join_err,
464 "writer task join failed while closing after connection error"
465 );
466 Err(loop_err)
467 }
468 (Ok(ConnectionLoopExit::PeerClosed), Ok(Ok(()))) => Ok(()),
469 (Ok(ConnectionLoopExit::PeerClosed), Ok(Err(writer_err))) => {
470 Err(ConnectionError::FrameIo(writer_err))
471 }
472 (Ok(ConnectionLoopExit::PeerClosed), Err(join_err)) => Err(join_err),
473 (Ok(ConnectionLoopExit::CloseRequested(_)), _) => {
474 unreachable!("close requests are handled before normal writer result matching")
475 }
476 }
477 };
478
479 match &result {
480 Ok(()) => {
481 if let Some(reason) = close_reason.as_deref() {
482 debug!(
483 connection_id = connection_id.get(),
484 close_reason = reason,
485 "subc connection closed by request"
486 );
487 } else {
488 debug!(
489 connection_id = connection_id.get(),
490 "subc connection closed"
491 );
492 }
493 }
494 Err(err) => debug!(
495 connection_id = connection_id.get(),
496 error = %err,
497 "subc connection exited with error"
498 ),
499 }
500
501 result
502}
503
504async fn connection_loop<R>(
505 read_half: &mut R,
506 router: Arc<Router>,
507 ctx: RouteCtx,
508 mut close_receiver: ConnectionCloseReceiver,
509 route_open_tasks: &mut JoinSet<Result<(), RouterError>>,
510) -> Result<ConnectionLoopExit, ConnectionError>
511where
512 R: AsyncRead + Unpin,
513{
514 loop {
515 while let Some(result) = route_open_tasks.try_join_next() {
516 finish_route_open_task(result)?;
517 }
518
519 let read = read_frame(&mut *read_half);
522 tokio::pin!(read);
523 let frame = loop {
524 tokio::select! {
525 close = &mut close_receiver => {
526 return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
527 }
528 result = route_open_tasks.join_next(), if !route_open_tasks.is_empty() => {
529 finish_route_open_task(
530 result.expect("a non-empty route.open JoinSet has a next task")
531 )?;
532 }
533 result = &mut read => {
534 break match result.map_err(ConnectionError::FrameIo)? {
535 Some(frame) => frame,
536 None => return Ok(ConnectionLoopExit::PeerClosed),
537 };
538 }
539 }
540 };
541
542 if let Some(target_module_id) = router.route_open_target(&frame) {
543 while let Some(result) = route_open_tasks.try_join_next() {
547 finish_route_open_task(result)?;
548 }
549
550 if route_open_tasks.len() >= MAX_PENDING_ROUTE_OPENS_PER_CONNECTION {
551 let refusal = router
554 .route_open_capacity_refusal(
555 &ctx,
556 &frame,
557 &target_module_id,
558 route_open_tasks.len(),
559 MAX_PENDING_ROUTE_OPENS_PER_CONNECTION,
560 )
561 .map_err(ConnectionError::Router)?;
562 let send_result = tokio::select! {
563 close = &mut close_receiver => {
564 return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
565 }
566 result = ctx.egress.send(refusal) => result,
567 };
568 send_result.map_err(ConnectionError::Router)?;
569 continue;
570 }
571
572 let task_router = Arc::clone(&router);
573 let task_ctx = ctx.clone();
574 let dispatch_started_at = Instant::now();
577 route_open_tasks.spawn(async move {
578 route_open_tail(task_router, task_ctx, frame, dispatch_started_at).await
579 });
580 continue;
581 }
582
583 let route_result = tokio::select! {
586 close = &mut close_receiver => {
587 return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
588 }
589 result = router.route_for_connection(&ctx, frame) => result,
590 };
591
592 if let Err(err) = route_result {
593 if let Some(error_frame) = err.to_error_frame() {
594 warn!(
595 connection_id = ctx.connection_id.get(),
596 error = %err,
597 "routing failure recovered with ERROR frame"
598 );
599 let send_result = tokio::select! {
600 close = &mut close_receiver => {
601 return Ok(ConnectionLoopExit::CloseRequested(close_reason(close)));
602 }
603 result = ctx.egress.send(error_frame) => result,
604 };
605 send_result.map_err(ConnectionError::Router)?;
606 } else {
607 debug!(
608 connection_id = ctx.connection_id.get(),
609 error = %err,
610 "fatal routing failure"
611 );
612 return Err(ConnectionError::Router(err));
613 }
614 }
615 }
616}
617
618async fn route_open_tail(
619 router: Arc<Router>,
620 ctx: RouteCtx,
621 frame: crate::Frame,
622 dispatch_started_at: Instant,
623) -> Result<(), RouterError> {
624 match router
625 .route_for_connection_started(&ctx, frame, Some(dispatch_started_at))
626 .await
627 {
628 Ok(()) => Ok(()),
629 Err(err) => {
630 let Some(error_frame) = err.to_error_frame() else {
631 return Err(err);
632 };
633 warn!(
634 connection_id = ctx.connection_id.get(),
635 error = %err,
636 "routing failure recovered with ERROR frame"
637 );
638 ctx.egress.send(error_frame).await
639 }
640 }
641}
642
643fn finish_route_open_task(
644 result: Result<Result<(), RouterError>, tokio::task::JoinError>,
645) -> Result<(), ConnectionError> {
646 match result {
647 Ok(Ok(())) => Ok(()),
648 Ok(Err(err)) => Err(ConnectionError::Router(err)),
649 Err(err) => Err(ConnectionError::Router(RouterError::backend(
650 0,
651 0,
652 format!("route.open task failed: {err}"),
653 ))),
654 }
655}
656
657fn close_reason(
658 result: Result<CloseReason, tokio::sync::oneshot::error::RecvError>,
659) -> CloseReason {
660 result.unwrap_or_else(|_| {
661 CloseReason::new(
662 "close_registry_dropped",
663 "connection close registration was dropped without a reason",
664 )
665 })
666}
667
668pub(crate) fn connection_egress() -> (FrameSink, mpsc::Receiver<crate::router::OutboundFrame>) {
672 let (tx, rx) = mpsc::channel::<crate::router::OutboundFrame>(CONNECTION_EGRESS_FRAME_CAP);
673 (
674 FrameSink::with_byte_budget(tx, CONNECTION_EGRESS_BYTE_BUDGET),
675 rx,
676 )
677}
678
679async fn drain_writer<W>(
680 write_half: W,
681 mut rx: mpsc::Receiver<crate::router::OutboundFrame>,
682) -> Result<(), FrameIoError>
683where
684 W: AsyncWrite + Unpin,
685{
686 let mut writer = BufWriter::new(write_half);
687 while let Some(outbound) = rx.recv().await {
688 write_outbound(&mut writer, outbound).await?;
689 while let Ok(outbound) = rx.try_recv() {
690 write_outbound(&mut writer, outbound).await?;
691 }
692 writer.flush().await.map_err(FrameIoError::Io)?;
693 }
694 writer.flush().await.map_err(FrameIoError::Io)?;
695 Ok(())
696}
697
698async fn write_outbound<W>(
705 writer: &mut BufWriter<W>,
706 outbound: crate::router::OutboundFrame,
707) -> Result<(), FrameIoError>
708where
709 W: AsyncWrite + Unpin,
710{
711 const SLOW_REPLY_QUEUE: Duration = Duration::from_millis(1000);
712 let queued = outbound.enqueued_at.elapsed();
713 let charge = outbound.charge;
716 if let Some(charge) = &charge {
717 charge.taken_by_writer();
718 }
719 let frame = outbound.frame;
720 if frame.header.channel == 0 && queued >= SLOW_REPLY_QUEUE {
721 let write_started = std::time::Instant::now();
722 let result = write_frame(writer, &frame).await;
723 tracing::warn!(
724 corr = frame.header.corr,
725 queued_ms = queued.as_millis() as u64,
726 write_ms = write_started.elapsed().as_millis() as u64,
727 "slow control reply write"
728 );
729 result?;
730 } else {
731 write_frame(writer, &frame).await?;
732 }
733 if let Some(flushed) = outbound.flushed {
734 writer.flush().await.map_err(FrameIoError::Io)?;
735 let _ = flushed.send(());
736 }
737 Ok(())
738}
739
740#[cfg(all(test, unix))]
741#[tokio::test]
742async fn shutdown_notice_ack_waits_for_socket_flush() {
743 let (socket, mut peer) = tokio::io::duplex(1);
744 let (tx, rx) = mpsc::channel(1);
745 let sink = FrameSink::new(tx);
746 let writer = tokio::spawn(drain_writer(socket, rx));
747 let frame = crate::Frame::build(
748 subc_protocol::FrameType::Push,
749 subc_protocol::Flags::new(false, subc_protocol::Priority::Interactive, false),
750 0,
751 0,
752 0,
753 b"notice".to_vec(),
754 )
755 .unwrap();
756 let send = sink.send_flushed(frame);
757 tokio::pin!(send);
758 assert!(
759 timeout(Duration::from_millis(20), &mut send).await.is_err(),
760 "queueing bytes is not a socket flush acknowledgement"
761 );
762 let (sent, received) = timeout(Duration::from_secs(1), async {
763 tokio::join!(&mut send, read_frame(&mut peer))
764 })
765 .await
766 .unwrap();
767 sent.unwrap();
768 assert_eq!(received.unwrap().unwrap().body, b"notice");
769 writer.abort();
770}
771
772#[derive(Debug)]
773pub enum ServerError {
774 NoListeners,
775 Accept {
776 local_addr: Option<SocketAddr>,
777 source: io::Error,
778 },
779}
780
781impl fmt::Display for ServerError {
782 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
783 match self {
784 Self::NoListeners => write!(f, "no TCP listeners were provided"),
785 Self::Accept { local_addr, source } => match local_addr {
786 Some(addr) => write!(f, "failed to accept TCP connection on {addr}: {source}"),
787 None => write!(f, "failed to accept TCP connection: {source}"),
788 },
789 }
790 }
791}
792
793impl Error for ServerError {
794 fn source(&self) -> Option<&(dyn Error + 'static)> {
795 match self {
796 Self::Accept { source, .. } => Some(source),
797 Self::NoListeners => None,
798 }
799 }
800}
801
802#[derive(Debug)]
803pub enum ConnectionError {
804 Auth(AuthError),
805 UnauthenticatedCapacity,
806 FrameIo(FrameIoError),
807 Router(RouterError),
808 WriterTask(tokio::task::JoinError),
809}
810
811impl ConnectionError {
812 fn is_quiet_reject(&self) -> bool {
813 matches!(self, Self::Auth(_) | Self::UnauthenticatedCapacity)
814 }
815}
816
817impl fmt::Display for ConnectionError {
818 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
819 match self {
820 Self::Auth(err) => write!(f, "connection auth failed: {err}"),
821 Self::UnauthenticatedCapacity => write!(
822 f,
823 "too many concurrent unauthenticated subc TCP connections"
824 ),
825 Self::FrameIo(err) => write!(f, "frame connection error: {err}"),
826 Self::Router(err) => write!(f, "router connection error: {err}"),
827 Self::WriterTask(err) => write!(f, "connection writer task failed: {err}"),
828 }
829 }
830}
831
832impl Error for ConnectionError {
833 fn source(&self) -> Option<&(dyn Error + 'static)> {
834 match self {
835 Self::Auth(err) => Some(err),
836 Self::FrameIo(err) => Some(err),
837 Self::Router(err) => Some(err),
838 Self::WriterTask(err) => Some(err),
839 Self::UnauthenticatedCapacity => None,
840 }
841 }
842}
843
844#[cfg(test)]
845mod tests {
846 use std::{
847 pin::Pin,
848 sync::atomic::{AtomicUsize, Ordering},
849 task::{Context, Poll},
850 };
851
852 use super::*;
853 use subc_protocol::{
854 DecodeError, ErrorBody, Flags, FrameType, Priority, HEADER_LEN, PROTOCOL_VERSION,
855 };
856 use tokio::io::{duplex, AsyncReadExt, AsyncWriteExt, ReadBuf};
857
858 use subc_transport::{authenticate_client, ConnectionInfo, Endpoint, SCHEMA_VERSION};
859
860 use crate::{ControlHandler, EchoBackend, Frame, ReadStage, Registry};
861
862 const TEST_DEADLINE: Duration = Duration::from_secs(2);
863 const TEST_DAEMON_VER: &str = "test-subc-server";
864
865 #[tokio::test]
870 async fn stale_queued_control_reply_logs_slow_reply_write() {
871 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::WARN);
872 let (tx, rx) = mpsc::channel::<crate::router::OutboundFrame>(4);
873 let reply = Frame::build_with_version(
874 PROTOCOL_VERSION,
875 FrameType::Error,
876 Flags::new(false, Priority::Interactive, false),
877 0,
878 0,
879 7,
880 serde_json::to_vec(&ErrorBody {
881 code: "test".into(),
882 message: "reply".into(),
883 detail: None,
884 })
885 .expect("body encodes"),
886 )
887 .expect("frame builds");
888 tx.send(crate::router::OutboundFrame {
889 frame: reply,
890 enqueued_at: std::time::Instant::now() - Duration::from_millis(1500),
891 flushed: None,
892 charge: None,
893 })
894 .await
895 .expect("queued");
896 drop(tx);
897 let (write_half, mut read_half) = duplex(64 * 1024);
898 drain_writer(write_half, rx).await.expect("writer drains");
899 let mut sink = Vec::new();
900 read_half.read_to_end(&mut sink).await.expect("read");
901 assert!(!sink.is_empty(), "frame reached the socket");
902 let captured = crate::router::test_log::captured_logs(&logs);
903 assert!(
904 captured.contains("slow control reply write") && captured.contains("corr=7"),
905 "expected slow reply WARN naming corr, got: {captured}"
906 );
907 let queued_ms: u64 = captured
908 .split("queued_ms=")
909 .nth(1)
910 .and_then(|s| s.split_whitespace().next())
911 .and_then(|s| s.parse().ok())
912 .expect("queued_ms present");
913 assert!(
914 queued_ms >= 1500,
915 "queued_ms reflects residency: {queued_ms}"
916 );
917 }
918
919 #[tokio::test]
923 async fn fresh_control_and_stale_data_frames_log_nothing() {
924 let (logs, _guard) = crate::router::test_log::log_capture(tracing::Level::WARN);
925 let (tx, rx) = mpsc::channel::<crate::router::OutboundFrame>(4);
926 let control = Frame::build_with_version(
927 PROTOCOL_VERSION,
928 FrameType::Error,
929 Flags::new(false, Priority::Interactive, false),
930 0,
931 0,
932 8,
933 serde_json::to_vec(&ErrorBody {
934 code: "test".into(),
935 message: "fresh".into(),
936 detail: None,
937 })
938 .expect("body encodes"),
939 )
940 .expect("frame builds");
941 tx.send(crate::router::OutboundFrame {
942 frame: control,
943 enqueued_at: std::time::Instant::now(),
944 flushed: None,
945 charge: None,
946 })
947 .await
948 .expect("queued");
949 let data = Frame::build_with_version(
950 PROTOCOL_VERSION,
951 FrameType::Error,
952 Flags::new(false, Priority::Interactive, false),
953 9,
954 1,
955 9,
956 serde_json::to_vec(&ErrorBody {
957 code: "test".into(),
958 message: "data".into(),
959 detail: None,
960 })
961 .expect("body encodes"),
962 )
963 .expect("frame builds");
964 tx.send(crate::router::OutboundFrame {
965 frame: data,
966 enqueued_at: std::time::Instant::now() - Duration::from_millis(5000),
967 flushed: None,
968 charge: None,
969 })
970 .await
971 .expect("queued");
972 drop(tx);
973 let (write_half, _read_half) = duplex(64 * 1024);
974 drain_writer(write_half, rx).await.expect("writer drains");
975 let captured = crate::router::test_log::captured_logs(&logs);
976 assert!(
977 !captured.contains("slow control reply write"),
978 "no WARN for fresh control or stale data frames, got: {captured}"
979 );
980 }
981
982 struct CountingReader {
983 bytes: Vec<u8>,
984 offset: usize,
985 first_read_end: Option<usize>,
986 reads: Arc<AtomicUsize>,
987 }
988
989 impl CountingReader {
990 fn new(bytes: Vec<u8>, first_read_end: Option<usize>) -> (Self, Arc<AtomicUsize>) {
991 let reads = Arc::new(AtomicUsize::new(0));
992 (
993 Self {
994 bytes,
995 offset: 0,
996 first_read_end,
997 reads: Arc::clone(&reads),
998 },
999 reads,
1000 )
1001 }
1002 }
1003
1004 impl AsyncRead for CountingReader {
1005 fn poll_read(
1006 mut self: Pin<&mut Self>,
1007 _cx: &mut Context<'_>,
1008 buf: &mut ReadBuf<'_>,
1009 ) -> Poll<io::Result<()>> {
1010 self.reads.fetch_add(1, Ordering::Relaxed);
1011 let available = self.bytes.len().saturating_sub(self.offset);
1012 let first_read_remaining = self
1013 .first_read_end
1014 .filter(|end| self.offset < *end)
1015 .map_or(available, |end| end - self.offset);
1016 let count = available.min(first_read_remaining).min(buf.remaining());
1017 let end = self.offset + count;
1018 buf.put_slice(&self.bytes[self.offset..end]);
1019 self.offset = end;
1020 Poll::Ready(Ok(()))
1021 }
1022 }
1023
1024 fn encode_frames(frames: &[Frame]) -> Vec<u8> {
1025 let mut bytes = Vec::new();
1026 for frame in frames {
1027 bytes.extend_from_slice(&frame.header.encode());
1028 bytes.extend_from_slice(&frame.body);
1029 }
1030 bytes
1031 }
1032
1033 async fn read_frames<R>(reader: &mut R, count: usize) -> Vec<Frame>
1034 where
1035 R: AsyncRead + Unpin,
1036 {
1037 let mut frames = Vec::with_capacity(count);
1038 for _ in 0..count {
1039 frames.push(read_frame(reader).await.unwrap().unwrap());
1040 }
1041 frames
1042 }
1043
1044 fn request(channel: u16, corr: u64, body: &[u8]) -> Frame {
1045 Frame::build(
1046 FrameType::Request,
1047 Flags::new(true, Priority::Interactive, false),
1048 channel,
1049 0,
1050 corr,
1051 body.to_vec(),
1052 )
1053 .unwrap()
1054 }
1055
1056 fn echo_router() -> Arc<Router> {
1057 let mut router = Router::with_default_self_handler();
1058 router.register_backend(7, EchoBackend).unwrap();
1059 router.register_backend(9, EchoBackend).unwrap();
1060 Arc::new(router)
1061 }
1062
1063 fn test_auth() -> (ServerAuth, ConnectionInfo) {
1064 test_auth_with_limit(4)
1065 }
1066
1067 fn test_auth_with_limit(max_unauthenticated: usize) -> (ServerAuth, ConnectionInfo) {
1068 let key = vec![0x42; 32];
1069 let daemon_id = [0x24; 16];
1070 let conn = ConnectionInfo {
1071 schema: SCHEMA_VERSION,
1072 wire_version: None,
1073 endpoints: vec![Endpoint {
1074 host: "127.0.0.1".to_owned(),
1075 port: 1,
1076 }],
1077 key: key.clone(),
1078 daemon_id,
1079 pid: std::process::id(),
1080 daemon_ver: TEST_DAEMON_VER.to_owned(),
1081 };
1082 (
1083 ServerAuth::with_limits(
1084 key,
1085 daemon_id,
1086 TEST_DAEMON_VER,
1087 TEST_DEADLINE,
1088 max_unauthenticated,
1089 ),
1090 conn,
1091 )
1092 }
1093
1094 async fn authenticate<S>(stream: &mut S, conn: &ConnectionInfo)
1095 where
1096 S: AsyncRead + AsyncWrite + Unpin,
1097 {
1098 authenticate_client(stream, conn, TEST_DEADLINE)
1099 .await
1100 .expect("test client should authenticate")
1101 }
1102
1103 #[tokio::test]
1104 async fn buffered_frame_reader_coalesces_reads_and_preserves_short_reads() {
1105 let frames = vec![
1106 request(7, 1, b"first"),
1107 request(9, 2, b"second"),
1108 request(7, 3, b"third"),
1109 request(9, 4, b"fourth"),
1110 ];
1111 let bytes = encode_frames(&frames);
1112
1113 let (mut direct, direct_reads) = CountingReader::new(bytes.clone(), None);
1114 assert_eq!(read_frames(&mut direct, frames.len()).await, frames);
1115 assert_eq!(direct_reads.load(Ordering::Relaxed), frames.len() * 3);
1116
1117 let (buffered_source, buffered_reads) = CountingReader::new(bytes, None);
1118 let mut buffered = BufReader::new(buffered_source);
1119 assert_eq!(read_frames(&mut buffered, frames.len()).await, frames);
1120 assert_eq!(buffered_reads.load(Ordering::Relaxed), 1);
1121
1122 let split_frame = request(7, 5, b"split-body");
1123 let split_bytes = encode_frames(std::slice::from_ref(&split_frame));
1124 let (split_source, split_reads) = CountingReader::new(split_bytes, Some(10));
1125 let mut split_reader = BufReader::new(split_source);
1126 assert_eq!(
1127 read_frame(&mut split_reader).await.unwrap(),
1128 Some(split_frame)
1129 );
1130 assert_eq!(split_reads.load(Ordering::Relaxed), 2);
1131 }
1132
1133 #[tokio::test]
1134 async fn interleaved_channels_on_one_stream_demux_byte_identically_after_auth() {
1135 let (mut client, server_stream) = duplex(4096);
1136 let (auth, conn) = test_auth();
1137 let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
1138 authenticate(&mut client, &conn).await;
1139 let frames = [
1140 request(7, 1, b"chan7-first\0opaque"),
1141 request(9, 2, b"chan9-middle-{json?}"),
1142 request(7, 3, b"chan7-second\xffbytes"),
1143 ];
1144
1145 for frame in &frames {
1146 crate::write_frame(&mut client, frame).await.unwrap();
1147 }
1148
1149 for expected in &frames {
1150 let response = crate::read_frame(&mut client).await.unwrap().unwrap();
1151 assert_eq!(response.header.ty, FrameType::Response);
1152 assert_eq!(response.header.channel, expected.header.channel);
1153 assert_eq!(response.header.corr, expected.header.corr);
1154 assert_eq!(response.body, expected.body);
1155 }
1156
1157 drop(client);
1158 server.await.unwrap().unwrap();
1159 }
1160
1161 #[tokio::test]
1162 async fn channel_zero_goes_to_subc_self_handler_after_auth() {
1163 let (mut client, server_stream) = duplex(512);
1164 let (auth, conn) = test_auth();
1165 let server = tokio::spawn(handle_connection(
1166 server_stream,
1167 Arc::new(Router::with_default_self_handler()),
1168 auth,
1169 ));
1170 authenticate(&mut client, &conn).await;
1171 let ping = Frame::build(
1172 FrameType::Ping,
1173 Flags::new(false, Priority::Passive, false),
1174 0,
1175 0,
1176 55,
1177 Vec::new(),
1178 )
1179 .unwrap();
1180
1181 crate::write_frame(&mut client, &ping).await.unwrap();
1182 let response = crate::read_frame(&mut client).await.unwrap().unwrap();
1183
1184 assert_eq!(response.header.ty, FrameType::Pong);
1185 assert_eq!(response.header.channel, 0);
1186 assert_eq!(response.header.corr, 55);
1187 assert!(response.body.is_empty());
1188
1189 drop(client);
1190 server.await.unwrap().unwrap();
1191 }
1192
1193 #[tokio::test]
1194 async fn unauthenticated_connection_is_rejected_before_routing() {
1195 let (mut client, server_stream) = duplex(512);
1196 let (auth, _conn) = test_auth();
1197 let registry = Arc::new(Registry::default());
1198 let router = Arc::new(Router::with_control_handler(Arc::new(ControlHandler::new(
1199 Arc::clone(®istry),
1200 ))));
1201 let server = tokio::spawn(handle_connection(server_stream, router, auth));
1202 let ping = Frame::build(
1203 FrameType::Ping,
1204 Flags::new(false, Priority::Passive, false),
1205 0,
1206 0,
1207 66,
1208 Vec::new(),
1209 )
1210 .unwrap();
1211
1212 crate::write_frame(&mut client, &ping).await.unwrap();
1213 if let Ok(Ok(Some(frame))) =
1214 tokio::time::timeout(Duration::from_millis(200), crate::read_frame(&mut client)).await
1215 {
1216 panic!("unauthenticated frame reached router: {frame:?}");
1217 }
1218
1219 let err = server.await.unwrap().unwrap_err();
1220 assert!(matches!(err, ConnectionError::Auth(_)));
1221 assert_eq!(registry.active_registration_count().unwrap(), 0);
1222 }
1223
1224 #[tokio::test]
1225 async fn over_cap_peer_queues_for_a_slot_and_authenticates_when_one_frees() {
1226 let (mut first_client, first_server_stream) = duplex(2048);
1231 let (mut second_client, second_server_stream) = duplex(2048);
1232 let (auth, conn) = test_auth_with_limit(1);
1233 let registry = Arc::new(Registry::default());
1234 let router = Arc::new(Router::with_control_handler(Arc::new(ControlHandler::new(
1235 Arc::clone(®istry),
1236 ))));
1237
1238 let first_server = tokio::spawn(handle_connection(
1239 first_server_stream,
1240 Arc::clone(&router),
1241 auth.clone(),
1242 ));
1243
1244 let second_server = tokio::spawn(handle_connection(
1245 second_server_stream,
1246 Arc::clone(&router),
1247 auth.clone(),
1248 ));
1249
1250 tokio::time::sleep(Duration::from_millis(50)).await;
1252 assert!(
1253 !second_server.is_finished(),
1254 "queued peer must not be reset"
1255 );
1256
1257 authenticate(&mut first_client, &conn).await;
1260 authenticate(&mut second_client, &conn).await;
1261
1262 drop(first_client);
1263 drop(second_client);
1264 let _ = first_server.await;
1265 let _ = second_server.await;
1266 assert_eq!(registry.active_registration_count().unwrap(), 0);
1267 }
1268
1269 #[tokio::test]
1270 async fn over_cap_peer_is_rejected_when_no_slot_frees_within_deadline() {
1271 let (mut second_client, second_server_stream) = duplex(512);
1277 let (auth, conn) = test_auth_with_limit(1);
1278 let registry = Arc::new(Registry::default());
1279 let router = Arc::new(Router::with_control_handler(Arc::new(ControlHandler::new(
1280 Arc::clone(®istry),
1281 ))));
1282
1283 let held_slot = auth
1284 .unauthenticated
1285 .clone()
1286 .try_acquire_owned()
1287 .expect("sole pre-auth slot");
1288
1289 let second_server = tokio::spawn(handle_connection(
1290 second_server_stream,
1291 Arc::clone(&router),
1292 auth.clone(),
1293 ));
1294 let second_err = tokio::time::timeout(TEST_DEADLINE * 2, second_server)
1295 .await
1296 .expect("capacity reject should settle at the deadline")
1297 .expect("second connection task should not panic")
1298 .expect_err("queued peer must be rejected when no slot frees");
1299 drop(held_slot);
1300 assert!(matches!(
1301 second_err,
1302 ConnectionError::UnauthenticatedCapacity
1303 ));
1304 let mut closed = [0u8; 1];
1305 assert_eq!(
1306 second_client.read(&mut closed).await.unwrap(),
1307 0,
1308 "capacity-rejected peer should observe a closed stream"
1309 );
1310 assert_eq!(registry.active_registration_count().unwrap(), 0);
1311
1312 let (mut authed_client, authed_server_stream) = duplex(2048);
1313 let authed_server = tokio::spawn(handle_connection(
1314 authed_server_stream,
1315 Arc::clone(&router),
1316 auth,
1317 ));
1318 authenticate(&mut authed_client, &conn).await;
1319 let ping = Frame::build(
1320 FrameType::Ping,
1321 Flags::new(false, Priority::Passive, false),
1322 0,
1323 0,
1324 77,
1325 Vec::new(),
1326 )
1327 .unwrap();
1328 crate::write_frame(&mut authed_client, &ping).await.unwrap();
1329 let pong = crate::read_frame(&mut authed_client)
1330 .await
1331 .unwrap()
1332 .unwrap();
1333 assert_eq!(pong.header.ty, FrameType::Pong);
1334 assert_eq!(pong.header.channel, 0);
1335 assert_eq!(pong.header.corr, 77);
1336 assert_eq!(registry.active_registration_count().unwrap(), 0);
1337
1338 drop(authed_client);
1339 authed_server.await.unwrap().unwrap();
1340 }
1341
1342 #[tokio::test]
1343 async fn bind_ack_during_partial_frame_preserves_next_request() {
1344 use subc_control::ClientControlRequest;
1345 use subc_protocol::{
1346 manifest::{
1347 Concurrency, ExecutionMode, IdentityScope, ModuleManifest, ProviderRole, Tool,
1348 },
1349 session::{ModuleControlRequest, ModuleControlResponse},
1350 BindIdentity, ModuleHelloBody, RouteTarget,
1351 };
1352
1353 let mut configured_router = Router::with_default_self_handler();
1354 configured_router.register_backend(7, EchoBackend).unwrap();
1355 let router = Arc::new(configured_router);
1356 let (auth, conn) = test_auth();
1357 let (mut module, module_stream) = duplex(4096);
1358 let module_server = tokio::spawn(handle_connection(
1359 module_stream,
1360 Arc::clone(&router),
1361 auth.clone(),
1362 ));
1363 authenticate(&mut module, &conn).await;
1364 let manifest = ModuleManifest::builder("frame-test", "0.1.0")
1365 .protocol_ver(PROTOCOL_VERSION)
1366 .provides(vec![ProviderRole::ToolProvider {
1367 tools: vec![Tool {
1368 name: "read".into(),
1369 description: None,
1370 execution_mode: ExecutionMode::Pure,
1371 schema: serde_json::json!({"type": "object"}),
1372 }],
1373 identity_scope: vec![IdentityScope::Project, IdentityScope::Session],
1374 concurrency: Concurrency::ModuleManaged,
1375 emits_push: true,
1376 sub_supervises: true,
1377 }])
1378 .build();
1379 let hello = Frame::build(
1380 FrameType::Hello,
1381 Flags::new(false, Priority::Passive, false),
1382 0,
1383 0,
1384 1,
1385 serde_json::to_vec(&ModuleHelloBody {
1386 manifest,
1387 protocol_ver: PROTOCOL_VERSION,
1388 control_ops: None,
1389 launch_nonce: None,
1390 })
1391 .unwrap(),
1392 )
1393 .unwrap();
1394 crate::write_frame(&mut module, &hello).await.unwrap();
1395 assert_eq!(
1396 read_frame(&mut module).await.unwrap().unwrap().header.ty,
1397 FrameType::HelloAck
1398 );
1399
1400 let (mut client, client_stream) = duplex(4096);
1401 let client_server = tokio::spawn(handle_connection(client_stream, router, auth));
1402 authenticate(&mut client, &conn).await;
1403 let open = Frame::build(
1404 FrameType::Request,
1405 Flags::new(false, Priority::Passive, false),
1406 0,
1407 0,
1408 2,
1409 serde_json::to_vec(&ClientControlRequest::RouteOpen {
1410 target: RouteTarget::ToolProvider {
1411 module_id: "frame-test".into(),
1412 },
1413 identity: BindIdentity::new(std::env::current_dir().unwrap(), "unit", "session"),
1414 consumer_identity: None,
1415 consumer_capabilities: None,
1416 admission_facts: None,
1417 })
1418 .unwrap(),
1419 )
1420 .unwrap();
1421 crate::write_frame(&mut client, &open).await.unwrap();
1422 let bind = timeout(TEST_DEADLINE, read_frame(&mut module))
1423 .await
1424 .unwrap()
1425 .unwrap()
1426 .unwrap();
1427 assert!(matches!(
1428 serde_json::from_slice::<ModuleControlRequest>(&bind.body).unwrap(),
1429 ModuleControlRequest::RouteBind { .. }
1430 ));
1431
1432 let ping = request(7, 3, b"partial-frame-body");
1433 client.write_all(&ping.header.encode()).await.unwrap();
1434 tokio::time::sleep(Duration::from_millis(30)).await;
1436 let ack = Frame::build(
1437 FrameType::Response,
1438 Flags::new(false, Priority::Passive, false),
1439 0,
1440 0,
1441 bind.header.corr,
1442 serde_json::to_vec(&ModuleControlResponse::RouteBindAck {}).unwrap(),
1443 )
1444 .unwrap();
1445 crate::write_frame(&mut module, &ack).await.unwrap();
1446 let opened = timeout(TEST_DEADLINE, read_frame(&mut client))
1447 .await
1448 .unwrap()
1449 .unwrap();
1450 if opened.is_none() {
1451 panic!("client closed: {:?}", client_server.await);
1452 }
1453 let opened = opened.unwrap();
1454 assert_eq!(opened.header.corr, 2);
1455 client.write_all(&ping.body).await.unwrap();
1456 let pong = timeout(TEST_DEADLINE, read_frame(&mut client))
1457 .await
1458 .expect("partial frame must reach the router after the bind ack")
1459 .expect("frame must decode")
1460 .expect("connection must remain open");
1461 assert_eq!(pong.header.ty, FrameType::Response);
1462 assert_eq!(pong.header.corr, 3);
1463 assert_eq!(pong.body, ping.body);
1464 drop(client);
1465 drop(module);
1466 client_server.await.unwrap().unwrap();
1467 let _ = module_server.await.unwrap();
1468 }
1469
1470 #[tokio::test]
1471 async fn aborted_accept_does_not_end_listener() {
1472 let listener = Arc::new(TcpListener::bind("127.0.0.1:0").await.unwrap());
1473 let addr = listener.local_addr().unwrap();
1474 let (auth, conn) = test_auth();
1475 let mut attempts = 0;
1476 let server = tokio::spawn(serve_listener_with_accept(
1477 Some(addr),
1478 echo_router(),
1479 auth,
1480 move || {
1481 attempts += 1;
1482 let result = if attempts == 1 {
1483 Some(io::Error::from(io::ErrorKind::ConnectionAborted))
1484 } else {
1485 None
1486 };
1487 let listener = Arc::clone(&listener);
1488 async move {
1489 match result {
1490 Some(err) => Err(err),
1491 None => listener.accept().await,
1492 }
1493 }
1494 },
1495 ));
1496 let mut client = tokio::net::TcpStream::connect(addr).await.unwrap();
1497 authenticate(&mut client, &conn).await;
1498 let ping = Frame::build(
1499 FrameType::Ping,
1500 Flags::new(false, Priority::Passive, false),
1501 0,
1502 0,
1503 77,
1504 Vec::new(),
1505 )
1506 .unwrap();
1507 crate::write_frame(&mut client, &ping).await.unwrap();
1508 let pong = timeout(TEST_DEADLINE, read_frame(&mut client))
1509 .await
1510 .unwrap()
1511 .unwrap()
1512 .unwrap();
1513 assert_eq!(pong.header.ty, FrameType::Pong);
1514 assert!(
1515 !server.is_finished(),
1516 "a temporary accept failure must not stop the listener"
1517 );
1518 server.abort();
1519 }
1520
1521 #[tokio::test]
1522 async fn exhausted_accept_backs_off_then_serves_connection() {
1523 let listener = Arc::new(TcpListener::bind("127.0.0.1:0").await.unwrap());
1524 let addr = listener.local_addr().unwrap();
1525 let (auth, conn) = test_auth();
1526 let mut attempts = 0;
1527 let server = tokio::spawn(serve_listener_with_accept(
1528 Some(addr),
1529 echo_router(),
1530 auth,
1531 move || {
1532 attempts += 1;
1533 #[cfg(unix)]
1536 let emfile = rustix::io::Errno::MFILE.raw_os_error();
1537 #[cfg(windows)]
1538 let emfile = 10024;
1539 let error = (attempts == 1).then(|| io::Error::from_raw_os_error(emfile));
1540 let listener = Arc::clone(&listener);
1541 async move {
1542 match error {
1543 Some(err) => Err(err),
1544 None => listener.accept().await,
1545 }
1546 }
1547 },
1548 ));
1549 let started = Instant::now();
1550 let mut client = tokio::net::TcpStream::connect(addr).await.unwrap();
1551 authenticate(&mut client, &conn).await;
1552 assert!(
1553 started.elapsed() >= Duration::from_millis(50),
1554 "fd exhaustion must back off before retrying"
1555 );
1556 assert!(!server.is_finished());
1557 server.abort();
1558 }
1559
1560 #[tokio::test]
1561 async fn fatal_accept_error_still_stops_listener() {
1562 let (auth, _) = test_auth();
1563 let err = serve_listener_with_accept(None, echo_router(), auth, || async {
1564 Err(io::Error::from(io::ErrorKind::PermissionDenied))
1565 })
1566 .await
1567 .unwrap_err();
1568 assert!(
1569 matches!(err, ServerError::Accept { source, .. } if source.kind() == io::ErrorKind::PermissionDenied)
1570 );
1571 }
1572
1573 #[tokio::test]
1574 async fn serve_listeners_with_no_listeners_returns_typed_error() {
1575 let (auth, _conn) = test_auth();
1576 let err = serve_listeners(Vec::new(), echo_router(), auth)
1577 .await
1578 .expect_err("empty listener set must fail loudly");
1579 assert!(matches!(err, ServerError::NoListeners));
1580 }
1581
1582 #[tokio::test]
1583 async fn malformed_header_returns_typed_error_no_panic() {
1584 let (mut client, server_stream) = duplex(128);
1585 let (auth, conn) = test_auth();
1586 let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
1587 authenticate(&mut client, &conn).await;
1588 let mut header = [0u8; HEADER_LEN];
1589 header[4] = PROTOCOL_VERSION;
1590 header[5] = 250;
1591
1592 client.write_all(&header).await.unwrap();
1593 drop(client);
1594
1595 let err = server.await.unwrap().unwrap_err();
1596 assert!(matches!(
1597 err,
1598 ConnectionError::FrameIo(FrameIoError::DecodeHeader(DecodeError::UnknownFrameType {
1599 byte: 250
1600 }))
1601 ));
1602 }
1603
1604 #[tokio::test]
1605 async fn truncated_body_returns_typed_error_no_panic() {
1606 let (mut client, server_stream) = duplex(128);
1607 let (auth, conn) = test_auth();
1608 let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
1609 authenticate(&mut client, &conn).await;
1610 let frame = request(7, 8, b"abcd");
1611
1612 client.write_all(&frame.header.encode()).await.unwrap();
1613 client.write_all(b"ab").await.unwrap();
1614 drop(client);
1615
1616 let err = server.await.unwrap().unwrap_err();
1617 assert!(matches!(
1618 err,
1619 ConnectionError::FrameIo(FrameIoError::UnexpectedEof {
1620 stage: ReadStage::Body,
1621 expected: 4,
1622 actual: 2
1623 })
1624 ));
1625 }
1626
1627 #[tokio::test]
1628 async fn unknown_channel_is_returned_as_error_frame_and_connection_continues() {
1629 let (mut client, server_stream) = duplex(1024);
1630 let (auth, conn) = test_auth();
1631 let server = tokio::spawn(handle_connection(server_stream, echo_router(), auth));
1632 authenticate(&mut client, &conn).await;
1633 let unknown = request(42, 10, b"lost");
1634 let known = request(7, 11, b"still-routes");
1635
1636 crate::write_frame(&mut client, &unknown).await.unwrap();
1637 crate::write_frame(&mut client, &known).await.unwrap();
1638
1639 let error = crate::read_frame(&mut client).await.unwrap().unwrap();
1640 assert_eq!(error.header.ty, FrameType::Error);
1641 assert_eq!(error.header.channel, 42);
1642 assert_eq!(error.header.corr, 10);
1643 let error_body: ErrorBody = serde_json::from_slice(&error.body).unwrap();
1644 assert_eq!(error_body.code, "unknown_channel");
1645
1646 let response = crate::read_frame(&mut client).await.unwrap().unwrap();
1647 assert_eq!(response.header.ty, FrameType::Response);
1648 assert_eq!(response.header.channel, 7);
1649 assert_eq!(response.header.corr, 11);
1650 assert_eq!(response.body, b"still-routes");
1651
1652 drop(client);
1653 server.await.unwrap().unwrap();
1654 }
1655}