use crate::krse::future::poll_fn;
use crate::krse::io::{AsyncRead, AsyncWrite, PollEvented};
use crate::krse::net::tcp::split::{split, ReadHalf, WriteHalf};
use crate::krse::net::ToSocketAddrs;
use bytes::Buf;
use iovec::IoVec;
use std::convert::TryFrom;
use std::fmt;
use std::io::{self, Read, Write};
use std::mem::MaybeUninit;
use std::net::{self, Shutdown, SocketAddr};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
macro_rules! ready {
($e:expr $(,)?) => {
match $e {
std::task::Poll::Ready(t) => t,
std::task::Poll::Pending => return std::task::Poll::Pending,
}
};
}
pub struct TcpStream {
io: PollEvented<mio::net::TcpStream>,
}
impl TcpStream {
pub async fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> {
let addrs = addr.to_socket_addrs().await?;
let mut last_err = None;
for addr in addrs {
match TcpStream::connect_addr(addr).await {
Ok(stream) => return Ok(stream),
Err(e) => last_err = Some(e),
}
}
Err(last_err.unwrap_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"could not resolve to any addresses",
)
}))
}
async fn connect_addr(addr: SocketAddr) -> io::Result<TcpStream> {
let sys = mio::net::TcpStream::connect(&addr)?;
let stream = TcpStream::new(sys)?;
poll_fn(|cx| stream.io.poll_write_ready(cx)).await?;
if let Some(e) = stream.io.get_ref().take_error()? {
return Err(e);
}
Ok(stream)
}
pub(crate) fn new(connected: mio::net::TcpStream) -> io::Result<TcpStream> {
let io = PollEvented::new(connected)?;
Ok(TcpStream { io })
}
pub fn from_std(stream: net::TcpStream) -> io::Result<TcpStream> {
let io = mio::net::TcpStream::from_stream(stream)?;
let io = PollEvented::new(io)?;
Ok(TcpStream { io })
}
#[doc(hidden)]
pub async fn connect_std(stream: net::TcpStream, addr: &SocketAddr) -> io::Result<TcpStream> {
let io = mio::net::TcpStream::connect_stream(stream, addr)?;
let io = PollEvented::new(io)?;
let stream = TcpStream { io };
poll_fn(|cx| stream.io.poll_write_ready(cx)).await?;
if let Some(e) = stream.io.get_ref().take_error()? {
return Err(e);
}
Ok(stream)
}
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().local_addr()
}
pub fn peer_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().peer_addr()
}
pub fn poll_peek(&mut self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().peek(buf) {
Ok(ret) => Poll::Ready(Ok(ret)),
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
Err(e) => Poll::Ready(Err(e)),
}
}
pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result<usize> {
poll_fn(|cx| self.poll_peek(cx, buf)).await
}
pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
self.io.get_ref().shutdown(how)
}
pub fn nodelay(&self) -> io::Result<bool> {
self.io.get_ref().nodelay()
}
pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
self.io.get_ref().set_nodelay(nodelay)
}
pub fn recv_buffer_size(&self) -> io::Result<usize> {
self.io.get_ref().recv_buffer_size()
}
pub fn set_recv_buffer_size(&self, size: usize) -> io::Result<()> {
self.io.get_ref().set_recv_buffer_size(size)
}
pub fn send_buffer_size(&self) -> io::Result<usize> {
self.io.get_ref().send_buffer_size()
}
pub fn set_send_buffer_size(&self, size: usize) -> io::Result<()> {
self.io.get_ref().set_send_buffer_size(size)
}
pub fn keepalive(&self) -> io::Result<Option<Duration>> {
self.io.get_ref().keepalive()
}
pub fn set_keepalive(&self, keepalive: Option<Duration>) -> io::Result<()> {
self.io.get_ref().set_keepalive(keepalive)
}
pub fn ttl(&self) -> io::Result<u32> {
self.io.get_ref().ttl()
}
pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
self.io.get_ref().set_ttl(ttl)
}
pub fn linger(&self) -> io::Result<Option<Duration>> {
self.io.get_ref().linger()
}
pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> {
self.io.get_ref().set_linger(dur)
}
pub fn split(&mut self) -> (ReadHalf<'_>, WriteHalf<'_>) {
split(self)
}
pub(crate) fn poll_read_priv(
&self,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().read(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
pub(super) fn poll_write_priv(
&self,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
match self.io.get_ref().write(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready(cx)?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
pub(super) fn poll_write_buf_priv<B: Buf>(
&self,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
use std::io::IoSlice;
ready!(self.io.poll_write_ready(cx))?;
static S: &[u8] = &[0];
const MAX_BUFS: usize = 64;
let mut slices: [IoSlice<'_>; MAX_BUFS] = [
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
IoSlice::new(S), IoSlice::new(S), IoSlice::new(S), IoSlice::new(S),
];
let cnt = buf.bytes_vectored(&mut slices);
let iovec = <&IoVec>::from(S);
let mut vecs = [iovec; MAX_BUFS];
for i in 0..cnt {
vecs[i] = (*slices[i]).into();
}
match self.io.get_ref().write_bufs(&vecs[..cnt]) {
Ok(n) => {
buf.advance(n);
Poll::Ready(Ok(n))
},
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready(cx)?;
Poll::Pending
},
Err(e) => Poll::Ready(Err(e)),
}
}
}
impl TryFrom<TcpStream> for mio::net::TcpStream {
type Error = io::Error;
fn try_from(value: TcpStream) -> Result<Self, Self::Error> {
value.io.into_inner()
}
}
impl TryFrom<net::TcpStream> for TcpStream {
type Error = io::Error;
fn try_from(stream: net::TcpStream) -> Result<Self, Self::Error> {
Self::from_std(stream)
}
}
impl AsyncRead for TcpStream {
unsafe fn prepare_uninitialized_buffer(&self, _: &mut [MaybeUninit<u8>]) -> bool {
false
}
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
self.poll_read_priv(cx, buf)
}
}
impl AsyncWrite for TcpStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.poll_write_priv(cx, buf)
}
fn poll_write_buf<B: Buf>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
self.poll_write_buf_priv(cx, buf)
}
#[inline]
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
self.shutdown(std::net::Shutdown::Write)?;
Poll::Ready(Ok(()))
}
}
impl fmt::Debug for TcpStream {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.io.get_ref().fmt(f)
}
}
#[cfg(unix)]
mod sys {
use super::TcpStream;
use std::os::unix::prelude::*;
impl AsRawFd for TcpStream {
fn as_raw_fd(&self) -> RawFd {
self.io.get_ref().as_raw_fd()
}
}
}
#[cfg(windows)]
mod sys {
}