use std::io;
use std::mem;
use std::net::SocketAddr;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use async_trait::async_trait;
use futures::channel::mpsc::{unbounded, UnboundedReceiver};
use futures::stream::{Fuse, Peekable, Stream, StreamExt};
use futures::{ready, Future, FutureExt, TryFutureExt};
use log::debug;
use crate::error::*;
use crate::xfer::{BufStreamHandle, SerialMessage};
#[async_trait]
pub trait Connect
where
Self: Sized,
{
type Transport: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send;
async fn connect(addr: &SocketAddr) -> io::Result<Self::Transport>;
}
enum WriteTcpState {
LenBytes {
pos: usize,
length: [u8; 2],
bytes: Vec<u8>,
},
Bytes {
pos: usize,
bytes: Vec<u8>,
},
Flushing,
}
pub enum ReadTcpState {
LenBytes {
pos: usize,
bytes: [u8; 2],
},
Bytes {
pos: usize,
bytes: Vec<u8>,
},
}
#[must_use = "futures do nothing unless polled"]
pub struct TcpStream<S> {
socket: S,
outbound_messages: Peekable<Fuse<UnboundedReceiver<SerialMessage>>>,
send_state: Option<WriteTcpState>,
read_state: ReadTcpState,
peer_addr: SocketAddr,
}
impl<S> TcpStream<S> {
pub fn peer_addr(&self) -> SocketAddr {
self.peer_addr
}
fn pollable_split(
&mut self,
) -> (
&mut S,
&mut Peekable<Fuse<UnboundedReceiver<SerialMessage>>>,
&mut Option<WriteTcpState>,
&mut ReadTcpState,
) {
(
&mut self.socket,
&mut self.outbound_messages,
&mut self.send_state,
&mut self.read_state,
)
}
}
impl<S: Connect + 'static> TcpStream<S> {
#[allow(clippy::new_ret_no_self, clippy::type_complexity)]
pub fn new<E>(
name_server: SocketAddr,
) -> (
impl Future<Output = Result<TcpStream<S::Transport>, io::Error>> + Send,
BufStreamHandle,
)
where
E: FromProtoError,
{
Self::with_timeout(name_server, Duration::from_secs(5))
}
#[allow(clippy::type_complexity)]
pub fn with_timeout(
name_server: SocketAddr,
timeout: Duration,
) -> (
impl Future<Output = Result<TcpStream<S::Transport>, io::Error>> + Send,
BufStreamHandle,
) {
let (message_sender, outbound_messages) = unbounded();
let message_sender = BufStreamHandle::new(message_sender);
let stream_fut = Self::connect(name_server, timeout, outbound_messages);
(stream_fut, message_sender)
}
async fn connect(
name_server: SocketAddr,
timeout: Duration,
outbound_messages: UnboundedReceiver<SerialMessage>,
) -> Result<TcpStream<S::Transport>, io::Error> {
let tcp = S::connect(&name_server);
tokio::time::timeout(timeout, tcp)
.map_err(move |_| {
debug!("timed out connecting to: {}", name_server);
io::Error::new(
io::ErrorKind::TimedOut,
format!("timed out connecting to: {}", name_server),
)
})
.map(
move |tcp_stream: Result<Result<S::Transport, io::Error>, _>| {
tcp_stream
.and_then(|tcp_stream| tcp_stream)
.map(|tcp_stream| {
debug!("TCP connection established to: {}", name_server);
TcpStream {
socket: tcp_stream,
outbound_messages: outbound_messages.fuse().peekable(),
send_state: None,
read_state: ReadTcpState::LenBytes {
pos: 0,
bytes: [0u8; 2],
},
peer_addr: name_server,
}
})
},
)
.await
}
}
impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite> TcpStream<S> {
pub fn from_stream(stream: S, peer_addr: SocketAddr) -> (Self, BufStreamHandle) {
let (message_sender, outbound_messages) = unbounded();
let message_sender = BufStreamHandle::new(message_sender);
let stream = Self::from_stream_with_receiver(stream, peer_addr, outbound_messages);
(stream, message_sender)
}
pub fn from_stream_with_receiver(
stream: S,
peer_addr: SocketAddr,
receiver: UnboundedReceiver<SerialMessage>,
) -> Self {
TcpStream {
socket: stream,
outbound_messages: receiver.fuse().peekable(),
send_state: None,
read_state: ReadTcpState::LenBytes {
pos: 0,
bytes: [0u8; 2],
},
peer_addr,
}
}
}
impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin> Stream for TcpStream<S> {
type Item = io::Result<SerialMessage>;
#[allow(clippy::cognitive_complexity)]
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
let peer = self.peer_addr;
let (socket, outbound_messages, send_state, read_state) = self.pollable_split();
let mut socket = Pin::new(socket);
let mut outbound_messages = Pin::new(outbound_messages);
loop {
if send_state.is_some() {
match send_state {
Some(WriteTcpState::LenBytes {
ref mut pos,
ref length,
..
}) => {
let wrote = ready!(socket.as_mut().poll_write(cx, &length[*pos..]))?;
*pos += wrote;
}
Some(WriteTcpState::Bytes {
ref mut pos,
ref bytes,
}) => {
let wrote = ready!(socket.as_mut().poll_write(cx, &bytes[*pos..]))?;
*pos += wrote;
}
Some(WriteTcpState::Flushing) => {
ready!(socket.as_mut().poll_flush(cx))?;
}
_ => (),
}
let current_state = send_state.take();
match current_state {
Some(WriteTcpState::LenBytes { pos, length, bytes }) => {
if pos < length.len() {
mem::replace(
send_state,
Some(WriteTcpState::LenBytes { pos, length, bytes }),
);
} else {
mem::replace(send_state, Some(WriteTcpState::Bytes { pos: 0, bytes }));
}
}
Some(WriteTcpState::Bytes { pos, bytes }) => {
if pos < bytes.len() {
mem::replace(send_state, Some(WriteTcpState::Bytes { pos, bytes }));
} else {
mem::replace(send_state, Some(WriteTcpState::Flushing));
}
}
Some(WriteTcpState::Flushing) => {
send_state.take();
}
None => (),
};
} else {
match outbound_messages.as_mut().poll_next(cx)
{
Poll::Ready(Some(message)) => {
let (buffer, dst) = message.unwrap();
if peer != dst {
return Poll::Ready(Some(Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("mismatched peer: {} and dst: {}", peer, dst),
))));
}
let len: [u8; 2] = [
(buffer.len() >> 8 & 0xFF) as u8,
(buffer.len() & 0xFF) as u8,
];
debug!("sending message len: {} to: {}", buffer.len(), dst);
*send_state = Some(WriteTcpState::LenBytes {
pos: 0,
length: len,
bytes: buffer,
});
}
Poll::Pending => break,
Poll::Ready(None) => {
debug!("no messages to send");
break;
}
}
}
}
let mut ret_buf: Option<Vec<u8>> = None;
while ret_buf.is_none() {
let new_state: Option<ReadTcpState> = match read_state {
ReadTcpState::LenBytes {
ref mut pos,
ref mut bytes,
} => {
let read = ready!(socket.as_mut().poll_read(cx, &mut bytes[*pos..]))?;
if read == 0 {
debug!("zero bytes read, stream closed?");
if *pos == 0 {
return Poll::Ready(None);
} else {
return Poll::Ready(Some(Err(io::Error::new(
io::ErrorKind::BrokenPipe,
"closed while reading length",
))));
}
}
debug!("in ReadTcpState::LenBytes: {}", pos);
*pos += read;
if *pos < bytes.len() {
debug!("remain ReadTcpState::LenBytes: {}", pos);
None
} else {
let length =
u16::from(bytes[0]) << 8 & 0xFF00 | u16::from(bytes[1]) & 0x00FF;
debug!("got length: {}", length);
let mut bytes = vec![0; length as usize];
bytes.resize(length as usize, 0);
debug!("move ReadTcpState::Bytes: {}", bytes.len());
Some(ReadTcpState::Bytes { pos: 0, bytes })
}
}
ReadTcpState::Bytes {
ref mut pos,
ref mut bytes,
} => {
let read = ready!(socket.as_mut().poll_read(cx, &mut bytes[*pos..]))?;
if read == 0 {
debug!("zero bytes read for message, stream closed?");
return Poll::Ready(Some(Err(io::Error::new(
io::ErrorKind::BrokenPipe,
"closed while reading message",
))));
}
debug!("in ReadTcpState::Bytes: {}", bytes.len());
*pos += read;
if *pos < bytes.len() {
debug!("remain ReadTcpState::Bytes: {}", bytes.len());
None
} else {
debug!("reset ReadTcpState::LenBytes: {}", 0);
Some(ReadTcpState::LenBytes {
pos: 0,
bytes: [0u8; 2],
})
}
}
};
if let Some(state) = new_state {
if let ReadTcpState::Bytes { pos, bytes } = mem::replace(read_state, state) {
debug!("returning bytes");
assert_eq!(pos, bytes.len());
ret_buf = Some(bytes);
}
}
}
if let Some(buffer) = ret_buf {
debug!("returning buffer");
let src_addr = self.peer_addr;
Poll::Ready(Some(Ok(SerialMessage::new(buffer, src_addr))))
} else {
debug!("bottomed out");
Poll::Pending
}
}
}
#[cfg(test)]
mod tests {
#[cfg(not(target_os = "linux"))]
use std::net::Ipv6Addr;
use std::net::{IpAddr, Ipv4Addr};
use tokio;
use super::*;
#[test]
fn test_tcp_client_stream_ipv4() {
tcp_client_stream_test(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)))
}
#[test]
#[cfg(not(target_os = "linux"))]
fn test_tcp_client_stream_ipv6() {
tcp_client_stream_test(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)))
}
#[cfg(test)]
const TEST_BYTES: &[u8; 8] = b"DEADBEEF";
#[cfg(test)]
const TEST_BYTES_LEN: usize = 8;
#[cfg(test)]
fn tcp_client_stream_test(server_addr: IpAddr) {
use std::io::{Read, Write};
use tokio::runtime;
let succeeded = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let succeeded_clone = succeeded.clone();
std::thread::Builder::new()
.name("thread_killer".to_string())
.spawn(move || {
let succeeded = succeeded_clone;
for _ in 0..15 {
std::thread::sleep(std::time::Duration::from_secs(1));
if succeeded.load(std::sync::atomic::Ordering::Relaxed) {
return;
}
}
panic!("timeout");
})
.unwrap();
let server = std::net::TcpListener::bind(SocketAddr::new(server_addr, 0)).unwrap();
let server_addr = server.local_addr().unwrap();
let send_recv_times = 4;
let server_handle = std::thread::Builder::new()
.name("test_tcp_client_stream:server".to_string())
.spawn(move || {
let (mut socket, _) = server.accept().expect("accept failed");
socket
.set_read_timeout(Some(std::time::Duration::from_secs(5)))
.unwrap();
socket
.set_write_timeout(Some(std::time::Duration::from_secs(5)))
.unwrap();
for _ in 0..send_recv_times {
let mut len_bytes = [0_u8; 2];
socket
.read_exact(&mut len_bytes)
.expect("SERVER: receive failed");
let length =
u16::from(len_bytes[0]) << 8 & 0xFF00 | u16::from(len_bytes[1]) & 0x00FF;
assert_eq!(length as usize, TEST_BYTES_LEN);
let mut buffer = [0_u8; TEST_BYTES_LEN];
socket.read_exact(&mut buffer).unwrap();
assert_eq!(&buffer, TEST_BYTES);
socket
.write_all(&len_bytes)
.expect("SERVER: send length failed");
socket
.write_all(&buffer)
.expect("SERVER: send buffer failed");
std::thread::yield_now();
}
})
.unwrap();
let mut io_loop = runtime::Runtime::new().unwrap();
let (stream, sender) = TcpStream::<tokio::net::TcpStream>::new::<ProtoError>(server_addr);
let mut stream = io_loop.block_on(stream).expect("run failed to get stream");
for _ in 0..send_recv_times {
sender
.unbounded_send(SerialMessage::new(TEST_BYTES.to_vec(), server_addr))
.expect("send failed");
let (buffer, stream_tmp) = io_loop.block_on(stream.into_future());
stream = stream_tmp;
let message = buffer
.expect("no buffer received")
.expect("error receiving buffer");
assert_eq!(message.bytes(), TEST_BYTES);
}
succeeded.store(true, std::sync::atomic::Ordering::Relaxed);
server_handle.join().expect("server thread failed");
}
}