h3x 0.2.0

pub mod stream;

use std::{
    pin::{Pin, pin},
    task::{Context, Poll},
};

use bytes::{Buf, Bytes};
use futures::{Sink, SinkExt, Stream, TryStream, TryStreamExt};
use tokio::io::{self, AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};

use crate::{
    buflist::BufList,
    codec::{
        DecodeExt, DecodeFrom, DecodeStreamError, EncodeError, EncodeExt, EncodeInto,
        EncodeStreamError, FixedLengthReader, StreamReader,
    },
    connection::StreamError,
    error::H3FrameDecodeError,
    quic::{self, GetStreamId, StopStream},
    varint::{self, VARINT_MAX, VarInt},
};

pin_project_lite::pin_project! {
    /// All frames have the following format:
    ///
    /// ``` plaintext
    /// HTTP/3 Frame Format {
    ///   Type (i),
    ///   Length (i),
    ///   Frame Payload (..),
    /// }
    /// ```
    ///
    /// Frame Payload has a generic type `P`, it could be:
    /// ### A buffer that implements [`Buf`]
    ///
    /// Crate such a frame with [`Frame::new`], encode it to the stream with [`Frame::encode`].
    ///
    /// ### A [Stream] of bytes
    ///
    /// You can decode a frame on the stream by calling [`Frame::decode`],
    /// and then you can asynchronously read the frame payload using [`AsyncRead`], [`AsyncBufRead`] or [`Stream`]
    ///
    /// A maximum of the frame length in bytes can be retrieved from the frame;
    /// if the amount of data in the stream is less than the frame length, reading the frame will result in an error.
    ///
    /// Once the payload of the current frame has been fully read,
    /// calling [`Frame::decode_next`] will retrieve the next new frame in the stream.
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct Frame<P: ?Sized> {
        r#type: VarInt,
        length: VarInt,
        #[pin]
        payload: P,
    }
}

impl Frame<()> {
    pub const DATA_FRAME_TYPE: VarInt = VarInt::from_u32(0x00);
    pub const HEADERS_FRAME_TYPE: VarInt = VarInt::from_u32(0x01);
    pub const CANCEL_PUSH_FRAME_TYPE: VarInt = VarInt::from_u32(0x03);
    pub const SETTINGS_FRAME_TYPE: VarInt = VarInt::from_u32(0x04);
    pub const PUSH_PROMISE_FRAME_TYPE: VarInt = VarInt::from_u32(0x05);
    pub const GOAWAY_FRAME_TYPE: VarInt = VarInt::from_u32(0x07);
    pub const MAX_PUSH_ID_FRAME_TYPE: VarInt = VarInt::from_u32(0x0d);
}

impl<P: ?Sized> Frame<P> {
    /// Create a new frame.
    ///
    /// Length is calculated from the payload
    pub fn new(r#type: VarInt, payload: P) -> Result<Self, varint::err::Overflow>
    where
        P: Buf + Sized,
    {
        let length = VarInt::try_from(payload.remaining())?;
        Ok(Self {
            r#type,
            length,
            payload,
        })
    }

    pub const fn r#type(&self) -> VarInt {
        self.r#type
    }

