1use std::future::Future;
2use std::net::SocketAddr;
3use std::sync::Arc;
4use std::time::Duration;
5
6use hyper::body::Bytes;
7use hyper::body::Incoming;
8use hyper::server::conn::http1;
9use hyper::{Method, Request, Response, StatusCode};
10use hyper::service::service_fn;
11use hyper_util::rt::TokioIo;
12use hyper_util::rt::TokioTimer;
13use http_body_util::combinators::BoxBody;
14use http_body_util::{BodyExt, Empty, Full};
15use tokio::net::{TcpListener, TcpStream};
16#[cfg(feature = "tls")]
17use tokio_rustls::TlsAcceptor;
18
19use crate::body::{MaxBodySize, DEFAULT_MAX_BODY_SIZE};
20use crate::cors::CorsConfig;
21use crate::error::ServeError;
22use crate::handler::{Handler, ResponseBody};
23use crate::router::{QueryParams, Router};
24use crate::state::State;
25
26const MAX_PATH_LEN: usize = 8_192;
27const MAX_QUERY_LEN: usize = 4_096;
28const DEFAULT_HEADER_READ_TIMEOUT: Duration = Duration::from_secs(30);
29const DEFAULT_MAX_CONNECTIONS: usize = 1024;
30#[cfg(feature = "tls")]
31const DEFAULT_TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
32const ACCEPT_BACKOFF_INITIAL: Duration = Duration::from_millis(10);
33const ACCEPT_BACKOFF_MAX: Duration = Duration::from_secs(1);
34
35#[cfg(test)]
36thread_local! {
37 static ERROR_LOG: std::cell::RefCell<Vec<(u16, String)>> = std::cell::RefCell::new(Vec::new());
38}
39
40#[cfg(test)]
41fn capture_error(code: u16, message: String) {
42 ERROR_LOG.with(|log| log.borrow_mut().push((code, message)));
43}
44
45#[cfg(test)]
46fn take_error_log() -> Vec<(u16, String)> {
47 ERROR_LOG.with(|log| log.borrow_mut().drain(..).collect())
48}
49
50trait TcpAccept {
55 async fn accept(&self) -> std::io::Result<(TcpStream, SocketAddr)>;
56}
57
58impl TcpAccept for TcpListener {
59 async fn accept(&self) -> std::io::Result<(TcpStream, SocketAddr)> {
60 TcpListener::accept(self).await
61 }
62}
63
64struct Backoff {
69 delay: Duration,
70}
71
72impl Backoff {
73 fn new() -> Self {
74 Backoff { delay: ACCEPT_BACKOFF_INITIAL }
75 }
76
77 fn next_delay(&mut self) -> Duration {
78 let delay = self.delay;
79 self.delay = (self.delay * 2).min(ACCEPT_BACKOFF_MAX);
80 delay
81 }
82
83 fn reset(&mut self) {
84 self.delay = ACCEPT_BACKOFF_INITIAL;
85 }
86}
87
88async fn accept_with_backoff<L: TcpAccept>(
89 listener: &L,
90 backoff: &mut Backoff,
91) -> (TcpStream, SocketAddr) {
92 loop {
93 match listener.accept().await {
94 Ok(conn) => {
95 backoff.reset();
96 return conn;
97 }
98 Err(_) => {
99 tokio::time::sleep(backoff.next_delay()).await;
100 }
101 }
102 }
103}
104
105pub type ErrorHandler =
106 Arc<dyn Fn(StatusCode, &str) -> Response<ResponseBody> + Send + Sync>;
107
108pub struct App<S> {
136 state: Arc<S>,
137 router: Arc<Router<S>>,
138 max_body_size: usize,
139 pub(crate) header_read_timeout: Duration,
140 pub(crate) max_connections: usize,
141 #[cfg(feature = "tls")]
142 pub(crate) tls_handshake_timeout: Duration,
143 error_handler: ErrorHandler,
144 cors_config: Option<CorsConfig>,
145}
146
147fn parse_query(query: Option<&str>) -> QueryParams {
148 let mut map = std::collections::HashMap::new();
149 if let Some(query) = query {
150 for pair in query.split('&').filter(|s| !s.is_empty()) {
151 if let Some((key, value)) = pair.split_once('=') {
152 let key = decode_query_component(key);
153 let value = decode_query_component(value);
154 map.insert(key, value);
155 } else {
156 let pair = decode_query_component(pair);
157 map.insert(pair, String::new());
158 }
159 }
160 }
161 QueryParams(map)
162}
163
164fn decode_query_component(s: &str) -> String {
165 let with_spaces = s.replace('+', " ");
166 percent_encoding::percent_decode_str(&with_spaces)
167 .decode_utf8_lossy()
168 .into_owned()
169}
170
171fn error_response(status: StatusCode, message: &str) -> Response<ResponseBody> {
172 #[cfg(test)]
173 capture_error(status.as_u16(), message.to_string());
174
175 let client_message = if status.is_server_error() {
176 "internal server error"
177 } else {
178 message
179 };
180
181 let body = serde_json::json!({ "message": client_message });
182 let json = serde_json::to_string(&body)
183 .unwrap_or_else(|_| r#"{"message":"internal server error"}"#.to_string());
184 Response::builder()
185 .status(status)
186 .header("content-type", "application/json")
187 .body(BoxBody::new(Full::new(Bytes::from(json)).map_err(|never: std::convert::Infallible| match never {})))
188 .expect("status is valid and headers are static ASCII")
189}
190
191fn default_error_handler() -> ErrorHandler {
192 Arc::new(error_response)
193}
194
195fn ephemeral_bind_addr() -> SocketAddr {
199 (std::net::Ipv4Addr::LOCALHOST, 0).into()
200}
201
202impl<S: Send + Sync + 'static> App<S> {
203 pub fn new(state: S) -> Self {
208 App {
209 state: Arc::new(state),
210 router: Arc::new(Router::new()),
211 max_body_size: DEFAULT_MAX_BODY_SIZE,
212 header_read_timeout: DEFAULT_HEADER_READ_TIMEOUT,
213 max_connections: DEFAULT_MAX_CONNECTIONS,
214 #[cfg(feature = "tls")]
215 tls_handshake_timeout: DEFAULT_TLS_HANDSHAKE_TIMEOUT,
216 error_handler: default_error_handler(),
217 cors_config: None,
218 }
219 }
220
221 pub fn state_arc(&self) -> Arc<S> {
226 Arc::clone(&self.state)
227 }
228
229 pub async fn route(&self, req: Request<Incoming>) -> Response<ResponseBody> {
230 let method = req.method().clone();
231 let resp = self.route_inner(req).await;
232
233 if method == Method::HEAD {
240 let (parts, _) = resp.into_parts();
241 Response::from_parts(parts, BoxBody::new(Empty::new().map_err(|never: std::convert::Infallible| match never {})))
242 } else {
243 resp
244 }
245 }
246
247 async fn route_inner(&self, req: Request<Incoming>) -> Response<ResponseBody> {
248 if req.uri().path().len() > MAX_PATH_LEN {
249 return (self.error_handler)(StatusCode::BAD_REQUEST, "path too long");
250 }
251 if req.uri().query().map(|q| q.len()).unwrap_or(0) > MAX_QUERY_LEN {
252 return (self.error_handler)(StatusCode::BAD_REQUEST, "query string too long");
253 }
254
255 let method = req.method().clone();
256 let path = req.uri().path().to_string();
257 let state = State::from_arc(Arc::clone(&self.state));
258 let query_params = parse_query(req.uri().query());
259
260 let req_origin = req
262 .headers()
263 .get("origin")
264 .and_then(|v| v.to_str().ok())
265 .map(|s| s.to_string());
266
267 if method == Method::OPTIONS && req_origin.is_some() {
269 if let Some(cfg) = &self.cors_config {
270 if self.router.path_exists(&path) {
271 let requested_headers = req
272 .headers()
273 .get("access-control-request-headers")
274 .and_then(|v| v.to_str().ok());
275 let allowed = self.allowed_methods_with_head(&path);
276 return cfg.preflight_response(req_origin.as_deref(), requested_headers, &allowed);
277 }
278 }
279 }
280
281 let method_to_match = if method == Method::HEAD {
282 Method::GET
283 } else {
284 method.clone()
285 };
286
287 match self.router.match_route(&method_to_match, &path) {
288 Some((handler, params)) => {
289 let mut req = req;
290 req.extensions_mut().insert(query_params);
291 req.extensions_mut().insert(params);
292 req.extensions_mut().insert(MaxBodySize(self.max_body_size));
293 match handler(req, state).await {
294 Ok(mut resp) => {
295 if let Some(cfg) = &self.cors_config {
296 cfg.apply_to_response(&mut resp, req_origin.as_deref());
297 }
298 resp
299 }
300 Err(e) => (self.error_handler)(
301 StatusCode::from_u16(e.code)
302 .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR),
303 &e.message,
304 ),
305 }
306 }
307 None => {
308 let allowed = self.allowed_methods_with_head(&path);
309 if !allowed.is_empty() {
310 let mut method_strs: Vec<&str> = allowed.iter().map(|m| m.as_str()).collect();
311 method_strs.sort();
312 method_strs.dedup();
313 let allow_header = method_strs.join(", ");
314 let mut resp = (self.error_handler)(StatusCode::METHOD_NOT_ALLOWED, "method not allowed");
315 if let Ok(val) = allow_header.parse() {
316 resp.headers_mut().insert("allow", val);
317 }
318 resp
319 } else {
320 (self.error_handler)(StatusCode::NOT_FOUND, "not found")
321 }
322 }
323 }
324 }
325
326 fn allowed_methods_with_head(&self, path: &str) -> Vec<Method> {
333 let mut allowed = self.router.allowed_methods(path);
334 if allowed.contains(&Method::GET) {
335 allowed.push(Method::HEAD);
336 }
337 allowed
338 }
339
340 pub async fn bind_ephemeral(self) -> Result<u16, ServeError> {
346 let listener = TcpListener::bind(ephemeral_bind_addr())
347 .await
348 .map_err(|e| ServeError::new(500, format!("failed to bind to ephemeral port: {e}")))?;
349 let port = listener
350 .local_addr()
351 .map_err(|e| ServeError::new(500, format!("failed to get assigned port: {e}")))?
352 .port();
353 let app = Arc::new(self);
354 tokio::spawn(async move {
355 serve_inner(listener, app).await;
356 });
357 Ok(port)
358 }
359
360 #[cfg(feature = "tls")]
366 pub async fn bind_tls_ephemeral(
367 self,
368 config: Arc<rustls::ServerConfig>,
369 ) -> Result<u16, ServeError> {
370 let listener = TcpListener::bind(ephemeral_bind_addr())
371 .await
372 .map_err(|e| ServeError::new(500, format!("failed to bind to ephemeral port: {e}")))?;
373 let port = listener
374 .local_addr()
375 .map_err(|e| ServeError::new(500, format!("failed to get assigned port: {e}")))?
376 .port();
377 let acceptor = TlsAcceptor::from(config);
378 let app = Arc::new(self);
379 tokio::spawn(async move {
380 serve_tls_inner(listener, app, acceptor).await;
381 });
382 Ok(port)
383 }
384
385 pub async fn run<F>(self, listener: TcpListener, shutdown: F) -> Result<(), ServeError>
390 where
391 F: Future<Output = ()> + Send + 'static,
392 {
393 let app = Arc::new(self);
394 serve_with_shutdown(listener, app, shutdown).await;
395 Ok(())
396 }
397
398 pub async fn bind(self, addr: SocketAddr) -> Result<(), ServeError> {
403 let listener = TcpListener::bind(addr)
404 .await
405 .map_err(|e| ServeError::new(500, format!("failed to bind to {addr}: {e}")))?;
406 self.run(listener, signal_shutdown()).await
407 }
408
409 #[cfg(feature = "tls")]
411 pub async fn run_tls<F>(
412 self,
413 listener: TcpListener,
414 config: Arc<rustls::ServerConfig>,
415 shutdown: F,
416 ) -> Result<(), ServeError>
417 where
418 F: Future<Output = ()> + Send + 'static,
419 {
420 let acceptor = TlsAcceptor::from(config);
421 let app = Arc::new(self);
422 serve_tls_with_shutdown(listener, app, acceptor, shutdown).await;
423 Ok(())
424 }
425
426 #[cfg(feature = "tls")]
428 pub async fn bind_tls(
429 self,
430 addr: SocketAddr,
431 config: Arc<rustls::ServerConfig>,
432 ) -> Result<(), ServeError> {
433 let listener = TcpListener::bind(addr)
434 .await
435 .map_err(|e| ServeError::new(500, format!("failed to bind to {addr}: {e}")))?;
436 self.run_tls(listener, config, signal_shutdown()).await
437 }
438}
439
440impl App<()> {
441 pub fn stateless() -> Self {
442 App::new(())
443 }
444}
445
446async fn serve_connection<S, IO>(io: IO, app: Arc<App<S>>, header_read_timeout: Duration)
450where
451 S: Send + Sync + 'static,
452 IO: hyper::rt::Read + hyper::rt::Write + Unpin + 'static,
453{
454 let svc = service_fn(move |req: Request<Incoming>| {
455 let app = app.clone();
456 async move {
457 Ok::<_, hyper::Error>(app.route(req).await)
458 }
459 });
460 let mut builder = http1::Builder::new();
461 builder.timer(TokioTimer::new());
462 builder.header_read_timeout(header_read_timeout);
463 let conn = builder.serve_connection(io, svc);
464 let _ = conn.await;
465}
466
467async fn serve_inner<S: Send + Sync + 'static>(
468 listener: TcpListener,
469 app: Arc<App<S>>,
470) {
471 let header_read_timeout = app.header_read_timeout;
472 let semaphore = Arc::new(tokio::sync::Semaphore::new(app.max_connections));
473 let mut backoff = Backoff::new();
474 loop {
475 let (stream, _) = accept_with_backoff(&listener, &mut backoff).await;
476 let sem = semaphore.clone();
477 let permit = match sem.acquire_owned().await {
478 Ok(p) => p,
479 Err(_) => continue,
480 };
481 let app = app.clone();
482 tokio::spawn(async move {
483 let _permit = permit;
484 serve_connection(TokioIo::new(stream), app, header_read_timeout).await;
485 });
486 }
487}
488
489#[cfg(feature = "tls")]
494async fn serve_tls_inner<S: Send + Sync + 'static>(
495 listener: TcpListener,
496 app: Arc<App<S>>,
497 acceptor: TlsAcceptor,
498) {
499 let header_read_timeout = app.header_read_timeout;
500 let handshake_timeout = app.tls_handshake_timeout;
501 let semaphore = Arc::new(tokio::sync::Semaphore::new(app.max_connections));
502 let mut backoff = Backoff::new();
503 loop {
504 let (stream, _) = accept_with_backoff(&listener, &mut backoff).await;
505 let sem = semaphore.clone();
506 let permit = match sem.acquire_owned().await {
507 Ok(p) => p,
508 Err(_) => continue,
509 };
510 let app = app.clone();
511 let acceptor = acceptor.clone();
512 tokio::spawn(async move {
513 let _permit = permit;
514 let tls_stream = match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
515 Ok(Ok(s)) => s,
516 Ok(Err(_)) | Err(_) => return,
517 };
518 serve_connection(TokioIo::new(tls_stream), app, header_read_timeout).await;
519 });
520 }
521}
522
523async fn accept_and_permit<L: TcpAccept>(
537 listener: &L,
538 backoff: &mut Backoff,
539 semaphore: &Arc<tokio::sync::Semaphore>,
540) -> Option<(TcpStream, tokio::sync::OwnedSemaphorePermit)> {
541 loop {
542 let (stream, _) = match listener.accept().await {
543 Ok(conn) => {
544 backoff.reset();
545 conn
546 }
547 Err(_) => {
548 tokio::time::sleep(backoff.next_delay()).await;
549 continue;
550 }
551 };
552 return match semaphore.clone().acquire_owned().await {
553 Ok(permit) => Some((stream, permit)),
554 Err(_) => None,
555 };
556 }
557}
558
559async fn signal_shutdown() {
561 let mut sigint = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())
562 .expect("failed to install SIGINT handler");
563 let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
564 .expect("failed to install SIGTERM handler");
565
566 tokio::select! {
567 _ = sigint.recv() => {}
568 _ = sigterm.recv() => {}
569 }
570}
571
572async fn serve_with_shutdown<S, F>(listener: TcpListener, app: Arc<App<S>>, shutdown: F)
581where
582 S: Send + Sync + 'static,
583 F: Future<Output = ()> + Send + 'static,
584{
585 let header_read_timeout = app.header_read_timeout;
586 let semaphore = Arc::new(tokio::sync::Semaphore::new(app.max_connections));
587 let mut backoff = Backoff::new();
588 let mut join_set: tokio::task::JoinSet<()> = tokio::task::JoinSet::new();
589 let mut shutdown_pin = std::pin::pin!(shutdown);
590 let mut shutting_down = false;
591
592 loop {
593 if !shutting_down {
594 tokio::select! {
595 accepted = accept_and_permit(&listener, &mut backoff, &semaphore) => {
596 match accepted {
597 Some((stream, permit)) => {
598 let app = app.clone();
599 join_set.spawn(async move {
600 let _permit = permit;
601 serve_connection(TokioIo::new(stream), app, header_read_timeout).await;
602 });
603 }
604 None => shutting_down = true,
605 }
606 }
607 _ = shutdown_pin.as_mut() => {
608 shutting_down = true;
609 }
610 }
611 continue;
612 }
613
614 match join_set.join_next().await {
615 Some(_) => continue,
616 None => break,
617 }
618 }
619}
620
621#[cfg(feature = "tls")]
626async fn serve_tls_with_shutdown<S, F>(
627 listener: TcpListener,
628 app: Arc<App<S>>,
629 acceptor: TlsAcceptor,
630 shutdown: F,
631) where
632 S: Send + Sync + 'static,
633 F: Future<Output = ()> + Send + 'static,
634{
635 let header_read_timeout = app.header_read_timeout;
636 let handshake_timeout = app.tls_handshake_timeout;
637 let semaphore = Arc::new(tokio::sync::Semaphore::new(app.max_connections));
638 let mut backoff = Backoff::new();
639 let mut join_set: tokio::task::JoinSet<()> = tokio::task::JoinSet::new();
640 let mut shutdown_pin = std::pin::pin!(shutdown);
641 let mut shutting_down = false;
642
643 loop {
644 if !shutting_down {
645 tokio::select! {
646 accepted = accept_and_permit(&listener, &mut backoff, &semaphore) => {
647 match accepted {
648 Some((stream, permit)) => {
649 let app = app.clone();
650 let acceptor = acceptor.clone();
651 join_set.spawn(async move {
652 let _permit = permit;
653 let tls_stream = match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
654 Ok(Ok(s)) => s,
655 Ok(Err(_)) | Err(_) => return,
656 };
657 serve_connection(TokioIo::new(tls_stream), app, header_read_timeout).await;
658 });
659 }
660 None => shutting_down = true,
661 }
662 }
663 _ = shutdown_pin.as_mut() => {
664 shutting_down = true;
665 }
666 }
667 continue;
668 }
669
670 match join_set.join_next().await {
671 Some(_) => continue,
672 None => break,
673 }
674 }
675}
676
677#[must_use = "RouteBuilder does nothing until .seal() is called"]
702pub struct RouteBuilder<S> {
703 state: Arc<S>,
704 router: Router<S>,
705 max_body_size: usize,
706 header_read_timeout: Duration,
707 max_connections: usize,
708 #[cfg(feature = "tls")]
709 tls_handshake_timeout: Duration,
710 error_handler: ErrorHandler,
711 cors_config: Option<CorsConfig>,
712}
713
714impl<S: Send + Sync + 'static> RouteBuilder<S> {
715 pub fn new(state: S) -> Self {
717 RouteBuilder {
718 state: Arc::new(state),
719 router: Router::new(),
720 max_body_size: DEFAULT_MAX_BODY_SIZE,
721 header_read_timeout: DEFAULT_HEADER_READ_TIMEOUT,
722 max_connections: DEFAULT_MAX_CONNECTIONS,
723 #[cfg(feature = "tls")]
724 tls_handshake_timeout: DEFAULT_TLS_HANDSHAKE_TIMEOUT,
725 error_handler: default_error_handler(),
726 cors_config: None,
727 }
728 }
729
730 pub fn with_max_body_size(mut self, max: usize) -> Self {
732 self.max_body_size = max;
733 self
734 }
735
736 pub fn with_header_read_timeout(mut self, d: Duration) -> Self {
738 self.header_read_timeout = d;
739 self
740 }
741
742 #[cfg(feature = "tls")]
745 pub fn with_tls_handshake_timeout(mut self, d: Duration) -> Self {
746 self.tls_handshake_timeout = d;
747 self
748 }
749
750 pub fn with_max_connections(mut self, max: usize) -> Self {
752 self.max_connections = max;
753 self
754 }
755
756 pub fn with_error_handler(
757 mut self,
758 f: impl Fn(StatusCode, &str) -> Response<ResponseBody> + Send + Sync + 'static,
759 ) -> Self {
760 self.error_handler = Arc::new(f);
761 self
762 }
763
764 pub fn with_cors(mut self, config: CorsConfig) -> Self {
765 self.cors_config = Some(config);
766 self
767 }
768
769 pub fn get(mut self, path: &str, handler: Handler<S>) -> Self {
770 self.router.insert(Method::GET, path, handler);
771 self
772 }
773
774 pub fn post(mut self, path: &str, handler: Handler<S>) -> Self {
775 self.router.insert(Method::POST, path, handler);
776 self
777 }
778
779 pub fn put(mut self, path: &str, handler: Handler<S>) -> Self {
780 self.router.insert(Method::PUT, path, handler);
781 self
782 }
783
784 pub fn delete(mut self, path: &str, handler: Handler<S>) -> Self {
785 self.router.insert(Method::DELETE, path, handler);
786 self
787 }
788
789 pub fn seal(self) -> App<S> {
790 App {
791 state: self.state,
792 router: Arc::new(self.router),
793 max_body_size: self.max_body_size,
794 header_read_timeout: self.header_read_timeout,
795 max_connections: self.max_connections,
796 #[cfg(feature = "tls")]
797 tls_handshake_timeout: self.tls_handshake_timeout,
798 error_handler: self.error_handler,
799 cors_config: self.cors_config,
800 }
801 }
802}
803
804impl RouteBuilder<()> {
805 pub fn stateless() -> Self {
806 RouteBuilder::new(())
807 }
808}
809
810#[cfg(test)]
811mod tests {
812 use super::*;
813 use std::sync::Mutex;
814 use std::sync::atomic::{AtomicUsize, Ordering};
815 use http_body_util::BodyExt;
816
817 #[test]
818 fn ephemeral_bind_addr_is_loopback_only() {
819 assert_eq!(
820 ephemeral_bind_addr().ip(),
821 std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)
822 );
823 }
824
825 #[test]
826 fn backoff_delays_double_up_to_a_cap() {
827 let mut backoff = Backoff::new();
828
829 assert_eq!(backoff.next_delay(), ACCEPT_BACKOFF_INITIAL);
830 assert_eq!(backoff.next_delay(), ACCEPT_BACKOFF_INITIAL * 2);
831 assert_eq!(backoff.next_delay(), ACCEPT_BACKOFF_INITIAL * 4);
832
833 let mut last = Duration::ZERO;
835 for _ in 0..20 {
836 last = backoff.next_delay();
837 }
838 assert_eq!(last, ACCEPT_BACKOFF_MAX);
839 }
840
841 #[test]
842 fn backoff_reset_returns_to_initial_delay() {
843 let mut backoff = Backoff::new();
844 backoff.next_delay();
845 backoff.next_delay();
846 backoff.reset();
847 assert_eq!(backoff.next_delay(), ACCEPT_BACKOFF_INITIAL);
848 }
849
850 struct FlakyListener {
854 inner: TcpListener,
855 remaining_failures: AtomicUsize,
856 attempts: Mutex<Vec<tokio::time::Instant>>,
857 }
858
859 impl TcpAccept for FlakyListener {
860 async fn accept(&self) -> std::io::Result<(TcpStream, SocketAddr)> {
861 self.attempts.lock().unwrap().push(tokio::time::Instant::now());
862 if self.remaining_failures.fetch_sub(1, Ordering::SeqCst) > 0 {
863 Err(std::io::Error::other("simulated accept error"))
864 } else {
865 TcpAccept::accept(&self.inner).await
866 }
867 }
868 }
869
870 #[tokio::test(start_paused = true)]
871 async fn accept_loop_backs_off_between_repeated_errors_instead_of_busy_spinning() {
872 let inner = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
873 let addr = inner.local_addr().unwrap();
874
875 let flaky = FlakyListener {
876 inner,
877 remaining_failures: AtomicUsize::new(5),
878 attempts: Mutex::new(Vec::new()),
879 };
880
881 tokio::spawn(async move {
882 let _ = TcpStream::connect(addr).await;
883 });
884
885 let mut backoff = Backoff::new();
886 accept_with_backoff(&flaky, &mut backoff).await;
887
888 let recorded = flaky.attempts.lock().unwrap();
889 assert_eq!(recorded.len(), 6, "5 failures then 1 success");
890
891 let expected_gaps = [
892 ACCEPT_BACKOFF_INITIAL,
893 ACCEPT_BACKOFF_INITIAL * 2,
894 ACCEPT_BACKOFF_INITIAL * 4,
895 ACCEPT_BACKOFF_INITIAL * 8,
896 ACCEPT_BACKOFF_INITIAL * 16,
897 ];
898 for (i, expected) in expected_gaps.iter().enumerate() {
899 let gap = recorded[i + 1] - recorded[i];
900 assert_eq!(
901 gap, *expected,
902 "gap between attempt {i} and {} should reflect the backoff delay, not a busy spin",
903 i + 1
904 );
905 }
906 }
907
908 #[tokio::test]
909 async fn error_handler_sanitizes_5xx_in_response_body() {
910 take_error_log(); let resp = error_response(StatusCode::INTERNAL_SERVER_ERROR, "raw db connection string leaked");
912
913 let (parts, body) = resp.into_parts();
914 assert_eq!(parts.status, StatusCode::INTERNAL_SERVER_ERROR);
915
916 let collected = body.collect().await.unwrap();
917 let bytes = collected.to_bytes();
918 let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
919 let msg = json.get("message").and_then(|v| v.as_str()).unwrap();
920 assert_eq!(msg, "internal server error", "5xx message should be sanitized");
921 }
922
923 #[test]
924 fn error_handler_captures_5xx_message_in_log() {
925 take_error_log(); let sensitive_msg = "raw db connection string leaked";
927 error_response(StatusCode::INTERNAL_SERVER_ERROR, sensitive_msg);
928
929 let log = take_error_log();
930 assert_eq!(log.len(), 1);
931 assert_eq!(log[0].0, 500);
932 assert_eq!(log[0].1, sensitive_msg);
933 }
934
935 #[tokio::test]
936 async fn error_handler_passes_through_4xx_in_response_body() {
937 take_error_log(); let msg = "bad request";
939 let resp = error_response(StatusCode::BAD_REQUEST, msg);
940
941 let (parts, body) = resp.into_parts();
942 assert_eq!(parts.status, StatusCode::BAD_REQUEST);
943
944 let collected = body.collect().await.unwrap();
945 let bytes = collected.to_bytes();
946 let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
947 let response_msg = json.get("message").and_then(|v| v.as_str()).unwrap();
948 assert_eq!(response_msg, msg, "4xx message should pass through");
949 }
950
951 #[test]
952 fn error_handler_logs_4xx_messages() {
953 take_error_log(); let msg = "bad request";
955 error_response(StatusCode::BAD_REQUEST, msg);
956
957 let log = take_error_log();
958 assert_eq!(log.len(), 1);
959 assert_eq!(log[0].0, 400);
960 assert_eq!(log[0].1, msg);
961 }
962}