use std::convert::TryFrom;
use std::error::Error;
use std::net::IpAddr;
use async_channel::Sender;
use async_trait::async_trait;
use tokio::net;
use tokio::time::{timeout as tokio_timeout, Duration};
use crate::io::{AsyncReader, AsyncWriter, Stream};
#[async_trait]
pub trait Listener {
async fn accept_clients(&mut self, new_clients: Sender<Stream>) -> Result<(), Box<dyn Error>>;
}
pub struct TcpListener {
listener: net::TcpListener,
}
impl TcpListener {
pub async fn new(local_address: IpAddr, local_port: u16) -> Result<Self, Box<dyn Error>> {
let listener_address = format!("{}:{}", local_address, local_port);
log::info!("start listening on {}", listener_address);
let listener = net::TcpListener::bind(listener_address).await?;
Ok(Self { listener })
}
}
#[async_trait]
impl Listener for TcpListener {
async fn accept_clients(&mut self, new_clients: Sender<Stream>) -> Result<(), Box<dyn Error>> {
while let Ok((client_stream, client_address)) = self.listener.accept().await {
log::debug!("got connection from {}", client_address);
let (client_reader, client_writer) = client_stream.into_split();
new_clients
.send(Stream::new(client_reader, client_writer))
.await?;
}
Ok(())
}
}
pub const MAX_UDP_PACKET_SIZE: usize = u16::MAX as usize;
pub const STREAMED_UDP_PACKET_HEADER_SIZE: usize = 2;
pub async fn stream_udp_packet(payload: &[u8], size: usize, writer: &mut Box<dyn AsyncWriter>) {
if payload.len() < size {
log::error!(
"payload {:?} is too small (expecting size {})",
payload,
size
);
return;
}
let size_u16 = match u16::try_from(size) {
Ok(s) => s,
Err(e) => {
log::error!("size {} can't fit in a u16: {}", size, e);
return;
}
};
if let Err(e) = writer.write(&size_u16.to_be_bytes()).await {
log::error!("failed to write header: {}", e);
return;
};
if let Err(e) = writer.write(&payload[..size]).await {
log::error!("failed to write payload: {}", e);
return;
};
}
#[derive(PartialEq, Debug)]
pub enum UnstreamPacketResult {
Error,
Timeout,
Payload(Vec<u8>),
}
pub async fn unstream_udp_packet(
reader: &mut Box<dyn AsyncReader>,
timeout: Option<Duration>,
) -> UnstreamPacketResult {
let mut header_bytes = [0; STREAMED_UDP_PACKET_HEADER_SIZE];
let read_header_future = reader.read_exact(&mut header_bytes);
let header_size_result = match timeout {
None => read_header_future.await,
Some(duration) => match tokio_timeout(duration, read_header_future).await {
Ok(size_result) => size_result,
Err(_) => return UnstreamPacketResult::Timeout,
},
};
let header_size = match header_size_result {
Ok(size) => size,
Err(e) => {
log::error!("failed to read header: {}", e);
return UnstreamPacketResult::Error;
}
};
if header_size != STREAMED_UDP_PACKET_HEADER_SIZE {
log::error!("got unexpected header size in bytes {}", header_size);
return UnstreamPacketResult::Error;
}
let header = u16::from_be_bytes(header_bytes);
let header_usize = header as usize;
let mut payload = vec![0; header_usize];
let size = match reader.read_exact(&mut payload).await {
Ok(size) => size,
Err(e) => {
log::error!("failed to read payload: {}", e);
return UnstreamPacketResult::Error;
}
};
if size != header_usize {
log::error!("got unexpected data size in bytes {}", header_size);
return UnstreamPacketResult::Error;
}
UnstreamPacketResult::Payload(payload)
}
#[cfg(test)]
mod tests {
use std::io::ErrorKind;
use tokio::io;
use tokio_test::io::Builder;
use crate::io::{AsyncReadWrapper, AsyncWriteWrapper};
use super::*;
#[tokio::test]
async fn stream_udp_packet_payload_too_small() -> Result<(), Box<dyn Error>> {
let payload = vec![1, 2, 3];
let mut writer: Box<dyn AsyncWriter> =
Box::new(AsyncWriteWrapper::new(Builder::new().build()));
stream_udp_packet(&payload, 7, &mut writer).await;
Ok(())
}
#[tokio::test]
async fn stream_udp_packet_size_not_fit_in_u16() -> Result<(), Box<dyn Error>> {
let payload = vec![0; u16::MAX as usize + 7];
let mut writer: Box<dyn AsyncWriter> =
Box::new(AsyncWriteWrapper::new(Builder::new().build()));
stream_udp_packet(&payload, payload.len(), &mut writer).await;
Ok(())
}
#[tokio::test]
async fn stream_udp_packet_write_header_failed() -> Result<(), Box<dyn Error>> {
let payload = vec![1, 2, 3];
let mut writer: Box<dyn AsyncWriter> = Box::new(AsyncWriteWrapper::new(
Builder::new()
.write_error(io::Error::new(ErrorKind::Other, "oh no!"))
.build(),
));
stream_udp_packet(&payload, payload.len(), &mut writer).await;
Ok(())
}
#[tokio::test]
async fn stream_udp_packet_write_payload_failed() -> Result<(), Box<dyn Error>> {
let payload = vec![1, 2, 3];
let mut writer: Box<dyn AsyncWriter> = Box::new(AsyncWriteWrapper::new(
Builder::new()
.write(vec![0u8, 3].as_slice())
.write_error(io::Error::new(ErrorKind::Other, "oh no!"))
.build(),
));
stream_udp_packet(&payload, payload.len(), &mut writer).await;
Ok(())
}
#[tokio::test]
async fn stream_udp_packet_success() -> Result<(), Box<dyn Error>> {
let payload = vec![1, 2, 3];
let mut writer: Box<dyn AsyncWriter> = Box::new(AsyncWriteWrapper::new(
Builder::new()
.write(vec![0u8, 3].as_slice())
.write(payload.as_slice())
.build(),
));
stream_udp_packet(&payload, payload.len(), &mut writer).await;
Ok(())
}
#[tokio::test]
async fn stream_udp_packet_timeout() -> Result<(), Box<dyn Error>> {
let mut reader: Box<dyn AsyncReader> = Box::new(AsyncReadWrapper::new(
Builder::new().wait(Duration::from_secs(5)).build(),
));
let res = unstream_udp_packet(&mut reader, Some(Duration::from_millis(1))).await;
assert_eq!(res, UnstreamPacketResult::Timeout);
Ok(())
}
#[tokio::test]
async fn stream_udp_packet_read_header_failed() -> Result<(), Box<dyn Error>> {
let mut reader: Box<dyn AsyncReader> =
Box::new(AsyncReadWrapper::new(Builder::new().build()));
let res = unstream_udp_packet(&mut reader, None).await;
assert_eq!(res, UnstreamPacketResult::Error);
Ok(())
}
#[tokio::test]
async fn stream_udp_packet_read_payload_failed() -> Result<(), Box<dyn Error>> {
let mut reader: Box<dyn AsyncReader> = Box::new(AsyncReadWrapper::new(
Builder::new().read(vec![0u8, 3].as_slice()).build(),
));
let res = unstream_udp_packet(&mut reader, None).await;
assert_eq!(res, UnstreamPacketResult::Error);
Ok(())
}
#[tokio::test]
async fn stream_udp_packet_read_payload_success() -> Result<(), Box<dyn Error>> {
let payload = vec![1u8, 2, 3];
let mut reader: Box<dyn AsyncReader> = Box::new(AsyncReadWrapper::new(
Builder::new()
.read(vec![0u8, 3].as_slice())
.read(payload.as_slice())
.build(),
));
let res = unstream_udp_packet(&mut reader, None).await;
assert_eq!(res, UnstreamPacketResult::Payload(payload));
Ok(())
}
}