use crate::network::adapter::{
Resource, Remote, Local, Adapter, SendStatus, AcceptedType, ReadStatus, ConnectionInfo,
ListeningInfo, PendingStatus,
};
use crate::network::{RemoteAddr, Readiness, TransportConnect, TransportListen};
use mio::net::{TcpListener, TcpStream};
use mio::event::{Source};
use socket2::{Socket, TcpKeepalive};
use std::net::{SocketAddr};
use std::io::{self, ErrorKind, Read, Write};
use std::ops::{Deref};
use std::mem::{forget, MaybeUninit};
#[cfg(target_os = "windows")]
use std::os::windows::io::{FromRawSocket, AsRawSocket};
#[cfg(not(target_os = "windows"))]
use std::os::{fd::AsRawFd, unix::io::FromRawFd};
pub const INPUT_BUFFER_SIZE: usize = u16::MAX as usize; #[derive(Clone, Debug, Default)]
pub struct TcpConnectConfig {
pub keepalive: Option<TcpKeepalive>,
}
#[derive(Clone, Debug, Default)]
pub struct TcpListenConfig {
pub keepalive: Option<TcpKeepalive>,
}
pub(crate) struct TcpAdapter;
impl Adapter for TcpAdapter {
type Remote = RemoteResource;
type Local = LocalResource;
}
pub(crate) struct RemoteResource {
stream: TcpStream,
keepalive: Option<TcpKeepalive>,
}
impl Resource for RemoteResource {
fn source(&mut self) -> &mut dyn Source {
&mut self.stream
}
}
impl Remote for RemoteResource {
fn connect_with(
config: TransportConnect,
remote_addr: RemoteAddr,
) -> io::Result<ConnectionInfo<Self>> {
let config = match config {
TransportConnect::Tcp(config) => config,
_ => panic!("Internal error: Got wrong config"),
};
let peer_addr = *remote_addr.socket_addr();
let stream = TcpStream::connect(peer_addr)?;
let local_addr = stream.local_addr()?;
Ok(ConnectionInfo {
remote: Self { stream, keepalive: config.keepalive },
local_addr,
peer_addr,
})
}
fn receive(&self, mut process_data: impl FnMut(&[u8])) -> ReadStatus {
let buffer: MaybeUninit<[u8; INPUT_BUFFER_SIZE]> = MaybeUninit::uninit();
let mut input_buffer = unsafe { buffer.assume_init() }; loop {
let stream = &self.stream;
match stream.deref().read(&mut input_buffer) {
Ok(0) => break ReadStatus::Disconnected,
Ok(size) => process_data(&input_buffer[..size]),
Err(ref err) if err.kind() == ErrorKind::Interrupted => continue,
Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
break ReadStatus::WaitNextEvent
}
Err(ref err) if err.kind() == ErrorKind::ConnectionReset => {
break ReadStatus::Disconnected
}
Err(err) => {
log::error!("TCP receive error: {}", err);
break ReadStatus::Disconnected }
}
}
}
fn send(&self, data: &[u8]) -> SendStatus {
let mut total_bytes_sent = 0;
loop {
let stream = &self.stream;
match stream.deref().write(&data[total_bytes_sent..]) {
Ok(bytes_sent) => {
total_bytes_sent += bytes_sent;
if total_bytes_sent == data.len() {
break SendStatus::Sent
}
}
Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => continue,
Err(err) => {
log::error!("TCP receive error: {}", err);
break SendStatus::ResourceNotFound }
}
}
}
fn pending(&self, _readiness: Readiness) -> PendingStatus {
let status = check_stream_ready(&self.stream);
if status == PendingStatus::Ready {
if let Some(keepalive) = &self.keepalive {
#[cfg(target_os = "windows")]
let socket = unsafe { Socket::from_raw_socket(self.stream.as_raw_socket()) };
#[cfg(not(target_os = "windows"))]
let socket = unsafe { Socket::from_raw_fd(self.stream.as_raw_fd()) };
if let Err(e) = socket.set_tcp_keepalive(keepalive) {
log::warn!("TCP set keepalive error: {}", e);
}
forget(socket);
}
}
status
}
}
pub fn check_stream_ready(stream: &TcpStream) -> PendingStatus {
if let Ok(Some(_)) = stream.take_error() {
return PendingStatus::Disconnected
}
match stream.peer_addr() {
Ok(_) => PendingStatus::Ready,
Err(err) if err.kind() == io::ErrorKind::NotConnected => PendingStatus::Incomplete,
Err(err) if err.kind() == io::ErrorKind::InvalidInput => PendingStatus::Incomplete,
Err(_) => PendingStatus::Disconnected,
}
}
pub(crate) struct LocalResource {
listener: TcpListener,
keepalive: Option<TcpKeepalive>,
}
impl Resource for LocalResource {
fn source(&mut self) -> &mut dyn Source {
&mut self.listener
}
}
impl Local for LocalResource {
type Remote = RemoteResource;
fn listen_with(config: TransportListen, addr: SocketAddr) -> io::Result<ListeningInfo<Self>> {
let config = match config {
TransportListen::Tcp(config) => config,
_ => panic!("Internal error: Got wrong config"),
};
let listener = TcpListener::bind(addr)?;
let local_addr = listener.local_addr().unwrap();
Ok(ListeningInfo {
local: { LocalResource { listener, keepalive: config.keepalive } },
local_addr,
})
}
fn accept(&self, mut accept_remote: impl FnMut(AcceptedType<'_, Self::Remote>)) {
loop {
match self.listener.accept() {
Ok((stream, addr)) => accept_remote(AcceptedType::Remote(
addr,
RemoteResource { stream, keepalive: self.keepalive.clone() },
)),
Err(ref err) if err.kind() == ErrorKind::WouldBlock => break,
Err(ref err) if err.kind() == ErrorKind::Interrupted => continue,
Err(err) => break log::error!("TCP accept error: {}", err), }
}
}
}