salvo_core 0.94.0

use std::cmp;
use std::error::Error as StdError;
use std::fmt::{self, Debug, Formatter};
#[cfg(any(feature = "http1", feature = "http2"))]
use std::future::pending;
use std::io::{Error as IoError, ErrorKind, IoSlice, Result as IoResult};
use std::marker::PhantomPinned;
use std::pin::Pin;
use std::task::{self, Context, Poll, ready};

use bytes::{Buf, Bytes};
use http::{Request, Response, Version};
use hyper::service::Service;
use pin_project::pin_project;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio_util::sync::CancellationToken;

use crate::ConnCtrl;
#[cfg(any(feature = "http1", feature = "http2"))]
use crate::conn::ctrl::ConnState;
use crate::fuse::FuseConfig;
use crate::http::body::{Body, HyperBody};
#[cfg(any(feature = "http1", feature = "http2"))]
use crate::rt::tokio::TokioIo;

type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;

#[cfg(feature = "http1")]
use hyper::server::conn::http1;
#[cfg(feature = "http2")]
use hyper::server::conn::http2;

#[cfg(feature = "quinn")]
use crate::conn::quinn;
#[cfg(feature = "http2")]
use crate::rt::tokio::TokioExecutor;

const H2_PREFACE: &[u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n";

#[doc(hidden)]
pub struct HttpBuilder {
    #[cfg(feature = "http1")]
    pub(crate) http1: http1::Builder,
    #[cfg(feature = "http2")]
    pub(crate) http2: http2::Builder<TokioExecutor>,
    #[cfg(feature = "quinn")]
    pub(crate) quinn: quinn::Builder,
}
impl Default for HttpBuilder {
    fn default() -> Self {
        Self::new()
    }
}
impl Debug for HttpBuilder {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("HttpBuilder").finish()
    }
}

impl HttpBuilder {
    #[must_use]
    pub fn new() -> Self {
        Self {
            // The timer is connection-independent, so set it once here rather than on
            // every accepted connection. Only the per-connection header-read timeout,
            // which varies with the fuse config, is applied later.
            #[cfg(feature = "http1")]
            http1: {
                let mut builder = http1::Builder::new();
                builder.timer(crate::rt::tokio::TokioTimer::new());
                builder
            },
            #[cfg(feature = "http2")]
            http2: http2::Builder::new(crate::rt::tokio::TokioExecutor::new()),
            #[cfg(feature = "quinn")]
            quinn: crate::conn::quinn::Builder::new(),
        }
    }

