rama_core/io/

read.rs

use pin_project_lite::pin_project;
use std::{
    io::{Cursor, Read},
    pin::Pin,
    task::{Context, Poll, ready},
};
use tokio::io::{self, AsyncBufRead, AsyncRead, ReadBuf};

use crate::bytes::{Buf, Bytes};

/// Read and discard exactly `n` bytes from `reader`.
///
/// Equivalent to `reader.read_exact(&mut vec![0u8; n])` but does not allocate
/// a buffer sized to `n`; it pipes the bytes through [`tokio::io::sink`] via
/// a `take` adaptor, reusing a small internal scratch buffer instead.
///
/// Returns an `UnexpectedEof` error if the reader yields fewer than `n` bytes
/// before EOF.
pub async fn discard<R>(reader: &mut R, n: u64) -> io::Result<()>
where
    R: AsyncRead + Unpin,
{
    if n == 0 {
        return Ok(());
    }
    let mut limited = tokio::io::AsyncReadExt::take(reader, n);
    let copied = tokio::io::copy(&mut limited, &mut tokio::io::sink()).await?;
    if copied < n {
        return Err(io::Error::new(
            io::ErrorKind::UnexpectedEof,
            "io::discard: reader EOF before requested byte count was consumed",
        ));
    }
    Ok(())
}

pin_project! {
    /// Reader for reading from a heap-allocated bytes buffer.
    #[derive(Debug, Clone)]
    pub struct HeapReader {
        #[pin]
        inner: Cursor<Vec<u8>>,
    }
}

impl HeapReader {
    /// Creates a new `HeapReader` with the specified bytes data.
    #[must_use]
    pub const fn new(data: Vec<u8>) -> Self {
        Self {
            inner: Cursor::new(data),
        }
    }

    /// How many bytes are there remaining
    #[must_use]
    pub fn remaining(&self) -> usize {
        self.inner.remaining()
    }

    /// Returns true if there are any more bytes to consume
    #[must_use]
    pub fn has_remaining(&self) -> bool {
        self.inner.has_remaining()
    }
}

impl From<Vec<u8>> for HeapReader {
    fn from(data: Vec<u8>) -> Self {
        Self::new(data)
    }
}

impl From<&[u8]> for HeapReader {
    fn from(data: &[u8]) -> Self {
        Self::new(data.to_vec())
    }
}

impl From<&str> for HeapReader {
    fn from(data: &str) -> Self {
        Self::new(data.as_bytes().to_vec())
    }
}

impl Default for HeapReader {
    fn default() -> Self {
        Self::new(Vec::new())
    }
}

impl From<Bytes> for HeapReader {
    fn from(data: Bytes) -> Self {
        Self::new(data.to_vec())
    }
}

#[warn(clippy::missing_trait_methods)]
impl AsyncRead for HeapReader {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        self.project().inner.poll_read(cx, buf)
    }
}

#[warn(clippy::missing_trait_methods)]
impl AsyncBufRead for HeapReader {
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        self.project().inner.poll_fill_buf(cx)
    }

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

impl Read for HeapReader {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.inner.read(buf)
    }

    #[inline]
    fn read_exact(&mut self, buf: &mut [u8]) -> std::io::Result<()> {
        self.inner.read_exact(buf)
    }

    #[inline]
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
        self.inner.read_to_end(buf)
    }

    #[inline]
    fn read_to_string(&mut self, buf: &mut String) -> std::io::Result<usize> {
        self.inner.read_to_string(buf)
    }

    #[inline]
    fn read_vectored(&mut self, bufs: &mut [std::io::IoSliceMut<'_>]) -> std::io::Result<usize> {
        self.inner.read_vectored(bufs)
    }
}

/// Reader for reading from a stack buffer
#[derive(Debug, Clone)]
pub struct StackReader<const N: usize> {
    data: [u8; N],
    offset: usize,
}

