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 return cfg.preflight_response(req_origin.as_deref());
272 }
273 }
274 }
275
276 let method_to_match = if method == Method::HEAD {
277 Method::GET
278 } else {
279 method.clone()
280 };
281
282 match self.router.match_route(&method_to_match, &path) {
283 Some((handler, params)) => {
284 let mut req = req;
285 req.extensions_mut().insert(query_params);
286 req.extensions_mut().insert(params);
287 req.extensions_mut().insert(MaxBodySize(self.max_body_size));
288 match handler(req, state).await {
289 Ok(mut resp) => {
290 if let Some(cfg) = &self.cors_config {
291 cfg.apply_to_response(&mut resp, req_origin.as_deref());
292 }
293 resp
294 }
295 Err(e) => (self.error_handler)(
296 StatusCode::from_u16(e.code)
297 .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR),
298 &e.message,
299 ),
300 }
301 }
302 None => {
303 let mut allowed = self.router.allowed_methods(&path);
304 if !allowed.is_empty() {
305 if allowed.contains(&Method::GET) {
306 allowed.push(Method::HEAD);
307 }
308 let mut method_strs: Vec<&str> = allowed.iter().map(|m| m.as_str()).collect();
309 method_strs.sort();
310 method_strs.dedup();
311 let allow_header = method_strs.join(", ");
312 let mut resp = (self.error_handler)(StatusCode::METHOD_NOT_ALLOWED, "method not allowed");
313 if let Ok(val) = allow_header.parse() {
314 resp.headers_mut().insert("allow", val);
315 }
316 resp
317 } else {
318 (self.error_handler)(StatusCode::NOT_FOUND, "not found")
319 }
320 }
321 }
322 }
323
324 pub async fn bind_ephemeral(self) -> Result<u16, ServeError> {
330 let listener = TcpListener::bind(ephemeral_bind_addr())
331 .await
332 .map_err(|e| ServeError::new(500, format!("failed to bind to ephemeral port: {e}")))?;
333 let port = listener
334 .local_addr()
335 .map_err(|e| ServeError::new(500, format!("failed to get assigned port: {e}")))?
336 .port();
337 let app = Arc::new(self);
338 tokio::spawn(async move {
339 serve_inner(listener, app).await;
340 });
341 Ok(port)
342 }
343
344 #[cfg(feature = "tls")]
350 pub async fn bind_tls_ephemeral(
351 self,
352 config: Arc<rustls::ServerConfig>,
353 ) -> Result<u16, ServeError> {
354 let listener = TcpListener::bind(ephemeral_bind_addr())
355 .await
356 .map_err(|e| ServeError::new(500, format!("failed to bind to ephemeral port: {e}")))?;
357 let port = listener
358 .local_addr()
359 .map_err(|e| ServeError::new(500, format!("failed to get assigned port: {e}")))?
360 .port();
361 let acceptor = TlsAcceptor::from(config);
362 let app = Arc::new(self);
363 tokio::spawn(async move {
364 serve_tls_inner(listener, app, acceptor).await;
365 });
366 Ok(port)
367 }
368
369 pub async fn run<F>(self, listener: TcpListener, shutdown: F) -> Result<(), ServeError>
374 where
375 F: Future<Output = ()> + Send + 'static,
376 {
377 let app = Arc::new(self);
378 serve_with_shutdown(listener, app, shutdown).await;
379 Ok(())
380 }
381
382 pub async fn bind(self, addr: SocketAddr) -> Result<(), ServeError> {
387 let listener = TcpListener::bind(addr)
388 .await
389 .map_err(|e| ServeError::new(500, format!("failed to bind to {addr}: {e}")))?;
390 self.run(listener, signal_shutdown()).await
391 }
392
393 #[cfg(feature = "tls")]
395 pub async fn run_tls<F>(
396 self,
397 listener: TcpListener,
398 config: Arc<rustls::ServerConfig>,
399 shutdown: F,
400 ) -> Result<(), ServeError>
401 where
402 F: Future<Output = ()> + Send + 'static,
403 {
404 let acceptor = TlsAcceptor::from(config);
405 let app = Arc::new(self);
406 serve_tls_with_shutdown(listener, app, acceptor, shutdown).await;
407 Ok(())
408 }
409
410 #[cfg(feature = "tls")]
412 pub async fn bind_tls(
413 self,
414 addr: SocketAddr,
415 config: Arc<rustls::ServerConfig>,
416 ) -> Result<(), ServeError> {
417 let listener = TcpListener::bind(addr)
418 .await
419 .map_err(|e| ServeError::new(500, format!("failed to bind to {addr}: {e}")))?;
420 self.run_tls(listener, config, signal_shutdown()).await
421 }
422}
423
424impl App<()> {
425 pub fn stateless() -> Self {
426 App::new(())
427 }
428}
429
430async fn serve_connection<S, IO>(io: IO, app: Arc<App<S>>, header_read_timeout: Duration)
434where
435 S: Send + Sync + 'static,
436 IO: hyper::rt::Read + hyper::rt::Write + Unpin + 'static,
437{
438 let svc = service_fn(move |req: Request<Incoming>| {
439 let app = app.clone();
440 async move {
441 Ok::<_, hyper::Error>(app.route(req).await)
442 }
443 });
444 let mut builder = http1::Builder::new();
445 builder.timer(TokioTimer::new());
446 builder.header_read_timeout(header_read_timeout);
447 let conn = builder.serve_connection(io, svc);
448 let _ = conn.await;
449}
450
451async fn serve_inner<S: Send + Sync + 'static>(
452 listener: TcpListener,
453 app: Arc<App<S>>,
454) {
455 let header_read_timeout = app.header_read_timeout;
456 let semaphore = Arc::new(tokio::sync::Semaphore::new(app.max_connections));
457 let mut backoff = Backoff::new();
458 loop {
459 let (stream, _) = accept_with_backoff(&listener, &mut backoff).await;
460 let sem = semaphore.clone();
461 let permit = match sem.acquire_owned().await {
462 Ok(p) => p,
463 Err(_) => continue,
464 };
465 let app = app.clone();
466 tokio::spawn(async move {
467 let _permit = permit;
468 serve_connection(TokioIo::new(stream), app, header_read_timeout).await;
469 });
470 }
471}
472
473#[cfg(feature = "tls")]
478async fn serve_tls_inner<S: Send + Sync + 'static>(
479 listener: TcpListener,
480 app: Arc<App<S>>,
481 acceptor: TlsAcceptor,
482) {
483 let header_read_timeout = app.header_read_timeout;
484 let handshake_timeout = app.tls_handshake_timeout;
485 let semaphore = Arc::new(tokio::sync::Semaphore::new(app.max_connections));
486 let mut backoff = Backoff::new();
487 loop {
488 let (stream, _) = accept_with_backoff(&listener, &mut backoff).await;
489 let sem = semaphore.clone();
490 let permit = match sem.acquire_owned().await {
491 Ok(p) => p,
492 Err(_) => continue,
493 };
494 let app = app.clone();
495 let acceptor = acceptor.clone();
496 tokio::spawn(async move {
497 let _permit = permit;
498 let tls_stream = match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
499 Ok(Ok(s)) => s,
500 Ok(Err(_)) | Err(_) => return,
501 };
502 serve_connection(TokioIo::new(tls_stream), app, header_read_timeout).await;
503 });
504 }
505}
506
507async fn accept_and_permit<L: TcpAccept>(
521 listener: &L,
522 backoff: &mut Backoff,
523 semaphore: &Arc<tokio::sync::Semaphore>,
524) -> Option<(TcpStream, tokio::sync::OwnedSemaphorePermit)> {
525 loop {
526 let (stream, _) = match listener.accept().await {
527 Ok(conn) => {
528 backoff.reset();
529 conn
530 }
531 Err(_) => {
532 tokio::time::sleep(backoff.next_delay()).await;
533 continue;
534 }
535 };
536 return match semaphore.clone().acquire_owned().await {
537 Ok(permit) => Some((stream, permit)),
538 Err(_) => None,
539 };
540 }
541}
542
543async fn signal_shutdown() {
545 let mut sigint = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())
546 .expect("failed to install SIGINT handler");
547 let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
548 .expect("failed to install SIGTERM handler");
549
550 tokio::select! {
551 _ = sigint.recv() => {}
552 _ = sigterm.recv() => {}
553 }
554}
555
556async fn serve_with_shutdown<S, F>(listener: TcpListener, app: Arc<App<S>>, shutdown: F)
565where
566 S: Send + Sync + 'static,
567 F: Future<Output = ()> + Send + 'static,
568{
569 let header_read_timeout = app.header_read_timeout;
570 let semaphore = Arc::new(tokio::sync::Semaphore::new(app.max_connections));
571 let mut backoff = Backoff::new();
572 let mut join_set: tokio::task::JoinSet<()> = tokio::task::JoinSet::new();
573 let mut shutdown_pin = std::pin::pin!(shutdown);
574 let mut shutting_down = false;
575
576 loop {
577 if !shutting_down {
578 tokio::select! {
579 accepted = accept_and_permit(&listener, &mut backoff, &semaphore) => {
580 match accepted {
581 Some((stream, permit)) => {
582 let app = app.clone();
583 join_set.spawn(async move {
584 let _permit = permit;
585 serve_connection(TokioIo::new(stream), app, header_read_timeout).await;
586 });
587 }
588 None => shutting_down = true,
589 }
590 }
591 _ = shutdown_pin.as_mut() => {
592 shutting_down = true;
593 }
594 }
595 continue;
596 }
597
598 match join_set.join_next().await {
599 Some(_) => continue,
600 None => break,
601 }
602 }
603}
604
605#[cfg(feature = "tls")]
610async fn serve_tls_with_shutdown<S, F>(
611 listener: TcpListener,
612 app: Arc<App<S>>,
613 acceptor: TlsAcceptor,
614 shutdown: F,
615) where
616 S: Send + Sync + 'static,
617 F: Future<Output = ()> + Send + 'static,
618{
619 let header_read_timeout = app.header_read_timeout;
620 let handshake_timeout = app.tls_handshake_timeout;
621 let semaphore = Arc::new(tokio::sync::Semaphore::new(app.max_connections));
622 let mut backoff = Backoff::new();
623 let mut join_set: tokio::task::JoinSet<()> = tokio::task::JoinSet::new();
624 let mut shutdown_pin = std::pin::pin!(shutdown);
625 let mut shutting_down = false;
626
627 loop {
628 if !shutting_down {
629 tokio::select! {
630 accepted = accept_and_permit(&listener, &mut backoff, &semaphore) => {
631 match accepted {
632 Some((stream, permit)) => {
633 let app = app.clone();
634 let acceptor = acceptor.clone();
635 join_set.spawn(async move {
636 let _permit = permit;
637 let tls_stream = match tokio::time::timeout(handshake_timeout, acceptor.accept(stream)).await {
638 Ok(Ok(s)) => s,
639 Ok(Err(_)) | Err(_) => return,
640 };
641 serve_connection(TokioIo::new(tls_stream), app, header_read_timeout).await;
642 });
643 }
644 None => shutting_down = true,
645 }
646 }
647 _ = shutdown_pin.as_mut() => {
648 shutting_down = true;
649 }
650 }
651 continue;
652 }
653
654 match join_set.join_next().await {
655 Some(_) => continue,
656 None => break,
657 }
658 }
659}
660
661#[must_use = "RouteBuilder does nothing until .seal() is called"]
686pub struct RouteBuilder<S> {
687 state: Arc<S>,
688 router: Router<S>,
689 max_body_size: usize,
690 header_read_timeout: Duration,
691 max_connections: usize,
692 #[cfg(feature = "tls")]
693 tls_handshake_timeout: Duration,
694 error_handler: ErrorHandler,
695 cors_config: Option<CorsConfig>,
696}
697
698impl<S: Send + Sync + 'static> RouteBuilder<S> {
699 pub fn new(state: S) -> Self {
701 RouteBuilder {
702 state: Arc::new(state),
703 router: Router::new(),
704 max_body_size: DEFAULT_MAX_BODY_SIZE,
705 header_read_timeout: DEFAULT_HEADER_READ_TIMEOUT,
706 max_connections: DEFAULT_MAX_CONNECTIONS,
707 #[cfg(feature = "tls")]
708 tls_handshake_timeout: DEFAULT_TLS_HANDSHAKE_TIMEOUT,
709 error_handler: default_error_handler(),
710 cors_config: None,
711 }
712 }
713
714 pub fn with_max_body_size(mut self, max: usize) -> Self {
716 self.max_body_size = max;
717 self
718 }
719
720 pub fn with_header_read_timeout(mut self, d: Duration) -> Self {
722 self.header_read_timeout = d;
723 self
724 }
725
726 #[cfg(feature = "tls")]
729 pub fn with_tls_handshake_timeout(mut self, d: Duration) -> Self {
730 self.tls_handshake_timeout = d;
731 self
732 }
733
734 pub fn with_max_connections(mut self, max: usize) -> Self {
736 self.max_connections = max;
737 self
738 }
739
740 pub fn with_error_handler(
741 mut self,
742 f: impl Fn(StatusCode, &str) -> Response<ResponseBody> + Send + Sync + 'static,
743 ) -> Self {
744 self.error_handler = Arc::new(f);
745 self
746 }
747
748 pub fn with_cors(mut self, config: CorsConfig) -> Self {
749 self.cors_config = Some(config);
750 self
751 }
752
753 pub fn get(mut self, path: &str, handler: Handler<S>) -> Self {
754 self.router.insert(Method::GET, path, handler);
755 self
756 }
757
758 pub fn post(mut self, path: &str, handler: Handler<S>) -> Self {
759 self.router.insert(Method::POST, path, handler);
760 self
761 }
762
763 pub fn put(mut self, path: &str, handler: Handler<S>) -> Self {
764 self.router.insert(Method::PUT, path, handler);
765 self
766 }
767
768 pub fn delete(mut self, path: &str, handler: Handler<S>) -> Self {
769 self.router.insert(Method::DELETE, path, handler);
770 self
771 }
772
773 pub fn seal(self) -> App<S> {
774 App {
775 state: self.state,
776 router: Arc::new(self.router),
777 max_body_size: self.max_body_size,
778 header_read_timeout: self.header_read_timeout,
779 max_connections: self.max_connections,
780 #[cfg(feature = "tls")]
781 tls_handshake_timeout: self.tls_handshake_timeout,
782 error_handler: self.error_handler,
783 cors_config: self.cors_config,
784 }
785 }
786}
787
788impl RouteBuilder<()> {
789 pub fn stateless() -> Self {
790 RouteBuilder::new(())
791 }
792}
793
794#[cfg(test)]
795mod tests {
796 use super::*;
797 use std::sync::Mutex;
798 use std::sync::atomic::{AtomicUsize, Ordering};
799 use http_body_util::BodyExt;
800
801 #[test]
802 fn ephemeral_bind_addr_is_loopback_only() {
803 assert_eq!(
804 ephemeral_bind_addr().ip(),
805 std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)
806 );
807 }
808
809 #[test]
810 fn backoff_delays_double_up_to_a_cap() {
811 let mut backoff = Backoff::new();
812
813 assert_eq!(backoff.next_delay(), ACCEPT_BACKOFF_INITIAL);
814 assert_eq!(backoff.next_delay(), ACCEPT_BACKOFF_INITIAL * 2);
815 assert_eq!(backoff.next_delay(), ACCEPT_BACKOFF_INITIAL * 4);
816
817 let mut last = Duration::ZERO;
819 for _ in 0..20 {
820 last = backoff.next_delay();
821 }
822 assert_eq!(last, ACCEPT_BACKOFF_MAX);
823 }
824
825 #[test]
826 fn backoff_reset_returns_to_initial_delay() {
827 let mut backoff = Backoff::new();
828 backoff.next_delay();
829 backoff.next_delay();
830 backoff.reset();
831 assert_eq!(backoff.next_delay(), ACCEPT_BACKOFF_INITIAL);
832 }
833
834 struct FlakyListener {
838 inner: TcpListener,
839 remaining_failures: AtomicUsize,
840 attempts: Mutex<Vec<tokio::time::Instant>>,
841 }
842
843 impl TcpAccept for FlakyListener {
844 async fn accept(&self) -> std::io::Result<(TcpStream, SocketAddr)> {
845 self.attempts.lock().unwrap().push(tokio::time::Instant::now());
846 if self.remaining_failures.fetch_sub(1, Ordering::SeqCst) > 0 {
847 Err(std::io::Error::other("simulated accept error"))
848 } else {
849 TcpAccept::accept(&self.inner).await
850 }
851 }
852 }
853
854 #[tokio::test(start_paused = true)]
855 async fn accept_loop_backs_off_between_repeated_errors_instead_of_busy_spinning() {
856 let inner = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
857 let addr = inner.local_addr().unwrap();
858
859 let flaky = FlakyListener {
860 inner,
861 remaining_failures: AtomicUsize::new(5),
862 attempts: Mutex::new(Vec::new()),
863 };
864
865 tokio::spawn(async move {
866 let _ = TcpStream::connect(addr).await;
867 });
868
869 let mut backoff = Backoff::new();
870 accept_with_backoff(&flaky, &mut backoff).await;
871
872 let recorded = flaky.attempts.lock().unwrap();
873 assert_eq!(recorded.len(), 6, "5 failures then 1 success");
874
875 let expected_gaps = [
876 ACCEPT_BACKOFF_INITIAL,
877 ACCEPT_BACKOFF_INITIAL * 2,
878 ACCEPT_BACKOFF_INITIAL * 4,
879 ACCEPT_BACKOFF_INITIAL * 8,
880 ACCEPT_BACKOFF_INITIAL * 16,
881 ];
882 for (i, expected) in expected_gaps.iter().enumerate() {
883 let gap = recorded[i + 1] - recorded[i];
884 assert_eq!(
885 gap, *expected,
886 "gap between attempt {i} and {} should reflect the backoff delay, not a busy spin",
887 i + 1
888 );
889 }
890 }
891
892 #[tokio::test]
893 async fn error_handler_sanitizes_5xx_in_response_body() {
894 take_error_log(); let resp = error_response(StatusCode::INTERNAL_SERVER_ERROR, "raw db connection string leaked");
896
897 let (parts, body) = resp.into_parts();
898 assert_eq!(parts.status, StatusCode::INTERNAL_SERVER_ERROR);
899
900 let collected = body.collect().await.unwrap();
901 let bytes = collected.to_bytes();
902 let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
903 let msg = json.get("message").and_then(|v| v.as_str()).unwrap();
904 assert_eq!(msg, "internal server error", "5xx message should be sanitized");
905 }
906
907 #[test]
908 fn error_handler_captures_5xx_message_in_log() {
909 take_error_log(); let sensitive_msg = "raw db connection string leaked";
911 error_response(StatusCode::INTERNAL_SERVER_ERROR, sensitive_msg);
912
913 let log = take_error_log();
914 assert_eq!(log.len(), 1);
915 assert_eq!(log[0].0, 500);
916 assert_eq!(log[0].1, sensitive_msg);
917 }
918
919 #[tokio::test]
920 async fn error_handler_passes_through_4xx_in_response_body() {
921 take_error_log(); let msg = "bad request";
923 let resp = error_response(StatusCode::BAD_REQUEST, msg);
924
925 let (parts, body) = resp.into_parts();
926 assert_eq!(parts.status, StatusCode::BAD_REQUEST);
927
928 let collected = body.collect().await.unwrap();
929 let bytes = collected.to_bytes();
930 let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
931 let response_msg = json.get("message").and_then(|v| v.as_str()).unwrap();
932 assert_eq!(response_msg, msg, "4xx message should pass through");
933 }
934
935 #[test]
936 fn error_handler_logs_4xx_messages() {
937 take_error_log(); let msg = "bad request";
939 error_response(StatusCode::BAD_REQUEST, msg);
940
941 let log = take_error_log();
942 assert_eq!(log.len(), 1);
943 assert_eq!(log[0].0, 400);
944 assert_eq!(log[0].1, msg);
945 }
946}