    /// Frame types of the format 0x1f * N + 0x21 for non-negative integer
    /// values of N are reserved to exercise the requirement that unknown
    /// types be ignored (Section 9). These frames have no semantics, and
    /// they MAY be sent on any stream where frames are allowed to be sent.
    /// This enables their use for application-layer padding. Endpoints MUST
    /// NOT consider these frames to have any meaning upon receipt.
    ///
    /// The payload and length of the frames are selected in any manner the
    /// implementation chooses.
    ///
    /// <https://datatracker.ietf.org/doc/html/rfc9114#name-reserved-frame-types>
    pub const fn is_reserved_frame(&self) -> bool {
        (self.r#type.into_inner() >= 0x21) && (self.r#type.into_inner() - 0x21).is_multiple_of(0x1f)
    }

    pub const fn length(&self) -> VarInt {
        self.length
    }

    pub const fn payload(&self) -> &P {
        &self.payload
    }

    pub fn into_payload(self) -> P
    where
        P: Sized,
    {
        self.payload
    }

    pub fn map<U>(self, map: impl FnOnce(P) -> U) -> Frame<U>
    where
        P: Sized,
    {
        Frame {
            r#type: self.r#type,
            length: self.length,
            payload: map(self.payload),
        }
    }
}

impl<P: AsyncWrite + ?Sized> AsyncWrite for Frame<P> {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        let project = self.project();
        match project.length.into_inner().checked_add(buf.len() as u64) {
            Some(new_length) if new_length > VARINT_MAX => {
                return Poll::Ready(Err(EncodeError::FramePayloadTooLarge.into()));
            }
            Some(new_length) => {
                *project.length =
                    VarInt::from_u64(new_length).expect("length checked to be within VarInt range");
            }
            None => return Poll::Ready(Err(EncodeError::FramePayloadTooLarge.into())),
        }

        project.payload.poll_write(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        self.project().payload.poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        self.project().payload.poll_shutdown(cx)
    }
}

impl<P: Sink<B> + ?Sized, B: Buf> Sink<B> for Frame<P>
where
    EncodeStreamError: From<P::Error>,
{
    type Error = EncodeStreamError;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.project()
            .payload
            .poll_ready(cx)
            .map_err(EncodeStreamError::from)
    }

    fn start_send(self: Pin<&mut Self>, item: B) -> Result<(), Self::Error> {
        let project = self.project();
        let len = item.remaining() as u64;
        match project.length.into_inner().checked_add(len) {
            Some(new_length) if new_length > VARINT_MAX => {
                return Err(EncodeError::FramePayloadTooLarge.into());
            }
            Some(new_length) => {
                *project.length =
                    VarInt::from_u64(new_length).expect("length checked to be within VarInt range");
            }
            None => return Err(EncodeError::FramePayloadTooLarge.into()),
        }
        project.payload.start_send(item)?;
        Ok(())
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.project()
            .payload
            .poll_flush(cx)
            .map_err(EncodeStreamError::from)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.project()
            .payload
            .poll_close(cx)
            .map_err(EncodeStreamError::from)
    }
}

impl<P: Buf + Send, S: AsyncWrite + Sink<Bytes, Error = quic::StreamError> + Send> EncodeInto<S>
    for Frame<P>
{
    type Output = ();

    type Error = quic::StreamError;

    async fn encode_into(self, stream: S) -> Result<Self::Output, Self::Error> {
        let Frame {
            r#type,
            length,
            mut payload,
        } = self;
        let mut stream = pin!(stream);
        stream.as_mut().encode(r#type).await?;
        stream.as_mut().encode(length).await?;
        while payload.has_remaining() {
            let bytes = payload.copy_to_bytes(payload.chunk().len());
            stream.as_mut().feed(bytes).await?;
        }
        Ok(())
    }
}

impl<P1, P> DecodeFrom<Frame<P>> for Frame<P1>
where
    P1: DecodeFrom<P> + Send,
    P: Send,
{
    type Error = <P1 as DecodeFrom<P>>::Error;

    async fn decode_from(stream: Frame<P>) -> Result<Self, Self::Error> {
        let Frame {
            r#type,
            length,
            payload,
        } = stream;
        let payload = P1::decode_from(payload).await?;
        Ok(Frame {
            r#type,
            length,
            payload,
        })
    }
}

impl<S> DecodeFrom<&mut StreamReader<S>> for Frame<BufList>
where
    S: TryStream<Ok = Bytes, Error = quic::StreamError> + Unpin + Send,
{
    type Error = StreamError;

    async fn decode_from(stream: &mut StreamReader<S>) -> Result<Self, Self::Error> {
        let decode = async move {
            let r#type = stream.decode_one::<VarInt>().await?;
            let length = stream.decode_one::<VarInt>().await?;

            let mut payload = BufList::new();
            let mut reader = FixedLengthReader::new(stream, length.into_inner());
            while let Some(bytes) = reader.try_next().await? {
                payload.write(bytes);
            }
            Ok(Frame {
                r#type,
                length,
                payload,
            })
        };

        decode.await.map_err(|error: DecodeStreamError| {
            error.map_decode_error(|decode_error| {
                H3FrameDecodeError {
                    source: decode_error,
                }
                .into()
            })
        })
    }
}

impl<P: AsyncRead + ?Sized> AsyncRead for Frame<P> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        self.project().payload.poll_read(cx, buf)
    }
}