impl<const N: usize> StackReader<N> {
    /// Creates a new `StackReader` with the specified bytes data.
    #[must_use]
    pub const fn new(data: [u8; N]) -> Self {
        Self { data, offset: 0 }
    }

    /// Skip up to n bytes, less if n < m
    pub fn skip(&mut self, n: usize) {
        self.offset = (self.offset + n).min(N);
    }

    /// How many bytes are there remaining
    #[must_use]
    pub fn remaining(&self) -> usize {
        N - self.offset
    }

    /// Returns true if there are any more bytes to consume
    #[must_use]
    pub fn has_remaining(&self) -> bool {
        self.remaining() > 0
    }
}

impl<const N: usize> From<[u8; N]> for StackReader<N> {
    #[inline]
    fn from(data: [u8; N]) -> Self {
        Self::new(data)
    }
}

impl<const N: usize> AsyncRead for StackReader<N> {
    fn poll_read(
        mut self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        if self.offset < N {
            let remaining = &self.data[self.offset..];
            let to_copy = remaining.len().min(buf.remaining());

            if to_copy > 0 {
                buf.put_slice(&remaining[..to_copy]);
                self.offset += to_copy;
            }
        }

        // done
        Poll::Ready(Ok(()))
    }
}

impl<const N: usize> AsyncBufRead for StackReader<N> {
    fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        let me = self.get_mut();
        Poll::Ready(Ok(if me.offset < N {
            &me.data[me.offset..]
        } else {
            &[]
        }))
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        self.get_mut().skip(amt)
    }
}

impl<const N: usize> Read for StackReader<N> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.offset < N {
            let remaining = &self.data[self.offset..];
            let to_copy = remaining.len().min(buf.len());

            if to_copy > 0 {
                buf[..to_copy].copy_from_slice(&remaining[..to_copy]);
                self.offset += to_copy;
                return Ok(to_copy);
            }
        }

        // done
        Ok(0)
    }
}

/// Reader replaying a [`Bytes`] buffer, releasing it once fully consumed.
///
/// Unlike [`HeapReader`] the backing allocation is dropped at EOF, so the
/// reader can be held for a connection's lifetime (e.g. as the prefix of a
/// [`PrefixedIo`]) without retaining the replayed bytes.
///
/// [`PrefixedIo`]: super::PrefixedIo
#[derive(Debug, Clone, Default)]
pub struct ReplayReader {
    data: Bytes,
}

impl ReplayReader {
    /// Creates a new `ReplayReader` with the specified bytes data.
    #[must_use]
    pub const fn new(data: Bytes) -> Self {
        Self { data }
    }

    /// How many bytes are there remaining
    #[must_use]
    pub fn remaining(&self) -> usize {
        self.data.len()
    }

    /// Returns true if there are any more bytes to consume
    #[must_use]
    pub fn has_remaining(&self) -> bool {
        !self.data.is_empty()
    }

    fn advance(&mut self, n: usize) {
        if n >= self.data.len() {
            // EOF: release the backing allocation
            self.data = Bytes::new();
        } else {
            self.data.advance(n);
        }
    }
}

impl From<Bytes> for ReplayReader {
    #[inline]
    fn from(data: Bytes) -> Self {
        Self::new(data)
    }
}

impl AsyncRead for ReplayReader {
    fn poll_read(
        mut self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let to_copy = self.data.len().min(buf.remaining());
        if to_copy > 0 {
            buf.put_slice(&self.data[..to_copy]);
            self.advance(to_copy);
        }

        // done
        Poll::Ready(Ok(()))
    }
}

impl AsyncBufRead for ReplayReader {
    fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        Poll::Ready(Ok(&self.get_mut().data))
    }

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

impl Read for ReplayReader {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let to_copy = self.data.len().min(buf.len());
        if to_copy > 0 {
            buf[..to_copy].copy_from_slice(&self.data[..to_copy]);
            self.advance(to_copy);
        }

        // done
        Ok(to_copy)
    }
}

