salvo_core 0.94.0

use std::fmt::{self, Debug, Formatter};
use std::io::{Error as IoError, ErrorKind, Result as IoResult};
use std::pin::Pin;
use std::task::{Context, Poll};

use bytes::Bytes;
use futures_util::stream::Stream;
use hyper::body::{Body, Frame, Incoming, SizeHint};

use crate::BoxedError;
use crate::fuse::FuseConfig;

#[doc(hidden)]
#[derive(Debug)]
pub struct BodyTimeout {
    duration: std::time::Duration,
    sleep: Pin<Box<tokio::time::Sleep>>,
    armed: bool,
}

impl BodyTimeout {
    fn new(duration: std::time::Duration) -> Self {
        Self {
            duration,
            sleep: Box::pin(tokio::time::sleep(std::time::Duration::from_secs(
                86400 * 365,
            ))),
            armed: false,
        }
    }

    fn apply(&mut self, poll: PollFrame, cx: &mut Context<'_>) -> PollFrame {
        match poll {
            Poll::Pending => {
                if !self.armed {
                    self.armed = true;
                    self.sleep
                        .as_mut()
                        .reset(tokio::time::Instant::now() + self.duration);
                }
                if self.sleep.as_mut().poll(cx).is_ready() {
                    // Fail only this request body. On HTTP/2 and HTTP/3 the error ends the
                    // offending stream (RST_STREAM) while sibling streams multiplexed over
                    // the same connection keep running; on HTTP/1 the connection carries a
                    // single request, so it is torn down as before. Deliberately does not
                    // abort the whole transport connection, which would take down unrelated
                    // in-flight streams.
                    Poll::Ready(Some(Err(IoError::new(
                        ErrorKind::TimedOut,
                        "request body timeout",
                    ))))
                } else {
                    Poll::Pending
                }
            }
            ready => {
                // A frame that arrives after the deadline already elapsed must still time out,
                // or a client dribbling each chunk just past the deadline evades the limit.
                if self.armed && self.sleep.as_mut().poll(cx).is_ready() {
                    self.armed = false;
                    return Poll::Ready(Some(Err(IoError::new(
                        ErrorKind::TimedOut,
                        "request body timeout",
                    ))));
                }
                self.armed = false;
                ready
            }
        }
    }
}

pub(crate) type BoxedBody =
    Pin<Box<dyn Body<Data = Bytes, Error = BoxedError> + Send + Sync + 'static>>;
pub(crate) type PollFrame = Poll<Option<Result<Frame<Bytes>, IoError>>>;

/// Body for HTTP request.
#[non_exhaustive]
#[derive(Default)]
pub enum ReqBody {
    /// None body.
    #[default]
    None,
    /// Once bytes body.
    Once(Bytes),
    /// Hyper default body.
    Hyper {
        /// Inner body.
        inner: Incoming,
        /// Request-body timeout state.
        fuse_config: Option<BodyTimeout>,
    },
    /// Boxed body.
    Boxed {
        /// Inner body.
        inner: BoxedBody,
        /// Request-body timeout state.
        fuse_config: Option<BodyTimeout>,
    },
}
impl ReqBody {
    #[doc(hidden)]
    pub fn set_fuse_config(&mut self, value: Option<FuseConfig>) {
        match self {
            Self::None | Self::Once(_) => {}
            Self::Hyper { fuse_config, .. } | Self::Boxed { fuse_config, .. } => {
                *fuse_config = value
                    .and_then(|config| config.request_body_timeout)
                    .map(BodyTimeout::new);
            }
        }
    }
    /// Returns true if the body is not set.
    #[inline]
    pub fn is_none(&self) -> bool {
        matches!(*self, Self::None)
    }
    /// Returns true if the body contains one byte buffer.
    #[inline]
    pub fn is_once(&self) -> bool {
        matches!(*self, Self::Once(_))
    }
    /// Returns true if the body is Hyper's default incoming body type.
    #[inline]
    pub fn is_hyper(&self) -> bool {
        matches!(*self, Self::Hyper { .. })
    }
    /// Returns true if the body is boxed.
    #[inline]
    pub fn is_boxed(&self) -> bool {
        matches!(*self, Self::Boxed { .. })
    }

    /// Set body to none and returns current body.
    #[inline]
    #[must_use]
    pub fn take(&mut self) -> Self {
        std::mem::replace(self, Self::None)
    }
}

impl Body for ReqBody {
    type Data = Bytes;
    type Error = IoError;

