1use std::convert::Infallible;
42use std::fmt::Debug;
43use std::io;
44use std::pin::Pin;
45use std::sync::{Arc, Mutex};
46use std::task::{Context, Poll};
47use std::time::Duration;
48
49use axum::serve::Listener;
50use hyper_util::rt::{TokioExecutor, TokioIo};
51use hyper_util::server::conn::auto::{Builder, UpgradeableConnection};
52use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
53use tokio::sync::{Notify, Semaphore, watch};
54use tokio::time::Instant;
55
56use crate::config::HttpServerConfig;
57
58#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
60pub struct HttpLimits {
61 pub header_read_timeout: Option<Duration>,
63 pub keep_alive_timeout: Option<Duration>,
65 pub max_header_bytes: Option<usize>,
68 pub http2_max_concurrent_streams: Option<u32>,
70 pub max_connections: Option<usize>,
72}
73
74impl From<&HttpServerConfig> for HttpLimits {
75 fn from(config: &HttpServerConfig) -> Self {
76 let ms = |value: Option<u64>| value.filter(|v| *v > 0).map(Duration::from_millis);
77 Self {
78 header_read_timeout: ms(config.header_read_timeout_ms),
79 keep_alive_timeout: ms(config.keep_alive_timeout_ms),
80 max_header_bytes: config
83 .max_header_bytes
84 .filter(|v| *v > 0)
85 .map(|v| v.max(HttpServerConfig::MIN_HEADER_BYTES)),
86 http2_max_concurrent_streams: config.http2_max_concurrent_streams,
87 max_connections: config
90 .max_connections
91 .filter(|v| *v > 0)
92 .map(|v| v.min(HttpServerConfig::MAX_CONNECTIONS)),
93 }
94 }
95}
96
97impl HttpLimits {
98 const fn has_timers(&self) -> bool {
99 self.header_read_timeout.is_some() || self.keep_alive_timeout.is_some()
100 }
101
102 fn builder(&self) -> Builder<TokioExecutor> {
103 let mut builder = Builder::new(TokioExecutor::new());
104 builder.http2().enable_connect_protocol();
106 if let Some(bytes) = self.max_header_bytes {
107 let bytes = bytes.max(HttpServerConfig::MIN_HEADER_BYTES);
110 builder.http1().max_buf_size(bytes);
111 builder
112 .http2()
113 .max_header_list_size(u32::try_from(bytes).unwrap_or(u32::MAX));
114 }
115 if let Some(streams) = self.http2_max_concurrent_streams {
116 builder.http2().max_concurrent_streams(streams);
117 }
118 builder
119 }
120}
121
122pub struct IncomingStream<'a, L: Listener> {
127 io: &'a L::Io,
128 remote_addr: L::Addr,
129}
130
131impl<L: Listener> IncomingStream<'_, L> {
132 #[must_use]
134 pub const fn io(&self) -> &L::Io {
135 self.io
136 }
137
138 #[must_use]
140 pub const fn remote_addr(&self) -> &L::Addr {
141 &self.remote_addr
142 }
143}
144
145pub async fn serve<L, M, S, F>(
158 mut listener: L,
159 make_service: M,
160 limits: HttpLimits,
161 signal: F,
162) -> io::Result<()>
163where
164 L: Listener,
165 L::Addr: Debug,
166 M: for<'a> tower::Service<IncomingStream<'a, L>, Error = Infallible, Response = S>
167 + Send
168 + 'static,
169 for<'a> <M as tower::Service<IncomingStream<'a, L>>>::Future: Send,
170 S: tower::Service<
171 axum::extract::Request,
172 Response = axum::response::Response,
173 Error = Infallible,
174 > + Clone
175 + Send
176 + 'static,
177 S::Future: Send,
178 F: std::future::Future<Output = ()> + Send + 'static,
179{
180 let handoff = tls_handoff(&mut listener, &limits);
181 serve_with_handoff(listener, make_service, limits, handoff, signal).await
182}
183
184#[cfg(feature = "tls")]
188fn tls_handoff<L: Listener>(listener: &mut L, limits: &HttpLimits) -> Option<SlotHandoff> {
189 let max = limits.max_connections?;
190 (listener as &mut dyn std::any::Any)
191 .downcast_mut::<crate::tls::TlsListener>()?
192 .limit_connections(max)
193}
194
195#[cfg(not(feature = "tls"))]
196const fn tls_handoff<L>(_listener: &mut L, _limits: &HttpLimits) -> Option<SlotHandoff> {
197 None
198}
199
200pub(crate) type SlotHandoff = Arc<Mutex<Option<Slot>>>;
204
205pub(crate) type Slot = Box<dyn Send + Sync>;
209
210impl HttpLimits {
211 fn slots(&self, listener_takes_slots: bool) -> Option<Arc<Semaphore>> {
213 self.max_connections
214 .filter(|_| !listener_takes_slots)
215 .map(|n| Arc::new(Semaphore::new(n.min(Semaphore::MAX_PERMITS))))
216 }
217}
218
219async fn wait_for_slot(
221 permits: Option<&Arc<Semaphore>>,
222 closed: impl std::future::Future<Output = ()>,
223) -> Option<Option<Slot>> {
224 let Some(permits) = permits else {
225 return Some(None);
226 };
227 tokio::select! {
228 permit = Arc::clone(permits).acquire_owned() => {
229 Some(permit.ok().map(|permit| Box::new(permit) as Slot))
230 }
231 () = closed => None,
232 }
233}
234
235enum Step<T> {
237 Done(T),
238 TimedOut,
239}
240
241async fn race_step<T>(
246 step: impl std::future::Future<Output = T>,
247 deadline: Option<Instant>,
248 shutdown: &watch::Sender<()>,
249) -> Step<T> {
250 tokio::select! {
251 value = step => Step::Done(value),
252 () = sleep_until(deadline) => Step::TimedOut,
253 () = async {
254 shutdown.closed().await;
255 tokio::time::sleep(FRESH_CONN_GRACE).await;
256 } => Step::TimedOut,
257 }
258}
259
260fn take_handoff(handoff: Option<&SlotHandoff>) -> Option<Slot> {
262 handoff?
263 .lock()
264 .unwrap_or_else(std::sync::PoisonError::into_inner)
265 .take()
266}
267
268pub(crate) async fn serve_with_handoff<L, M, S, F>(
271 mut listener: L,
272 mut make_service: M,
273 limits: HttpLimits,
274 handoff: Option<SlotHandoff>,
275 signal: F,
276) -> io::Result<()>
277where
278 L: Listener,
279 L::Addr: Debug,
280 M: for<'a> tower::Service<IncomingStream<'a, L>, Error = Infallible, Response = S>
281 + Send
282 + 'static,
283 for<'a> <M as tower::Service<IncomingStream<'a, L>>>::Future: Send,
284 S: tower::Service<
285 axum::extract::Request,
286 Response = axum::response::Response,
287 Error = Infallible,
288 > + Clone
289 + Send
290 + 'static,
291 S::Future: Send,
292 F: std::future::Future<Output = ()> + Send + 'static,
293{
294 use tower::ServiceExt as _;
295
296 let permits = limits.slots(handoff.is_some());
297 let builder = limits.builder();
298 let (signal_tx, signal_rx) = watch::channel(());
299 tokio::spawn(async move {
300 signal.await;
301 drop(signal_rx);
302 });
303 let (close_tx, close_rx) = watch::channel(());
304
305 loop {
306 #[allow(
309 clippy::significant_drop_tightening,
310 reason = "the slot must live as long as the connection"
311 )]
312 let Some(permit) = wait_for_slot(permits.as_ref(), signal_tx.closed()).await else {
313 break;
314 };
315 let (io, remote_addr) = tokio::select! {
317 biased;
318 () = signal_tx.closed() => break,
319 conn = listener.accept() => conn,
320 };
321 let permit = permit.or_else(|| take_handoff(handoff.as_ref()));
322
323 let timers = limits
326 .has_timers()
327 .then(|| Arc::new(ConnTimers::new(&limits)));
328 let head = timers.as_ref().and_then(|t| t.deadline(None));
329 let Step::Done(_) = race_step(make_service.ready(), head, &signal_tx).await else {
330 continue;
331 };
332 let call = make_service.call(IncomingStream {
333 io: &io,
334 remote_addr,
335 });
336 let Step::Done(service) = race_step(call, head, &signal_tx).await else {
337 continue;
338 };
339 let service = service.unwrap_or_else(|err| match err {});
340 let first_request = Arc::new(FirstRequest::default());
341 let service = Tracked {
342 inner: service,
343 timers: timers.clone(),
344 first_request: Arc::clone(&first_request),
345 };
346 let service = service.map_request(|req: axum::http::Request<hyper::body::Incoming>| {
347 req.map(axum::body::Body::new)
348 });
349 let hyper_service = hyper_util::service::TowerToHyperService::new(service);
350 let io = TokioIo::new(ConnIo {
351 inner: io,
352 timers: timers.clone(),
353 _permit: permit,
354 });
355
356 let builder = builder.clone();
357 let signal_tx = signal_tx.clone();
358 let close_rx = close_rx.clone();
359 #[allow(
360 clippy::significant_drop_tightening,
361 reason = "the connection must live until the task ends"
362 )]
363 let connection = async move {
364 let _close_timers = CloseOnDrop(timers.clone());
365 let conn = builder.serve_connection_with_upgrades(io, hyper_service);
366 let conn = std::pin::pin!(conn);
367 let state = ConnState {
368 timers,
369 first_request,
370 };
371 drive_connection(
372 conn,
373 UpgradeableConnection::graceful_shutdown,
374 &signal_tx,
375 &state,
376 )
377 .await;
378 drop(close_rx);
379 };
380 tokio::spawn(connection);
381 }
382
383 drop(close_rx);
384 drop(listener);
385 close_tx.closed().await;
386 Ok(())
387}
388
389struct ConnState {
391 timers: Option<Arc<ConnTimers>>,
392 first_request: Arc<FirstRequest>,
393}
394
395async fn drive_connection<C: std::future::Future>(
398 mut conn: Pin<&mut C>,
399 drain: impl Fn(Pin<&mut C>),
400 signal_tx: &watch::Sender<()>,
401 state: &ConnState,
402) {
403 let ConnState {
404 timers,
405 first_request,
406 } = state;
407 let mut signal_closed = std::pin::pin!(signal_tx.closed());
408 let mut grace_until: Option<Instant> = None;
411 let mut fresh_until: Option<Instant> = None;
414 loop {
415 let shutting_down = grace_until.is_some();
416 let waiting_fresh = fresh_until.is_some();
417 let deadline = timers.as_ref().and_then(|t| t.deadline(grace_until));
418 tokio::select! {
420 biased;
421 _ = conn.as_mut() => break,
422 () = &mut signal_closed, if !shutting_down && !waiting_fresh => {
423 if first_request.is_seen() {
424 grace_until = Some(Instant::now() + SHUTDOWN_GRACE);
425 drain(conn.as_mut());
426 } else {
427 let until = Instant::now() + FRESH_CONN_GRACE;
429 let timer = timers.as_ref().and_then(|t| t.deadline(None));
430 fresh_until = Some(timer.map_or(until, |timer| timer.min(until)));
431 }
432 }
433 () = fresh_wait(first_request, fresh_until), if waiting_fresh && !shutting_down => {
435 if !first_request.is_seen() {
438 break;
439 }
440 fresh_until = None;
441 grace_until = Some(Instant::now() + SHUTDOWN_GRACE);
442 drain(conn.as_mut());
443 }
444 () = sleep_until(deadline), if deadline.is_some() => {
447 match timers.as_ref().and_then(|t| t.expired()) {
448 Some(Expiry::Head) => break,
450 Some(Expiry::Idle) if shutting_down => break,
452 Some(Expiry::Idle) => {
453 grace_until = Some(Instant::now() + SHUTDOWN_GRACE);
454 drain(conn.as_mut());
455 }
456 None => {}
457 }
458 }
459 () = changed(timers.as_deref()) => {}
460 }
461 }
462}
463
464struct CloseOnDrop(Option<Arc<ConnTimers>>);
466
467impl Drop for CloseOnDrop {
468 fn drop(&mut self) {
469 if let Some(timers) = &self.0 {
470 timers.close();
471 }
472 }
473}
474
475async fn changed(timers: Option<&ConnTimers>) {
477 match timers {
478 Some(timers) => timers.changed.notified().await,
479 None => std::future::pending().await,
480 }
481}
482
483const SHUTDOWN_GRACE: Duration = Duration::from_secs(5);
485
486const FRESH_CONN_GRACE: Duration = Duration::from_secs(2);
490
491async fn fresh_wait(first_request: &FirstRequest, until: Option<Instant>) {
493 let Some(until) = until else {
494 return std::future::pending().await;
495 };
496 if first_request.is_seen() {
497 return;
498 }
499 tokio::select! {
500 () = first_request.notify.notified() => {}
501 () = tokio::time::sleep_until(until) => {}
502 }
503}
504
505async fn sleep_until(deadline: Option<Instant>) {
506 match deadline {
507 Some(deadline) => tokio::time::sleep_until(deadline).await,
508 None => std::future::pending().await,
509 }
510}
511
512#[derive(Debug, Clone, Copy)]
515enum Phase {
516 Head { since: Instant },
517 Busy { in_flight: usize },
518 Idle { since: Instant },
519}
520
521#[derive(Debug, Clone, Copy, PartialEq, Eq)]
522enum Expiry {
523 Head,
524 Idle,
525}
526
527#[derive(Debug)]
528struct PhaseState {
529 phase: Phase,
530 scan: Scan,
531 header_block_since: Option<Instant>,
533}
534
535const H2_PREFACE: &[u8; 24] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n";
537const H2_FRAME_HEADERS: u8 = 0x1;
538const H2_FRAME_CONTINUATION: u8 = 0x9;
539const H2_FLAG_END_HEADERS: u8 = 0x4;
540
541#[derive(Debug)]
543enum Scan {
544 Preface {
546 matched: usize,
547 },
548 Http1(H1Scan),
549 Frames {
553 header: [u8; 9],
554 have: usize,
555 skip: usize,
556 ends_block: bool,
557 provisional: bool,
560 },
561}
562
563#[derive(Debug, Default)]
568struct H1Scan {
569 state: H1State,
570 line: LineScan,
572 content_length: u64,
573 chunked: bool,
574 position: HeadPosition,
575}
576
577#[derive(Debug, Default)]
581struct LineScan {
582 len: usize,
584 name: Vec<u8>,
587 in_value: bool,
588 field: Field,
589 number: Number,
591 in_extension: bool,
593 tail: Vec<u8>,
596}
597
598#[derive(Debug, Default, Clone, Copy)]
600enum Field {
601 #[default]
602 Other,
603 ContentLength,
604 TransferEncoding,
605}
606
607#[derive(Debug, Default)]
610struct Number {
611 item: Option<u64>,
612 first: Option<u64>,
613 bad: bool,
614}
615
616impl Number {
617 fn push(&mut self, byte: u8, radix: u32) {
618 match char::from(byte).to_digit(radix) {
619 Some(digit) => {
620 let next = self
621 .item
622 .unwrap_or(0)
623 .checked_mul(u64::from(radix))
624 .and_then(|v| v.checked_add(u64::from(digit)));
625 self.bad |= next.is_none();
626 self.item = next;
627 }
628 None if byte == b',' && radix == 10 => self.end_item(),
629 None => self.bad = true,
630 }
631 }
632
633 const fn end_item(&mut self) {
634 match (self.item.take(), self.first) {
635 (Some(item), None) => self.first = Some(item),
636 (Some(item), Some(first)) if item == first => {}
637 (Some(_), Some(_)) => self.bad = true,
638 (None, _) => {}
639 }
640 }
641
642 const fn finish(mut self) -> Option<u64> {
644 self.end_item();
645 if self.bad { None } else { self.first }
646 }
647}
648
649#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
651enum HeadPosition {
652 #[default]
655 Before,
656 RequestLine,
657 Fields,
658}
659
660const H1_NAME_KEEP: usize = 20;
662
663#[derive(Debug, Default, Clone, Copy)]
664enum H1State {
665 #[default]
666 Head,
667 Body(u64),
668 ChunkSize,
669 ChunkData(u64),
670 ChunkDataEnd(u8),
671 Trailers,
672 Stopped,
674}
675
676impl H1Scan {
677 fn scan(&mut self, mut bytes: &[u8], head_since: &mut Option<Instant>) {
678 while !bytes.is_empty() {
679 match self.state {
680 H1State::Stopped => return,
681 H1State::Body(n) | H1State::ChunkData(n) => {
682 let k = usize::try_from(n).unwrap_or(usize::MAX).min(bytes.len());
683 bytes = &bytes[k..];
684 let left = n - k as u64;
685 self.state = match (self.state, left) {
686 (H1State::Body(_), 0) => self.next_message(),
687 (H1State::Body(_), _) => H1State::Body(left),
688 (_, 0) => H1State::ChunkDataEnd(2),
689 _ => H1State::ChunkData(left),
690 };
691 }
692 H1State::ChunkDataEnd(n) => {
693 let k = usize::from(n).min(bytes.len());
694 bytes = &bytes[k..];
695 self.state = match n.saturating_sub(u8::try_from(k).unwrap_or(u8::MAX)) {
696 0 => H1State::ChunkSize,
697 left => H1State::ChunkDataEnd(left),
698 };
699 }
700 H1State::Head | H1State::ChunkSize | H1State::Trailers => {
701 let byte = bytes[0];
702 bytes = &bytes[1..];
703 if matches!(self.state, H1State::Head) && self.position == HeadPosition::Before
704 {
705 head_since.get_or_insert_with(Instant::now);
709 if byte == b'\r' || byte == b'\n' {
710 continue;
711 }
712 self.position = HeadPosition::RequestLine;
713 }
714 if byte == b'\n' {
715 self.on_line(head_since);
716 } else if byte != b'\r' {
717 self.on_line_byte(byte);
718 }
719 }
720 }
721 }
722 }
723
724 const fn stop(&mut self, head_since: &mut Option<Instant>) {
727 self.state = H1State::Stopped;
728 *head_since = None;
729 }
730
731 const fn next_message(&mut self) -> H1State {
732 self.content_length = 0;
733 self.chunked = false;
734 self.position = HeadPosition::Before;
735 H1State::Head
736 }
737
738 fn on_line_byte(&mut self, byte: u8) {
739 let line = &mut self.line;
740 line.len = line.len.saturating_add(1);
741 let blank = byte == b' ' || byte == b'\t';
742 match self.state {
743 H1State::Head if self.position == HeadPosition::Fields => {
744 if line.in_value {
745 match line.field {
746 _ if blank => {}
748 Field::ContentLength => line.number.push(byte, 10),
749 Field::TransferEncoding => {
750 if line.tail.len() == b"chunked".len() {
751 line.tail.remove(0);
752 }
753 line.tail.push(byte.to_ascii_lowercase());
754 }
755 Field::Other => {}
756 }
757 } else if byte == b':' {
758 line.in_value = true;
759 line.field = match line.name.as_slice() {
760 b"content-length" => Field::ContentLength,
761 b"transfer-encoding" => Field::TransferEncoding,
762 _ => Field::Other,
763 };
764 } else if line.name.len() <= H1_NAME_KEEP {
765 line.name.push(byte.to_ascii_lowercase());
766 }
767 }
768 H1State::ChunkSize if !line.in_extension && !blank => {
769 if byte == b';' {
770 line.in_extension = true;
771 } else {
772 line.number.push(byte, 16);
773 }
774 }
775 _ => {}
776 }
777 }
778
779 fn on_line(&mut self, head_since: &mut Option<Instant>) {
780 let line = std::mem::take(&mut self.line);
781 let empty = line.len == 0;
782 match self.state {
783 H1State::Head if empty => {
784 *head_since = None;
785 self.state = if self.chunked {
789 H1State::ChunkSize
790 } else if self.content_length > 0 {
791 H1State::Body(self.content_length)
792 } else {
793 self.next_message()
794 };
795 }
796 H1State::Head if self.position == HeadPosition::RequestLine => {
797 self.position = HeadPosition::Fields;
798 }
799 H1State::Head => match line.field {
800 Field::ContentLength => match line.number.finish() {
801 Some(len) => self.content_length = len,
802 None => self.stop(head_since),
804 },
805 Field::TransferEncoding => self.chunked = line.tail == b"chunked",
807 Field::Other => {}
808 },
809 H1State::ChunkSize => match line.number.finish() {
810 Some(0) => self.state = H1State::Trailers,
811 Some(n) => self.state = H1State::ChunkData(n),
812 None => self.stop(head_since),
813 },
814 H1State::Trailers if empty => self.state = self.next_message(),
815 _ => {}
816 }
817 }
818}
819
820impl PhaseState {
821 fn scan(&mut self, mut bytes: &[u8]) {
823 while !bytes.is_empty() {
824 match &mut self.scan {
825 Scan::Http1(h1) => {
826 h1.scan(bytes, &mut self.header_block_since);
827 return;
828 }
829 Scan::Preface { matched } => {
830 let rest = &H2_PREFACE[*matched..];
831 let n = rest.len().min(bytes.len());
832 if bytes[..n] != rest[..n] {
833 let mut h1 = H1Scan::default();
835 h1.scan(&H2_PREFACE[..*matched], &mut self.header_block_since);
836 h1.scan(bytes, &mut self.header_block_since);
837 self.scan = Scan::Http1(h1);
838 return;
839 }
840 *matched += n;
841 bytes = &bytes[n..];
842 if *matched == H2_PREFACE.len() {
843 self.scan = Scan::Frames {
844 header: [0; 9],
845 have: 0,
846 skip: 0,
847 ends_block: false,
848 provisional: false,
849 };
850 }
851 }
852 Scan::Frames {
853 header,
854 have,
855 skip,
856 ends_block,
857 provisional,
858 } => {
859 if *skip > 0 {
860 let n = (*skip).min(bytes.len());
861 *skip -= n;
862 bytes = &bytes[n..];
863 if *skip == 0 && std::mem::take(ends_block) {
865 self.header_block_since = None;
866 }
867 continue;
868 }
869 let n = (9 - *have).min(bytes.len());
870 header[*have..*have + n].copy_from_slice(&bytes[..n]);
871 *have += n;
872 bytes = &bytes[n..];
873 if *have < 9 {
874 let head = *have < 4
876 || header[3] == H2_FRAME_HEADERS
877 || header[3] == H2_FRAME_CONTINUATION;
878 if head && self.header_block_since.is_none() {
879 self.header_block_since = Some(Instant::now());
880 *provisional = true;
881 } else if !head && std::mem::take(provisional) {
882 self.header_block_since = None;
883 }
884 return;
885 }
886 *have = 0;
887 *skip = usize::from(header[0]) << 16
888 | usize::from(header[1]) << 8
889 | usize::from(header[2]);
890 let (kind, flags) = (header[3], header[4]);
891 if kind == H2_FRAME_HEADERS || kind == H2_FRAME_CONTINUATION {
892 *provisional = false;
893 self.header_block_since.get_or_insert_with(Instant::now);
894 *ends_block = flags & H2_FLAG_END_HEADERS != 0;
895 if *ends_block && *skip == 0 {
896 *ends_block = false;
897 self.header_block_since = None;
898 }
899 } else if std::mem::take(provisional) {
900 self.header_block_since = None;
901 }
902 }
903 }
904 }
905 }
906}
907
908#[derive(Debug)]
909struct ConnTimers {
910 closed: std::sync::atomic::AtomicBool,
913 state: Mutex<PhaseState>,
914 changed: Notify,
915 header_read_timeout: Option<Duration>,
916 keep_alive_timeout: Option<Duration>,
917}
918
919impl ConnTimers {
920 fn new(limits: &HttpLimits) -> Self {
921 Self {
922 closed: std::sync::atomic::AtomicBool::new(false),
923 state: Mutex::new(PhaseState {
924 phase: Phase::Head {
925 since: Instant::now(),
926 },
927 scan: Scan::Preface { matched: 0 },
928 header_block_since: None,
929 }),
930 changed: Notify::new(),
931 header_read_timeout: limits.header_read_timeout,
932 keep_alive_timeout: limits.keep_alive_timeout,
933 }
934 }
935
936 fn update(&self, f: impl FnOnce(&mut PhaseState)) {
937 {
938 let mut state = self
939 .state
940 .lock()
941 .unwrap_or_else(std::sync::PoisonError::into_inner);
942 f(&mut state);
943 }
944 self.changed.notify_one();
945 }
946
947 fn on_read(&self, bytes: &[u8]) {
950 if self.closed.load(std::sync::atomic::Ordering::Relaxed) {
951 return;
952 }
953 let mut state = self
954 .state
955 .lock()
956 .unwrap_or_else(std::sync::PoisonError::into_inner);
957 let block_before = state.header_block_since.is_some();
958 state.scan(bytes);
959 let changed = block_before != state.header_block_since.is_some();
960 drop(state);
961 if changed {
962 self.changed.notify_one();
963 }
964 }
965
966 fn close(&self) {
968 self.closed
969 .store(true, std::sync::atomic::Ordering::Relaxed);
970 }
971
972 fn request_started(&self) {
973 self.update(|state| {
974 state.phase = match state.phase {
975 Phase::Busy { in_flight } => Phase::Busy {
976 in_flight: in_flight + 1,
977 },
978 Phase::Head { .. } | Phase::Idle { .. } => Phase::Busy { in_flight: 1 },
979 };
980 });
981 }
982
983 fn request_finished(&self) {
984 self.update(|state| {
985 if let Phase::Busy { in_flight } = state.phase {
986 state.phase = if in_flight > 1 {
987 Phase::Busy {
988 in_flight: in_flight - 1,
989 }
990 } else {
991 Phase::Idle {
992 since: Instant::now(),
993 }
994 };
995 }
996 });
997 }
998
999 fn deadline(&self, grace_until: Option<Instant>) -> Option<Instant> {
1002 self.deadline_and_expiry()
1003 .map(|(deadline, expiry)| match (expiry, grace_until) {
1004 (Expiry::Idle, Some(grace)) => deadline.max(grace),
1005 _ => deadline,
1006 })
1007 }
1008
1009 fn deadline_and_expiry(&self) -> Option<(Instant, Expiry)> {
1010 let state = self
1011 .state
1012 .lock()
1013 .unwrap_or_else(std::sync::PoisonError::into_inner);
1014 let phase = match state.phase {
1015 Phase::Head { since } => self
1017 .header_read_timeout
1018 .or(self.keep_alive_timeout)
1019 .map(|timeout| (since + timeout, Expiry::Head)),
1020 Phase::Idle { .. }
1023 if state.header_block_since.is_some() && self.header_read_timeout.is_some() =>
1024 {
1025 None
1026 }
1027 Phase::Idle { since } => self
1028 .keep_alive_timeout
1029 .map(|timeout| (since + timeout, Expiry::Idle)),
1030 Phase::Busy { .. } => None,
1031 };
1032 let block = state
1034 .header_block_since
1035 .zip(self.header_read_timeout)
1036 .map(|(since, timeout)| (since + timeout, Expiry::Head));
1037 match (phase, block) {
1038 (Some(a), Some(b)) => Some(if b.0 < a.0 { b } else { a }),
1039 (a, b) => a.or(b),
1040 }
1041 }
1042
1043 fn expired(&self) -> Option<Expiry> {
1046 self.deadline_and_expiry()
1047 .filter(|(deadline, _)| *deadline <= Instant::now())
1048 .map(|(_, expiry)| expiry)
1049 }
1050}
1051
1052struct ConnIo<I> {
1054 inner: I,
1055 timers: Option<Arc<ConnTimers>>,
1056 _permit: Option<Slot>,
1057}
1058
1059#[derive(Default)]
1061struct FirstRequest {
1062 seen: std::sync::atomic::AtomicBool,
1063 notify: Notify,
1064}
1065
1066impl FirstRequest {
1067 fn mark(&self) {
1068 if !self.seen.swap(true, std::sync::atomic::Ordering::Relaxed) {
1069 self.notify.notify_one();
1070 }
1071 }
1072
1073 fn is_seen(&self) -> bool {
1074 self.seen.load(std::sync::atomic::Ordering::Relaxed)
1075 }
1076}
1077
1078impl<I: AsyncRead + Unpin> AsyncRead for ConnIo<I> {
1079 fn poll_read(
1080 mut self: Pin<&mut Self>,
1081 cx: &mut Context<'_>,
1082 buf: &mut ReadBuf<'_>,
1083 ) -> Poll<io::Result<()>> {
1084 let before = buf.filled().len();
1085 let result = Pin::new(&mut self.inner).poll_read(cx, buf);
1086 if buf.filled().len() > before
1087 && let Some(timers) = &self.timers
1088 {
1089 timers.on_read(&buf.filled()[before..]);
1090 }
1091 result
1092 }
1093}
1094
1095impl<I: AsyncWrite + Unpin> AsyncWrite for ConnIo<I> {
1096 fn poll_write(
1097 mut self: Pin<&mut Self>,
1098 cx: &mut Context<'_>,
1099 buf: &[u8],
1100 ) -> Poll<io::Result<usize>> {
1101 Pin::new(&mut self.inner).poll_write(cx, buf)
1102 }
1103
1104 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
1105 Pin::new(&mut self.inner).poll_flush(cx)
1106 }
1107
1108 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
1109 Pin::new(&mut self.inner).poll_shutdown(cx)
1110 }
1111
1112 fn poll_write_vectored(
1113 mut self: Pin<&mut Self>,
1114 cx: &mut Context<'_>,
1115 bufs: &[io::IoSlice<'_>],
1116 ) -> Poll<io::Result<usize>> {
1117 Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
1118 }
1119
1120 fn is_write_vectored(&self) -> bool {
1121 self.inner.is_write_vectored()
1122 }
1123}
1124
1125#[derive(Clone)]
1127struct Tracked<S> {
1128 inner: S,
1129 timers: Option<Arc<ConnTimers>>,
1130 first_request: Arc<FirstRequest>,
1131}
1132
1133struct InFlight(Arc<ConnTimers>);
1134
1135impl Drop for InFlight {
1136 fn drop(&mut self) {
1137 self.0.request_finished();
1138 }
1139}
1140
1141#[derive(Clone)]
1148pub(crate) struct TunnelGuard {
1149 _mark: Arc<InFlight>,
1150}
1151
1152pub struct KeepTunnel {
1177 guard: Option<TunnelGuard>,
1178}
1179
1180impl std::fmt::Debug for KeepTunnel {
1181 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1182 f.debug_struct("KeepTunnel")
1183 .field("holds_tunnel", &self.guard.is_some())
1184 .finish()
1185 }
1186}
1187
1188impl<S: Send + Sync> axum::extract::FromRequestParts<S> for KeepTunnel {
1189 type Rejection = Infallible;
1190
1191 async fn from_request_parts(
1192 parts: &mut axum::http::request::Parts,
1193 _state: &S,
1194 ) -> Result<Self, Infallible> {
1195 Ok(Self {
1196 guard: parts.extensions.remove::<TunnelGuard>(),
1197 })
1198 }
1199}
1200
1201impl<S> tower::Service<axum::extract::Request> for Tracked<S>
1202where
1203 S: tower::Service<
1204 axum::extract::Request,
1205 Response = axum::response::Response,
1206 Error = Infallible,
1207 >,
1208{
1209 type Response = axum::response::Response;
1210 type Error = Infallible;
1211 type Future = TrackedFuture<S::Future>;
1212
1213 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
1214 self.first_request.mark();
1217 self.inner.poll_ready(cx)
1218 }
1219
1220 fn call(&mut self, mut req: axum::extract::Request) -> Self::Future {
1221 self.first_request.mark();
1222 let guard = self.timers.as_ref().map(|timers| {
1223 if req.method() == axum::http::Method::CONNECT {
1227 timers.request_started();
1228 req.extensions_mut().insert(TunnelGuard {
1229 _mark: Arc::new(InFlight(Arc::clone(timers))),
1230 });
1231 }
1232 timers.request_started();
1233 InFlight(Arc::clone(timers))
1234 });
1235 TrackedFuture {
1236 inner: self.inner.call(req),
1237 guard,
1238 }
1239 }
1240}
1241
1242pin_project_lite::pin_project! {
1243 struct TrackedFuture<F> {
1246 #[pin]
1247 inner: F,
1248 guard: Option<InFlight>,
1249 }
1250}
1251
1252impl<F> std::future::Future for TrackedFuture<F>
1253where
1254 F: std::future::Future<Output = Result<axum::response::Response, Infallible>>,
1255{
1256 type Output = Result<axum::response::Response, Infallible>;
1257
1258 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1259 let this = self.project();
1260 let response = std::task::ready!(this.inner.poll(cx))?;
1261 Poll::Ready(Ok(match this.guard.take() {
1262 Some(guard) => response.map(|body| {
1263 axum::body::Body::new(TrackedBody {
1264 inner: body,
1265 _guard: guard,
1266 })
1267 }),
1268 None => response,
1269 }))
1270 }
1271}
1272
1273struct TrackedBody {
1275 inner: axum::body::Body,
1276 _guard: InFlight,
1277}
1278
1279impl http_body::Body for TrackedBody {
1280 type Data = axum::body::Bytes;
1281 type Error = axum::Error;
1282
1283 fn poll_frame(
1284 mut self: Pin<&mut Self>,
1285 cx: &mut Context<'_>,
1286 ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
1287 Pin::new(&mut self.inner).poll_frame(cx)
1288 }
1289
1290 fn is_end_stream(&self) -> bool {
1291 self.inner.is_end_stream()
1292 }
1293
1294 fn size_hint(&self) -> http_body::SizeHint {
1295 self.inner.size_hint()
1296 }
1297}
1298
1299#[cfg(test)]
1300mod tests {
1301 use super::*;
1302
1303 fn state() -> PhaseState {
1304 PhaseState {
1305 phase: Phase::Busy { in_flight: 0 },
1306 scan: Scan::Preface { matched: 0 },
1307 header_block_since: None,
1308 }
1309 }
1310
1311 fn frame(kind: u8, flags: u8, payload: &[u8]) -> Vec<u8> {
1312 let len = u8::try_from(payload.len()).expect("small test payload");
1313 let mut out = vec![0, 0, len, kind, flags, 0, 0, 0, 1];
1314 out.extend_from_slice(payload);
1315 out
1316 }
1317
1318 #[test]
1319 fn http1_bytes_end_the_scan() {
1320 let mut state = state();
1321 state.scan(b"GET / HTTP/1.1\r\n");
1322 assert!(matches!(state.scan, Scan::Http1(_)));
1323 assert!(!matches!(state.scan, Scan::Frames { .. }));
1324 }
1325
1326 #[test]
1327 fn http1_head_opens_at_the_first_request_and_closes_at_the_blank_line() {
1328 let mut state = state();
1329 state.scan(b"POST /x HTTP/1.1\r\nHost: t\r\n");
1330 assert!(state.header_block_since.is_some());
1331 state.scan(b"Content-Length: 12\r\n\r\n");
1332 assert!(state.header_block_since.is_none(), "head complete");
1333 state.scan(b"GET / HTTP/1"); assert!(state.header_block_since.is_none());
1335 state.scan(b"GET /next");
1336 assert!(state.header_block_since.is_some(), "a pipelined head");
1337 }
1338
1339 #[test]
1340 fn http1_a_crlf_between_messages_starts_the_head_timer() {
1341 let mut state = state();
1342 state.scan(b"GET / HTTP/1.1\r\nHost: t\r\n\r\n");
1343 assert!(state.header_block_since.is_none(), "head complete");
1344 state.scan(b"\r\n");
1347 let since = state
1348 .header_block_since
1349 .expect("the CRLF starts the head timer");
1350 state.scan(b"GET /2 HTTP/1.1\r\n");
1351 assert_eq!(
1352 state.header_block_since,
1353 Some(since),
1354 "more bytes do not restart it"
1355 );
1356 state.scan(b"Host: t\r\n\r\n");
1357 assert!(state.header_block_since.is_none(), "head complete");
1358 }
1359
1360 #[test]
1361 fn http1_pipelined_partial_head_in_one_read_is_open() {
1362 let mut state = state();
1363 state.scan(b"GET / HTTP/1.1\r\nHost: t\r\n\r\nGET /2 HTTP/1.1\r\nHo");
1364 assert!(state.header_block_since.is_some());
1365 }
1366
1367 #[test]
1368 fn http1_chunked_body_is_skipped() {
1369 let mut state = state();
1370 state.scan(b"POST / HTTP/1.1\r\nTransfer-Encoding: chunked\r\n\r\n");
1371 state.scan(b"5\r\nGET /\r\n0\r\n\r\n");
1372 assert!(
1373 state.header_block_since.is_none(),
1374 "chunk data is not a head"
1375 );
1376 state.scan(b"G");
1377 assert!(state.header_block_since.is_some());
1378 }
1379
1380 #[test]
1381 fn http1_long_whitespace_in_a_field_keeps_the_length() {
1382 let mut state = state();
1383 let ows = " ".repeat(300);
1384 state.scan(format!("POST / HTTP/1.1\r\nContent-Length:{ows}5\r\n\r\n").as_bytes());
1385 assert!(state.header_block_since.is_none(), "head complete");
1386 state.scan(b"GET /"); assert!(state.header_block_since.is_none());
1388 state.scan(b"G");
1389 assert!(state.header_block_since.is_some(), "the next head is timed");
1390 }
1391
1392 fn assert_frames(head: &[u8], body: &[u8]) {
1395 let mut state = state();
1396 state.scan(head);
1397 assert!(state.header_block_since.is_none(), "head complete");
1398 state.scan(body);
1399 assert!(state.header_block_since.is_none(), "the body is not a head");
1400 state.scan(b"G");
1401 assert!(state.header_block_since.is_some(), "the next head is timed");
1402 }
1403
1404 #[test]
1405 fn http1_long_valid_framing_values_still_frame() {
1406 let zeros = "0".repeat(300);
1407 assert_frames(
1408 format!("POST / HTTP/1.1\r\nContent-Length: {zeros}5\r\n\r\n").as_bytes(),
1409 b"GET /",
1410 );
1411 let codings = "gzip, ".repeat(60);
1412 assert_frames(
1413 format!("POST / HTTP/1.1\r\nTransfer-Encoding: {codings}chunked\r\n\r\n").as_bytes(),
1414 b"5\r\nGET /\r\n0\r\n\r\n",
1415 );
1416 let ext = "x".repeat(300);
1417 assert_frames(
1418 b"POST / HTTP/1.1\r\nTransfer-Encoding: chunked\r\n\r\n",
1419 format!("5;{ext}\r\nGET /\r\n0\r\n\r\n").as_bytes(),
1420 );
1421 assert_frames(b"POST / HTTP/1.1\r\nContent-Length: 5, 5\r\n\r\n", b"GET /");
1422 }
1423
1424 #[test]
1425 fn http1_chunked_must_be_the_last_coding() {
1426 let mut scan = H1Scan::default();
1428 let mut since = None;
1429 scan.scan(
1430 b"POST / HTTP/1.1\r\nTransfer-Encoding: chunked, gzip\r\n\r\n",
1431 &mut since,
1432 );
1433 assert!(!scan.chunked);
1434 }
1435
1436 #[test]
1437 fn http1_a_value_the_scan_cannot_read_stops_head_timing() {
1438 let mut state = state();
1439 state.scan(b"POST / HTTP/1.1\r\nContent-Length: abc\r\n");
1440 assert!(state.header_block_since.is_none());
1442 state.scan(b"\r\nGET /");
1443 assert!(state.header_block_since.is_none());
1444 }
1445
1446 #[test]
1447 fn a_rejected_http1_upgrade_keeps_the_scan() {
1448 let mut state = state();
1449 state.scan(b"GET /ws HTTP/1.1\r\nUpgrade: websocket\r\n\r\n");
1450 state.scan(b"GET /next");
1452 assert!(state.header_block_since.is_some(), "the next head is timed");
1453 }
1454
1455 #[test]
1456 fn a_closed_connection_stops_reading_bytes() {
1457 let timers = ConnTimers::new(&HttpLimits {
1458 header_read_timeout: Some(Duration::from_secs(10)),
1459 ..HttpLimits::default()
1460 });
1461 timers.request_started();
1462 timers.close();
1463 timers.on_read(b"GET /not-a-head");
1465 assert!(
1466 timers.state.lock().unwrap().header_block_since.is_none(),
1467 "a closed connection ignores the bytes"
1468 );
1469 }
1470
1471 #[tokio::test]
1472 async fn open_header_block_is_tracked_across_split_reads() {
1473 let mut state = state();
1474 let mut bytes = H2_PREFACE.to_vec();
1475 bytes.extend(frame(0x4, 0, &[0; 6])); bytes.extend(frame(H2_FRAME_HEADERS, 0, &[0x82, 0x86])); for chunk in bytes.chunks(5) {
1478 state.scan(chunk);
1479 }
1480 assert!(matches!(state.scan, Scan::Frames { .. }));
1481 assert!(state.header_block_since.is_some(), "block is open");
1482
1483 state.scan(&frame(H2_FRAME_CONTINUATION, H2_FLAG_END_HEADERS, &[0x84]));
1484 assert!(state.header_block_since.is_none(), "END_HEADERS closes it");
1485
1486 state.scan(&frame(H2_FRAME_HEADERS, H2_FLAG_END_HEADERS | 0x1, &[0x82]));
1487 assert!(
1488 state.header_block_since.is_none(),
1489 "a full block stays closed"
1490 );
1491 }
1492
1493 #[tokio::test(start_paused = true)]
1494 async fn an_open_header_block_suspends_the_idle_deadline() {
1495 let timers = ConnTimers::new(&HttpLimits {
1496 header_read_timeout: Some(Duration::from_secs(10)),
1497 keep_alive_timeout: Some(Duration::from_secs(75)),
1498 ..HttpLimits::default()
1499 });
1500 timers.on_read(H2_PREFACE);
1501 timers.request_started();
1502 timers.request_finished();
1503 tokio::time::advance(Duration::from_secs(74)).await;
1504
1505 let full = frame(H2_FRAME_HEADERS, H2_FLAG_END_HEADERS, &[0x82, 0x86, 0x84]);
1508 timers.on_read(&full[..9]);
1509 let (deadline, expiry) = timers.deadline_and_expiry().expect("a deadline");
1510 assert_eq!(expiry, Expiry::Head);
1511 assert_eq!(deadline, Instant::now() + Duration::from_secs(10));
1512
1513 timers.on_read(&full[9..]);
1514 let (_, expiry) = timers.deadline_and_expiry().expect("idle again");
1515 assert_eq!(expiry, Expiry::Idle);
1516 }
1517
1518 fn connect_request() -> axum::extract::Request {
1519 axum::http::Request::builder()
1520 .method(axum::http::Method::CONNECT)
1521 .uri("/ws")
1522 .body(axum::body::Body::empty())
1523 .unwrap()
1524 }
1525
1526 #[test]
1527 fn a_huge_connection_cap_does_not_panic() {
1528 let limits = HttpLimits {
1531 max_connections: Some(usize::MAX),
1532 ..HttpLimits::default()
1533 };
1534 let slots = limits.slots(false).expect("a cap is set");
1535 assert_eq!(slots.available_permits(), Semaphore::MAX_PERMITS);
1536 }
1537
1538 #[test]
1539 fn a_small_header_limit_does_not_panic() {
1540 let limits = HttpLimits {
1544 max_header_bytes: Some(4096),
1545 ..HttpLimits::default()
1546 };
1547 let _ = limits.builder();
1548 }
1549
1550 fn idle_timers() -> Arc<ConnTimers> {
1551 let timers = Arc::new(ConnTimers::new(&HttpLimits {
1552 keep_alive_timeout: Some(Duration::from_secs(60)),
1553 ..HttpLimits::default()
1554 }));
1555 timers.request_started();
1556 timers.request_finished();
1557 timers
1558 }
1559
1560 fn expiry(timers: &ConnTimers) -> Option<Expiry> {
1561 timers.deadline_and_expiry().map(|(_, expiry)| expiry)
1562 }
1563
1564 #[tokio::test]
1565 async fn a_request_waiting_for_service_readiness_counts_as_a_request() {
1566 use tower::Service as _;
1567 let first_request = Arc::new(FirstRequest::default());
1568 let mut service = Tracked {
1569 first_request: Arc::clone(&first_request),
1570 inner: tower::service_fn(|_req: axum::extract::Request| async {
1571 Ok::<_, Infallible>(axum::response::Response::new(axum::body::Body::empty()))
1572 }),
1573 timers: None,
1574 };
1575 std::future::poll_fn(|cx| service.poll_ready(cx))
1577 .await
1578 .unwrap();
1579 assert!(first_request.is_seen());
1580 }
1581
1582 #[tokio::test]
1583 async fn a_rejected_connect_releases_its_in_flight_mark() {
1584 use tower::Service as _;
1585 let timers = idle_timers();
1586 let mut service = Tracked {
1587 first_request: Arc::default(),
1588 inner: tower::service_fn(|_req: axum::extract::Request| async {
1589 Ok::<_, Infallible>(
1590 axum::response::Response::builder()
1591 .status(404)
1592 .body(axum::body::Body::empty())
1593 .unwrap(),
1594 )
1595 }),
1596 timers: Some(Arc::clone(&timers)),
1597 };
1598 drop(service.call(connect_request()).await.unwrap());
1599 assert_eq!(expiry(&timers), Some(Expiry::Idle));
1600 }
1601
1602 #[tokio::test]
1603 async fn a_kept_tunnel_guard_holds_the_connection_until_dropped() {
1604 use tower::Service as _;
1605 let timers = idle_timers();
1606 let (tx, rx) = std::sync::mpsc::channel::<TunnelGuard>();
1607 let mut service = Tracked {
1608 first_request: Arc::default(),
1609 inner: tower::service_fn(move |mut req: axum::extract::Request| {
1610 let guard = req.extensions_mut().remove::<TunnelGuard>();
1612 tx.send(guard.expect("CONNECT carries a guard")).unwrap();
1613 async {
1614 Ok::<_, Infallible>(axum::response::Response::new(axum::body::Body::empty()))
1615 }
1616 }),
1617 timers: Some(Arc::clone(&timers)),
1618 };
1619 drop(service.call(connect_request()).await.unwrap());
1620 let guard = rx.recv().unwrap();
1621 assert_eq!(expiry(&timers), None, "the tunnel is still open");
1622 drop(guard);
1623 assert_eq!(expiry(&timers), Some(Expiry::Idle));
1624 }
1625
1626 #[tokio::test]
1627 async fn keep_tunnel_holds_the_connection_for_a_raw_upgrade() {
1628 use axum::extract::FromRequestParts as _;
1629 use tower::Service as _;
1630 let timers = idle_timers();
1631 let (tx, rx) = std::sync::mpsc::channel::<KeepTunnel>();
1632 let mut service = Tracked {
1633 first_request: Arc::default(),
1634 inner: tower::service_fn(move |req: axum::extract::Request| {
1635 let tx = tx.clone();
1636 async move {
1637 let (mut parts, _) = req.into_parts();
1638 let keep = KeepTunnel::from_request_parts(&mut parts, &())
1639 .await
1640 .unwrap();
1641 tx.send(keep).unwrap();
1642 Ok::<_, Infallible>(axum::response::Response::new(axum::body::Body::empty()))
1643 }
1644 }),
1645 timers: Some(Arc::clone(&timers)),
1646 };
1647 drop(service.call(connect_request()).await.unwrap());
1648 let keep = rx.recv().unwrap();
1649 assert_eq!(expiry(&timers), None, "the raw tunnel is still open");
1650 drop(keep);
1651 assert_eq!(expiry(&timers), Some(Expiry::Idle));
1652 }
1653
1654 #[tokio::test]
1655 async fn a_connect_response_body_stays_in_flight() {
1656 use tower::Service as _;
1657 let timers = idle_timers();
1658 let mut service = Tracked {
1659 first_request: Arc::default(),
1660 inner: tower::service_fn(|_req: axum::extract::Request| async {
1663 Ok::<_, Infallible>(axum::response::Response::new(axum::body::Body::from(
1664 "tunnel bytes",
1665 )))
1666 }),
1667 timers: Some(Arc::clone(&timers)),
1668 };
1669 let response = service.call(connect_request()).await.unwrap();
1670 assert_eq!(expiry(&timers), None, "the response body is still open");
1671 drop(response);
1672 assert_eq!(expiry(&timers), Some(Expiry::Idle));
1673 }
1674
1675 #[tokio::test]
1676 async fn a_partial_frame_header_starts_the_head_timer() {
1677 let mut state = state();
1678 state.scan(H2_PREFACE);
1679 let headers = frame(H2_FRAME_HEADERS, 0, &[0x82]);
1680 state.scan(&headers[..3]);
1681 let since = state
1682 .header_block_since
1683 .expect("timer starts at the first byte");
1684 state.scan(&headers[3..]);
1685 assert_eq!(
1686 state.header_block_since,
1687 Some(since),
1688 "a HEADERS frame keeps the time of its first byte"
1689 );
1690
1691 let mut state = state_after_preface();
1692 let ping = frame(0x6, 0, &[0; 8]);
1693 state.scan(&ping[..3]); assert!(state.header_block_since.is_some());
1695 state.scan(&ping[3..]);
1696 assert!(
1697 state.header_block_since.is_none(),
1698 "control traffic cancels the timer"
1699 );
1700 }
1701
1702 #[test]
1703 fn a_known_control_frame_type_cancels_the_provisional_timer() {
1704 let mut state = state_after_preface();
1705 let ping = frame(0x6, 0, &[0; 8]); state.scan(&ping[..4]); assert!(
1708 state.header_block_since.is_none(),
1709 "a PING is not a head, so no head timer"
1710 );
1711 let mut state = state_after_preface();
1712 let headers = frame(H2_FRAME_HEADERS, H2_FLAG_END_HEADERS, &[0x82]);
1713 state.scan(&headers[..4]);
1714 assert!(
1715 state.header_block_since.is_some(),
1716 "a HEADERS frame is a head"
1717 );
1718 }
1719
1720 fn state_after_preface() -> PhaseState {
1721 let mut state = state();
1722 state.scan(H2_PREFACE);
1723 state
1724 }
1725
1726 #[tokio::test]
1727 async fn end_headers_closes_the_block_only_after_its_payload() {
1728 let mut state = state();
1729 state.scan(H2_PREFACE);
1730 let full = frame(H2_FRAME_HEADERS, H2_FLAG_END_HEADERS, &[0x82, 0x86, 0x84]);
1731 state.scan(&full[..10]);
1733 assert!(
1734 state.header_block_since.is_some(),
1735 "the head is not complete until the payload is in"
1736 );
1737 state.scan(&full[10..]);
1738 assert!(state.header_block_since.is_none());
1739 }
1740
1741 async fn start_server() -> (
1744 std::net::SocketAddr,
1745 tokio_util::sync::CancellationToken,
1746 tokio::task::JoinHandle<io::Result<()>>,
1747 ) {
1748 let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
1749 .await
1750 .expect("bind");
1751 let addr = listener.local_addr().expect("addr");
1752 let stop = tokio_util::sync::CancellationToken::new();
1753 let router = axum::Router::new().route("/", axum::routing::get(|| async { "ok" }));
1754 let signal = {
1755 let stop = stop.clone();
1756 async move { stop.cancelled().await }
1757 };
1758 let server = tokio::spawn(serve(
1759 listener,
1760 router.into_make_service(),
1761 HttpLimits::default(),
1762 signal,
1763 ));
1764 (addr, stop, server)
1765 }
1766
1767 #[tokio::test]
1768 async fn a_fresh_connection_is_served_when_its_request_is_late() {
1769 use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
1770 let (addr, stop, server) = start_server().await;
1771 let mut client = tokio::net::TcpStream::connect(addr).await.expect("connect");
1772 tokio::time::sleep(Duration::from_millis(30)).await;
1773 stop.cancel();
1774 tokio::time::sleep(Duration::from_millis(400)).await;
1776 let _ = client
1777 .write_all(b"GET / HTTP/1.1\r\nHost: t\r\nConnection: close\r\n\r\n")
1778 .await;
1779 let mut raw = Vec::new();
1780 let _ = client.read_to_end(&mut raw).await;
1781 server.await.expect("task").expect("serve");
1782 let text = String::from_utf8_lossy(&raw);
1783 assert!(text.starts_with("HTTP/1.1 200"), "got {text:?}");
1784 }
1785
1786 #[tokio::test]
1787 async fn a_partial_request_does_not_hold_the_drain_past_the_grace() {
1788 use tokio::io::AsyncWriteExt as _;
1789 let (addr, stop, server) = start_server().await;
1790 let mut client = tokio::net::TcpStream::connect(addr).await.expect("connect");
1791 tokio::time::sleep(Duration::from_millis(30)).await;
1792 stop.cancel();
1793 client.write_all(b"GET / HT").await.expect("write");
1795 let done = tokio::time::timeout(FRESH_CONN_GRACE + Duration::from_secs(3), server).await;
1796 assert!(
1797 done.is_ok(),
1798 "an incomplete request must not block the drain"
1799 );
1800 }
1801
1802 #[tokio::test]
1803 async fn a_silent_connection_does_not_hold_the_drain_past_the_grace() {
1804 let (addr, stop, server) = start_server().await;
1805 let _client = tokio::net::TcpStream::connect(addr).await.expect("connect");
1806 tokio::time::sleep(Duration::from_millis(30)).await;
1807 stop.cancel();
1808 let done = tokio::time::timeout(FRESH_CONN_GRACE + Duration::from_secs(3), server).await;
1809 assert!(
1810 done.is_ok(),
1811 "the drain must end after the fresh-connection grace"
1812 );
1813 }
1814
1815 fn start_with(
1817 listener: tokio::net::TcpListener,
1818 limits: HttpLimits,
1819 ) -> (
1820 tokio::sync::oneshot::Sender<()>,
1821 tokio::task::JoinHandle<io::Result<()>>,
1822 ) {
1823 let (stop, stopped) = tokio::sync::oneshot::channel::<()>();
1824 let router = axum::Router::new().route("/", axum::routing::get(|| async { "ok" }));
1825 let server = tokio::spawn(serve(
1826 listener,
1827 router.into_make_service(),
1828 limits,
1829 async move {
1830 let _ = stopped.await;
1831 },
1832 ));
1833 (stop, server)
1834 }
1835
1836 #[tokio::test]
1838 async fn the_head_timer_bounds_the_fresh_wait() {
1839 let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
1840 .await
1841 .expect("bind");
1842 let addr = listener.local_addr().expect("addr");
1843 let limits = HttpLimits {
1844 header_read_timeout: Some(Duration::from_millis(150)),
1845 ..HttpLimits::default()
1846 };
1847 let (stop, server) = start_with(listener, limits);
1848 let _client = tokio::net::TcpStream::connect(addr).await.expect("connect");
1849 tokio::time::sleep(Duration::from_millis(30)).await;
1850 let started = Instant::now();
1851 let _ = stop.send(());
1852 tokio::time::timeout(Duration::from_secs(3), server)
1853 .await
1854 .expect("the drain must end")
1855 .expect("task")
1856 .expect("serve");
1857 assert!(
1858 started.elapsed() < FRESH_CONN_GRACE,
1859 "the head timer must end the wait before the grace, took {:?}",
1860 started.elapsed()
1861 );
1862 }
1863
1864 #[tokio::test]
1866 async fn an_h2_preface_alone_does_not_start_the_drain() {
1867 use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
1868
1869 let (addr, stop, server) = start_server().await;
1870 let mut client = tokio::net::TcpStream::connect(addr).await.expect("connect");
1871 client.write_all(H2_PREFACE).await.expect("preface");
1872 tokio::time::sleep(Duration::from_millis(30)).await;
1873 stop.cancel();
1874 tokio::time::sleep(Duration::from_millis(400)).await;
1875 let settings = [0, 0, 0, 0x4, 0, 0, 0, 0, 0]; let headers = frame(
1878 H2_FRAME_HEADERS,
1879 H2_FLAG_END_HEADERS | 0x1,
1880 &[0x82, 0x86, 0x84, 0x41, 0x01, b't'],
1881 );
1882 let _ = client.write_all(&settings).await;
1883 let _ = client.write_all(&headers).await;
1884 let mut raw = Vec::new();
1885 let _ = tokio::time::timeout(Duration::from_secs(2), client.read_to_end(&mut raw)).await;
1886 let mut types = Vec::new();
1889 let mut i = 0;
1890 while i + 9 <= raw.len() {
1891 let len = (usize::from(raw[i]) << 16)
1892 | (usize::from(raw[i + 1]) << 8)
1893 | usize::from(raw[i + 2]);
1894 types.push(raw[i + 3]);
1895 i += 9 + len;
1896 }
1897 let _ = tokio::time::timeout(Duration::from_secs(3), server).await;
1898 assert!(
1899 types.contains(&H2_FRAME_HEADERS),
1900 "no response, frame types {types:?}"
1901 );
1902 }
1903
1904 #[tokio::test]
1906 async fn a_stuck_make_service_does_not_hang_the_drain() {
1907 let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
1908 .await
1909 .expect("bind");
1910 let addr = listener.local_addr().expect("addr");
1911 let (stop, stopped) = tokio::sync::oneshot::channel::<()>();
1912 let stuck = tower::service_fn(|_: IncomingStream<'_, tokio::net::TcpListener>| {
1913 std::future::pending::<Result<axum::Router, Infallible>>()
1914 });
1915 let server = tokio::spawn(serve(listener, stuck, HttpLimits::default(), async move {
1916 let _ = stopped.await;
1917 }));
1918 let _client = tokio::net::TcpStream::connect(addr).await.expect("connect");
1919 tokio::time::sleep(Duration::from_millis(30)).await;
1920 let _ = stop.send(());
1921 tokio::time::timeout(FRESH_CONN_GRACE + Duration::from_secs(3), server)
1922 .await
1923 .expect("the drain must end")
1924 .expect("task")
1925 .expect("serve");
1926 }
1927}