use futures::{prelude::*, try_ready};
use std::{cmp, error, fmt, io::Cursor, mem};
use tokio_io::{io, AsyncRead, AsyncWrite};
#[inline]
pub fn write_one<TSocket, TData>(socket: TSocket, data: TData) -> WriteOne<TSocket, TData>
where
TSocket: AsyncWrite,
TData: AsRef<[u8]>,
{
let len_data = build_int_buffer(data.as_ref().len());
WriteOne {
inner: WriteOneInner::WriteLen(io::write_all(socket, len_data), data),
}
}
fn build_int_buffer(num: usize) -> io::Window<[u8; 10]> {
let mut len_data = unsigned_varint::encode::u64_buffer();
let encoded_len = unsigned_varint::encode::u64(num as u64, &mut len_data).len();
let mut len_data = io::Window::new(len_data);
len_data.set_end(encoded_len);
len_data
}
pub struct WriteOne<TSocket, TData = Vec<u8>> {
inner: WriteOneInner<TSocket, TData>,
}
enum WriteOneInner<TSocket, TData> {
WriteLen(io::WriteAll<TSocket, io::Window<[u8; 10]>>, TData),
Write(io::WriteAll<TSocket, TData>),
Shutdown(io::Shutdown<TSocket>),
Poisoned,
}
impl<TSocket, TData> Future for WriteOne<TSocket, TData>
where
TSocket: AsyncWrite,
TData: AsRef<[u8]>,
{
type Item = ();
type Error = std::io::Error;
#[inline]
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
Ok(self.inner.poll()?.map(|_socket| ()))
}
}
impl<TSocket, TData> Future for WriteOneInner<TSocket, TData>
where
TSocket: AsyncWrite,
TData: AsRef<[u8]>,
{
type Item = TSocket;
type Error = std::io::Error;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
loop {
match mem::replace(self, WriteOneInner::Poisoned) {
WriteOneInner::WriteLen(mut inner, data) => match inner.poll()? {
Async::Ready((socket, _)) => {
*self = WriteOneInner::Write(io::write_all(socket, data));
}
Async::NotReady => {
*self = WriteOneInner::WriteLen(inner, data);
}
},
WriteOneInner::Write(mut inner) => match inner.poll()? {
Async::Ready((socket, _)) => {
*self = WriteOneInner::Shutdown(tokio_io::io::shutdown(socket));
}
Async::NotReady => {
*self = WriteOneInner::Write(inner);
}
},
WriteOneInner::Shutdown(ref mut inner) => {
let socket = try_ready!(inner.poll());
return Ok(Async::Ready(socket));
}
WriteOneInner::Poisoned => panic!(),
}
}
}
}
#[inline]
pub fn read_one<TSocket>(
socket: TSocket,
max_size: usize,
) -> ReadOne<TSocket>
{
ReadOne {
inner: ReadOneInner::ReadLen {
socket,
len_buf: Cursor::new([0; 10]),
max_size,
},
}
}
pub struct ReadOne<TSocket> {
inner: ReadOneInner<TSocket>,
}
enum ReadOneInner<TSocket> {
ReadLen {
socket: TSocket,
len_buf: Cursor<[u8; 10]>,
max_size: usize,
},
ReadRest(io::ReadExact<TSocket, io::Window<Vec<u8>>>),
Poisoned,
}
impl<TSocket> Future for ReadOne<TSocket>
where
TSocket: AsyncRead,
{
type Item = Vec<u8>;
type Error = ReadOneError;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
Ok(self.inner.poll()?.map(|(_, out)| out))
}
}
impl<TSocket> Future for ReadOneInner<TSocket>
where
TSocket: AsyncRead,
{
type Item = (TSocket, Vec<u8>);
type Error = ReadOneError;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
loop {
match mem::replace(self, ReadOneInner::Poisoned) {
ReadOneInner::ReadLen {
mut socket,
mut len_buf,
max_size,
} => {
match socket.read_buf(&mut len_buf)? {
Async::Ready(num_read) => {
if num_read == 0 {
if len_buf.position() == 0 {
return Ok(Async::Ready((socket, Vec::new())));
} else {
return Err(ReadOneError::Io(
std::io::ErrorKind::UnexpectedEof.into(),
));
}
}
let len_buf_with_data =
&len_buf.get_ref()[..len_buf.position() as usize];
if let Ok((len, data_start)) =
unsigned_varint::decode::usize(len_buf_with_data)
{
if len >= max_size {
return Err(ReadOneError::TooLarge {
requested: len,
max: max_size,
});
}
let n = cmp::min(data_start.len(), len);
let mut data_buf = vec![0; len];
data_buf[.. n].copy_from_slice(&data_start[.. n]);
let mut data_buf = io::Window::new(data_buf);
data_buf.set_start(data_start.len());
*self = ReadOneInner::ReadRest(io::read_exact(socket, data_buf));
} else {
*self = ReadOneInner::ReadLen {
socket,
len_buf,
max_size,
};
}
}
Async::NotReady => {
*self = ReadOneInner::ReadLen {
socket,
len_buf,
max_size,
};
return Ok(Async::NotReady);
}
}
}
ReadOneInner::ReadRest(mut inner) => {
match inner.poll()? {
Async::Ready((socket, data)) => {
return Ok(Async::Ready((socket, data.into_inner())));
}
Async::NotReady => {
*self = ReadOneInner::ReadRest(inner);
return Ok(Async::NotReady);
}
}
}
ReadOneInner::Poisoned => panic!(),
}
}
}
}
#[derive(Debug)]
pub enum ReadOneError {
Io(std::io::Error),
TooLarge {
requested: usize,
max: usize,
},
}
impl From<std::io::Error> for ReadOneError {
#[inline]
fn from(err: std::io::Error) -> ReadOneError {
ReadOneError::Io(err)
}
}
impl fmt::Display for ReadOneError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match *self {
ReadOneError::Io(ref err) => write!(f, "{}", err),
ReadOneError::TooLarge { .. } => write!(f, "Received data size over maximum"),
}
}
}
impl error::Error for ReadOneError {
fn source(&self) -> Option<&(dyn error::Error + 'static)> {
match *self {
ReadOneError::Io(ref err) => Some(err),
ReadOneError::TooLarge { .. } => None,
}
}
}
#[inline]
pub fn read_one_then<TSocket, TParam, TThen, TOut, TErr>(
socket: TSocket,
max_size: usize,
param: TParam,
then: TThen,
) -> ReadOneThen<TSocket, TParam, TThen>
where
TSocket: AsyncRead,
TThen: FnOnce(Vec<u8>, TParam) -> Result<TOut, TErr>,
TErr: From<ReadOneError>,
{
ReadOneThen {
inner: read_one(socket, max_size),
then: Some((param, then)),
}
}
pub struct ReadOneThen<TSocket, TParam, TThen> {
inner: ReadOne<TSocket>,
then: Option<(TParam, TThen)>,
}
impl<TSocket, TParam, TThen, TOut, TErr> Future for ReadOneThen<TSocket, TParam, TThen>
where
TSocket: AsyncRead,
TThen: FnOnce(Vec<u8>, TParam) -> Result<TOut, TErr>,
TErr: From<ReadOneError>,
{
type Item = TOut;
type Error = TErr;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
match self.inner.poll()? {
Async::Ready(buffer) => {
let (param, then) = self.then.take()
.expect("Future was polled after it was finished");
Ok(Async::Ready(then(buffer, param)?))
},
Async::NotReady => Ok(Async::NotReady),
}
}
}
#[inline]
pub fn read_respond<TSocket, TThen, TParam, TOut, TErr>(
socket: TSocket,
max_size: usize,
param: TParam,
then: TThen,
) -> ReadRespond<TSocket, TParam, TThen>
where
TSocket: AsyncRead,
TThen: FnOnce(TSocket, Vec<u8>, TParam) -> Result<TOut, TErr>,
TErr: From<ReadOneError>,
{
ReadRespond {
inner: read_one(socket, max_size).inner,
then: Some((then, param)),
}
}
pub struct ReadRespond<TSocket, TParam, TThen> {
inner: ReadOneInner<TSocket>,
then: Option<(TThen, TParam)>,
}
impl<TSocket, TThen, TParam, TOut, TErr> Future for ReadRespond<TSocket, TParam, TThen>
where
TSocket: AsyncRead,
TThen: FnOnce(TSocket, Vec<u8>, TParam) -> Result<TOut, TErr>,
TErr: From<ReadOneError>,
{
type Item = TOut;
type Error = TErr;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
match self.inner.poll()? {
Async::Ready((socket, buffer)) => {
let (then, param) = self.then.take().expect("Future was polled after it was finished");
Ok(Async::Ready(then(socket, buffer, param)?))
},
Async::NotReady => Ok(Async::NotReady),
}
}
}
#[inline]
pub fn request_response<TSocket, TData, TParam, TThen, TOut, TErr>(
socket: TSocket,
data: TData,
max_size: usize,
param: TParam,
then: TThen,
) -> RequestResponse<TSocket, TParam, TThen, TData>
where
TSocket: AsyncRead + AsyncWrite,
TData: AsRef<[u8]>,
TThen: FnOnce(Vec<u8>, TParam) -> Result<TOut, TErr>,
{
RequestResponse {
inner: RequestResponseInner::Write(write_one(socket, data).inner, max_size, param, then),
}
}
pub struct RequestResponse<TSocket, TParam, TThen, TData = Vec<u8>> {
inner: RequestResponseInner<TSocket, TData, TParam, TThen>,
}
enum RequestResponseInner<TSocket, TData, TParam, TThen> {
Write(WriteOneInner<TSocket, TData>, usize, TParam, TThen),
Read(ReadOneThen<TSocket, TParam, TThen>),
Poisoned,
}
impl<TSocket, TData, TParam, TThen, TOut, TErr> Future for RequestResponse<TSocket, TParam, TThen, TData>
where
TSocket: AsyncRead + AsyncWrite,
TData: AsRef<[u8]>,
TThen: FnOnce(Vec<u8>, TParam) -> Result<TOut, TErr>,
TErr: From<ReadOneError>,
{
type Item = TOut;
type Error = TErr;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
loop {
match mem::replace(&mut self.inner, RequestResponseInner::Poisoned) {
RequestResponseInner::Write(mut inner, max_size, param, then) => {
match inner.poll().map_err(ReadOneError::Io)? {
Async::Ready(socket) => {
self.inner =
RequestResponseInner::Read(read_one_then(socket, max_size, param, then));
}
Async::NotReady => {
self.inner = RequestResponseInner::Write(inner, max_size, param, then);
return Ok(Async::NotReady);
}
}
}
RequestResponseInner::Read(mut inner) => match inner.poll()? {
Async::Ready(packet) => return Ok(Async::Ready(packet)),
Async::NotReady => {
self.inner = RequestResponseInner::Read(inner);
return Ok(Async::NotReady);
}
},
RequestResponseInner::Poisoned => panic!(),
};
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{self, Cursor};
use tokio::runtime::current_thread::Runtime;
#[test]
fn write_one_works() {
let data = (0..rand::random::<usize>() % 10_000)
.map(|_| rand::random::<u8>())
.collect::<Vec<_>>();
let mut out = vec![0; 10_000];
let future = write_one(Cursor::new(&mut out[..]), data.clone());
Runtime::new().unwrap().block_on(future).unwrap();
let (out_len, out_data) = unsigned_varint::decode::usize(&out).unwrap();
assert_eq!(out_len, data.len());
assert_eq!(&out_data[..out_len], &data[..]);
}
#[test]
fn read_one_works() {
let original_data = (0..rand::random::<usize>() % 10_000)
.map(|_| rand::random::<u8>())
.collect::<Vec<_>>();
let mut len_buf = unsigned_varint::encode::usize_buffer();
let len_buf = unsigned_varint::encode::usize(original_data.len(), &mut len_buf);
let mut in_buffer = len_buf.to_vec();
in_buffer.extend_from_slice(&original_data);
let future = read_one_then(Cursor::new(in_buffer), 10_000, (), move |out, ()| -> Result<_, ReadOneError> {
assert_eq!(out, original_data);
Ok(())
});
Runtime::new().unwrap().block_on(future).unwrap();
}
#[test]
fn read_one_zero_len() {
let future = read_one_then(Cursor::new(vec![0]), 10_000, (), move |out, ()| -> Result<_, ReadOneError> {
assert!(out.is_empty());
Ok(())
});
Runtime::new().unwrap().block_on(future).unwrap();
}
#[test]
fn read_checks_length() {
let mut len_buf = unsigned_varint::encode::u64_buffer();
let len_buf = unsigned_varint::encode::u64(5_000, &mut len_buf);
let mut in_buffer = len_buf.to_vec();
in_buffer.extend((0..5000).map(|_| 0));
let future = read_one_then(Cursor::new(in_buffer), 100, (), move |_, ()| -> Result<_, ReadOneError> {
Ok(())
});
match Runtime::new().unwrap().block_on(future) {
Err(ReadOneError::TooLarge { .. }) => (),
_ => panic!(),
}
}
#[test]
fn read_one_accepts_empty() {
let future = read_one_then(Cursor::new([]), 10_000, (), move |out, ()| -> Result<_, ReadOneError> {
assert!(out.is_empty());
Ok(())
});
Runtime::new().unwrap().block_on(future).unwrap();
}
#[test]
fn read_one_eof_before_len() {
let future = read_one_then(Cursor::new([0x80]), 10_000, (), move |_, ()| -> Result<(), ReadOneError> {
unreachable!()
});
match Runtime::new().unwrap().block_on(future) {
Err(ReadOneError::Io(ref err)) if err.kind() == io::ErrorKind::UnexpectedEof => (),
_ => panic!()
}
}
}