pin_project! {
    /// Reader that can be used to chain two readers together.
    #[must_use = "streams do nothing unless polled"]
    #[derive(Debug, Clone)]
    pub struct ChainReader<T, U> {
        #[pin]
        first: T,
        #[pin]
        second: U,
        done_first: bool,
    }
}

impl<T, U> ChainReader<T, U>
where
    T: AsyncRead,
    U: AsyncRead,
{
    /// Creates a new `ChainReader` with the specified readers.
    pub const fn new(first: T, second: U) -> Self {
        Self {
            first,
            second,
            done_first: false,
        }
    }

    /// Gets references to the underlying readers in this `ChainReader`.
    pub fn get_ref(&self) -> (&T, &U) {
        (&self.first, &self.second)
    }

    /// Gets mutable references to the underlying readers in this `ChainReader`.
    ///
    /// Care should be taken to avoid modifying the internal I/O state of the
    /// underlying readers as doing so may corrupt the internal state of this
    /// `ChainReader`.
    pub fn get_mut(&mut self) -> (&mut T, &mut U) {
        (&mut self.first, &mut self.second)
    }

    /// Gets pinned mutable references to the underlying readers in this `ChainReader`.
    ///
    /// Care should be taken to avoid modifying the internal I/O state of the
    /// underlying readers as doing so may corrupt the internal state of this
    /// `ChainReader`.
    #[must_use]
    pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut T>, Pin<&mut U>) {
        let me = self.project();
        (me.first, me.second)
    }

    /// Consumes the `ChainReader`, returning the wrapped readers.
    pub fn into_inner(self) -> (T, U) {
        (self.first, self.second)
    }
}

impl<T, U> AsyncRead for ChainReader<T, U>
where
    T: AsyncRead,
    U: AsyncRead,
{
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let me = self.project();

        if !*me.done_first {
            let rem = buf.remaining();
            ready!(me.first.poll_read(cx, buf))?;
            if buf.remaining() == rem {
                *me.done_first = true;
            } else {
                return Poll::Ready(Ok(()));
            }
        }
        me.second.poll_read(cx, buf)
    }
}

impl<T, U> Read for ChainReader<T, U>
where
    T: Read,
    U: Read,
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if !self.done_first {
            let n = self.first.read(buf)?;
            if n == 0 {
                self.done_first = true;
            } else {
                return Ok(n);
            }
        }
        self.second.read(buf)
    }
}

