axum-server 0.9.1

use crate::{
    accept::{Accept, DefaultAcceptor},
    connection::{ConnectionShutdown, ShutdownKind},
    handle::Handle,
    service::{MakeService, SendService},
};
use either::Either;
use http::Request;
use hyper::body::Incoming;
use hyper_util::{
    rt::{TokioExecutor, TokioIo, TokioTimer},
    server::conn::auto::Builder,
    service::TowerToHyperService,
};
use std::{
    fmt,
    future::poll_fn,
    io::{self, ErrorKind},
    net::SocketAddr as IpSocketAddr,
    sync::{Arc, Mutex},
    task::{Context, Poll},
    time::Duration,
};
use tokio::{
    io::{AsyncRead, AsyncWrite},
    net::{TcpListener, TcpStream},
    sync::{oneshot, watch},
};
use tower_service::Service;

#[cfg(unix)]
use {
    std::os::unix::net::SocketAddr as UnixSocketAddr,
    tokio::net::{UnixListener, UnixStream},
};

/// HTTP server.
pub struct Server<Addr: Address, A = DefaultAcceptor> {
    acceptor: A,
    builder: Builder<TokioExecutor>,
    listener: Listener<Addr>,
    handle: Handle<Addr>,
    http_version: Option<HttpVersion>,
    first_request_timeout: Option<Duration>,
    accept_error_handler: Option<AcceptErrorHandler<Addr>>,
}

type AcceptErrorHandler<A> = Arc<dyn Fn(Option<&A>, &io::Error) + Send + Sync>;

// Builder doesn't implement Debug or Clone right now
impl<A: Address, B> fmt::Debug for Server<A, B>
where
    Listener<A>: fmt::Debug,
    Handle<A>: fmt::Debug,
    B: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Server")
            .field("acceptor", &self.acceptor)
            .field("listener", &self.listener)
            .field("handle", &self.handle)
            .finish_non_exhaustive()
    }
}

#[derive(Debug)]
enum Listener<A: Address> {
    Bind(A),
    Ready(A::Listener),
}

/// Create a [`Server`] that will bind to provided address.
pub fn bind<A: Address>(addr: A) -> Server<A> {
    Server::bind(addr)
}

/// Create a [`Server`] from existing `std::net::TcpListener`.
///
/// The listener is put into non-blocking mode.
pub fn from_tcp(listener: std::net::TcpListener) -> io::Result<Server<IpSocketAddr>> {
    // `tokio::net::TcpListener::from_std` requires the listener to be in
    // non-blocking mode; callers of this function should not have to know that.
    listener.set_nonblocking(true)?;

    Ok(Server::from_listener(TcpListener::from_std(listener)?))
}

/// Create a [`Server`] from existing `std::os::unix::net::UnixListener`.
///
/// The listener is put into non-blocking mode.
#[cfg(unix)]
pub fn from_unix(listener: std::os::unix::net::UnixListener) -> io::Result<Server<UnixSocketAddr>> {
    // `tokio::net::UnixListener::from_std` requires the listener to be in
    // non-blocking mode; callers of this function should not have to know that.
    listener.set_nonblocking(true)?;

    Ok(Server::from_listener(UnixListener::from_std(listener)?))
}

/// A trait to abstract over a listener that can accept an incoming connection of different types
/// and can be created from an address.
pub trait AddrListener<Stream, Addr: Address>: std::marker::Sized {
    /// Create a new `AddrListener` bound to the given address.
    fn bind_to(addr: Addr) -> impl std::future::Future<Output = io::Result<Self>> + Send;

    /// Accept an incoming connection, returning the stream and the remote address.
    fn accept_stream(&self)
        -> impl std::future::Future<Output = io::Result<(Stream, Addr)>> + Send;

    /// Get the address we're listening on.
    fn get_local_addr(&self) -> io::Result<Addr>;
}

#[cfg(unix)]
impl AddrListener<UnixStream, UnixSocketAddr> for UnixListener {
    async fn bind_to(addr: UnixSocketAddr) -> io::Result<Self> {
        UnixListener::bind(addr.as_pathname().ok_or_else(|| {
            io::Error::new(
                ErrorKind::InvalidInput,
                "A UnixListener can only be bound to a path address!",
            )
        })?)
    }

    async fn accept_stream(&self) -> io::Result<(UnixStream, UnixSocketAddr)> {
        let (stream, tokio_addr) = self.accept().await?;
        Ok((stream, tokio_addr.into()))
    }

    fn get_local_addr(&self) -> io::Result<UnixSocketAddr> {
        self.local_addr().map(tokio::net::unix::SocketAddr::into)
    }
}

impl AddrListener<TcpStream, IpSocketAddr> for TcpListener {
    async fn bind_to(addr: IpSocketAddr) -> io::Result<Self> {
        TcpListener::bind(addr).await
    }

    async fn accept_stream(&self) -> io::Result<(TcpStream, IpSocketAddr)> {
        self.accept().await
    }

    fn get_local_addr(&self) -> io::Result<IpSocketAddr> {
        self.local_addr()
    }
}

/// An address which an associated `AddrListener` can bind to.
pub trait Address: std::marker::Sized + Clone {
    /// The underlying stream we will obtain once an `AddrListener` is bound to this address.
    type Stream;

    /// An `AddrListener` that can bind to this `Address`.
    type Listener: AddrListener<Self::Stream, Self>;
}

#[cfg(unix)]
impl Address for UnixSocketAddr {
    type Stream = UnixStream;
    type Listener = UnixListener;
}

impl Address for IpSocketAddr {
    type Stream = TcpStream;
    type Listener = TcpListener;
}

