| use std::cmp; |
| use std::io::{self, Read, Write}; |
| |
| use bytes::{Buf, BufMut, Bytes, IntoBuf}; |
| use futures::{Async, Poll}; |
| use tokio_io::{AsyncRead, AsyncWrite}; |
| |
| /// Combine a buffer with an IO, rewinding reads to use the buffer. |
| #[derive(Debug)] |
| pub(crate) struct Rewind<T> { |
| pre: Option<Bytes>, |
| inner: T, |
| } |
| |
| impl<T> Rewind<T> { |
| pub(crate) fn new(io: T) -> Self { |
| Rewind { |
| pre: None, |
| inner: io, |
| } |
| } |
| |
| pub(crate) fn new_buffered(io: T, buf: Bytes) -> Self { |
| Rewind { |
| pre: Some(buf), |
| inner: io, |
| } |
| } |
| |
| pub(crate) fn rewind(&mut self, bs: Bytes) { |
| debug_assert!(self.pre.is_none()); |
| self.pre = Some(bs); |
| } |
| |
| pub(crate) fn into_inner(self) -> (T, Bytes) { |
| (self.inner, self.pre.unwrap_or_else(Bytes::new)) |
| } |
| } |
| |
| impl<T> Read for Rewind<T> |
| where |
| T: Read, |
| { |
| #[inline] |
| fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { |
| if let Some(pre_bs) = self.pre.take() { |
| // If there are no remaining bytes, let the bytes get dropped. |
| if pre_bs.len() > 0 { |
| let mut pre_reader = pre_bs.into_buf().reader(); |
| let read_cnt = pre_reader.read(buf)?; |
| |
| let mut new_pre = pre_reader.into_inner().into_inner(); |
| new_pre.advance(read_cnt); |
| |
| // Put back whats left |
| if new_pre.len() > 0 { |
| self.pre = Some(new_pre); |
| } |
| |
| return Ok(read_cnt); |
| } |
| } |
| self.inner.read(buf) |
| } |
| } |
| |
| impl<T> Write for Rewind<T> |
| where |
| T: Write, |
| { |
| #[inline] |
| fn write(&mut self, buf: &[u8]) -> io::Result<usize> { |
| self.inner.write(buf) |
| } |
| |
| #[inline] |
| fn flush(&mut self) -> io::Result<()> { |
| self.inner.flush() |
| } |
| } |
| |
| impl<T> AsyncRead for Rewind<T> |
| where |
| T: AsyncRead, |
| { |
| #[inline] |
| unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool { |
| self.inner.prepare_uninitialized_buffer(buf) |
| } |
| |
| #[inline] |
| fn read_buf<B: BufMut>(&mut self, buf: &mut B) -> Poll<usize, io::Error> { |
| if let Some(bs) = self.pre.take() { |
| let pre_len = bs.len(); |
| // If there are no remaining bytes, let the bytes get dropped. |
| if pre_len > 0 { |
| let cnt = cmp::min(buf.remaining_mut(), pre_len); |
| let pre_buf = bs.into_buf(); |
| let mut xfer = Buf::take(pre_buf, cnt); |
| buf.put(&mut xfer); |
| |
| let mut new_pre = xfer.into_inner().into_inner(); |
| new_pre.advance(cnt); |
| |
| // Put back whats left |
| if new_pre.len() > 0 { |
| self.pre = Some(new_pre); |
| } |
| |
| return Ok(Async::Ready(cnt)); |
| } |
| } |
| self.inner.read_buf(buf) |
| } |
| } |
| |
| impl<T> AsyncWrite for Rewind<T> |
| where |
| T: AsyncWrite, |
| { |
| #[inline] |
| fn shutdown(&mut self) -> Poll<(), io::Error> { |
| AsyncWrite::shutdown(&mut self.inner) |
| } |
| |
| #[inline] |
| fn write_buf<B: Buf>(&mut self, buf: &mut B) -> Poll<usize, io::Error> { |
| self.inner.write_buf(buf) |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| extern crate tokio_mockstream; |
| use self::tokio_mockstream::MockStream; |
| use std::io::Cursor; |
| |
| // Test a partial rewind |
| #[test] |
| fn async_partial_rewind() { |
| let bs = &mut [104, 101, 108, 108, 111]; |
| let o1 = &mut [0, 0]; |
| let o2 = &mut [0, 0, 0, 0, 0]; |
| |
| let mut stream = Rewind::new(MockStream::new(bs)); |
| let mut o1_cursor = Cursor::new(o1); |
| // Read off some bytes, ensure we filled o1 |
| match stream.read_buf(&mut o1_cursor).unwrap() { |
| Async::NotReady => panic!("should be ready"), |
| Async::Ready(cnt) => assert_eq!(2, cnt), |
| } |
| |
| // Rewind the stream so that it is as if we never read in the first place. |
| let read_buf = Bytes::from(&o1_cursor.into_inner()[..]); |
| stream.rewind(read_buf); |
| |
| // We poll 2x here since the first time we'll only get what is in the |
| // prefix (the rewinded part) of the Rewind.\ |
| let mut o2_cursor = Cursor::new(o2); |
| stream.read_buf(&mut o2_cursor).unwrap(); |
| stream.read_buf(&mut o2_cursor).unwrap(); |
| let o2_final = o2_cursor.into_inner(); |
| |
| // At this point we should have read everything that was in the MockStream |
| assert_eq!(&o2_final, &bs); |
| } |
| // Test a full rewind |
| #[test] |
| fn async_full_rewind() { |
| let bs = &mut [104, 101, 108, 108, 111]; |
| let o1 = &mut [0, 0, 0, 0, 0]; |
| let o2 = &mut [0, 0, 0, 0, 0]; |
| |
| let mut stream = Rewind::new(MockStream::new(bs)); |
| let mut o1_cursor = Cursor::new(o1); |
| match stream.read_buf(&mut o1_cursor).unwrap() { |
| Async::NotReady => panic!("should be ready"), |
| Async::Ready(cnt) => assert_eq!(5, cnt), |
| } |
| |
| let read_buf = Bytes::from(&o1_cursor.into_inner()[..]); |
| stream.rewind(read_buf); |
| |
| let mut o2_cursor = Cursor::new(o2); |
| stream.read_buf(&mut o2_cursor).unwrap(); |
| stream.read_buf(&mut o2_cursor).unwrap(); |
| let o2_final = o2_cursor.into_inner(); |
| |
| assert_eq!(&o2_final, &bs); |
| } |
| #[test] |
| fn partial_rewind() { |
| let bs = &mut [104, 101, 108, 108, 111]; |
| let o1 = &mut [0, 0]; |
| let o2 = &mut [0, 0, 0, 0, 0]; |
| |
| let mut stream = Rewind::new(MockStream::new(bs)); |
| stream.read(o1).unwrap(); |
| |
| let read_buf = Bytes::from(&o1[..]); |
| stream.rewind(read_buf); |
| let cnt = stream.read(o2).unwrap(); |
| stream.read(&mut o2[cnt..]).unwrap(); |
| assert_eq!(&o2, &bs); |
| } |
| #[test] |
| fn full_rewind() { |
| let bs = &mut [104, 101, 108, 108, 111]; |
| let o1 = &mut [0, 0, 0, 0, 0]; |
| let o2 = &mut [0, 0, 0, 0, 0]; |
| |
| let mut stream = Rewind::new(MockStream::new(bs)); |
| stream.read(o1).unwrap(); |
| |
| let read_buf = Bytes::from(&o1[..]); |
| stream.rewind(read_buf); |
| let cnt = stream.read(o2).unwrap(); |
| stream.read(&mut o2[cnt..]).unwrap(); |
| assert_eq!(&o2, &bs); |
| } |
| } |