84 lines
2.1 KiB
Rust
84 lines
2.1 KiB
Rust
use std::net::SocketAddr;
|
|
use std::pin::Pin;
|
|
use std::task::{self, Poll};
|
|
|
|
use async_compat::Compat as AsyncCompat;
|
|
use futures::io::{AsyncRead, AsyncWrite};
|
|
use tokio::net::{TcpListener, TcpStream, ToSocketAddrs};
|
|
|
|
use super::Accept;
|
|
|
|
#[pin_project::pin_project]
|
|
pub struct AddrIncoming {
|
|
pub local_addr: SocketAddr,
|
|
#[pin]
|
|
listener: TcpListener,
|
|
}
|
|
|
|
impl AddrIncoming {
|
|
pub async fn new(addr: impl ToSocketAddrs) -> std::io::Result<Self> {
|
|
let listener = TcpListener::bind(addr).await?;
|
|
let local_addr = listener.local_addr()?;
|
|
|
|
Ok(Self {
|
|
local_addr,
|
|
listener,
|
|
})
|
|
}
|
|
}
|
|
|
|
impl Accept for AddrIncoming {
|
|
type Conn = AddrStream;
|
|
type Error = std::io::Error;
|
|
|
|
fn poll_accept(
|
|
self: Pin<&mut Self>,
|
|
cx: &mut task::Context,
|
|
) -> task::Poll<Option<Result<Self::Conn, Self::Error>>> {
|
|
let this = self.project();
|
|
|
|
let (stream, remote_addr) = futures::ready!(this.listener.poll_accept(cx))?;
|
|
Poll::Ready(Some(Ok(AddrStream {
|
|
local_addr: *this.local_addr,
|
|
remote_addr,
|
|
stream: AsyncCompat::new(stream),
|
|
})))
|
|
}
|
|
}
|
|
|
|
#[pin_project::pin_project]
|
|
pub struct AddrStream {
|
|
pub local_addr: SocketAddr,
|
|
pub remote_addr: SocketAddr,
|
|
#[pin]
|
|
stream: AsyncCompat<TcpStream>,
|
|
}
|
|
|
|
impl AsyncRead for AddrStream {
|
|
fn poll_read(
|
|
self: Pin<&mut Self>,
|
|
cx: &mut task::Context<'_>,
|
|
buf: &mut [u8],
|
|
) -> Poll<std::io::Result<usize>> {
|
|
self.project().stream.poll_read(cx, buf)
|
|
}
|
|
}
|
|
|
|
impl AsyncWrite for AddrStream {
|
|
fn poll_write(
|
|
self: Pin<&mut Self>,
|
|
cx: &mut task::Context<'_>,
|
|
buf: &[u8],
|
|
) -> Poll<std::io::Result<usize>> {
|
|
self.project().stream.poll_write(cx, buf)
|
|
}
|
|
|
|
fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<std::io::Result<()>> {
|
|
self.project().stream.poll_flush(cx)
|
|
}
|
|
|
|
fn poll_close(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<std::io::Result<()>> {
|
|
self.project().stream.poll_close(cx)
|
|
}
|
|
}
|