    /// Serve a connection with the given service.
    #[allow(unused_variables)]
    pub async fn serve_connection<I, S, B>(
        &self,
        socket: I,
        service: S,
        fuse_config: Option<FuseConfig>,
        conn_ctrl: ConnCtrl,
        graceful_stop_token: Option<CancellationToken>,
    ) -> Result<()>
    where
        S: Service<Request<HyperBody>, Response = Response<B>> + Send,
        S::Future: Send + 'static,
        S::Error: Into<Box<dyn StdError + Send + Sync>>,
        B: Body + Send + 'static,
        B::Data: Send,
        B::Error: Into<Box<dyn StdError + Send + Sync>>,
        I: AsyncRead + AsyncWrite + Unpin + Send + 'static,
    {
        // When both HTTP/1 and HTTP/2 are enabled, consume only enough bytes to
        // distinguish the HTTP/2 prior-knowledge preface. `Rewind` preserves
        // those bytes so the selected Hyper connection receives the full input.
        //
        // Detection and Hyper's own header-read timeout share one `http1_header_timeout`
        // budget: track when detection began so the remainder can be handed to Hyper,
        // rather than letting a client that stalls detection then get a fresh full timeout.
        #[cfg(all(feature = "http1", feature = "http2"))]
        let detect_started = tokio::time::Instant::now();
        #[cfg(all(feature = "http1", feature = "http2"))]
        let (version, socket) = {
            // The initial protocol-detection read is bounded by the HTTP/1 header timeout.
            // A disabled fuse (`disable_fuse()` → no config, or a config with this timeout
            // set to `None`) is honored and leaves the read unbounded; the on-by-default
            // config keeps its header timeout, so a silent connection is still bounded.
            let detect_timeout = fuse_config.and_then(|config| config.http1_header_timeout);
            tokio::select! {
                result = read_version(socket) => result?,
                _ = async {
                    match detect_timeout {
                        Some(timeout) => tokio::time::sleep(timeout).await,
                        None => pending::<()>().await,
                    }
                } => {
                    tracing::info!("closing connection: protocol-detection read timed out");
                    return Ok(());
                }
                state = conn_ctrl.notified() => {
                    tracing::info!(?state, "closing connection during protocol detection");
                    return Ok(());
                },
            }
        };
        #[cfg(all(not(feature = "http1"), not(feature = "http2")))]
        let version = Version::HTTP_11; // Just make the compiler happy.
        #[cfg(all(feature = "http1", not(feature = "http2")))]
        let version = Version::HTTP_11;
        #[cfg(all(not(feature = "http1"), feature = "http2"))]
        let version = Version::HTTP_2;

        match version {
            Version::HTTP_10 | Version::HTTP_11 => {
                #[cfg(not(feature = "http1"))]
                return Err(std::io::Error::other("http1 feature not enabled").into());
                #[cfg(feature = "http1")]
                {
                    let mut http1 = self.http1.clone();
                    // Only override Hyper's header-read timeout when the fuse actually asks for
                    // one. A disabled fuse (or a config with this timeout unset) leaves whatever
                    // the caller configured via `HttpBuilder` / `http1_mut()` intact instead of
                    // silently clearing it.
                    if let Some(timeout) =
                        fuse_config.and_then(|config| config.http1_header_timeout)
                    {
                        // When both protocols are enabled, the detection read already spent part
                        // of this budget; subtract it so the whole header-read window stays a
                        // single `http1_header_timeout` instead of restarting after detection.
                        #[cfg(feature = "http2")]
                        let timeout = timeout.saturating_sub(detect_started.elapsed());
                        http1.header_read_timeout(Some(timeout));
                    }
                    let mut conn = http1
                        .serve_connection(TokioIo::new(socket), service)
                        .with_upgrades();

                    // The connection future, server shutdown, and handler-level
                    // connection control are driven by the same task. This avoids
                    // spawning a supervisor task for every accepted connection.
                    tokio::select! {
                        result = &mut conn => {
                            if let Err(error) = result {
                                tracing::debug!(?error, "HTTP/1 connection ended with an error");
                            }
                            return Ok(());
                        },
                        // Server-wide graceful shutdown stops HTTP keep-alive but
                        // allows the currently accepted request to complete.
                        _ = async {
                            if let Some(token) = &graceful_stop_token {
                                token.cancelled().await;
                            } else {
                                pending::<()>().await;
                            }
                        } => {
                            tracing::info!("gracefully shutting down HTTP/1 connection");
                            Pin::new(&mut conn).graceful_shutdown();
                            tokio::select! {
                                result = &mut conn => {
                                    if let Err(error) = result {
                                        tracing::debug!(?error, "HTTP/1 connection ended during server graceful shutdown");
                                    }
                                }
                                _ = conn_ctrl.aborted() => {
                                    tracing::info!("handler aborted HTTP/1 connection during server graceful shutdown");
                                }
                            }
                        }
                        // `ConnCtrl` is shared with handlers. Abort drops the
                        // Hyper connection immediately; graceful shutdown first
                        // disables keep-alive and remains abortable.
                        state = conn_ctrl.notified() => {
                            if state == ConnState::GracefulShutdown {
                                tracing::info!("handler requested graceful HTTP/1 shutdown");
                                Pin::new(&mut conn).graceful_shutdown();
                                tokio::select! {
                                    result = &mut conn => {
                                        if let Err(error) = result {
                                            tracing::debug!(?error, "HTTP/1 connection ended during graceful shutdown");
                                        }
                                    }
                                    _ = conn_ctrl.aborted() => {
                                        tracing::info!("handler escalated HTTP/1 shutdown to abort");
                                    }
                                }
                            } else {
                                tracing::info!("handler aborted HTTP/1 connection");
                            }
                        }
                    }
                }
            }
            Version::HTTP_2 => {
                #[cfg(not(feature = "http2"))]
                return Err(std::io::Error::other("http2 feature not enabled").into());
                #[cfg(feature = "http2")]
                {
                    let mut conn = self.http2.serve_connection(TokioIo::new(socket), service);

                    // HTTP/2 uses the same lifecycle arbitration as HTTP/1.
                    // Hyper translates graceful shutdown into a GOAWAY frame.
                    tokio::select! {
                        result = &mut conn => {
                            if let Err(error) = result {
                                tracing::debug!(?error, "HTTP/2 connection ended with an error");
                            }
                            return Ok(());
                        },
                        _ = async {
                            if let Some(token) = &graceful_stop_token {
                                token.cancelled().await;
                            } else {
                                pending::<()>().await;
                            }
                        } => {
                            tracing::info!("gracefully shutting down HTTP/2 connection");
                            Pin::new(&mut conn).graceful_shutdown();
                            tokio::select! {
                                result = &mut conn => {
                                    if let Err(error) = result {
                                        tracing::debug!(?error, "HTTP/2 connection ended during server graceful shutdown");
                                    }
                                }
                                _ = conn_ctrl.aborted() => {
                                    tracing::info!("handler aborted HTTP/2 connection during server graceful shutdown");
                                }
                            }
                        }
                        state = conn_ctrl.notified() => {
                            if state == ConnState::GracefulShutdown {
                                tracing::info!("handler requested graceful HTTP/2 shutdown");
                                Pin::new(&mut conn).graceful_shutdown();
                                tokio::select! {
                                    result = &mut conn => {
                                        if let Err(error) = result {
                                            tracing::debug!(?error, "HTTP/2 connection ended during graceful shutdown");
                                        }
                                    }
                                    _ = conn_ctrl.aborted() => {
                                        tracing::info!("handler escalated HTTP/2 shutdown to abort");
                                    }
                                }
                            } else {
                                tracing::info!("handler aborted HTTP/2 connection");
                            }
                        }
                    }
                }
            }
            _ => {
                tracing::info!("unsupported protocol version: {:?}", version);
            }
        }

        Ok(())
    }
}