/// Create the HTTP protocol builder used by [`Server`].
///
/// A timer is installed by default so that `header_read_timeout` and
/// `keep_alive_timeout` work without extra setup.
fn default_builder() -> Builder<TokioExecutor> {
    let mut builder = Builder::new(TokioExecutor::new());

    // Hyper defaults to a 30 second header read timeout. Without a timer that
    // default was silently ignored, so keep it disabled to preserve the
    // previous behavior; callers can still opt in through `http_builder`.
    builder
        .http1()
        .timer(TokioTimer::new())
        .header_read_timeout(None);
    builder.http2().timer(TokioTimer::new());

    builder
}

impl<A: Address> Server<A> {
    /// Create a server that will bind to provided address.
    pub fn bind(addr: A) -> Self {
        let acceptor = DefaultAcceptor::new();
        let builder = default_builder();
        let handle = Handle::new();

        Self {
            acceptor,
            builder,
            listener: Listener::Bind(addr),
            handle,
            http_version: None,
            first_request_timeout: None,
            accept_error_handler: None,
        }
    }

    /// Create a server from an existing `AddrListener`.
    pub fn from_listener(listener: A::Listener) -> Self {
        let acceptor = DefaultAcceptor::new();
        let builder = default_builder();
        let handle = Handle::new();

        Self {
            acceptor,
            builder,
            listener: Listener::Ready(listener),
            handle,
            http_version: None,
            first_request_timeout: None,
            accept_error_handler: None,
        }
    }
}

#[derive(Clone, Copy, Eq, PartialEq)]
enum HttpVersion {
    Http1,
    Http2,
}

impl<A: Address, Acc> Server<A, Acc> {
    /// Overwrite acceptor.
    pub fn acceptor<Acceptor>(self, acceptor: Acceptor) -> Server<A, Acceptor> {
        Server {
            acceptor,
            builder: self.builder,
            listener: self.listener,
            handle: self.handle,
            http_version: self.http_version,
            first_request_timeout: self.first_request_timeout,
            accept_error_handler: self.accept_error_handler,
        }
    }

    /// Map acceptor.
    pub fn map<Acceptor, F>(self, acceptor: F) -> Server<A, Acceptor>
    where
        F: FnOnce(Acc) -> Acceptor,
    {
        Server {
            acceptor: acceptor(self.acceptor),
            builder: self.builder,
            listener: self.listener,
            handle: self.handle,
            http_version: self.http_version,
            first_request_timeout: self.first_request_timeout,
            accept_error_handler: self.accept_error_handler,
        }
    }

    /// Returns a reference to the acceptor.
    pub fn get_ref(&self) -> &Acc {
        &self.acceptor
    }

    /// Returns a mutable reference to the acceptor.
    pub fn get_mut(&mut self) -> &mut Acc {
        &mut self.acceptor
    }

    /// Returns a mutable reference to the Http builder.
    pub fn http_builder(&mut self) -> &mut Builder<TokioExecutor> {
        &mut self.builder
    }

    /// Only accepts HTTP/1.
    ///
    /// Note that HTTP/1 protocol upgrades (for example `WebSockets`) are not
    /// supported when this is used; the default protocol-detecting mode does
    /// support them.
    pub fn http1_only(mut self) -> Self {
        self.http_version = Some(HttpVersion::Http1);
        self.builder = self.builder.http1_only();
        self
    }

    /// Only accepts HTTP/2.
    pub fn http2_only(mut self) -> Self {
        self.http_version = Some(HttpVersion::Http2);
        self.builder = self.builder.http2_only();
        self
    }

    /// Provide a handle for additional utilities.
    ///
    /// To share shutdown and connection counting between servers listening on
    /// different address types, build each server's handle from one
    /// [`Control`](crate::Control) with [`Handle::with_control`].
    pub fn handle(mut self, handle: Handle<A>) -> Self {
        self.handle = handle;
        self
    }

    /// Set a timeout for the first request on a connection.
    ///
    /// If the client does not send a complete first request within this
    /// duration, the connection is closed. This bounds the time a newly
    /// accepted connection can occupy resources before the HTTP protocol is
    /// determined, which is otherwise not covered by
    /// [`http_builder`](Self::http_builder) timeouts.
    ///
    /// The timeout starts after the connection acceptor finishes. For a TLS
    /// server this means the TLS handshake is bounded by the acceptor's own
    /// handshake timeout instead.
    ///
    /// `None` disables the timeout. The default is `None`.
    pub fn first_request_timeout(mut self, timeout: impl Into<Option<Duration>>) -> Self {
        self.first_request_timeout = timeout.into();
        self
    }

    /// Set a handler called when accepting a connection fails.
    ///
    /// The handler receives the remote address when it is known and the error
    /// that occurred. Accept errors do not stop the server; it backs off and
    /// keeps accepting.
    ///
    /// The handler is called from the accept loop, so keep it cheap. Use it for
    /// logging or metrics rather than blocking work.
    pub fn accept_error_handler<F>(mut self, handler: F) -> Self
    where
        F: Fn(Option<&A>, &io::Error) + Send + Sync + 'static,
    {
        self.accept_error_handler = Some(Arc::new(handler));
        self
    }

