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::*;
pub struct Framed<S, C> {
stream: S,
codec: C,
read_buf: BytesMut,
write_buf: BytesMut,
}
impl<S, C: Encode + Decode> Framed<S, C> {
pub fn new(stream: S, codec: C) -> Self {
Framed {
stream,
codec,
read_buf: BytesMut::default(),
write_buf: BytesMut::default(),
}
}
}
#[derive(Debug)]
pub enum ReadFrameError<E> {
Io(io::Error),
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 => {}
}
}
read_buf.reserve(1);
let eof = unsafe {
let n = {
let b = zero_buf(read_buf.bytes_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(),
))));
}
}
}
}
}
#[derive(Debug)]
pub enum WriteFrameError<E> {
Io(io::Error),
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.bytes_mut());
match codec.encode(&item, b) {
EncodeResult::Ok(len) => {
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);
}
}