Skip to main content

actus_server/
server.rs

1//! [`Server`] — the hyper-based HTTP server that owns the request lifecycle:
2//! routing, body limiting, the middleware chain, CORS, compression, and the
3//! HTTP-correctness stamping (`Allow`, `Vary`), with graceful shutdown.
4
5#[cfg(feature = "compression")]
6use crate::compression::CompressionLayer;
7use crate::cors::CorsLayer;
8use crate::error::ServerError;
9use crate::middleware::{Middleware, MiddlewareChain, Outcome};
10use crate::request::{ContentCoding, Request, content_coding, has_undecoded_transfer_coding};
11use crate::router::Router;
12#[cfg(feature = "websocket")]
13use crate::websocket;
14use actus_reply::{Finalizer, ProblemDetails, ReplyData, WebError};
15use bytes::Bytes;
16use http::{HeaderMap, StatusCode};
17#[cfg(feature = "websocket")]
18use http::{HeaderValue, header};
19use http_body_util::combinators::BoxBody;
20use hyper::body::{Body, Incoming};
21use hyper::service::service_fn;
22use hyper::{Request as HyperRequest, Response as HyperResponse};
23use std::future::Future;
24use std::io;
25use std::net::SocketAddr;
26use std::pin::Pin;
27use std::sync::Arc;
28use std::sync::atomic::{AtomicU64, Ordering};
29use std::task::{Context, Poll};
30use std::time::Duration;
31use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
32use tokio::net::{TcpListener, TcpStream};
33use tokio::sync::Semaphore;
34use tokio::task::JoinSet;
35use tracing::{Instrument, Level, error, info, span, warn};
36
37type ResponseBody = BoxBody<Bytes, WebError>;
38
39/// One kibibyte (1024 bytes). For readable byte-size limits, e.g.
40/// `#[controller(max_body_bytes = 4 * KIB)]`.
41pub const KIB: usize = 1024;
42/// One mebibyte (1024 × 1024 bytes), e.g. `Server::with_max_body_bytes(2 * MIB)`.
43pub const MIB: usize = 1024 * KIB;
44/// One gibibyte (1024 × 1024 × 1024 bytes).
45pub const GIB: usize = 1024 * MIB;
46
47/// Default cap on the request body Actus will buffer: **2 MiB** — a safe
48/// ceiling for the common case (JSON APIs, forms). Endpoints that accept larger
49/// bodies (uploads) opt in via [`Server::with_max_body_bytes`] or a
50/// per-controller `#[controller(max_body_bytes = …)]`. Matches axum's default.
51pub const DEFAULT_MAX_BODY_BYTES: usize = 2 * MIB;
52
53/// Default grace period for in-flight connections to finish after a
54/// shutdown signal: 30 seconds. Override with [`Server::with_drain_deadline`].
55pub const DEFAULT_DRAIN_DEADLINE: Duration = Duration::from_secs(30);
56
57/// The main Actus server.
58pub struct Server {
59    router: Arc<Router>,
60    middleware_chain: Arc<MiddlewareChain>,
61    finalizer: Arc<Finalizer>,
62    max_body_bytes: usize,
63    cors: Option<Arc<CorsLayer>>,
64    #[cfg(feature = "compression")]
65    compression: Option<Arc<CompressionLayer>>,
66    /// `Some(d)` caps each request's total time (parse → middleware →
67    /// handler → after-chain → finalize) at `d`; an over-budget request is
68    /// aborted and replied with `504 Gateway Timeout`. `None` disables the
69    /// per-request timer (the default).
70    request_timeout: Option<Duration>,
71    /// Grace period for in-flight connections to drain after shutdown.
72    drain_deadline: Duration,
73    /// Cap on concurrent connection tasks. `Some(n)` installs an
74    /// `Arc<Semaphore>` of `n` permits in the accept loop; while at
75    /// capacity, the loop pauses on permit acquisition and new SYNs queue
76    /// in the kernel's accept backlog (`SOMAXCONN`), at which point the
77    /// kernel drops them. `None` is unbounded.
78    max_connections: Option<usize>,
79    /// Cap on the total bytes being buffered across all in-flight body
80    /// reads. `Some(n)` installs a byte-permit semaphore; each
81    /// `collect_body_capped` reserves its per-request cap upfront and
82    /// releases the permits when the body is buffered or rejected. Refuses
83    /// excess requests with `503 Service Unavailable` (via `WebError::Busy`).
84    /// `None` is unbounded.
85    max_inflight_body_bytes: Option<Arc<Semaphore>>,
86    /// `Some(d)` is forwarded to hyper's `http1::Builder::header_read_timeout`
87    /// — bounds how long after starting to read request headers we'll wait
88    /// before dropping the connection. Catches slowloris and clients that
89    /// TCP-connect-and-send-nothing. `None` leaves hyper's default (none).
90    header_read_timeout: Option<Duration>,
91    /// `Some(counter)`: every byte written to a connection is added to
92    /// `counter` ([`Server::with_sent_bytes`]). `None` counts nothing (the
93    /// default).
94    sent_bytes: Option<Arc<AtomicU64>>,
95    /// `true`: a `gzip`-coded request body is decoded rather than refused
96    /// ([`Server::with_gzip_request_bodies`]). `false` by default.
97    #[cfg(feature = "compression")]
98    gzip_request_bodies: bool,
99}
100
101impl Server {
102    /// Create a server for `router` with default settings: no middleware, no
103    /// CORS, the default body cap, and no DoS limits. Configure it with the
104    /// `with_*` builder methods, then call [`run`](Self::run).
105    pub fn new(router: Router) -> Self {
106        Self {
107            router: Arc::new(router),
108            middleware_chain: Arc::new(MiddlewareChain::new()),
109            finalizer: Arc::new(Finalizer::new()),
110            max_body_bytes: DEFAULT_MAX_BODY_BYTES,
111            cors: None,
112            #[cfg(feature = "compression")]
113            compression: None,
114            request_timeout: None,
115            drain_deadline: DEFAULT_DRAIN_DEADLINE,
116            max_connections: None,
117            max_inflight_body_bytes: None,
118            header_read_timeout: None,
119            sent_bytes: None,
120            #[cfg(feature = "compression")]
121            gzip_request_bodies: false,
122        }
123    }
124
125    /// The router this server serves, shared.
126    ///
127    /// Lets application code keep using the route tree after `Server::new`
128    /// has consumed it — most usefully from a `Middleware` that wants the
129    /// framework's own longest-prefix matcher instead of re-deriving it. A
130    /// declaration-keyed gate is the motivating case:
131    ///
132    /// ```ignore
133    /// let server = Server::new(router);
134    /// let router = server.router(); // Arc<Router>
135    /// let server = server.with_middleware(FloorGate { router });
136    /// // …in the gate's `before`:
137    /// //   if let Some(rm) = self.router.match_controller(&req.path_parts) {
138    /// //       if rm.controller.actus_expects() == Some("credential") && … { … }
139    /// //   }
140    /// ```
141    ///
142    /// The clone is an `Arc` clone; a gate's extra `match_controller` call is
143    /// one map lookup per path segment on top of the one the server already
144    /// does. See the README's "Route families" section for the full pattern.
145    /// A request routed by a host alias needs nothing extra: its `path_parts`
146    /// already hold the aliased path by the time any middleware runs.
147    pub fn router(&self) -> Arc<Router> {
148        Arc::clone(&self.router)
149    }
150
151    /// Adds a middleware to the server's request processing chain.
152    pub fn with_middleware(mut self, middleware: impl Middleware + 'static) -> Self {
153        let mut chain = Arc::try_unwrap(self.middleware_chain).unwrap_or_else(|arc| (*arc).clone());
154        chain.add(middleware);
155        self.middleware_chain = Arc::new(chain);
156        self
157    }
158
159    /// Enables CORS with the given policy. The server then answers preflight
160    /// (`OPTIONS`) requests itself and adds the `Access-Control-*` headers to
161    /// every cross-origin response (including error responses). See
162    /// [`CorsLayer`].
163    pub fn with_cors(mut self, cors: CorsLayer) -> Self {
164        self.cors = Some(Arc::new(cors));
165        self
166    }
167
168    /// Enables response compression (gzip / brotli). For each response Actus
169    /// picks an encoding from the request's `Accept-Encoding` and compresses
170    /// buffered, compressible bodies above the layer's size threshold. See
171    /// [`CompressionLayer`]. *(Requires the `compression` feature.)*
172    #[cfg(feature = "compression")]
173    pub fn with_compression(mut self, layer: CompressionLayer) -> Self {
174        self.compression = Some(Arc::new(layer));
175        self
176    }
177
178    /// Decodes `gzip` request bodies: a request sent with
179    /// `Content-Encoding: gzip` reaches middleware and the handler as the
180    /// content it encodes, with `Content-Encoding` and `Content-Length`
181    /// removed. Without this, such a request is refused with `415`.
182    ///
183    /// **The body cap applies to the decoded size** — the controller's
184    /// `max_body_bytes`, else [`Server::with_max_body_bytes`] — as well as to
185    /// the coded bytes read, so a few KB of gzip that would decode to
186    /// gigabytes is a `413` as soon as its output passes the cap. Decoding is
187    /// streamed: memory stays at the cap, as for an unencoded body. A corrupt
188    /// or truncated stream is a `400`.
189    ///
190    /// Only a single `gzip` is decoded (`x-gzip` is its alias). Any other
191    /// coding, or codings stacked, is still a `415`, whose
192    /// `Accept-Encoding: gzip, identity` names what would have been
193    /// accepted. Middleware sees the decoded body, so a check over the bytes
194    /// as sent — a webhook signature — needs a sender that does not compress.
195    /// *(Requires the `compression` feature.)*
196    #[cfg(feature = "compression")]
197    pub fn with_gzip_request_bodies(mut self) -> Self {
198        self.gzip_request_bodies = true;
199        self
200    }
201
202    /// Counts every byte the server writes to its connections into `counter`
203    /// — what actually went out: status lines and headers, bodies after
204    /// compression, chunked framing, and the frames of an upgraded
205    /// (WebSocket) connection. Each write is counted as the operating system
206    /// accepts it, so the count is live, and a client that disconnects
207    /// mid-response leaves only what was written before the connection
208    /// failed.
209    ///
210    /// Request bytes are not counted, nor is anything outside the connection:
211    /// TCP/IP overhead, or what a proxy in front of the server adds (TLS) or
212    /// changes.
213    ///
214    /// The counter is the caller's: read it, or swap it to zero, from
215    /// anywhere. Typical use is metering a deployment's egress.
216    pub fn with_sent_bytes(mut self, counter: Arc<AtomicU64>) -> Self {
217        self.sent_bytes = Some(counter);
218        self
219    }
220
221    /// Caps the request body Actus will buffer (default
222    /// [`DEFAULT_MAX_BODY_BYTES`] = 2 MiB). A larger body is rejected with
223    /// `413 Payload Too Large` before it can exhaust memory — the limit
224    /// bounds buffered bytes, so it also covers chunked bodies that lie about
225    /// (or omit) `Content-Length`.
226    ///
227    /// `0` is accepted and means "reject every non-empty body" — typically
228    /// only useful on a strictly-GET surface that should never see a body.
229    pub fn with_max_body_bytes(mut self, max: usize) -> Self {
230        self.max_body_bytes = max;
231        self
232    }
233
234    /// Cap the total time any single request may take — body parse,
235    /// middleware `before`, handler, middleware `after`, and finalization
236    /// combined. An over-budget request is aborted (the handler's future
237    /// is dropped) and the client gets `504 Gateway Timeout`. No timeout
238    /// is set by default.
239    ///
240    /// **Scope.** The timer covers the request/response exchange. A
241    /// WebSocket upgrade succeeds inside the timer (the `101` is the
242    /// response); the post-upgrade conversation runs in its own task and
243    /// is not bound by this timeout.
244    ///
245    /// **Effect of an over-budget request.** When the timer elapses the
246    /// in-flight future is dropped, which cancels whatever the handler
247    /// was awaiting (DB query, channel recv, etc.). The 504 reply is
248    /// one-shot — the after-chain doesn't run on it (by definition,
249    /// some component upstream was unresponsive; running more risks
250    /// hanging again).
251    pub fn with_request_timeout(mut self, d: Duration) -> Self {
252        self.request_timeout = Some(d);
253        self
254    }
255
256    /// Override the grace period for in-flight connections after a
257    /// shutdown signal (default [`DEFAULT_DRAIN_DEADLINE`] = 30 s).
258    /// Anything still running at the deadline is hard-aborted via
259    /// `JoinSet::shutdown`. Use a longer value for surfaces that hold
260    /// long-lived connections (large file downloads, WebSockets);
261    /// a shorter value for fast-iteration dev workflows. `Duration::ZERO`
262    /// aborts every in-flight task immediately.
263    pub fn with_drain_deadline(mut self, d: Duration) -> Self {
264        self.drain_deadline = d;
265        self
266    }
267
268    /// Cap concurrent connection tasks at `n`. While the cap is held, the
269    /// accept loop pauses on permit acquisition; new SYNs queue in the
270    /// kernel's accept backlog and (once that fills, governed by
271    /// `SOMAXCONN`) get dropped at the OS level. No userland reject /
272    /// no-503-per-conn cost — the kernel handles the spillover.
273    ///
274    /// Each connection task holds its permit until it ends, including the
275    /// post-handshake WebSocket conversation. Size accordingly: a
276    /// `with_max_connections(N)` server can hold `N` long-lived WebSockets
277    /// *before* it stops accepting new connections of any kind.
278    ///
279    /// Unbounded by default (no semaphore installed).
280    pub fn with_max_connections(mut self, n: usize) -> Self {
281        self.max_connections = Some(n);
282        self
283    }
284
285    /// Cap the total bytes being buffered across all in-flight body reads
286    /// at `n`. Each request reserves its per-request cap (see
287    /// `with_max_body_bytes`) from this global budget upfront; if the
288    /// budget is exhausted, the request is refused with `503 Service
289    /// Unavailable` (via [`WebError::Busy`]) and a short `Retry-After`.
290    ///
291    /// Together with [`Self::with_max_connections`] this puts a hard
292    /// ceiling on the framework's memory under adversarial load:
293    /// `with_max_connections(C) * with_max_body_bytes(B)` is the *worst*
294    /// case absent this knob; with it, the ceiling is `min(C * B, this
295    /// value)`.
296    ///
297    /// Pre-reserving the per-request cap over-counts (a 1 KB request
298    /// reserves up to its full cap); the alternative — incremental
299    /// per-chunk byte accounting — is more code for the same effective
300    /// ceiling, and a request that has already started buffering can't
301    /// be sensibly aborted partway through anyway.
302    ///
303    /// `n` is clamped to `u32::MAX` internally (Tokio's `Semaphore`
304    /// permit count uses `u32`); for practical deployments this is no
305    /// limit (4 GiB).
306    ///
307    /// Unbounded by default.
308    pub fn with_max_inflight_body_bytes(mut self, n: usize) -> Self {
309        // u32 cap is a tokio Semaphore constraint, not a design choice.
310        let n_capped = n.min(u32::MAX as usize);
311        self.max_inflight_body_bytes = Some(Arc::new(Semaphore::new(n_capped)));
312        self
313    }
314
315    /// Bound how long after starting to read request headers we'll wait
316    /// before dropping the connection. Forwards to hyper's
317    /// `http1::Builder::header_read_timeout`. Catches slowloris (sending
318    /// headers one byte at a time) and clients that TCP-connect-and-send-
319    /// nothing — the most common file-descriptor-exhaustion attack on a
320    /// keep-alive HTTP server.
321    ///
322    /// Note: hyper 1.x doesn't have a separate "idle between requests"
323    /// timeout (after a complete request, an idle keep-alive connection
324    /// stays open until either side closes or the OS-level TCP keep-alive
325    /// fires). If that matters for your deployment, either disable
326    /// keep-alive entirely upstream of Actus or rely on the OS knobs.
327    ///
328    /// No timeout by default (hyper's default).
329    pub fn with_header_read_timeout(mut self, d: Duration) -> Self {
330        self.header_read_timeout = Some(d);
331        self
332    }
333
334    /// Runs the server on `127.0.0.1:port` (loopback only). For a different
335    /// bind address — e.g. `0.0.0.0:port` to accept connections from other
336    /// hosts in a container — use [`Server::run_on`].
337    ///
338    /// Listens for SIGTERM/SIGINT (Unix) or Ctrl-C (cross-platform) and
339    /// shuts down gracefully: stops accepting new connections, signals
340    /// in-flight connections to finish, and waits up to 30 seconds for
341    /// them to drain before returning.
342    pub async fn run(self, port: u16) -> Result<(), ServerError> {
343        self.run_on(SocketAddr::from(([127, 0, 0, 1], port))).await
344    }
345
346    /// Like [`Server::run`] but binds an arbitrary address. Pass
347    /// `0.0.0.0:port` (or `[::]:port`) to accept connections from other hosts.
348    pub async fn run_on(self, addr: SocketAddr) -> Result<(), ServerError> {
349        self.run_with_shutdown_on(addr, default_shutdown_signal())
350            .await
351    }
352
353    /// Like [`Server::run`] but with a custom shutdown trigger (a future that,
354    /// when it resolves, starts the graceful drain). Binds `127.0.0.1:port`;
355    /// see [`Server::run_with_shutdown_on`] for a custom bind address. Useful
356    /// for tests or for embedding the server in a larger supervision tree.
357    pub async fn run_with_shutdown(
358        self,
359        port: u16,
360        shutdown: impl Future<Output = ()> + Send + 'static,
361    ) -> Result<(), ServerError> {
362        self.run_with_shutdown_on(SocketAddr::from(([127, 0, 0, 1], port)), shutdown)
363            .await
364    }
365
366    /// Bind `addr`, then serve exactly as
367    /// [`Server::run_with_shutdown_listener`] — the general form — does.
368    /// [`Server::run`], [`Server::run_on`], and [`Server::run_with_shutdown`]
369    /// are thin wrappers over this; this is a thin bind over the listener
370    /// form.
371    ///
372    /// **Drain bound.** Once `shutdown` resolves the server stops accepting
373    /// and signals every in-flight connection to wind down. The drain
374    /// deadline defaults to [`DEFAULT_DRAIN_DEADLINE`] (30 s); override
375    /// with [`Server::with_drain_deadline`]. Anything still running at
376    /// the deadline is hard-aborted via `JoinSet::shutdown`. In particular,
377    /// long-lived connections (WebSockets, slow downloads, kept-alive idle
378    /// clients) and any connection task that raced the shutdown notification
379    /// and missed it both get aborted at the deadline rather than draining
380    /// gracefully.
381    pub async fn run_with_shutdown_on(
382        self,
383        addr: SocketAddr,
384        shutdown: impl Future<Output = ()> + Send + 'static,
385    ) -> Result<(), ServerError> {
386        let listener = TcpListener::bind(addr).await?;
387        self.run_with_shutdown_listener(listener, shutdown).await
388    }
389
390    /// Like [`Server::run_on`] but serves on a listener the caller already
391    /// bound (or inherited), with the default SIGTERM / SIGINT shutdown
392    /// trigger. See [`Server::run_with_shutdown_listener`] for what a
393    /// pre-bound listener buys and how to adopt an inherited one.
394    pub async fn run_listener(self, listener: TcpListener) -> Result<(), ServerError> {
395        self.run_with_shutdown_listener(listener, default_shutdown_signal())
396            .await
397    }
398
399    /// The most general form: serve on a listener the caller already bound
400    /// (or inherited) until `shutdown` resolves, then drain. Every other
401    /// `run*` method is a wrapper over this one.
402    ///
403    /// **Why hand the server a listener.** Two cases the bind-it-yourself
404    /// forms can't express:
405    ///
406    /// - **Socket activation** — a supervisor (systemd's `LISTEN_FDS`
407    ///   protocol, launchd, …) owns the listening socket and passes it to
408    ///   the process it spawns. Because the socket outlives the process,
409    ///   connections arriving during a restart queue in the kernel's accept
410    ///   backlog instead of being refused, and the next process serves them.
411    /// - **Race-free embedding and tests** — bind `127.0.0.1:0`, keep the
412    ///   listener, and pass it in; no bind-drop-rebind window in which the
413    ///   port can be lost, and requests may connect before the accept loop
414    ///   even starts (the kernel queues them).
415    ///
416    /// An inherited `std::net::TcpListener` must be set non-blocking before
417    /// conversion — tokio requires it, and a supervisor passes the fd
418    /// blocking:
419    ///
420    /// ```no_run
421    /// # async fn doc(server: actus_server::Server, inherited: std::net::TcpListener)
422    /// # -> Result<(), actus_server::ServerError> {
423    /// inherited.set_nonblocking(true)?;
424    /// let listener = tokio::net::TcpListener::from_std(inherited)?;
425    /// server.run_with_shutdown_listener(listener, std::future::pending()).await
426    /// # }
427    /// ```
428    ///
429    /// **Drain bound.** As [`Server::run_with_shutdown_on`]: once `shutdown`
430    /// resolves the server stops accepting (the listener is dropped — under
431    /// socket activation the *socket* stays open in the supervisor, which is
432    /// the point) and in-flight connections get [`Server::with_drain_deadline`]
433    /// (default [`DEFAULT_DRAIN_DEADLINE`] = 30 s) to finish before being
434    /// aborted. A streaming response that never completes (SSE) always rides
435    /// to the deadline.
436    pub async fn run_with_shutdown_listener(
437        self,
438        listener: TcpListener,
439        shutdown: impl Future<Output = ()> + Send + 'static,
440    ) -> Result<(), ServerError> {
441        match listener.local_addr() {
442            Ok(addr) => info!("Server listening on http://{}", addr),
443            // An inherited fd can decline local_addr (exotic socket family);
444            // serving still works, so log what we know and carry on.
445            Err(_) => info!("Server listening on an inherited socket"),
446        }
447
448        let app = Arc::new(self);
449        // Per-connection cancellation: once `Notify::notify_waiters` fires,
450        // every in-flight task wakes up and asks hyper to gracefully close
451        // its connection (finishing the current response, then exiting).
452        let notify = Arc::new(tokio::sync::Notify::new());
453        let mut tasks: JoinSet<()> = JoinSet::new();
454
455        // Optional cap on concurrent connections. When at-capacity the
456        // accept loop pauses on permit acquisition; new SYNs queue in the
457        // kernel accept backlog (SOMAXCONN) and get dropped at the OS
458        // level once that fills. Each spawned connection task moves its
459        // permit in; the permit releases when the task exits.
460        let conn_permits = app.max_connections.map(|n| Arc::new(Semaphore::new(n)));
461
462        tokio::pin!(shutdown);
463
464        loop {
465            tokio::select! {
466                // Accept-branch: acquire a connection permit first (when
467                // a cap is configured), then accept. The outer select
468                // races this against shutdown so a paused-at-capacity
469                // accept loop still notices the shutdown signal.
470                accept_with_permit = async {
471                    let permit = match &conn_permits {
472                        Some(s) => Some(s.clone().acquire_owned().await.expect("semaphore never closed")),
473                        None => None,
474                    };
475                    let result = listener.accept().await;
476                    (result, permit)
477                } => {
478                    let (accept_result, permit) = accept_with_permit;
479                    let (stream, _peer) = match accept_result {
480                        Ok(s) => s,
481                        Err(e) => {
482                            error!("accept error: {}", e);
483                            // permit released when this branch ends — fine
484                            continue;
485                        }
486                    };
487                    // Every byte the connection writes passes through here —
488                    // responses, hyper's own replies, an upgraded connection's
489                    // frames — so this is where `with_sent_bytes` counts.
490                    let io = hyper_util::rt::TokioIo::new(CountingStream {
491                        inner: stream,
492                        sent_bytes: app.sent_bytes.clone(),
493                    });
494                    let app = app.clone();
495                    let notify = notify.clone();
496                    let header_timeout = app.header_read_timeout;
497                    tasks.spawn(async move {
498                        // The permit (if any) lives for the connection's
499                        // lifetime; releasing happens at task drop.
500                        let _permit = permit;
501
502                        let mut builder = hyper::server::conn::http1::Builder::new();
503                        if let Some(d) = header_timeout {
504                            builder.header_read_timeout(d);
505                        }
506                        let conn = builder.serve_connection(
507                            io,
508                            service_fn(move |req| app.clone().handle_request(req)),
509                        );
510                        // With the `websocket` feature, allow `101 Switching
511                        // Protocols` responses to hand off the connection.
512                        #[cfg(feature = "websocket")]
513                        let conn = conn.with_upgrades();
514                        tokio::pin!(conn);
515                        tokio::select! {
516                            res = conn.as_mut() => {
517                                if let Err(err) = res {
518                                    error!("Error serving connection: {}", err);
519                                }
520                            }
521                            _ = notify.notified() => {
522                                conn.as_mut().graceful_shutdown();
523                                if let Err(err) = conn.await {
524                                    error!("Error during graceful shutdown: {}", err);
525                                }
526                            }
527                        }
528                    });
529                }
530                // Reap finished connection tasks so the `JoinSet` doesn't grow
531                // without bound over the server's lifetime — and so a panicked
532                // connection task is logged promptly, not only at shutdown.
533                joined = tasks.join_next(), if !tasks.is_empty() => {
534                    match joined {
535                        Some(Err(e)) if e.is_panic() => error!("Connection task panicked: {}", e),
536                        Some(Err(e)) => error!("Connection task failed: {}", e),
537                        Some(Ok(())) | None => {}
538                    }
539                }
540                _ = &mut shutdown => {
541                    info!("Shutdown signal received; draining in-flight requests");
542                    break;
543                }
544            }
545        }
546
547        // Stop accepting; signal connections to wind down.
548        drop(listener);
549        notify.notify_waiters();
550
551        // Drain. The grace period is configurable via
552        // `Server::with_drain_deadline` (default 30 s).
553        let drain_deadline = tokio::time::sleep(app.drain_deadline);
554        tokio::pin!(drain_deadline);
555        loop {
556            tokio::select! {
557                next = tasks.join_next() => {
558                    match next {
559                        Some(Ok(())) => {}
560                        Some(Err(e)) if e.is_panic() => {
561                            error!("Connection task panicked: {}", e);
562                        }
563                        Some(Err(e)) => {
564                            error!("Connection task failed: {}", e);
565                        }
566                        None => break,
567                    }
568                }
569                _ = &mut drain_deadline => {
570                    warn!("Drain deadline exceeded; aborting {} connection(s)", tasks.len());
571                    tasks.shutdown().await;
572                    break;
573                }
574            }
575        }
576
577        info!("Server shutdown complete");
578        Ok(())
579    }
580
581    /// Stamp the configured CORS response headers onto `response` (no-op when
582    /// CORS isn't enabled, or the request had no allowed `Origin`). Applied to
583    /// *every* outgoing response — success and error alike — so the browser
584    /// can read 4xx/5xx bodies.
585    fn with_cors_headers(
586        &self,
587        request: &Request,
588        mut response: HyperResponse<ResponseBody>,
589    ) -> HyperResponse<ResponseBody> {
590        if let Some(cors) = &self.cors {
591            cors.apply(&request.headers, response.headers_mut(), false);
592        }
593        response
594    }
595
596    /// Fulfil a `ReplyData::Upgrade` from a handler when the request was
597    /// genuinely a WebSocket handshake: send `101 Switching Protocols` and
598    /// spawn the handler on the upgraded connection. (The "handler returned
599    /// Upgrade but the request wasn't a handshake" case is rewritten to a
600    /// 426 reply in [`Self::finalize_reply`] before reaching this method, so
601    /// it can flow through the after-chain like any other error.)
602    ///
603    /// No CORS headers on the `101`: WebSocket handshakes are scoped by
604    /// browser origin checks (the handler inspects `Origin` itself before
605    /// calling `ws::upgrade`), not by the CORS protocol — `Access-Control-*`
606    /// on a `101` is meaningless to the browser.
607    #[cfg(feature = "websocket")]
608    async fn complete_ws_upgrade(
609        &self,
610        handler: Box<dyn std::any::Any + Send>,
611        ws_upgrade: (hyper::upgrade::OnUpgrade, HeaderValue),
612    ) -> HyperResponse<ResponseBody> {
613        // `ReplyData::Upgrade` is only constructible via `ws::upgrade(...)`,
614        // which always boxes an `UpgradeTask`. A failing downcast would mean
615        // a crate-internal invariant is broken; surface as a panic rather
616        // than silently producing a 500.
617        let task = handler
618            .downcast::<websocket::UpgradeTask>()
619            .expect("ReplyData::Upgrade always carries an UpgradeTask");
620        let (on_upgrade, accept) = ws_upgrade;
621        tokio::spawn(websocket::run_upgrade(on_upgrade, *task));
622        let mut resp = self.finalizer.build_response(ReplyData::Empty).await;
623        *resp.status_mut() = StatusCode::SWITCHING_PROTOCOLS;
624        let h = resp.headers_mut();
625        h.insert(header::CONNECTION, HeaderValue::from_static("upgrade"));
626        h.insert(header::UPGRADE, HeaderValue::from_static("websocket"));
627        h.insert(header::SEC_WEBSOCKET_ACCEPT, accept);
628        resp
629    }
630
631    /// Build the error reply for `error` and route it through
632    /// [`finalize_reply`](Self::finalize_reply), so the after-chain,
633    /// compression, and CORS apply to errors exactly as they do to handler
634    /// successes. This is the canonical way to produce a `WebError`
635    /// response anywhere a `Request` exists.
636    async fn finalize_error(
637        &self,
638        error: WebError,
639        request: &Request,
640        #[cfg(feature = "websocket")] ws_upgrade: Option<(hyper::upgrade::OnUpgrade, HeaderValue)>,
641    ) -> HyperResponse<ResponseBody> {
642        let data = self.finalizer.error_to_reply(error);
643        self.finalize_reply(
644            data,
645            request,
646            #[cfg(feature = "websocket")]
647            ws_upgrade,
648        )
649        .await
650    }
651
652    /// How to buffer a request's body, decided from its codings before it is
653    /// read — or the reply refusing it. A handler must never receive encoded
654    /// bytes as though they were the content: a JSON route would fail to
655    /// parse them, and a route storing bytes would store the encoding.
656    ///
657    /// - A transfer coding besides a lone `chunked` is one the server does not
658    ///   understand: `501 Not Implemented` (RFC 9112 §6.1).
659    /// - A `gzip` content coding is decoded, when
660    ///   `Server::with_gzip_request_bodies` turned that on.
661    /// - Any other content coding but `identity`: `415 Unsupported Media
662    ///   Type`, with an `Accept-Encoding` naming what would have been accepted
663    ///   (RFC 9110 §12.5.3).
664    ///
665    /// A request with no body is taken as sent: there is nothing to decode,
666    /// and it is not refused over the header alone.
667    fn body_coding(&self, headers: &HeaderMap, body: &Incoming) -> Result<BodyCoding, ReplyData> {
668        if has_undecoded_transfer_coding(headers) {
669            return Err(self.finalizer.error_to_reply(WebError::Problem(
670                ProblemDetails::new(StatusCode::NOT_IMPLEMENTED, "Not Implemented")
671                    .detail("send the request body with no transfer coding but chunked"),
672            )));
673        }
674        if body.is_end_stream() {
675            return Ok(BodyCoding::AsSent);
676        }
677        match content_coding(headers) {
678            ContentCoding::Identity => Ok(BodyCoding::AsSent),
679            #[cfg(feature = "compression")]
680            ContentCoding::Gzip if self.gzip_request_bodies => Ok(BodyCoding::Gzip),
681            _ => {
682                let mut data = self.finalizer.error_to_reply(WebError::Problem(
683                    ProblemDetails::new(
684                        StatusCode::UNSUPPORTED_MEDIA_TYPE,
685                        "Unsupported Media Type",
686                    )
687                    .detail("Content-Encoding not accepted; Accept-Encoding lists what is"),
688                ));
689                data.add_header("Accept-Encoding", self.accepted_request_codings());
690                Err(data)
691            }
692        }
693    }
694
695    /// The request content codings this server takes, as the `Accept-Encoding`
696    /// of a content-coding `415` lists them.
697    fn accepted_request_codings(&self) -> &'static str {
698        #[cfg(feature = "compression")]
699        if self.gzip_request_bodies {
700            return "gzip, identity";
701        }
702        "identity"
703    }
704
705    /// Run the after-middleware chain, then turn the reply into a `Response`
706    /// via [`dispatch_reply`](Self::dispatch_reply).
707    ///
708    /// **After-chain runs on every reply with a body and a `Request`.** That
709    /// includes handler successes, `Outcome::Respond` short-circuits, *and*
710    /// every error the application produced (404 / 405 / 401 / 400 / a
711    /// handler-returned `Err(WebError)`, etc.). The README's promise that a
712    /// request-id stamper "still fires on a short-circuit" generalizes to
713    /// every reply — that's the contract this method enforces.
714    ///
715    /// **Exceptions** (the after-chain *doesn't* run):
716    /// - **101 Switching Protocols** — a WebSocket-handshake success has no
717    ///   HTTP body to decorate, and the upgrade machinery consumes the
718    ///   connection.
719    /// - **Pre-parse failures** — a request that fails before
720    ///   [`Request::from_hyper`] returns a skeleton (e.g. malformed HTTP
721    ///   from hyper itself) has no `Request` to give the hook. The body-cap
722    ///   413 and truncated-body 400 are *not* exceptions here: `from_hyper`
723    ///   now returns a skeleton `Request` on those, so they do run through
724    ///   the after-chain.
725    /// - **CORS preflight 204** — synthesized before middleware or routing;
726    ///   not an application request (see [`Self::handle_request`]).
727    async fn finalize_reply(
728        &self,
729        #[allow(unused_mut)] mut data: ReplyData,
730        request: &Request,
731        #[cfg(feature = "websocket")] ws_upgrade: Option<(hyper::upgrade::OnUpgrade, HeaderValue)>,
732    ) -> HyperResponse<ResponseBody> {
733        // If the handler returned `ws::upgrade(...)` but the request isn't a
734        // real WebSocket handshake, rewrite to a 426 error reply *here* so
735        // it flows through the same after-chain / compression / CORS path as
736        // any other error. Only the success-handshake path (Upgrade reply +
737        // ws_upgrade present) keeps the after-chain bypass — a 101 has no
738        // HTTP body to decorate.
739        #[cfg(feature = "websocket")]
740        if matches!(data, ReplyData::Upgrade(_)) && ws_upgrade.is_none() {
741            data = self.finalizer.error_to_reply(WebError::Problem(
742                ProblemDetails::new(StatusCode::UPGRADE_REQUIRED, "WebSocket Upgrade Required")
743                    .detail("this endpoint expects a WebSocket handshake"),
744            ));
745        }
746
747        let needs_after_chain = !matches!(data, ReplyData::Upgrade(_));
748        if needs_after_chain
749            && let Err(e) = self
750                .middleware_chain
751                .process_response(request, &mut data)
752                .await
753        {
754            // After-chain itself errored. Build a plain error response
755            // (no further after-chain — recursion prevention) so a buggy
756            // hook can't infinite-loop the request.
757            return self.with_cors_headers(request, self.finalizer.build_error(e).await);
758        }
759        self.dispatch_reply(
760            data,
761            request,
762            #[cfg(feature = "websocket")]
763            ws_upgrade,
764        )
765        .await
766    }
767
768    /// Turn a handler's (or a short-circuiting middleware's) `ReplyData` into a
769    /// fully processed response: WebSocket upgrade if it's an `Upgrade` reply;
770    /// otherwise compress (if enabled), finalize, and stamp CORS / `Vary`.
771    async fn dispatch_reply(
772        &self,
773        #[allow(unused_mut)] mut data: ReplyData,
774        request: &Request,
775        #[cfg(feature = "websocket")] ws_upgrade: Option<(hyper::upgrade::OnUpgrade, HeaderValue)>,
776    ) -> HyperResponse<ResponseBody> {
777        // A handler that returned `ws::upgrade(...)`: complete the handshake
778        // instead of finalizing a body. `finalize_reply` only lets us reach
779        // here for an `Upgrade` reply when the request *was* a real
780        // handshake (otherwise it rewrote the reply to a 426 error), so
781        // `ws_upgrade` is guaranteed `Some` on this branch.
782        #[cfg(feature = "websocket")]
783        if matches!(data, ReplyData::Upgrade(_)) {
784            let ReplyData::Upgrade(handler) = data else {
785                unreachable!()
786            };
787            let ws_upgrade =
788                ws_upgrade.expect("finalize_reply rewrites Upgrade-without-handshake to 426");
789            return self.complete_ws_upgrade(handler, ws_upgrade).await;
790        }
791        // Compression is the last transform — after response middleware,
792        // before the bytes leave. (Only buffered, compressible bodies above
793        // the threshold are touched.)
794        #[cfg(feature = "compression")]
795        if let Some(c) = &self.compression {
796            data = c.compress_reply(
797                data,
798                request
799                    .headers
800                    .get("accept-encoding")
801                    .and_then(|v| v.to_str().ok()),
802            );
803        }
804        let response = self.finalizer.build_response(data).await;
805        let response = self.with_cors_headers(request, response);
806        #[cfg(feature = "compression")]
807        let response = crate::compression::tag_vary_if_encoded(response);
808        response
809    }
810
811    /// Handles an individual incoming `hyper::Request`.
812    ///
813    /// Wraps [`handle_request_inner`](Self::handle_request_inner) in a
814    /// per-request timeout when one is configured (see
815    /// [`Server::with_request_timeout`]); a timed-out request gets a
816    /// one-shot `504 Gateway Timeout` (no after-chain, since by definition
817    /// something upstream was unresponsive).
818    ///
819    /// Every reply with a `Request` flows through
820    /// [`finalize_reply`](Self::finalize_reply) — handler successes,
821    /// `Outcome::Respond` short-circuits, and *every* error (middleware
822    /// `Err`, body parse failure, 404 / 405 from the router, handler-returned
823    /// `Err`, the 415 / 501 for an encoded body, even the 413 / 400 from the
824    /// body-cap path). The after-chain, compression, and CORS apply
825    /// uniformly. CORS preflight is the one short-circuit that bypasses the
826    /// pipeline — it's HTTP-protocol traffic rather than an application
827    /// request.
828    async fn handle_request(
829        self: Arc<Self>,
830        req: HyperRequest<Incoming>,
831    ) -> Result<HyperResponse<ResponseBody>, hyper::Error> {
832        let timeout = self.request_timeout;
833        let app = self.clone();
834        let inner = app.handle_request_inner(req);
835        match timeout {
836            None => inner.await,
837            Some(d) => match tokio::time::timeout(d, inner).await {
838                Ok(r) => r,
839                Err(_) => {
840                    warn!(timeout = ?d, "request exceeded configured timeout");
841                    Ok(self.finalizer.build_error(WebError::Timeout).await)
842                }
843            },
844        }
845    }
846
847    /// The actual request pipeline. Split from
848    /// [`handle_request`](Self::handle_request) so the latter can wrap it
849    /// in a timeout when one is configured.
850    ///
851    /// **Lifecycle order:**
852    ///
853    /// 1. capture WS upgrade (if request looks like a handshake)
854    /// 2. build the `Request` skeleton (no body yet) — and, when the router
855    ///    has host aliases, route an aliased request as the path it names
856    /// 3. CORS preflight short-circuit (uses headers only)
857    /// 4. match controller — 404 short-circuits *without* buffering the
858    ///    body (efficiency win on adversarial bad-path requests); then stamp
859    ///    the matched controller's rate-limit class onto the request, so a
860    ///    `before` middleware (which only gets `&Request`) can read it; then
861    ///    decide from its codings how to buffer the body — as sent, decoded
862    ///    from gzip, or refused (`415` / `501`) — also without reading it
863    /// 5. buffer the body, capped per the resolved policy (today: server-
864    ///    wide; soon, per-controller / per-route)
865    /// 6. middleware `before`
866    /// 7. `to_params` (Content-Type-driven body parse)
867    /// 8. dispatch via the already-matched controller
868    /// 9. middleware `after` + finalize
869    ///
870    /// The "route before buffer" order is what lets the body cap depend on
871    /// the matched route — without it, the framework would have to commit
872    /// to a single cap before knowing where the request is headed.
873    async fn handle_request_inner(
874        self: Arc<Self>,
875        #[allow(unused_mut)] mut req: HyperRequest<Incoming>,
876    ) -> Result<HyperResponse<ResponseBody>, hyper::Error> {
877        let request_span = span!(Level::INFO, "request");
878        async move {
879            // 1. Capture the WS upgrade handshake (if any) before
880            //    `from_hyper_parts` consumes the request: the `OnUpgrade`
881            //    future and the derived `Sec-WebSocket-Accept`. (See
882            //    `websocket` module docs for why this happens up front.)
883            #[cfg(feature = "websocket")]
884            let ws_upgrade: Option<(hyper::upgrade::OnUpgrade, HeaderValue)> =
885                if websocket::is_upgrade_request(req.method(), req.headers()) {
886                    websocket::accept_key(req.headers())
887                        .map(|accept| (hyper::upgrade::on(&mut req), accept))
888                } else {
889                    None
890                };
891
892            // 2. Build the skeleton (method / path / query / headers); the
893            //    body stream is held aside for step 5.
894            //
895            //    A host alias is applied here, before anything else sees the
896            //    request: from now on, for the router, every middleware and
897            //    the handler, it *is* the request for the aliased path. With
898            //    no aliases registered this is the only host work there is —
899            //    one check; `Host` is not even read. The host is read before
900            //    the skeleton takes the request, because an absolute-form
901            //    request's authority is in the URI, which the skeleton does
902            //    not keep.
903            let (mut request, body_stream, alias_prefix) = if self.router.has_host_aliases() {
904                let hit = self.router.host_alias_hit(req.uri(), req.headers());
905                let (mut request, body_stream) = Request::from_hyper_parts(req);
906                let alias_prefix = hit.and_then(|hit| hit.apply(&mut request.path_parts));
907                (request, body_stream, alias_prefix)
908            } else {
909                let (request, body_stream) = Request::from_hyper_parts(req);
910                (request, body_stream, None)
911            };
912
913            // 3. CORS preflight: synthesize the 204 ourselves before any
914            //    application-layer work. Preflights are HTTP-protocol
915            //    traffic; neither `before` nor `after` middleware runs on
916            //    them (see `CLAUDE.md` principle 1).
917            if let Some(cors) = &self.cors
918                && CorsLayer::is_preflight(&request.method, &request.headers)
919            {
920                let mut resp = self.finalizer.build_response(ReplyData::Empty).await;
921                cors.apply(&request.headers, resp.headers_mut(), true);
922                return Ok(resp);
923            }
924
925            // 4. Match controller. A path that hits nothing is 404 *before*
926            //    body buffering — a 10 MiB POST to a non-existent URL no
927            //    longer wastes 10 MiB of memory.
928            let route_match = match self.router.match_controller(&request.path_parts) {
929                Some(rm) => rm,
930                None => {
931                    return Ok(self
932                        .finalize_error(
933                            WebError::NotFound,
934                            &request,
935                            #[cfg(feature = "websocket")]
936                            ws_upgrade,
937                        )
938                        .await);
939                }
940            };
941
942            // 4b. Stamp the matched controller's rate-limit class onto the
943            //     request (the skeleton predates routing, so it was `None`).
944            //     This is the one piece of routing context a `before`
945            //     middleware can't otherwise see — it gets `&mut Request`,
946            //     not the matched controller. An application rate-limit
947            //     middleware reads `request.rate_limit_class` and applies its
948            //     own per-class policy; the framework owns the label and the
949            //     `429` response, not the limiter. (Set before the body
950            //     buffer so it survives `collect_body` and is present for the
951            //     whole pipeline, including the after-chain on error replies.)
952            request.rate_limit_class = route_match.controller.actus_rate_limit();
953
954            // 4c. Decide from the body's codings how to buffer it — as sent,
955            //     decoded from gzip, or not at all (`415` / `501`, see
956            //     `body_coding`) — before reading it, so a refused body never
957            //     touches the body cap or the inflight budget.
958            let coding = match self.body_coding(&request.headers, &body_stream) {
959                Ok(coding) => coding,
960                Err(refusal) => {
961                    return Ok(self
962                        .finalize_reply(
963                            refusal,
964                            &request,
965                            #[cfg(feature = "websocket")]
966                            ws_upgrade,
967                        )
968                        .await);
969                }
970            };
971
972            // 5. Resolve the effective body cap: the matched controller's
973            //    `#[controller(max_body_bytes = …)]` if it set one, otherwise
974            //    the server-wide `with_max_body_bytes` cap (default 2 MiB). A
975            //    future Phase 2 adds per-route overrides at the top of this
976            //    fall-through.
977            //
978            //    The error path returns the same skeleton so the after-chain
979            //    still has a `Request`.
980            //
981            //    A gzip body is decoded as it is buffered, the cap holding its
982            //    decoded size as well as its coded one.
983            let effective_cap = route_match
984                .controller
985                .actus_max_body_bytes()
986                .unwrap_or(self.max_body_bytes);
987            let budget = self.max_inflight_body_bytes.as_ref();
988            let collected = match coding {
989                BodyCoding::AsSent => {
990                    request
991                        .collect_body(body_stream, effective_cap, budget)
992                        .await
993                }
994                #[cfg(feature = "compression")]
995                BodyCoding::Gzip => {
996                    request
997                        .collect_gzip_body(body_stream, effective_cap, budget)
998                        .await
999                }
1000            };
1001            request = match collected {
1002                Ok(r) => r,
1003                Err((request, e)) => {
1004                    warn!("rejecting request before parse: {}", e);
1005                    return Ok(self
1006                        .finalize_error(
1007                            e,
1008                            &request,
1009                            #[cfg(feature = "websocket")]
1010                            ws_upgrade,
1011                        )
1012                        .await);
1013                }
1014            };
1015
1016            // 6. Middleware `before` chain.
1017            let pre_data: Option<ReplyData> =
1018                match self.middleware_chain.process_request(&mut request).await {
1019                    Ok(Outcome::Continue) => None,
1020                    Ok(Outcome::Respond(data)) => Some(data),
1021                    Err(e) => {
1022                        return Ok(self
1023                            .finalize_error(
1024                                e,
1025                                &request,
1026                                #[cfg(feature = "websocket")]
1027                                ws_upgrade,
1028                            )
1029                            .await);
1030                    }
1031                };
1032
1033            // A `before` hook short-circuited with a reply — skip routing
1034            // and the handler, but still run the after-chain.
1035            if let Some(data) = pre_data {
1036                return Ok(self
1037                    .finalize_reply(
1038                        data,
1039                        &request,
1040                        #[cfg(feature = "websocket")]
1041                        ws_upgrade,
1042                    )
1043                    .await);
1044            }
1045
1046            // 7. Body parse (JSON / form / opaque, per Content-Type).
1047            //    Malformed body → 400 through the after-chain.
1048            let mut params = match request.to_params() {
1049                Ok(p) => p,
1050                Err(e) => {
1051                    return Ok(self
1052                        .finalize_error(
1053                            e,
1054                            &request,
1055                            #[cfg(feature = "websocket")]
1056                            ws_upgrade,
1057                        )
1058                        .await);
1059                }
1060            };
1061            // The router's answer to "did an alias route this?", for the
1062            // handler — so it never has to match `Host` a second time.
1063            if let Some(prefix) = alias_prefix {
1064                params.set_alias_prefix(prefix);
1065            }
1066
1067            // 8. Dispatch via the matched controller. 405 (verb mismatch
1068            //    inside the controller) and handler-returned errors both
1069            //    come back through here.
1070            match route_match
1071                .controller
1072                .actus_dispatch(&route_match.action, params)
1073                .await
1074            {
1075                Ok(data) => Ok(self
1076                    .finalize_reply(
1077                        data,
1078                        &request,
1079                        #[cfg(feature = "websocket")]
1080                        ws_upgrade,
1081                    )
1082                    .await),
1083                Err(e) => Ok(self
1084                    .finalize_error(
1085                        e,
1086                        &request,
1087                        #[cfg(feature = "websocket")]
1088                        ws_upgrade,
1089                    )
1090                    .await),
1091            }
1092        }
1093        .instrument(request_span)
1094        .await
1095    }
1096}
1097
1098/// Default shutdown trigger: resolves on SIGTERM, SIGINT (Unix), or Ctrl-C
1099/// (Windows). This is what [`Server::run`] uses; for tests or embedding,
1100/// see [`Server::run_with_shutdown`].
1101async fn default_shutdown_signal() {
1102    #[cfg(unix)]
1103    {
1104        use tokio::signal::unix::{SignalKind, signal};
1105        let mut sigterm = signal(SignalKind::terminate()).expect("install SIGTERM handler");
1106        let mut sigint = signal(SignalKind::interrupt()).expect("install SIGINT handler");
1107        tokio::select! {
1108            _ = sigterm.recv() => info!("Received SIGTERM"),
1109            _ = sigint.recv() => info!("Received SIGINT"),
1110        }
1111    }
1112    #[cfg(not(unix))]
1113    {
1114        tokio::signal::ctrl_c()
1115            .await
1116            .expect("install Ctrl-C handler");
1117        info!("Received Ctrl-C");
1118    }
1119}
1120
1121/// How a request body is buffered, as [`Server::body_coding`] decided from its
1122/// codings.
1123enum BodyCoding {
1124    /// As sent: there is no coding to undo.
1125    AsSent,
1126    /// Decoded from `gzip` ([`Server::with_gzip_request_bodies`]).
1127    #[cfg(feature = "compression")]
1128    Gzip,
1129}
1130
1131/// An accepted connection that adds every byte written to it to the
1132/// [`Server::with_sent_bytes`] counter, when one is configured. Reads pass
1133/// through untouched.
1134struct CountingStream {
1135    inner: TcpStream,
1136    sent_bytes: Option<Arc<AtomicU64>>,
1137}
1138
1139impl CountingStream {
1140    /// Count a write's outcome: the bytes the operating system accepted.
1141    fn count(&self, written: &Poll<io::Result<usize>>) {
1142        if let (Some(counter), Poll::Ready(Ok(n))) = (&self.sent_bytes, written) {
1143            counter.fetch_add(*n as u64, Ordering::Relaxed);
1144        }
1145    }
1146}
1147
1148impl AsyncRead for CountingStream {
1149    fn poll_read(
1150        self: Pin<&mut Self>,
1151        cx: &mut Context<'_>,
1152        buf: &mut ReadBuf<'_>,
1153    ) -> Poll<io::Result<()>> {
1154        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
1155    }
1156}
1157
1158impl AsyncWrite for CountingStream {
1159    fn poll_write(
1160        self: Pin<&mut Self>,
1161        cx: &mut Context<'_>,
1162        buf: &[u8],
1163    ) -> Poll<io::Result<usize>> {
1164        let this = self.get_mut();
1165        let written = Pin::new(&mut this.inner).poll_write(cx, buf);
1166        this.count(&written);
1167        written
1168    }
1169
1170    fn poll_write_vectored(
1171        self: Pin<&mut Self>,
1172        cx: &mut Context<'_>,
1173        bufs: &[io::IoSlice<'_>],
1174    ) -> Poll<io::Result<usize>> {
1175        let this = self.get_mut();
1176        let written = Pin::new(&mut this.inner).poll_write_vectored(cx, bufs);
1177        this.count(&written);
1178        written
1179    }
1180
1181    // hyper picks its write strategy from this; forwarding it keeps the
1182    // vectored writes a bare `TcpStream` gets.
1183    fn is_write_vectored(&self) -> bool {
1184        self.inner.is_write_vectored()
1185    }
1186
1187    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
1188        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
1189    }
1190
1191    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
1192        Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
1193    }
1194}