impl<P: AsyncBufRead + ?Sized> AsyncBufRead for Frame<P> {
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        self.project().payload.poll_fill_buf(cx)
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        self.project().payload.consume(amt);
    }
}

impl<P: TryStream<Ok = Bytes, Error = DecodeStreamError> + ?Sized> Stream for Frame<P> {
    type Item = Result<Bytes, DecodeStreamError>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.project().payload.try_poll_next(cx)
    }
}

impl<P: GetStreamId + ?Sized> GetStreamId for Frame<P> {
    fn poll_stream_id(
        self: Pin<&mut Self>,
        cx: &mut Context,
    ) -> Poll<Result<VarInt, quic::StreamError>> {
        self.project().payload.poll_stream_id(cx)
    }
}

impl<P: StopStream + ?Sized> StopStream for Frame<P> {
    fn poll_stop(
        self: Pin<&mut Self>,
        cx: &mut Context,
        code: VarInt,
    ) -> Poll<Result<(), quic::StreamError>> {
        self.project().payload.poll_stop(cx, code)
    }
}

#[cfg(test)]
mod tests {
    use futures::stream::{self, StreamExt};
    use tokio::io::AsyncReadExt;

    use super::*;
    use crate::{
        codec::{DecodeError, EncodeExt, SinkWriter},
        dhttp::frame::stream::FrameStream,
        error::Code,
    };

    fn to_pre_byte_stream(
        data: impl IntoIterator<Item = u8>,
    ) -> impl Stream<Item = Result<Bytes, quic::StreamError>> {
        stream::iter(
            data.into_iter()
                .map(|byte| vec![byte])
                .map(Bytes::from_owner)
                .map(Ok),
        )
    }

    #[test]
    fn new_frame() {
        Frame::new(Frame::DATA_FRAME_TYPE, Bytes::new()).unwrap();
    }