impl<T, U> AsyncBufRead for ChainReader<T, U>
where
    T: AsyncBufRead,
    U: AsyncBufRead,
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        let me = self.project();

        if !*me.done_first {
            match ready!(me.first.poll_fill_buf(cx)?) {
                [] => {
                    *me.done_first = true;
                }
                buf => return Poll::Ready(Ok(buf)),
            }
        }
        me.second.poll_fill_buf(cx)
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        let me = self.project();
        if !*me.done_first {
            me.first.consume(amt)
        } else {
            me.second.consume(amt)
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    use tokio::io::AsyncReadExt;

    #[tokio::test]
    async fn test_replay_reader_releases_buffer_at_eof() {
        let data = Bytes::from(vec![7u8; 64]);
        let mut reader = ReplayReader::new(data.clone());
        let mut out = Vec::new();
        tokio::io::AsyncReadExt::read_to_end(&mut reader, &mut out)
            .await
            .unwrap();
        assert_eq!(vec![7u8; 64], out);
        // the reader dropped its handle at EOF: ours is unique again
        data.try_into_mut().expect("replay buffer was released");
    }

    #[tokio::test]
    async fn test_discard_consumes_exact_bytes() {
        let mut reader = Cursor::new(b"abcdefghij".to_vec());
        discard(&mut reader, 4).await.unwrap();
        let mut rest = Vec::new();
        tokio::io::AsyncReadExt::read_to_end(&mut reader, &mut rest)
            .await
            .unwrap();
        assert_eq!(rest, b"efghij");
    }

    #[tokio::test]
    async fn test_discard_zero_is_noop() {
        let mut reader = Cursor::new(b"abc".to_vec());
        discard(&mut reader, 0).await.unwrap();
        let mut rest = Vec::new();
        tokio::io::AsyncReadExt::read_to_end(&mut reader, &mut rest)
            .await
            .unwrap();
        assert_eq!(rest, b"abc");
    }

    #[tokio::test]
    async fn test_discard_eof_before_n_errors() {
        let mut reader = Cursor::new(b"abc".to_vec());
        let err = discard(&mut reader, 10).await.unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
    }

    async fn test_multi_read_async<const N: usize>(
        mut stream: impl AsyncRead + Unpin,
        cases: &[&str],
    ) {
        let mut buf = [0u8; N];

        for (i, case) in cases.iter().enumerate() {
            let n = stream.read(&mut buf).await.unwrap();
            assert_eq!(
                n,
                case.len(),
                "[{N}][async] step #{} for cases: {:?}",
                i + 1,
                cases
            );
            assert_eq!(
                &buf[..n],
                case.as_bytes(),
                "[{N}][async] step #{} for cases: {:?}",
                i + 1,
                cases
            );
        }
    }

    fn test_multi_read_sync<const N: usize>(mut stream: impl Read, cases: &[&str]) {
        let mut buf = [0u8; N];

        for (i, case) in cases.iter().enumerate() {
            let n = stream.read(&mut buf).unwrap();
            assert_eq!(
                n,
                case.len(),
                "[{N}][sync] step #{} for cases: {:?}",
                i + 1,
                cases
            );
            assert_eq!(
                &buf[..n],
                case.as_bytes(),
                "[{N}][sync] step #{} for cases: {:?}",
                i + 1,
                cases
            );
        }
    }

    #[derive(Debug)]
    struct TestCase<const N: usize, R> {
        reader: R,
        expected_reads: &'static [&'static str],
    }

    impl<const N: usize, R: AsyncRead + Clone + Unpin + Read> TestCase<N, R> {
        async fn test_sync_and_async(&self) {
            let new_stream = || self.reader.clone();

            test_multi_read_async::<N>(&mut new_stream(), self.expected_reads).await;
            test_multi_read_sync::<N>(&mut new_stream(), self.expected_reads);
        }
    }

    #[tokio::test]
    async fn test_heap_reader() {
        TestCase::<5, _> {
            reader: HeapReader::from(""),
            expected_reads: &[""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<5, _> {
            reader: HeapReader::from("hello world"),
            expected_reads: &["hello", " worl", "d", ""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<10, _> {
            reader: HeapReader::from("hello world"),
            expected_reads: &["hello worl", "d", ""],
        }
        .test_sync_and_async()
        .await;
    }

    #[tokio::test]
    async fn test_stack_reader() {
        TestCase::<5, _> {
            reader: StackReader::new(*b""),
            expected_reads: &[""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<5, _> {
            reader: StackReader::new(*b"hello world"),
            expected_reads: &["hello", " worl", "d", ""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<10, _> {
            reader: StackReader::from(*b"hello world"),
            expected_reads: &["hello worl", "d", ""],
        }
        .test_sync_and_async()
        .await;
    }

    #[tokio::test]
    async fn test_chain_reader() {
        TestCase::<5, _> {
            reader: ChainReader::new(Cursor::new(""), Cursor::new("")),
            expected_reads: &[""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<5, _> {
            reader: ChainReader::new(Cursor::new("hello world"), Cursor::new("")),
            expected_reads: &["hello", " worl", "d", ""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<5, _> {
            reader: ChainReader::new(Cursor::new("hello "), Cursor::new("world")),
            expected_reads: &["hello", " ", "world", ""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<5, _> {
            reader: ChainReader::new(Cursor::new(""), Cursor::new("hello world")),
            expected_reads: &["hello", " worl", "d", ""],
        }
        .test_sync_and_async()
        .await;
    }
}