    fn poll_frame(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> PollFrame {
        #[inline]
        fn through_fuse_config(
            poll: PollFrame,
            fuse_config: Option<&mut BodyTimeout>,
            cx: &mut Context<'_>,
        ) -> PollFrame {
            match fuse_config {
                Some(timeout) => timeout.apply(poll, cx),
                None => poll,
            }
        }
        match &mut *self {
            Self::None => Poll::Ready(None),
            Self::Once(bytes) => {
                if bytes.is_empty() {
                    Poll::Ready(None)
                } else {
                    let bytes = std::mem::take(bytes);
                    Poll::Ready(Some(Ok(Frame::data(bytes))))
                }
            }
            Self::Hyper { inner, fuse_config } => {
                let poll = Pin::new(inner).poll_frame(cx).map_err(IoError::other);
                through_fuse_config(poll, fuse_config.as_mut(), cx)
            }
            Self::Boxed { inner, fuse_config } => {
                let poll = Pin::new(inner).poll_frame(cx).map_err(IoError::other);
                through_fuse_config(poll, fuse_config.as_mut(), cx)
            }
        }
    }

    fn is_end_stream(&self) -> bool {
        match self {
            Self::None => true,
            Self::Once(bytes) => bytes.is_empty(),
            Self::Hyper { inner, .. } => inner.is_end_stream(),
            Self::Boxed { inner, .. } => inner.is_end_stream(),
        }
    }

    fn size_hint(&self) -> SizeHint {
        match self {
            Self::None => SizeHint::with_exact(0),
            Self::Once(bytes) => SizeHint::with_exact(bytes.len() as u64),
            Self::Hyper { inner, .. } => inner.size_hint(),
            Self::Boxed { inner, .. } => inner.size_hint(),
        }
    }
}
impl Stream for ReqBody {
    type Item = IoResult<Frame<Bytes>>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match Body::poll_frame(self, cx) {
            Poll::Ready(Some(Ok(frame))) => Poll::Ready(Some(Ok(frame))),
            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(IoError::other(e)))),
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }
}

impl From<Bytes> for ReqBody {
    fn from(value: Bytes) -> Self {
        Self::Once(value)
    }
}
impl From<Incoming> for ReqBody {
    fn from(inner: Incoming) -> Self {
        Self::Hyper {
            inner,
            fuse_config: None,
        }
    }
}
impl From<String> for ReqBody {
    #[inline]
    fn from(value: String) -> Self {
        Self::Once(value.into())
    }
}
impl TryFrom<ReqBody> for Incoming {
    type Error = crate::Error;
    fn try_from(body: ReqBody) -> Result<Self, Self::Error> {
        match body {
            ReqBody::None => Err(crate::Error::other(
                "ReqBody::None cannot convert to Incoming",
            )),
            ReqBody::Once(_) => Err(crate::Error::other(
                "ReqBody::Once cannot convert to Incoming",
            )),
            ReqBody::Hyper { inner, .. } => Ok(inner),
            ReqBody::Boxed { .. } => Err(crate::Error::other(
                "ReqBody::Boxed cannot convert to Incoming",
            )),
        }
    }
}

impl From<&'static [u8]> for ReqBody {
    fn from(value: &'static [u8]) -> Self {
        Self::Once(Bytes::from_static(value))
    }
}

impl From<&'static str> for ReqBody {
    fn from(value: &'static str) -> Self {
        Self::Once(Bytes::from_static(value.as_bytes()))
    }
}

impl From<Vec<u8>> for ReqBody {
    fn from(value: Vec<u8>) -> Self {
        Self::Once(value.into())
    }
}

impl<T> From<Box<T>> for ReqBody
where
    T: Into<Self>,
{
    fn from(value: Box<T>) -> Self {
        (*value).into()
    }
}

