use std::io::{IoSlice, IoSliceMut, Read as _, Write as _};
use std::net::SocketAddr;
use std::pin::Pin;
use crate::future;
use crate::io::{self, Read, Write};
use crate::net::driver::Watcher;
use crate::net::ToSocketAddrs;
use crate::task::{spawn_blocking, Context, Poll};
#[derive(Debug)]
pub struct TcpStream {
pub(super) watcher: Watcher<mio::net::TcpStream>,
}
impl TcpStream {
pub async fn connect<A: ToSocketAddrs>(addrs: A) -> io::Result<TcpStream> {
let mut last_err = None;
for addr in addrs.to_socket_addrs().await? {
let res = spawn_blocking(move || {
let std_stream = std::net::TcpStream::connect(addr)?;
let mio_stream = mio::net::TcpStream::from_stream(std_stream)?;
Ok(TcpStream {
watcher: Watcher::new(mio_stream),
})
})
.await;
match res {
Ok(stream) => return Ok(stream),
Err(err) => last_err = Some(err),
}
}
Err(last_err.unwrap_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"could not resolve to any addresses",
)
}))
}
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.watcher.get_ref().local_addr()
}
pub fn peer_addr(&self) -> io::Result<SocketAddr> {
self.watcher.get_ref().peer_addr()
}
pub fn ttl(&self) -> io::Result<u32> {
self.watcher.get_ref().ttl()
}
pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
self.watcher.get_ref().set_ttl(ttl)
}
pub async fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
future::poll_fn(|cx| self.watcher.poll_read_with(cx, |inner| inner.peek(buf))).await
}
pub fn nodelay(&self) -> io::Result<bool> {
self.watcher.get_ref().nodelay()
}
pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
self.watcher.get_ref().set_nodelay(nodelay)
}
pub fn shutdown(&self, how: std::net::Shutdown) -> std::io::Result<()> {
self.watcher.get_ref().shutdown(how)
}
}
impl Read for TcpStream {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut &*self).poll_read(cx, buf)
}
fn poll_read_vectored(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &mut [IoSliceMut<'_>],
) -> Poll<io::Result<usize>> {
Pin::new(&mut &*self).poll_read_vectored(cx, bufs)
}
}
impl Read for &TcpStream {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
self.watcher.poll_read_with(cx, |mut inner| inner.read(buf))
}
}
impl Write for TcpStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut &*self).poll_write(cx, buf)
}
fn poll_write_vectored(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &[IoSlice<'_>],
) -> Poll<io::Result<usize>> {
Pin::new(&mut &*self).poll_write_vectored(cx, bufs)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut &*self).poll_flush(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut &*self).poll_close(cx)
}
}
impl Write for &TcpStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.watcher
.poll_write_with(cx, |mut inner| inner.write(buf))
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.watcher.poll_write_with(cx, |mut inner| inner.flush())
}
fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
}
impl From<std::net::TcpStream> for TcpStream {
fn from(stream: std::net::TcpStream) -> TcpStream {
let mio_stream = mio::net::TcpStream::from_stream(stream).unwrap();
TcpStream {
watcher: Watcher::new(mio_stream),
}
}
}
cfg_unix! {
use crate::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
impl AsRawFd for TcpStream {
fn as_raw_fd(&self) -> RawFd {
self.watcher.get_ref().as_raw_fd()
}
}
impl FromRawFd for TcpStream {
unsafe fn from_raw_fd(fd: RawFd) -> TcpStream {
std::net::TcpStream::from_raw_fd(fd).into()
}
}
impl IntoRawFd for TcpStream {
fn into_raw_fd(self) -> RawFd {
self.watcher.into_inner().into_raw_fd()
}
}
}
cfg_windows! {
}