Skip to main content

autumn_web/
http_server.rs

1//! HTTP serve loop with connection limits (issue #3065).
2//!
3//! [`serve`] replaces `axum::serve`. axum gives no access to the hyper
4//! builder. Thus axum cannot set a header-read timeout, a header size limit
5//! or an HTTP/2 stream limit. This loop applies [`HttpLimits`] from
6//! `[server.http]`. With default limits it behaves like `axum::serve`.
7//!
8//! # Timers
9//!
10//! Each connection is in one of three phases:
11//!
12//! - **Head**: the connection opened and no request arrived yet.
13//!   `header_read_timeout` applies (if it is not set, `keep_alive_timeout`).
14//! - **Busy**: at least one request is in flight. No timer applies. The
15//!   request timeout layer limits handlers.
16//! - **Idle**: no request in flight. `keep_alive_timeout` applies.
17//!
18//! In every phase, an open head also starts `header_read_timeout`, and
19//! suspends the idle timer. The timer starts at the first byte of the head.
20//! More bytes do not restart it, so a slowloris client cannot keep the
21//! connection open. To find heads, the read path scans the client bytes:
22//!
23//! - HTTP/1: a small framer skips bodies (`Content-Length`, chunked). A head
24//!   runs from its first byte to its blank line, also when it is pipelined
25//!   behind another request. CRLF between messages does not open a head.
26//! - HTTP/2: a head is a header block (`HEADERS` and `CONTINUATION` frames,
27//!   up to the payload of the `END_HEADERS` frame). Control frames (for
28//!   example `PING`) do not open one.
29//!
30//! The timers also run during a graceful drain.
31//!
32//! # Drain of a fresh connection
33//!
34//! A connection that has received no request is not closed at the drain signal. Its
35//! request can be in flight, and a close would reset it. The server waits up
36//! to `FRESH_CONN_GRACE` (2 s) for the first request, or until the head timer
37//! ends. Then it starts the drain. Setup of an accepted connection (the
38//! make-service) gets the same grace after a shutdown, so a stuck one cannot
39//! hang the drain.
40
41use 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/// Resolved `[server.http]` limits. `0` in the config means "off".
59#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
60pub struct HttpLimits {
61    /// Maximum time to receive a full request head.
62    pub header_read_timeout: Option<Duration>,
63    /// Maximum idle time with no request in flight.
64    pub keep_alive_timeout: Option<Duration>,
65    /// Maximum request head size, in bytes. A value below
66    /// [`HttpServerConfig::MIN_HEADER_BYTES`] is raised to it.
67    pub max_header_bytes: Option<usize>,
68    /// Maximum concurrent HTTP/2 streams per connection.
69    pub http2_max_concurrent_streams: Option<u32>,
70    /// Maximum open connections.
71    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            // hyper panics below the minimum. `validate` rejects it; clamp
81            // for a config that skipped `validate`.
82            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            // `validate` rejects a value above the semaphore limit; clamp
88            // for a config that skipped `validate`.
89            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        // CONNECT protocol for HTTP/2 WebSockets, as `axum::serve` sets it.
105        builder.http2().enable_connect_protocol();
106        if let Some(bytes) = self.max_header_bytes {
107            // hyper panics below the minimum. The config path clamps already;
108            // a value built by hand skipped that.
109            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
122/// One accepted connection, given to the make-service.
123///
124/// The same role as `axum::serve::IncomingStream`, which has no public
125/// constructor.
126pub struct IncomingStream<'a, L: Listener> {
127    io: &'a L::Io,
128    remote_addr: L::Addr,
129}
130
131impl<L: Listener> IncomingStream<'_, L> {
132    /// The accepted IO.
133    #[must_use]
134    pub const fn io(&self) -> &L::Io {
135        self.io
136    }
137
138    /// The peer address.
139    #[must_use]
140    pub const fn remote_addr(&self) -> &L::Addr {
141        &self.remote_addr
142    }
143}
144
145/// Serve `make_service` on `listener` until `signal` resolves, then drain.
146///
147/// After `signal`, the loop stops accepting, asks every connection to shut
148/// down gracefully, and returns when all connections have closed.
149///
150/// With `max_connections` set and a `tls::TlsListener` (the `tls` feature),
151/// a connection counts from TCP accept, so the cap also covers TLS
152/// handshakes in flight.
153///
154/// # Errors
155///
156/// Never returns an error today. The `io::Result` matches `axum::serve`.
157pub 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/// A `tls::TlsListener` takes the `max_connections` slots before TCP
185/// accept, as `App::run` sets it up, so a connection counts during its
186/// handshake too. Other listeners leave the slots to the serve loop.
187#[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
200/// Where a listener that accepts in the background (TLS) puts the
201/// `max_connections` slot of the stream it returns. The listener takes the
202/// slot before TCP accept, so the serve loop does not take one.
203pub(crate) type SlotHandoff = Arc<Mutex<Option<Slot>>>;
204
205/// A held `max_connections` slot. Dropping it frees the slot. The permit
206/// type is hidden, so a type that holds a slot is not a "significant drop"
207/// type for clippy in user code (for example `TlsListener`).
208pub(crate) type Slot = Box<dyn Send + Sync>;
209
210impl HttpLimits {
211    /// The `max_connections` slots, unless the listener takes them.
212    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
219/// Wait for a free connection slot. `None` when `closed` resolves first.
220async 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
235/// How a step of setting up a connection ended.
236enum Step<T> {
237    Done(T),
238    TimedOut,
239}
240
241/// Run `step`, unless `deadline` passes first. A shutdown does not stop it:
242/// the connection is accepted, so its request must be served. But after a
243/// shutdown, `step` gets only `FRESH_CONN_GRACE` more, so a stuck one cannot
244/// hang the drain.
245async 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
260/// The slot the listener put in `handoff` for the stream it returned.
261fn take_handoff(handoff: Option<&SlotHandoff>) -> Option<Slot> {
262    handoff?
263        .lock()
264        .unwrap_or_else(std::sync::PoisonError::into_inner)
265        .take()
266}
267
268/// [`serve`], with the connection slots taken by the listener when
269/// `handoff` is set.
270pub(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        // At the cap, wait for a free slot before `accept`. The kernel queues
307        // new connections meanwhile.
308        #[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        // `biased`: after a shutdown, never accept another connection.
316        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        // The head timer runs from accept, so a make-service that is slow to
324        // get ready cannot hold the connection (and its slot) past it.
325        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
389/// What the connection task knows about its connection.
390struct ConnState {
391    timers: Option<Arc<ConnTimers>>,
392    first_request: Arc<FirstRequest>,
393}
394
395/// Run one connection until it closes. After the server signal, `drain` asks
396/// hyper to shut the connection down gracefully (`graceful_shutdown`).
397async 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    // Set by a graceful shutdown. Until then hyper can finish its
409    // GOAWAY sequence and flush the last response bytes.
410    let mut grace_until: Option<Instant> = None;
411    // Set when the signal finds a connection with no request. Its
412    // request can be in flight. Hyper must not close it as idle.
413    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        // `biased`: read bytes already in the socket before a shutdown acts.
419        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                    // The head timer bounds the wait, as it bounds every head.
428                    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            // The first request came, or the wait ended: drain now.
434            () = fresh_wait(first_request, fresh_until), if waiting_fresh && !shutting_down => {
435                // No request in time: close. Else a partial head
436                // could hold the drain when no timer is set.
437                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            // The timers also run during a drain, so a slow client
445            // cannot hold the drain open.
446            () = sleep_until(deadline), if deadline.is_some() => {
447                match timers.as_ref().and_then(|t| t.expired()) {
448                    // A slow head: close now. No response is lost.
449                    Some(Expiry::Head) => break,
450                    // Still open after the grace period: close now.
451                    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
464/// Closes the timers when the connection task ends, also on abort.
465struct 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
475/// Resolves when the connection phase changes. Never without timers.
476async fn changed(timers: Option<&ConnTimers>) {
477    match timers {
478        Some(timers) => timers.changed.notified().await,
479        None => std::future::pending().await,
480    }
481}
482
483/// Time a connection gets to close by itself after a graceful shutdown.
484const SHUTDOWN_GRACE: Duration = Duration::from_secs(5);
485
486/// Time a connection with no request can wait for its first request after a
487/// graceful shutdown starts. Hyper closes an idle connection at once. That
488/// close resets a request that is still in flight.
489const FRESH_CONN_GRACE: Duration = Duration::from_secs(2);
490
491/// Resolves at the first request, or when the fresh-connection grace ends.
492async 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// ── Per-connection timers ─────────────────────────────────────────────
513
514#[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    /// Set while a request head (HTTP/1) or a header block (HTTP/2) is open.
532    header_block_since: Option<Instant>,
533}
534
535/// The HTTP/2 client preface.
536const 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/// A scan of the bytes the client sends: HTTP/1, or HTTP/2 frame headers.
542#[derive(Debug)]
543enum Scan {
544    /// Comparing the first bytes with the HTTP/2 preface.
545    Preface {
546        matched: usize,
547    },
548    Http1(H1Scan),
549    /// Between HTTP/2 frames: `have` bytes of the 9-byte frame header, then
550    /// `skip` payload bytes. `ends_block` is set while the payload of the
551    /// last frame of a header block (`END_HEADERS`) is still to come.
552    Frames {
553        header: [u8; 9],
554        have: usize,
555        skip: usize,
556        ends_block: bool,
557        /// The header timer started at the first byte of a frame header that
558        /// is not complete yet. Control traffic cancels it.
559        provisional: bool,
560    },
561}
562
563/// A minimal HTTP/1 framer. It finds where each request head starts and
564/// ends, and skips request bodies (`Content-Length` or chunked). hyper parses
565/// the same bytes; this scan only times the heads, so a pipelined partial
566/// head is timed too.
567#[derive(Debug, Default)]
568struct H1Scan {
569    state: H1State,
570    /// What the scan keeps of the current line.
571    line: LineScan,
572    content_length: u64,
573    chunked: bool,
574    position: HeadPosition,
575}
576
577/// The parts of the current line that frame a message. The scan reads
578/// values byte by byte and does not keep the line, so a long valid line
579/// (optional whitespace, leading zeros, chunk extensions) still frames.
580#[derive(Debug, Default)]
581struct LineScan {
582    /// Bytes in the line, without CR.
583    len: usize,
584    /// The field name, lowercase. Only the first bytes: a longer name is
585    /// not a name the scan reads.
586    name: Vec<u8>,
587    in_value: bool,
588    field: Field,
589    /// `Content-Length` (decimal) or the chunk size (hexadecimal).
590    number: Number,
591    /// The chunk size is complete; the rest is a chunk extension.
592    in_extension: bool,
593    /// The last bytes of a `Transfer-Encoding` value, lowercase, without
594    /// spaces and tabs.
595    tail: Vec<u8>,
596}
597
598/// A field this scan reads.
599#[derive(Debug, Default, Clone, Copy)]
600enum Field {
601    #[default]
602    Other,
603    ContentLength,
604    TransferEncoding,
605}
606
607/// A number read one digit at a time. A comma-separated list is one value
608/// if all items are equal, as hyper reads `Content-Length`.
609#[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    /// The value, or `None` when it is not a valid number.
643    const fn finish(mut self) -> Option<u64> {
644        self.end_item();
645        if self.bad { None } else { self.first }
646    }
647}
648
649/// Where the scan is inside the current head.
650#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
651enum HeadPosition {
652    /// No byte of the request line yet. A CRLF between messages starts the
653    /// head timer but not the request line.
654    #[default]
655    Before,
656    RequestLine,
657    Fields,
658}
659
660/// Field names longer than this are not names the scan reads.
661const 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    /// A value hyper rejects. Heads are not timed after it.
673    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                        // The first byte after an idle period starts the head
706                        // timer, also a CRLF: hyper skips it, but a client must
707                        // not stall after it with only the idle timeout.
708                        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    /// Stop the scan at a value hyper rejects. Fail open: the head timer
725    /// does not close the connection.
726    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                        // Optional whitespace can be long; values do not need it.
747                        _ 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                // An upgrade (`Upgrade`, `CONNECT`) needs no special case:
786                // a rejected one stays HTTP/1, and an accepted one ends the
787                // connection task, which closes the timers.
788                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                    // hyper rejects the head and closes the connection.
803                    None => self.stop(head_since),
804                },
805                // `chunked` must be the last coding.
806                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    /// Advance the scan over `bytes`. Opens or closes the HTTP/2 header block.
822    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                        // HTTP/1: replay the bytes that matched the preface.
834                        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                        // The block ends only when its last payload byte is in.
864                        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                        // Until the type byte is in, the frame can be a head.
875                        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    /// Set when the connection task ends. After an upgrade the IO lives on
911    /// (a WebSocket), but no timer applies, so reads skip the scan.
912    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    /// Bytes arrived. On HTTP/1, the first byte after an idle period starts
948    /// the head timer. On HTTP/2, an open header block starts it.
949    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    /// The connection task ended. Stop scanning reads.
967    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    /// When the connection must close. A slow head closes at once. During a
1000    /// shutdown grace (`grace_until`), an idle close waits for the grace.
1001    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            // With no head timeout, the idle timeout bounds the wait.
1016            Phase::Head { since } => self
1017                .header_read_timeout
1018                .or(self.keep_alive_timeout)
1019                .map(|timeout| (since + timeout, Expiry::Head)),
1020            // An open HTTP/2 header block suspends the idle timer: the new
1021            // request gets the full header timeout.
1022            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        // An open HTTP/2 header block, also next to other streams in flight.
1033        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    /// The expired timer, if any. The loop checks it again on wake, because
1044    /// the phase can change while the timer sleeps.
1045    fn expired(&self) -> Option<Expiry> {
1046        self.deadline_and_expiry()
1047            .filter(|(deadline, _)| *deadline <= Instant::now())
1048            .map(|(_, expiry)| expiry)
1049    }
1050}
1051
1052/// The accepted IO, with the connection permit and the read hook.
1053struct ConnIo<I> {
1054    inner: I,
1055    timers: Option<Arc<ConnTimers>>,
1056    _permit: Option<Slot>,
1057}
1058
1059/// Records the first request on a connection, and wakes a waiting task.
1060#[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/// Marks a request in flight from `call` until its response body is done.
1126#[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/// The in-flight mark of an HTTP/2 `CONNECT` stream (a WebSocket, RFC 8441).
1142///
1143/// The tunnel stays open after its response, so the mark rides in the
1144/// request extensions instead of the response body. The WebSocket upgrade
1145/// keeps it until the socket closes. When no handler keeps it (for example a
1146/// `404`), it drops with the request.
1147#[derive(Clone)]
1148pub(crate) struct TunnelGuard {
1149    _mark: Arc<InFlight>,
1150}
1151
1152/// Keeps an HTTP/2 WebSocket connection out of its idle timer.
1153///
1154/// Over HTTP/2, a WebSocket is a `CONNECT` stream. The server counts the
1155/// connection as busy while the request holds a tunnel mark. A handler on
1156/// axum's own `WebSocketUpgrade` drops that mark after the handshake, so
1157/// `keep_alive_timeout_ms` closes the connection while the socket is open.
1158/// Extract `KeepTunnel` and keep it for the life of the socket.
1159/// `autumn_web::ws::WebSocketUpgrade` does this for you.
1160///
1161/// It never rejects. On HTTP/1, or for a request that is not `CONNECT`, it
1162/// holds nothing.
1163///
1164/// ```rust,ignore
1165/// use autumn_web::http_server::KeepTunnel;
1166///
1167/// async fn raw(tunnel: KeepTunnel, ws: axum::extract::ws::WebSocketUpgrade)
1168///     -> axum::response::Response
1169/// {
1170///     ws.on_upgrade(move |mut socket| async move {
1171///         let _tunnel = tunnel; // drop it when the socket closes
1172///         while let Some(Ok(_)) = socket.recv().await {}
1173///     })
1174/// }
1175/// ```
1176pub 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        // hyper polls readiness only for a parsed request. A busy inner
1215        // service must not make that request look absent.
1216        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            // A CONNECT stream gets two marks: one on the response body, as
1224            // for every request, and one for the tunnel, which can outlive
1225            // the response.
1226            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    /// The response future of [`Tracked`]. Moves the in-flight mark into the
1244    /// response body.
1245    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
1273/// A response body that ends the in-flight mark when dropped.
1274struct 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"); // the 12-byte body, not a head
1334        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        // hyper skips the CRLF, but it is the first byte after the idle
1345        // period: a client that stalls after it gets the head timeout.
1346        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 /"); // the 5-byte body, not a head
1387        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    /// Scan `head`, then `body`, then one byte of the next head. The head
1393    /// must close, the body must not open a head, and the next head must.
1394    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        // hyper rejects this head; the body is not chunked.
1427        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        // Fail open: hyper rejects this head and closes the connection.
1441        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        // The server answered 400, so the next bytes are a new request.
1451        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        // After an upgrade, WebSocket frames are not scanned.
1464        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])); // SETTINGS with a payload
1476        bytes.extend(frame(H2_FRAME_HEADERS, 0, &[0x82, 0x86])); // no END_HEADERS
1477        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        // A request starts 1 s before the idle deadline: it gets the full
1506        // header timeout, not the 1 s that is left of the idle period.
1507        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        // `Semaphore::new` panics above `MAX_PERMITS`. `HttpLimits` is public,
1529        // so a value that skipped config validation must still be safe.
1530        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        // hyper's `max_buf_size` panics below its minimum. `HttpLimits` is
1541        // public, so a value that skipped config validation must still be
1542        // safe: it is raised to the minimum.
1543        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        // hyper polls readiness once a request is parsed, before `call`.
1576        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                // What the WebSocket upgrade does: keep the guard.
1611                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            // A CONNECT handler that is not a WebSocket: it drops the guard
1661            // and streams a body.
1662            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]); // the type byte is not in yet
1694        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]); // PING
1706        state.scan(&ping[..4]); // the type byte is in; the flags are not
1707        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        // Only the 9-byte frame header and one payload byte.
1732        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    /// Start `serve` on a fresh port. Returns the address, the stop token and
1742    /// the server task.
1743    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        // Longer than ACCEPT_SETTLE (100 ms).
1775        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        // A head that never completes. No timers are set.
1794        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    /// Spawn `serve` with the given limits and a stop token as the signal.
1816    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    /// The head timer still bounds a fresh connection during the drain.
1837    #[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    /// An HTTP/2 preface is not a request: the drain waits for the HEADERS.
1865    #[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]; // empty SETTINGS, stream 0
1876        // HPACK: :method GET, :scheme http, :path /, :authority t
1877        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        // A response HEADERS frame must be among the frames. A GOAWAY alone
1887        // means the request was refused.
1888        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    /// A stuck make-service must not hang the drain.
1905    #[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}