    #[tokio::test]
    async fn test_data_frames() {
        let stream = StreamReader::new(to_pre_byte_stream(vec![
            0x00, // Type: DATA
            0x05, // Length: 5
            b'H', b'e', b'l', b'l', b'o', // Payload: "Hello"
            0x00, // Type: DATA
            0x05, // Length: 5
            b'W', b'o', b'r', b'l', b'd', // Payload: "World"
        ]));

        let mut stream = pin!(FrameStream::new(stream));
        let mut frame1 = stream.as_mut().next_frame().await.unwrap().unwrap();
        assert_eq!(frame1.r#type().into_inner(), 0);
        assert_eq!(frame1.length().into_inner(), 5);
        let mut payload = vec![];
        frame1.read_to_end(&mut payload).await.unwrap();
        assert_eq!(&payload[..], b"Hello");

        let mut frame2 = stream.next_frame().await.unwrap().unwrap();
        let mut payload = vec![];
        frame2.read_to_end(&mut payload).await.unwrap();
        assert_eq!(&payload[..], b"World");
        assert_eq!(frame2.length().into_inner(), 5);
        assert_eq!(frame2.r#type().into_inner(), 0);
    }

    #[tokio::test]
    async fn incomplete_type() {
        let mut stream = StreamReader::new(to_pre_byte_stream(vec![
            // 0x00, // Type: DATA
            // 0x05, // Length: 5
            // b'H', b'e', b'l', b'l', b'o', // Payload: "Hello"
        ]));

        assert!(matches!(
            stream.decode_one::<Frame<_>>().await,
            Err(StreamError::Code { source }) if source.code() == Code::H3_FRAME_ERROR
        ));
    }

    #[tokio::test]
    async fn incomplete_length() {
        let mut stream = StreamReader::new(to_pre_byte_stream(vec![
            0x00, // Type: DATA
                 // 0x05, // Length: 5
                 // b'H', b'e', b'l', b'l', b'o', // Payload: "Hello"
        ]));

        assert!(matches!(
            stream.decode_one::<Frame<_>>().await,
            Err(StreamError::Code { source }) if source.code() == Code::H3_FRAME_ERROR
        ));
    }

    #[tokio::test]
    async fn incomplete_payload() {
        let stream = StreamReader::new(to_pre_byte_stream(vec![
            0x00, // Type: DATA
            0x05, // Length: 5
            b'H', b'e', b'l', b'l', /* b'o', */ // Payload: "Hello"
        ]));

        let stream = pin!(FrameStream::new(stream));
        let mut frame1 = stream.next_frame().await.unwrap().unwrap();
        assert_eq!(frame1.r#type().into_inner(), 0);
        assert_eq!(frame1.length().into_inner(), 5);
        let mut payload = vec![];
        assert!(matches!(
            DecodeStreamError::from(frame1.read_to_end(&mut payload).await.unwrap_err()),
            DecodeStreamError::Decode {
                source: DecodeError::Incomplete
            }
        ));
        assert_eq!(payload.as_slice(), b"Hell")
    }

    fn channel() -> (
        impl Sink<Bytes, Error = quic::StreamError>,
        impl Stream<Item = Result<Bytes, quic::StreamError>>,
    ) {
        let (sink, stream) = futures::channel::mpsc::channel(8);
        (sink.sink_map_err(|_e| unreachable!()), stream.map(Ok))
    }

    #[tokio::test]
    async fn encode_and_decode() {
        let (sink, stream) = channel();

        let decode = tokio::spawn(async move {
            let stream = StreamReader::new(stream);

            let mut stream = pin!(FrameStream::new(stream));
            let mut frame1 = stream.as_mut().next_frame().await.unwrap().unwrap();
            assert_eq!(frame1.r#type().into_inner(), 0);
            assert_eq!(frame1.length().into_inner(), 5);
            let mut payload = vec![];
            frame1.read_to_end(&mut payload).await.unwrap();
            assert_eq!(&payload[..], b"Hello");

            let mut frame2 = stream.next_frame().await.unwrap().unwrap();
            let mut payload = vec![];
            frame2.read_to_end(&mut payload).await.unwrap();
            assert_eq!(&payload[..], b"World");
            assert_eq!(frame2.length().into_inner(), 5);
            assert_eq!(frame2.r#type().into_inner(), 0);
        });
        let encode = tokio::spawn(async move {
            let mut sink = SinkWriter::new(sink);

            let frame1 = Frame::new(VarInt::from_u32(0), Bytes::from_static(b"Hello")).unwrap();
            assert!(frame1.r#type().into_inner() == 0);
            assert!(frame1.length().into_inner() == 5);
            assert!(frame1.payload().as_ref() == b"Hello");
            sink.encode_one(frame1).await.unwrap();

            let frame2 = Frame::new(VarInt::from_u32(0), Bytes::from_static(b"World")).unwrap();
            assert!(frame2.r#type().into_inner() == 0);
            assert!(frame2.length().into_inner() == 5);
            assert!(frame2.payload().as_ref() == b"World");
            sink.encode_one(frame2).await.unwrap();

            Pin::new(&mut sink).flush_buffer().await.unwrap();
        });
        tokio::try_join!(encode, decode).unwrap();
    }
}