    /// Serve provided [`MakeService`].
    ///
    /// To create [`MakeService`] easily, `Shared` from [`tower`] can be used.
    ///
    /// # Errors
    ///
    /// An error will be returned when:
    ///
    /// - Binding to an address fails.
    /// - `make_service` returns an error when `poll_ready` is called. This never happens on
    ///   [`axum`] make services.
    ///
    /// [`axum`]: https://docs.rs/axum/0.8
    /// [`tower`]: https://docs.rs/tower
    /// [`MakeService`]: https://docs.rs/tower/0.5/tower/make/trait.MakeService.html
    pub async fn serve<M>(self, mut make_service: M) -> io::Result<()>
    where
        M: MakeService<A, Request<Incoming>>,
        A: Send + 'static,
        A::Stream: Send,
        Acc: Accept<A::Stream, M::Service> + Clone + Send + Sync + 'static,
        Acc::Stream: AsyncRead + AsyncWrite + Unpin + Send,
        Acc::Service: SendService<Request<Incoming>> + Send,
        Acc::Future: Send,
    {
        let acceptor = self.acceptor;
        let handle = self.handle;
        let builder = std::sync::Arc::new(self.builder);
        let first_request_timeout = self.first_request_timeout;
        let accept_error_handler = self.accept_error_handler;

        let mut incoming = match bind_incoming(self.listener).await {
            Ok(v) => v,
            Err(e) => {
                handle.notify_listening(None);
                return Err(e);
            }
        };

        handle.notify_listening(incoming.get_local_addr().ok());

        let accept_loop_future = async {
            loop {
                let (tcp_stream, socket_addr) = tokio::select! {
                    biased;
                    result = accept(&mut incoming, accept_error_handler.as_ref()) => result,
                    _ = handle.wait_graceful_shutdown() => return Ok(()),
                };

                poll_fn(|cx| make_service.poll_ready(cx))
                    .await
                    .map_err(io_other)?;

                let remote_addr = socket_addr.clone();
                let service = match make_service.make_service(socket_addr).await {
                    Ok(service) => service,
                    Err(e) => {
                        if let Some(handler) = &accept_error_handler {
                            handler(Some(&remote_addr), &io_other(e));
                        }
                        continue;
                    }
                };

                let acceptor = acceptor.clone();
                let watcher = handle.watcher();
                let builder = builder.clone();
                let http_version = self.http_version;
                let accept_error_handler = accept_error_handler.clone();

                tokio::spawn(async move {
                    // Do not let a connection that is still in the TLS handshake
                    // keep a task (and its file descriptor) alive after shutdown.
                    let accepted = tokio::select! {
                        biased;
                        _ = watcher.wait_shutdown() => return,
                        _ = watcher.wait_graceful_shutdown() => return,
                        accepted = acceptor.accept(tcp_stream, service) => accepted,
                    };

                    let (stream, send_service) = match accepted {
                        Ok(accepted) => accepted,
                        Err(e) => {
                            if let Some(handler) = &accept_error_handler {
                                handler(Some(&remote_addr), &e);
                            }
                            return;
                        }
                    };

                    let io = TokioIo::new(stream);
                    let service = send_service.into_service();

                    let (shutdown_tx, mut shutdown_rx) = watch::channel(ShutdownKind::None);

                    let (first_request_rx, notify_first_request) = match first_request_timeout {
                        Some(_) => {
                            let (tx, rx) = oneshot::channel();
                            (Some(rx), Some(Arc::new(Mutex::new(Some(tx)))))
                        }
                        None => (None, None),
                    };

                    let service = FirstRequestService::new(service, notify_first_request);
                    let service = ConnectionService::new(
                        service,
                        ConnectionShutdown::new(shutdown_tx.clone()),
                    );
                    let service = TowerToHyperService::new(service);
                    let serve_future = match http_version {
                        Some(_) => Either::Left(builder.serve_connection(io, service)),
                        _ => Either::Right(builder.serve_connection_with_upgrades(io, service)),
                    };
                    tokio::pin!(serve_future);
                    let mut serve_future = serve_future.as_pin_mut();

                    // Bound the time before the first request is dispatched. This
                    // also covers hyper's un-timed protocol detection.
                    if let (Some(timeout), Some(mut first_request_rx)) =
                        (first_request_timeout, first_request_rx)
                    {
                        tokio::select! {
                            biased;
                            _ = watcher.wait_shutdown() => return,
                            _ = watcher.wait_graceful_shutdown() => {}
                            _ = tokio::time::sleep(timeout) => return,
                            result = &mut first_request_rx => {
                                if result.is_err() {
                                    return;
                                }
                            }
                            _ = &mut serve_future => return,
                        }
                    }

                    let mut graceful = false;
                    let mut apply_graceful = false;

                    loop {
                        if apply_graceful {
                            apply_graceful = false;

                            match &mut serve_future {
                                Either::Left(serve_future) => {
                                    serve_future.as_mut().graceful_shutdown()
                                }
                                Either::Right(serve_future) => {
                                    serve_future.as_mut().graceful_shutdown()
                                }
                            }
                        }

                        tokio::select! {
                            biased;
                            _ = watcher.wait_shutdown() => return,
                            _ = watcher.wait_graceful_shutdown(), if !graceful => {
                                graceful = true;
                                apply_graceful = true;
                            }
                            // Keeping `shutdown_tx` alive means this only
                            // resolves when a handler asks to close.
                            _ = shutdown_rx.changed() => {
                                match *shutdown_rx.borrow() {
                                    ShutdownKind::Abort => return,
                                    ShutdownKind::Graceful if !graceful => {
                                        graceful = true;
                                        apply_graceful = true;
                                    }
                                    ShutdownKind::None | ShutdownKind::Graceful => {}
                                }
                            }
                            _ = &mut serve_future => return,
                        }
                    }
                });
            }
        };

        let result = tokio::select! {
            biased;
            _ = handle.wait_shutdown() => return Ok(()),
            result = accept_loop_future => result,
        };

        // Tokio internally accepts TCP connections while the TCPListener is active;
        // drop the listener to immediately refuse connections rather than letting
        // them hang.
        drop(incoming);

        // attempting to do a "result?;" requires us to specify the type of result which is annoying
        #[allow(clippy::question_mark)]
        if let Err(e) = result {
            return Err(e);
        }

        handle.wait_connections_end().await;

        Ok(())
    }
}

async fn bind_incoming<A: Address>(listener: Listener<A>) -> io::Result<A::Listener> {
    match listener {
        Listener::Bind(addr) => A::Listener::bind_to(addr).await,
        Listener::Ready(listener) => Ok(listener),
    }
}

