async-codec 0.4.1

//! The `Framed` type for `futures::io::{AsyncRead, AsyncWrite}` streams

use std::{
    error::Error,
    fmt::{
        self,
        Display,
    },
    io,
    pin::Pin,
};

use bytes::{
    buf::UninitSlice,
    Buf,
    BufMut,
    BytesMut,
};

use futures::{
    io::{
        AsyncRead,
        AsyncWrite,
    },
    prelude::*,
    ready,
    task::{
        Context,
        Poll,
    },
};

use super::*;

/// A wrapper around a byte stream that uses an `Encode + Decode` type to
/// produce a `Sink + Stream`
pub struct Framed<S, C> {
    stream: S,
    codec: C,
    read_buf: BytesMut,
    write_buf: BytesMut,
}

impl<S, C: Encode + Decode> Framed<S, C> {
    /// Create a new framed stream from a byte stream and a codec.
    pub fn new(stream: S, codec: C) -> Self {
        Framed {
            stream,
            codec,
            read_buf: BytesMut::default(),
            write_buf: BytesMut::default(),
        }
    }
}

/// Errors arising from reading a frame
#[derive(Debug)]
pub enum ReadFrameError<E> {
    /// There was an error in the underlying stream
    Io(io::Error),
    /// There was an error decoding the frame
    Decode(E),
}

impl<E: Error + Send + Sync + 'static> From<ReadFrameError<E>> for io::Error {
    fn from(other: ReadFrameError<E>) -> io::Error {
        match other {
            ReadFrameError::Decode(err) => io::Error::new(io::ErrorKind::InvalidData, err),
            ReadFrameError::Io(err) => err,
        }
    }
}

impl<E: Display> Display for ReadFrameError<E> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            ReadFrameError::Io(e) => write!(f, "error reading from stream: {}", e),
            ReadFrameError::Decode(e) => write!(f, "error decoding frame: {}", e),
        }
    }
}

impl<E> Error for ReadFrameError<E>
where
    E: Error + 'static,
{
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            ReadFrameError::Io(ref e) => Some(e),
            ReadFrameError::Decode(ref e) => Some(e),
        }
    }
}

impl<S, C> Stream for Framed<S, C>
where
    C: Decode,
    S: AsyncRead,
{
    type Item = Result<C::Item, ReadFrameError<C::Error>>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let (mut stream, codec, read_buf) = unsafe {
            let this = self.get_unchecked_mut();
            (
                Pin::new_unchecked(&mut this.stream),
                &mut this.codec,
                &mut this.read_buf,
            )
        };
        loop {
            let empty = read_buf.is_empty();
            if !empty {
                let (consumed, decode_res) = codec.decode(read_buf);
                read_buf.advance(consumed);
                match decode_res {
                    DecodeResult::Ok(value) => {
                        return Poll::Ready(Some(Ok(value)));
                    }
                    DecodeResult::Err(e) => {
                        return Poll::Ready(Some(Err(ReadFrameError::Decode(e))));
                    }
                    DecodeResult::UnexpectedEnd => {}
                }
            }

            // Make sure there's at least one byte available to read into
            read_buf.reserve(1);

            // Safety: the buffer that we're reading into is immediately zeroed
            // without reading it.
            let eof = unsafe {
                let n = {
                    let b = zero_buf(read_buf.chunk_mut());
                    match ready!(stream.as_mut().poll_read(cx, b)).map_err(ReadFrameError::Io) {
                        Err(e) => return Poll::Ready(Some(Err(e))),
                        Ok(n) => n,
                    }
                };

                read_buf.advance_mut(n);

                n == 0
            };

            if eof {
                if empty {
                    return Poll::Ready(None);
                } else {
                    return Poll::Ready(Some(Err(ReadFrameError::Io(
                        io::ErrorKind::UnexpectedEof.into(),
                    ))));
                }
            }
        }
    }
}

/// Errors arising from writing a frame to a stream
#[derive(Debug)]
pub enum WriteFrameError<E> {
    /// An error in the underlying stream
    Io(io::Error),
    /// An error from encoding the frame
    Encode(E),
}

impl<E: Error + Send + Sync + 'static> From<WriteFrameError<E>> for io::Error {
    fn from(other: WriteFrameError<E>) -> io::Error {
        match other {
            WriteFrameError::Encode(err) => io::Error::new(io::ErrorKind::InvalidInput, err),
            WriteFrameError::Io(err) => err,
        }
    }
}

impl<E: Display> Display for WriteFrameError<E> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            WriteFrameError::Io(e) => write!(f, "error writing to stream: {}", e),
            WriteFrameError::Encode(e) => write!(f, "error encoding frame: {}", e),
        }
    }
}

impl<E> Error for WriteFrameError<E>
where
    E: Error + 'static,
{
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            WriteFrameError::Io(ref e) => Some(e),
            WriteFrameError::Encode(ref e) => Some(e),
        }
    }
}