cfg_feature! {
    #![feature = "quinn"]
    pub(crate) mod h3 {
        use std::boxed::Box;
        use std::pin::Pin;
        use std::task::{ready, Context, Poll};
        use std::fmt::{self, Debug, Formatter};

        use hyper::body::{Body, Frame, SizeHint};
        use salvo_http3::quic::RecvStream;
        use salvo_http3::error::Code;

        use bytes::{Buf, Bytes};

        use crate::BoxedError;
        use crate::http::ReqBody;

        /// HTTP/3 request body.
        pub struct H3ReqBody<S, B>
        where
            S: RecvStream + Send + Unpin,
            B: Buf + Send + Unpin,
        {
            inner: salvo_http3::server::RequestStream<S, B>,
        }
        impl<S, B> Debug for H3ReqBody<S, B>
        where
            S: RecvStream + Send + Unpin,
            B: Buf + Send + Unpin,
        {
            fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
                f.debug_struct("H3ReqBody").finish()
            }
        }
        impl<S, B> H3ReqBody<S, B>
        where
            S: RecvStream + Send + Unpin + 'static,
            B: Buf + Send + Unpin + 'static,
        {
            /// Creates a new `H3ReqBody` instance.
            pub fn new(inner: salvo_http3::server::RequestStream<S, B>) -> Self {
                Self { inner }
            }
        }

        impl<S, B> Body for H3ReqBody<S, B>
        where
            S: RecvStream + Send + Unpin,
            B: Buf + Send + Unpin,
        {
            type Data = Bytes;
            type Error = BoxedError;

            fn poll_frame(
                mut self: Pin<&mut Self>,
                cx: &mut Context<'_>,
            ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
                let this = &mut *self;
                match ready!(this.inner.poll_recv_data(cx)) {
                    Ok(Some(buf)) => {
                        Poll::Ready(Some(Ok(Frame::data(Bytes::copy_from_slice(buf.chunk())))))
                    }
                    Ok(None) => Poll::Ready(None),
                    Err(e) => Poll::Ready(Some(Err(e.into()))),
                }
            }

            fn is_end_stream(&self) -> bool {
                false // TODO: check this
            }

            fn size_hint(&self) -> SizeHint {
                SizeHint::default()
            }
        }

        impl<S, B> Drop for H3ReqBody<S, B>
        where
            S: RecvStream + Send + Unpin,
            B: Buf + Send + Unpin,
        {
            fn drop(&mut self) {
                self.inner.stop_sending(Code::H3_NO_ERROR);
            }
        }

        impl<S, B> From<H3ReqBody<S, B>> for ReqBody
        where
            S: RecvStream + Send + Sync + Unpin + 'static,
            B: Buf + Send + Sync + Unpin + 'static,
        {
            fn from(value: H3ReqBody<S, B>) -> Self {
                Self::Boxed {
                    inner: Box::pin(value),
                    fuse_config: None,
                }
            }
        }
    }
}

impl Debug for ReqBody {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Self::None => write!(f, "ReqBody::None"),
            Self::Once(value) => f.debug_tuple("ReqBody::Once").field(value).finish(),
            Self::Hyper { inner, .. } => f
                .debug_struct("ReqBody::Hyper")
                .field("inner", inner)
                .finish(),
            Self::Boxed { .. } => write!(f, "ReqBody::Boxed{{..}}"),
        }
    }
}

#[cfg(test)]
mod tests {
    use bytes::Bytes;

    use super::*;

    #[test]
    fn test_take() {
        let mut b = ReqBody::Once(Bytes::from("abc"));
        let old = b.take();
        assert!(matches!(old, ReqBody::Once(_)));
        assert!(b.is_none());
    }

    #[test]
    fn test_debug() {
        let b = ReqBody::None;
        let s = format!("{b:?}");
        assert!(s.contains("ReqBody::None"));
    }

    #[test]
    fn test_is_end_stream() {
        let b = ReqBody::None;
        assert!(b.is_end_stream());
        let b = ReqBody::Once(Bytes::new());
        assert!(b.is_end_stream());
    }

    #[tokio::test]
    async fn request_body_timeout_fails_only_the_body() {
        use std::future::poll_fn;
        use std::time::Duration;

        struct PendingBody;
        impl Body for PendingBody {
            type Data = Bytes;
            type Error = BoxedError;
            fn poll_frame(
                self: Pin<&mut Self>,
                _cx: &mut Context<'_>,
            ) -> Poll<Option<Result<Frame<Bytes>, Self::Error>>> {
                Poll::Pending
            }
        }

        let mut body = ReqBody::Boxed {
            inner: Box::pin(PendingBody),
            fuse_config: None,
        };
        body.set_fuse_config(Some(FuseConfig {
            request_body_timeout: Some(Duration::from_millis(10)),
            ..FuseConfig::disabled()
        }));

        // The stalled body must surface a `TimedOut` error to its reader. That error
        // ends only this request stream; no connection-wide abort is involved.
        let frame = poll_fn(|cx| Pin::new(&mut body).poll_frame(cx)).await;
        let error = frame.expect("body must yield a frame").unwrap_err();
        assert_eq!(error.kind(), ErrorKind::TimedOut);
    }

    #[tokio::test]
    async fn request_body_timeout_rejects_a_late_frame() {
        use std::future::poll_fn;
        use std::time::Duration;

        let mut timeout = BodyTimeout::new(Duration::from_millis(10));

        // Arm the timer as a pending frame poll would, then let the deadline lapse.
        poll_fn(|cx| {
            let _ = timeout.apply(Poll::Pending, cx);
            Poll::Ready(())
        })
        .await;
        tokio::time::sleep(Duration::from_millis(40)).await;

        // A frame that only arrives now — past the deadline — must be converted to a timeout,
        // not accepted as if the gap had been within bounds.
        let out = poll_fn(|cx| {
            let frame = Frame::data(Bytes::from_static(b"x"));
            Poll::Ready(timeout.apply(Poll::Ready(Some(Ok(frame))), cx))
        })
        .await;
        match out {
            Poll::Ready(Some(Err(error))) => assert_eq!(error.kind(), ErrorKind::TimedOut),
            _ => panic!("a late body frame must time out"),
        }
    }
}