pub(crate) async fn accept<L: AddrListener<S, A>, S, A: Address>(
    listener: &mut L,
    error_handler: Option<&AcceptErrorHandler<A>>,
) -> (S, A) {
    loop {
        match listener.accept_stream().await {
            Ok(value) => return value,
            Err(e) => {
                if let Some(handler) = error_handler {
                    handler(None, &e);
                }

                // Avoid a hot loop when the listener is failing continuously,
                // for example because of `EMFILE`.
                tokio::time::sleep(Duration::from_millis(50)).await;
            }
        }
    }
}

/// A service wrapper that reports when the first request is dispatched.
#[derive(Clone)]
struct FirstRequestService<S> {
    inner: S,
    notify: Option<Arc<Mutex<Option<oneshot::Sender<()>>>>>,
}

impl<S> FirstRequestService<S> {
    fn new(inner: S, notify: Option<Arc<Mutex<Option<oneshot::Sender<()>>>>>) -> Self {
        Self { inner, notify }
    }
}

impl<S> Service<Request<Incoming>> for FirstRequestService<S>
where
    S: Service<Request<Incoming>>,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request<Incoming>) -> Self::Future {
        if let Some(notify) = &self.notify {
            if let Some(tx) = notify.lock().unwrap().take() {
                let _ = tx.send(());
            }
        }

        self.inner.call(req)
    }
}

/// A service wrapper that adds the connection shutdown handle to request
/// extensions.
#[derive(Clone, Debug)]
struct ConnectionService<S> {
    inner: S,
    shutdown: ConnectionShutdown,
}

impl<S> ConnectionService<S> {
    fn new(inner: S, shutdown: ConnectionShutdown) -> Self {
        Self { inner, shutdown }
    }
}

impl<S> Service<Request<Incoming>> for ConnectionService<S>
where
    S: Service<Request<Incoming>>,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut req: Request<Incoming>) -> Self::Future {
        req.extensions_mut().insert(self.shutdown.clone());

        self.inner.call(req)
    }
}

type BoxError = Box<dyn std::error::Error + Send + Sync>;

pub(crate) fn io_other<E: Into<BoxError>>(error: E) -> io::Error {
    io::Error::other(error)
}

#[cfg(test)]
mod tests {
    use crate::{
        accept::{Accept, LocalAddr, LocalAddrAcceptor},
        handle::{Control, Handle},
        server::Server,
        ConnectionShutdown,
    };
    use axum::body::Body;
    use axum::extract::State;
    use axum::response::Response;
    use axum::routing::post;
    use axum::{routing::get, Extension, Router};
    use bytes::Bytes;
    use futures_util::{stream, StreamExt};
    use http::{Method, Request, Uri};
    use http_body::Frame;
    use http_body_util::{BodyExt, StreamBody};
    use hyper::client;
    use hyper::client::conn::{http1, http2};
    use hyper_util::rt::{TokioExecutor, TokioIo};
    use std::{
        io,
        net::SocketAddr as IpSocketAddr,
        sync::{Arc, Mutex},
        time::Duration,
    };
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::sync::oneshot;
    use tokio::{net::TcpStream, task::JoinHandle, time::timeout};

    #[derive(Clone, Copy)]
    enum TestHttpVersion {
        Http1,
        Http2,
    }

    #[tokio::test]
    async fn start_and_request() {
        let (_handle, _server_task, addr) = start_server().await;

        let (mut client, _conn) = connect(addr).await;

        // Client can send requests

        do_empty_request_h1(&mut client).await.unwrap();

        do_slow_request(&mut client, Duration::from_millis(50))
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_shutdown() {
        let (handle, _server_task, addr) = start_server().await;

        let (mut client, conn) = connect(addr).await;

        // Client can send request before shutdown.
        do_empty_request_h1(&mut client).await.unwrap();

        handle.shutdown();

        // After shutdown, all client requests should fail.
        do_empty_request_h1(&mut client).await.unwrap_err();

        // Connection should finish soon.
        let _ = timeout(Duration::from_secs(1), conn).await.unwrap();
    }

    #[tokio::test]
    async fn connection_graceful_shutdown_sends_response_then_closes() {
        let (handle, _server_task, addr) = start_server_with_app(Router::new().route(
            "/",
            get(
                |Extension(shutdown): Extension<ConnectionShutdown>| async move {
                    shutdown.graceful_shutdown();
                    "goodbye"
                },
            ),
        ))
        .await;

        let (mut client, conn) = connect(addr).await;

        client.ready().await.unwrap();
        let response = client
            .send_request(Request::new(Body::empty()))
            .await
            .unwrap();

        let body = response.into_body().collect().await.unwrap().to_bytes();
        assert_eq!(body.as_ref(), b"goodbye");

        timeout(Duration::from_secs(1), conn)
            .await
            .unwrap()
            .unwrap();
        wait_for_no_connections(&handle).await;

        do_empty_request_h1(&mut client).await.unwrap_err();
    }

    #[derive(Clone)]
    struct ShutdownHolder(Arc<Mutex<Option<ConnectionShutdown>>>);

    async fn stash_shutdown(
        State(holder): State<ShutdownHolder>,
        Extension(shutdown): Extension<ConnectionShutdown>,
    ) -> &'static str {
        *holder.0.lock().unwrap() = Some(shutdown);
        "ok"
    }

    #[tokio::test]
    async fn connection_shutdown_aborts_connection() {
        let holder = ShutdownHolder(Arc::new(Mutex::new(None)));

        let app = Router::new()
            .route("/", get(stash_shutdown))
            .with_state(holder.clone());

        let (handle, _server_task, addr) = start_server_with_app(app).await;

        let (mut client, _conn) = connect(addr).await;

        client.ready().await.unwrap();
        let response = client
            .send_request(Request::new(Body::empty()))
            .await
            .unwrap();

        let body = response.into_body().collect().await.unwrap().to_bytes();
        assert_eq!(body.as_ref(), b"ok");

        holder.0.lock().unwrap().take().unwrap().shutdown();

        wait_for_no_connections(&handle).await;

        do_empty_request_h1(&mut client).await.unwrap_err();
    }