impl<S, C> Sink<C::Item> for Framed<S, C>
where
    C: Encode,
    C::Error: std::fmt::Debug,
    S: AsyncWrite,
{
    type Error = WriteFrameError<C::Error>;
    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let (buffer, mut stream) = unsafe {
            let this = self.get_unchecked_mut();
            (&mut this.write_buf, Pin::new_unchecked(&mut this.stream))
        };
        loop {
            if buffer.len() == 0 {
                return Poll::Ready(Ok(()));
            }
            let written = ready!(stream
                .as_mut()
                .poll_write(cx, &buffer)
                .map_err(WriteFrameError::Io))?;
            if written == 0 {
                return Poll::Ready(Err(WriteFrameError::Io(io::ErrorKind::WriteZero.into())));
            }
            buffer.advance(written);
        }
    }
    fn start_send(self: Pin<&mut Self>, item: C::Item) -> Result<(), Self::Error> {
        let (buffer, codec) = unsafe {
            let this = self.get_unchecked_mut();
            (&mut this.write_buf, &mut this.codec)
        };
        codec.reset();
        loop {
            let b = zero_buf(buffer.chunk_mut());
            match codec.encode(&item, b) {
                EncodeResult::Ok(len) => {
                    // Safety: We made sure to zero the buffer above, so this is
                    // safe assuming the advance_mut implementation won't
                    // overflow the buffer
                    unsafe { buffer.advance_mut(len) };
                    return Ok(());
                }
                EncodeResult::Err(e) => return Err(WriteFrameError::Encode(e)),
                EncodeResult::Overflow(0) => buffer.reserve(buffer.remaining_mut() * 2),
                EncodeResult::Overflow(new_size) => buffer.reserve(new_size),
            }
        }
    }
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        ready!(self.as_mut().poll_ready(cx))?;

        unsafe { Pin::new_unchecked(&mut self.get_unchecked_mut().stream) }
            .poll_flush(cx)
            .map_err(WriteFrameError::Io)
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        ready!(self.as_mut().poll_flush(cx))?;

        unsafe { Pin::new_unchecked(&mut self.get_unchecked_mut().stream) }
            .poll_close(cx)
            .map_err(WriteFrameError::Io)
    }
}

fn zero_buf(b: &mut UninitSlice) -> &mut [u8] {
    for i in 0..b.len() {
        b.write_byte(i, 0)
    }
    unsafe { std::mem::transmute(b) }
}

#[cfg(test)]
mod test {
    use std::str::Utf8Error;

    use futures::{
        io::Cursor,
        prelude::*,
    };

    use super::*;

    struct LineCodec;

    impl Encode for LineCodec {
        type Item = String;
        type Error = ();
        fn encode(&mut self, item: &String, buf: &mut [u8]) -> EncodeResult<()> {
            let needed = item.as_bytes().len() + 1;
            if buf.len() < needed {
                return EncodeResult::Overflow(needed);
            }
            buf[..needed - 1].copy_from_slice(item.as_bytes());
            buf[needed - 1] = b'\n';
            Ok(needed).into()
        }
    }

    impl Decode for LineCodec {
        type Item = String;
        type Error = Utf8Error;

        fn decode(&mut self, buf: &mut [u8]) -> (usize, DecodeResult<String, Utf8Error>) {
            let newline = match buf.iter().position(|b| *b == b'\n') {
                Some(idx) => idx,
                None => return (0, DecodeResult::UnexpectedEnd),
            };
            let string_bytes = &buf[..newline];
            (
                newline + 1,
                std::str::from_utf8(string_bytes).map(String::from).into(),
            )
        }
    }

    const SHAKESPEARE: &str = r#"Now is the winter of our discontent
Made glorious summer by this sun of York.
Some are born great, some achieve greatness
And some have greatness thrust upon them.
Friends, Romans, countrymen - lend me your ears!
I come not to praise Caesar, but to bury him.
The evil that men do lives after them
The good is oft interred with their bones.
                    It is a tale
Told by an idiot, full of sound and fury
Signifying nothing.
Ay me! For aught that I could ever read,
Could ever hear by tale or history,
The course of true love never did run smooth.
I have full cause of weeping, but this heart
Shall break into a hundred thousand flaws,
Or ere I'll weep.-O Fool, I shall go mad!
                    Each your doing,
So singular in each particular,
Crowns what you are doing in the present deed,
That all your acts are queens.
"#;

    #[async_std::test]
    async fn test_framed_stream() {
        let reader = Cursor::new(Vec::from(SHAKESPEARE.as_bytes()));
        let mut framed = Framed::new(reader, LineCodec);
        let expected = SHAKESPEARE.lines().map(String::from).collect::<Vec<_>>();
        let mut actual = vec![];
        while let Some(frame) = framed.next().await.transpose().unwrap() {
            actual.push(frame);
        }
        assert_eq!(actual, expected);
    }

    #[async_std::test]
    async fn test_framed_sink() {
        let frames = SHAKESPEARE.lines().map(String::from).collect::<Vec<_>>();
        let mut actual = vec![0u8; SHAKESPEARE.as_bytes().len()];
        {
            let writer = Cursor::new(&mut actual);
            let mut framed = Framed::new(writer, LineCodec);
            for frame in frames {
                framed.send(frame).await.unwrap();
            }
        }
        assert_eq!(std::str::from_utf8(&actual).unwrap(), SHAKESPEARE);
    }
}