#[allow(dead_code)]
#[allow(clippy::future_not_send)]
pub(crate) async fn read_version<A>(mut reader: A) -> IoResult<(Version, Rewind<A>)>
where
    A: AsyncRead + Unpin,
{
    let mut buf = [0; 24];
    let (version, buf) = ReadVersion {
        reader: &mut reader,
        buf: ReadBuf::new(&mut buf),
        version: Version::HTTP_11,
        _pin: PhantomPinned,
    }
    .await?;
    Ok((version, Rewind::new_buffered(Bytes::from(buf), reader)))
}

#[derive(Debug)]
#[pin_project]
#[must_use = "futures do nothing unless you `.await` or poll them"]
struct ReadVersion<'a, A: ?Sized> {
    reader: &'a mut A,
    buf: ReadBuf<'a>,
    version: Version,
    // Make this future `!Unpin` for compatibility with async trait methods.
    #[pin]
    _pin: PhantomPinned,
}

impl<A> Future for ReadVersion<'_, A>
where
    A: AsyncRead + Unpin + ?Sized,
{
    type Output = IoResult<(Version, Vec<u8>)>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<(Version, Vec<u8>)>> {
        let this = self.project();

        while this.buf.remaining() != 0 {
            if this.buf.filled() != &H2_PREFACE[0..this.buf.filled().len()] {
                return Poll::Ready(Ok((*this.version, this.buf.filled().to_vec())));
            }
            // if our buffer is empty, then we need to read some data to continue.
            let rem = this.buf.remaining();
            ready!(Pin::new(&mut *this.reader).poll_read(cx, this.buf))?;
            if this.buf.remaining() == rem {
                return Err(IoError::new(ErrorKind::UnexpectedEof, "early eof")).into();
            }
        }
        if this.buf.filled() == H2_PREFACE {
            *this.version = Version::HTTP_2;
        }
        Poll::Ready(Ok((*this.version, this.buf.filled().to_vec())))
    }
}

// from https://github.com/hyperium/hyper-util/pull/11/files#diff-1bd3ef8e9a23396b76bdb4ec6ab5aba4c48dd0511d287e485148a90170c6b4fd
/// Combine a buffer with an IO, rewinding reads to use the buffer.
#[derive(Debug)]
pub(crate) struct Rewind<T> {
    pre: Option<Bytes>,
    inner: T,
}
#[allow(dead_code)]
impl<T> Rewind<T> {
    fn new_buffered(buf: Bytes, io: T) -> Self {
        Self {
            pre: Some(buf),
            inner: io,
        }
    }
}

impl<T> AsyncRead for Rewind<T>
where
    T: AsyncRead + Unpin,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut task::Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<IoResult<()>> {
        if let Some(mut prefix) = self.pre.take() {
            // If there are no remaining bytes, let the bytes get dropped.
            if !prefix.is_empty() {
                let copy_len = cmp::min(prefix.len(), buf.remaining());
                // TODO: There should be a way to do following two lines cleaner...
                buf.put_slice(&prefix[..copy_len]);
                prefix.advance(copy_len);
                // Put back what's left
                if !prefix.is_empty() {
                    self.pre = Some(prefix);
                }
                return Poll::Ready(Ok(()));
            }
        }
        Pin::new(&mut self.inner).poll_read(cx, buf)
    }
}

impl<T> AsyncWrite for Rewind<T>
where
    T: AsyncWrite + Unpin,
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut task::Context<'_>,
        buf: &[u8],
    ) -> Poll<IoResult<usize>> {
        Pin::new(&mut self.inner).poll_write(cx, buf)
    }

    fn poll_write_vectored(
        mut self: Pin<&mut Self>,
        cx: &mut task::Context<'_>,
        bufs: &[IoSlice<'_>],
    ) -> Poll<IoResult<usize>> {
        Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<IoResult<()>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<IoResult<()>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }

    fn is_write_vectored(&self) -> bool {
        self.inner.is_write_vectored()
    }
}