    #[tokio::test]
    async fn connection_graceful_shutdown_closes_http2_connection() {
        let (handle, _server_task, addr) = start_server_with_app(Router::new().route(
            "/",
            get(
                |Extension(shutdown): Extension<ConnectionShutdown>| async move {
                    shutdown.graceful_shutdown();
                    "goodbye"
                },
            ),
        ))
        .await;

        let (mut client, conn) = connect_h2(addr).await;

        client.ready().await.unwrap();
        let response = client
            .send_request(Request::new(Body::empty()))
            .await
            .unwrap();

        let body = response.into_body().collect().await.unwrap().to_bytes();
        assert_eq!(body.as_ref(), b"goodbye");

        timeout(Duration::from_secs(1), conn)
            .await
            .unwrap()
            .unwrap();
        wait_for_no_connections(&handle).await;
    }

    async fn wait_for_no_connections(handle: &Handle<IpSocketAddr>) {
        timeout(Duration::from_secs(1), async {
            while handle.connection_count() != 0 {
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .unwrap();
    }

    #[tokio::test]
    async fn control_shares_tcp_servers() {
        let control = Control::new();
        let app = Router::new().route("/", get(|| async { "Hello, world!" }));

        let handle_a = Handle::<IpSocketAddr>::with_control(control.clone());
        let handle_b = Handle::<IpSocketAddr>::with_control(control.clone());

        let server_a = tokio::spawn({
            let app = app.clone();
            let server =
                Server::bind(IpSocketAddr::from(([127, 0, 0, 1], 0))).handle(handle_a.clone());

            async move { server.serve(app.into_make_service()).await }
        });
        let server_b = tokio::spawn({
            let app = app.clone();
            let server =
                Server::bind(IpSocketAddr::from(([127, 0, 0, 1], 0))).handle(handle_b.clone());

            async move { server.serve(app.into_make_service()).await }
        });

        let addr_a = handle_a.listening().await.unwrap();
        let addr_b = handle_b.listening().await.unwrap();
        assert_ne!(addr_a, addr_b);

        let (mut client_a, _conn_a) = connect(addr_a).await;
        let (mut client_b, _conn_b) = connect(addr_b).await;
        do_empty_request_h1(&mut client_a).await.unwrap();
        do_empty_request_h1(&mut client_b).await.unwrap();

        assert_eq!(control.connection_count(), 2);

        control.graceful_shutdown(None);

        timeout(Duration::from_secs(1), server_a)
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        timeout(Duration::from_secs(1), server_b)
            .await
            .unwrap()
            .unwrap()
            .unwrap();

        assert_eq!(control.connection_count(), 0);
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn control_shares_tcp_and_unix_servers() {
        use std::os::unix::net::{SocketAddr as UnixSocketAddr, UnixListener as StdUnixListener};
        use tokio::net::UnixStream;

        let control = Control::new();
        let app = Router::new().route("/", get(|| async { "Hello, world!" }));

        let tcp_handle = Handle::<IpSocketAddr>::with_control(control.clone());
        let unix_handle = Handle::<UnixSocketAddr>::with_control(control.clone());

        let tcp_server = tokio::spawn({
            let app = app.clone();
            let server =
                Server::bind(IpSocketAddr::from(([127, 0, 0, 1], 0))).handle(tcp_handle.clone());

            async move { server.serve(app.into_make_service()).await }
        });

        let path =
            std::env::temp_dir().join(format!("axum-server-control-{}.sock", std::process::id()));
        let _ = std::fs::remove_file(&path);
        let listener = StdUnixListener::bind(&path).unwrap();

        let unix_server = tokio::spawn({
            let app = app.clone();
            let server = crate::from_unix(listener)
                .unwrap()
                .handle(unix_handle.clone());

            async move { server.serve(app.into_make_service()).await }
        });

        let tcp_addr = tcp_handle.listening().await.unwrap();
        let unix_addr = unix_handle.listening().await.unwrap();
        assert_eq!(unix_addr.as_pathname(), Some(path.as_path()));

        let (mut tcp_client, _tcp_conn) = connect(tcp_addr).await;
        do_empty_request_h1(&mut tcp_client).await.unwrap();

        let io = TokioIo::new(UnixStream::connect(&path).await.unwrap());
        let (mut unix_client, unix_conn) = client::conn::http1::handshake(io).await.unwrap();
        let unix_conn = tokio::spawn(async move {
            let _ = unix_conn.await;
        });
        do_empty_request_h1(&mut unix_client).await.unwrap();

        assert_eq!(control.connection_count(), 2);

        control.graceful_shutdown(None);

        timeout(Duration::from_secs(1), tcp_server)
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        timeout(Duration::from_secs(1), unix_server)
            .await
            .unwrap()
            .unwrap()
            .unwrap();

        assert_eq!(control.connection_count(), 0);

        let _ = timeout(Duration::from_secs(1), unix_conn).await;
        let _ = std::fs::remove_file(&path);
    }

    // Test graceful shutdown with no timeout.
    #[tokio::test]
    async fn test_graceful_shutdown_no_timeout() {
        let (handle, server_task, addr) = start_server().await;

        let (mut client1, _conn1) = connect(addr).await;
        let (mut client2, _conn2) = connect(addr).await;

        // Clients can send request before graceful shutdown.
        do_empty_request_h1(&mut client1).await.unwrap();
        do_empty_request_h1(&mut client2).await.unwrap();

        let start = tokio::time::Instant::now();

        let (hdr1_tx, hdr1_rx) = oneshot::channel::<()>();

        let fut1 = async {
            // A slow request made before graceful shutdown is handled.
            // Since there's no request timeout, this can take as long as it
            // needs.
            let hdr1 = send_slow_request(&mut client1, Duration::from_millis(500))
                .await
                .unwrap();
            hdr1_tx.send(()).unwrap();
            recv_slow_response_body(hdr1).await.unwrap();

            assert!(start.elapsed() >= Duration::from_millis(500));
        };
        let fut2 = async {
            // Graceful shutdown partway through
            tokio::time::sleep(Duration::from_millis(250)).await;
            hdr1_rx.await.unwrap();
            handle.graceful_shutdown(None);

            // Any new requests after graceful shutdown begins will fail
            do_empty_request_h1(&mut client2).await.unwrap_err();
            do_empty_request_h1(&mut client2).await.unwrap_err();
            do_empty_request_h1(&mut client2).await.unwrap_err();
        };

        tokio::join!(fut1, fut2);

        // At this point, graceful shutdown must have occurred, and the slow
        // request must have finished. Since there was no timeout, the elapsed
        // time should be at least 500 ms (slow request duration).
        assert!(start.elapsed() >= Duration::from_millis(500 + 100));

        // Server task should finish soon.
        timeout(Duration::from_secs(1), server_task)
            .await
            .unwrap()
            .unwrap()
            .unwrap();
    }

    // Test graceful shutdown with a timeout.
    #[tokio::test]
    async fn test_graceful_shutdown_timeout() {
        let (handle, server_task, addr) = start_server().await;

        let (mut client1, _conn1) = connect(addr).await;
        let (mut client2, _conn2) = connect(addr).await;

        // Clients can send request before graceful shutdown.
        do_empty_request_h1(&mut client1).await.unwrap();
        do_empty_request_h1(&mut client2).await.unwrap();

        let start = tokio::time::Instant::now();

        let (hdr1_tx, hdr1_rx) = oneshot::channel::<()>();

        let task1 = async {
            // A slow request made before graceful shutdown is handled.
            // This one is shorter than the timeout, so it should succeed.
            let hdr1 = send_slow_request(&mut client1, Duration::from_millis(222)).await;
            hdr1_tx.send(()).unwrap();

            let res1 = recv_slow_response_body(hdr1.unwrap()).await;
            res1.unwrap();
        };
        let task2 = async {
            // A slow request made before graceful shutdown is handled.
            // This one is much longer than the timeout; it should fail sometime
            // after the graceful shutdown timeout.

            let hdr2 = send_slow_request(&mut client2, Duration::from_millis(5_555)).await;
            hdr2.unwrap_err();
        };
        let task3 = async {
            // Begin graceful shutdown after we receive response headers for (1).
            hdr1_rx.await.unwrap();

            // Set a timeout on requests to finish before we drop them.
            handle.graceful_shutdown(Some(Duration::from_millis(333)));

            // Server task should finish soon.
            timeout(Duration::from_secs(1), server_task)
                .await
                .unwrap()
                .unwrap()
                .unwrap();

            // At this point, graceful shutdown must have occurred.
            assert!(start.elapsed() >= Duration::from_millis(222 + 333));
            assert!(start.elapsed() <= Duration::from_millis(5_555));
        };

        tokio::join!(task1, task2, task3);
    }

    #[tokio::test]
    async fn test_http1_only() {
        let (_handle, _server_task, addr) =
            start_server_with_http_version(Some(TestHttpVersion::Http1)).await;

        let (mut client, _conn) = connect_h1(addr).await;

        do_empty_request_h1(&mut client).await.unwrap();

        do_slow_request(&mut client, Duration::from_millis(50))
            .await
            .unwrap();

        let (mut client, _conn) = connect_h2(addr).await;
        do_empty_request_h2(&mut client).await.unwrap_err();
    }

    #[tokio::test]
    async fn test_http2_only() {
        let (_handle, _server_task, addr) =
            start_server_with_http_version(Some(TestHttpVersion::Http2)).await;

        let (mut client, _conn) = connect_h2(addr).await;

        do_empty_request_h2(&mut client).await.unwrap();

        do_slow_request_h2(&mut client, Duration::from_millis(50))
            .await
            .unwrap();

        let (mut client, _conn) = connect_h1(addr).await;
        do_empty_request_h1(&mut client).await.unwrap_err();
    }

    #[tokio::test]
    async fn test_listening_returns_none_on_bind_failure() {
        // Keep the port occupied so that the server fails to bind.
        let occupied = std::net::TcpListener::bind(("127.0.0.1", 0)).unwrap();
        let addr = occupied.local_addr().unwrap();

        let handle = Handle::new();
        let server_handle = handle.clone();

        let server_task = tokio::spawn(async move {
            Server::bind(addr)
                .handle(server_handle)
                .serve(Router::new().into_make_service())
                .await
        });

        let result = server_task.await.unwrap();
        assert!(result.is_err(), "expected bind failure");

        let listening = timeout(Duration::from_secs(1), handle.listening())
            .await
            .expect("listening() hung after bind failure instead of returning None");

        assert_eq!(listening, None);
    }

    #[tokio::test]
    async fn test_header_read_timeout_with_default_timer() {
        let handle = Handle::new();
        let server_handle = handle.clone();
        let addr = IpSocketAddr::from(([127, 0, 0, 1], 0));

        let mut server = Server::bind(addr);
        server
            .http_builder()
            .http1()
            .header_read_timeout(Duration::from_millis(100));

        let server_task = tokio::spawn(async move {
            server
                .handle(server_handle)
                .serve(
                    Router::new()
                        .route("/", get(|| async { "Hello, world!" }))
                        .into_make_service(),
                )
                .await
        });

        let addr = handle.listening().await.unwrap();

        // A normal request still works and does not panic because a timer is installed.
        let (mut client, _conn) = connect(addr).await;
        do_empty_request_h1(&mut client).await.unwrap();

        // A client that stalls after the request line is disconnected.
        let mut stalled = TcpStream::connect(addr).await.unwrap();
        stalled.write_all(b"GET / HTTP/1.1\r\n").await.unwrap();

        let mut buf = [0u8; 32];
        let read = timeout(Duration::from_secs(1), stalled.read(&mut buf))
            .await
            .expect("stalled connection was not closed by the header read timeout");

        match read {
            Ok(0) => {}
            Ok(n) => panic!("expected the connection to be closed, read {n} bytes"),
            Err(_) => {}
        }

        handle.shutdown();
        let result = timeout(Duration::from_secs(1), server_task)
            .await
            .unwrap()
            .unwrap();
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn first_request_timeout_closes_idle_connection() {
        let handle = Handle::new();
        let server_handle = handle.clone();
        let addr = IpSocketAddr::from(([127, 0, 0, 1], 0));

        let server_task = tokio::spawn(async move {
            Server::bind(addr)
                .first_request_timeout(Duration::from_millis(100))
                .handle(server_handle)
                .serve(
                    Router::new()
                        .route("/", get(|| async { "Hello, world!" }))
                        .into_make_service(),
                )
                .await
        });

        let addr = handle.listening().await.unwrap();
        let mut stream = TcpStream::connect(addr).await.unwrap();

        let mut buf = [0u8; 1];
        let read = timeout(Duration::from_secs(1), stream.read(&mut buf))
            .await
            .expect("idle connection was not closed by the first request timeout");

        match read {
            Ok(0) => {}
            Ok(n) => panic!("expected the connection to be closed, read {n} bytes"),
            Err(_) => {}
        }

        tokio::time::sleep(Duration::from_millis(100)).await;
        assert_eq!(handle.connection_count(), 0);

        handle.shutdown();
        let result = timeout(Duration::from_secs(1), server_task)
            .await
            .unwrap()
            .unwrap();
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn first_request_timeout_keeps_established_connection() {
        let handle = Handle::new();
        let server_handle = handle.clone();
        let addr = IpSocketAddr::from(([127, 0, 0, 1], 0));

        let server_task = tokio::spawn(async move {
            Server::bind(addr)
                .first_request_timeout(Duration::from_millis(100))
                .handle(server_handle)
                .serve(
                    Router::new()
                        .route("/", get(|| async { "Hello, world!" }))
                        .into_make_service(),
                )
                .await
        });

        let addr = handle.listening().await.unwrap();
        let (mut client, _conn) = connect(addr).await;

        do_empty_request_h1(&mut client).await.unwrap();

        // The connection has already served a request, so the first request
        // timeout no longer applies.
        tokio::time::sleep(Duration::from_millis(300)).await;

        do_empty_request_h1(&mut client).await.unwrap();

        handle.shutdown();
        let result = timeout(Duration::from_secs(1), server_task)
            .await
            .unwrap()
            .unwrap();
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn accept_error_handler_is_called() {
        use std::sync::{Arc, Mutex};

        #[derive(Clone)]
        struct FailingAcceptor;

        impl<S> Accept<TcpStream, S> for FailingAcceptor {
            type Stream = TcpStream;
            type Service = S;
            type Future = std::future::Ready<io::Result<(TcpStream, S)>>;

            fn accept(&self, _stream: TcpStream, _service: S) -> Self::Future {
                std::future::ready(Err(io::Error::other("acceptor failed")))
            }
        }

        let errors = Arc::new(Mutex::new(Vec::new()));
        let handler_errors = errors.clone();

        let handle = Handle::new();
        let server_handle = handle.clone();
        let addr = IpSocketAddr::from(([127, 0, 0, 1], 0));

        let server_task = tokio::spawn(async move {
            Server::bind(addr)
                .acceptor(FailingAcceptor)
                .accept_error_handler(move |addr, error| {
                    handler_errors
                        .lock()
                        .unwrap()
                        .push((addr.cloned(), error.to_string()));
                })
                .handle(server_handle)
                .serve(
                    Router::new()
                        .route("/", get(|| async { "Hello, world!" }))
                        .into_make_service(),
                )
                .await
        });

        let addr = handle.listening().await.unwrap();
        let _stream = TcpStream::connect(addr).await.unwrap();

        for _ in 0..100 {
            if !errors.lock().unwrap().is_empty() {
                break;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        {
            let errors = errors.lock().unwrap();
            assert_eq!(errors.len(), 1, "error handler was not called: {errors:?}");
            assert_eq!(errors[0].1, "acceptor failed");
            assert!(errors[0].0.is_some(), "remote address was not reported");
        }

        handle.shutdown();
        let result = timeout(Duration::from_secs(1), server_task)
            .await
            .unwrap()
            .unwrap();
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn local_addr_acceptor_adds_extension() {
        let handle = Handle::new();
        let server_handle = handle.clone();
        let addr = IpSocketAddr::from(([127, 0, 0, 1], 0));

        let server_task = tokio::spawn(async move {
            Server::bind(addr)
                .acceptor(LocalAddrAcceptor::new())
                .handle(server_handle)
                .serve(
                    Router::new()
                        .route(
                            "/",
                            get(|Extension(local): Extension<LocalAddr>| async move {
                                local.0.to_string()
                            }),
                        )
                        .into_make_service(),
                )
                .await
        });

        let addr = handle.listening().await.unwrap();
        let (mut client, _conn) = connect(addr).await;

        let body = client
            .send_request(Request::new(Body::empty()))
            .await
            .unwrap()
            .into_body()
            .collect()
            .await
            .unwrap()
            .to_bytes();

        assert_eq!(body.as_ref(), addr.to_string().as_bytes());

        handle.shutdown();
        let result = timeout(Duration::from_secs(1), server_task)
            .await
            .unwrap()
            .unwrap();
        assert!(result.is_ok());
    }

    async fn start_server_with_http_version(
        http_version: Option<TestHttpVersion>,
    ) -> (
        Handle<IpSocketAddr>,
        JoinHandle<io::Result<()>>,
        IpSocketAddr,
    ) {
        let handle = Handle::new();

        let server_handle = handle.clone();
        let server_task = tokio::spawn(async move {
            let app = Router::new()
                .route("/", get(|| async { "Hello, world!" }))
                .route(
                    "/echo_slowly",
                    post(|body: Bytes| async move {
                        // Stream a response slowly, byte-by-byte, over 100ms
                        Response::new(slow_body(body.len(), Duration::from_millis(100)))
                    }),
                );

            let addr = IpSocketAddr::from(([127, 0, 0, 1], 0));
            let server = Server::bind(addr);
            let server = match http_version {
                Some(TestHttpVersion::Http1) => server.http1_only(),
                Some(TestHttpVersion::Http2) => server.http2_only(),
                None => server,
            };

            server
                .handle(server_handle)
                .serve(app.into_make_service())
                .await
        });

        let addr = handle.listening().await.unwrap();

        (handle, server_task, addr)
    }

    async fn start_server() -> (
        Handle<IpSocketAddr>,
        JoinHandle<io::Result<()>>,
        IpSocketAddr,
    ) {
        start_server_with_http_version(None).await
    }

    async fn start_server_with_app(
        app: Router,
    ) -> (
        Handle<IpSocketAddr>,
        JoinHandle<io::Result<()>>,
        IpSocketAddr,
    ) {
        let handle = Handle::new();

        let server_handle = handle.clone();
        let server_task = tokio::spawn(async move {
            let addr = IpSocketAddr::from(([127, 0, 0, 1], 0));

            Server::bind(addr)
                .handle(server_handle)
                .serve(app.into_make_service())
                .await
        });

        let addr = handle.listening().await.unwrap();

        (handle, server_task, addr)
    }

    async fn connect(addr: IpSocketAddr) -> (http1::SendRequest<Body>, JoinHandle<()>) {
        connect_h1(addr).await
    }

    async fn connect_h1(addr: IpSocketAddr) -> (http1::SendRequest<Body>, JoinHandle<()>) {
        let stream = TokioIo::new(TcpStream::connect(addr).await.unwrap());
        let (send_request, connection) = client::conn::http1::handshake(stream).await.unwrap();

        let task = tokio::spawn(async move {
            let _ = connection.await;
        });

        (send_request, task)
    }

    async fn connect_h2(addr: IpSocketAddr) -> (http2::SendRequest<Body>, JoinHandle<()>) {
        let stream = TokioIo::new(TcpStream::connect(addr).await.unwrap());
        let (send_request, connection) =
            client::conn::http2::handshake(TokioExecutor::new(), stream)
                .await
                .unwrap();

        let task = tokio::spawn(async move {
            let _ = connection.await;
        });

        (send_request, task)
    }

    // Send a basic `GET /` request.
    async fn do_empty_request_h1(client: &mut http1::SendRequest<Body>) -> hyper::Result<()> {
        client.ready().await?;

        let body = client
            .send_request(Request::new(Body::empty()))
            .await?
            .into_body();

        let body = body.collect().await?.to_bytes();
        assert_eq!(body.as_ref(), b"Hello, world!");
        Ok(())
    }

    // Send a basic `GET /` request.
    async fn do_empty_request_h2(client: &mut http2::SendRequest<Body>) -> hyper::Result<()> {
        client.ready().await?;

        let body = client
            .send_request(Request::new(Body::empty()))
            .await?
            .into_body();

        let body = body.collect().await?.to_bytes();
        assert_eq!(body.as_ref(), b"Hello, world!");
        Ok(())
    }

    // Send a request with a body streamed byte-by-byte, over a given duration,
    // then wait for the full response.
    async fn do_slow_request(
        client: &mut http1::SendRequest<Body>,
        duration: Duration,
    ) -> hyper::Result<()> {
        let response = send_slow_request(client, duration).await?;
        recv_slow_response_body(response).await
    }

    async fn do_slow_request_h2(
        client: &mut http2::SendRequest<Body>,
        duration: Duration,
    ) -> hyper::Result<()> {
        let response = send_slow_request_h2(client, duration).await?;
        recv_slow_response_body(response).await
    }

    async fn send_slow_request(
        client: &mut http1::SendRequest<Body>,
        duration: Duration,
    ) -> hyper::Result<http::Response<hyper::body::Incoming>> {
        let req_body_len: usize = 10;
        let mut req = Request::new(slow_body(req_body_len, duration));
        *req.method_mut() = Method::POST;
        *req.uri_mut() = Uri::from_static("/echo_slowly");

        client.ready().await?;
        client.send_request(req).await
    }

    async fn send_slow_request_h2(
        client: &mut http2::SendRequest<Body>,
        duration: Duration,
    ) -> hyper::Result<http::Response<hyper::body::Incoming>> {
        let req_body_len: usize = 10;
        let mut req = Request::new(slow_body(req_body_len, duration));
        *req.method_mut() = Method::POST;
        *req.uri_mut() = Uri::from_static("/echo_slowly");

        client.ready().await?;
        client.send_request(req).await
    }

    async fn recv_slow_response_body(
        response: http::Response<hyper::body::Incoming>,
    ) -> hyper::Result<()> {
        let resp_body = response.into_body();
        let resp_body_bytes = resp_body.collect().await?.to_bytes();
        assert_eq!(10, resp_body_bytes.len());
        Ok(())
    }

    // A stream of n response data `Frame`s, where n = `length`, and each frame
    // consists of a single byte. The whole response is smeared out over
    // a `duration` length of time.
    fn slow_body(length: usize, duration: Duration) -> axum::body::Body {
        let frames =
            (0..length).map(move |_| Ok::<_, hyper::Error>(Frame::data(Bytes::from_static(b"X"))));

        let stream = stream::iter(frames).then(move |frame| async move {
            tokio::time::sleep(duration / (length as u32)).await;
            frame
        });

        axum::body::Body::new(StreamBody::new(stream))
    }
}