fire_stream/util/
listener.rs

1use super::ByteStream;
2
3use std::io;
4use std::net::{self, SocketAddrV4, SocketAddrV6};
5use std::task::{Poll, Context};
6use std::future::Future;
7use std::pin::Pin;
8
9use tokio::net::{TcpListener, TcpStream};
10
11/// SocketAddr struct equivalent to the one in stdlib but
12/// the unix::SocketAddr was added.
13// todo make this non_complete
14pub enum SocketAddr {
15	V4(SocketAddrV4),
16	V6(SocketAddrV6),
17	#[cfg(unix)]
18	Unix(tokio::net::unix::SocketAddr)
19}
20
21pub trait Listener {
22	type Stream: ByteStream;
23
24	fn poll_accept(
25		&self,
26		cx: &mut Context<'_>
27	) -> Poll<io::Result<(Self::Stream, SocketAddr)>>;
28}
29
30impl Listener for TcpListener {
31	type Stream = TcpStream;
32
33	fn poll_accept(
34		&self,
35		cx: &mut Context<'_>
36	) -> Poll<io::Result<(Self::Stream, SocketAddr)>> {
37		match self.poll_accept(cx) {
38			Poll::Ready(Ok((stream, addr))) => Poll::Ready(Ok((
39				stream,
40				match addr {
41					net::SocketAddr::V4(v4) => SocketAddr::V4(v4),
42					net::SocketAddr::V6(v6) => SocketAddr::V6(v6)
43				}
44			))),
45			Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
46			Poll::Pending => Poll::Pending
47		}
48	}
49}
50
51#[cfg(unix)]
52mod unix {
53
54	use super::*;
55
56	use tokio::net::{UnixListener, UnixStream};
57
58	impl Listener for UnixListener {
59		type Stream = UnixStream;
60
61		fn poll_accept(
62			&self,
63			cx: &mut Context<'_>
64		) -> Poll<io::Result<(Self::Stream, SocketAddr)>> {
65			match self.poll_accept(cx) {
66				Poll::Ready(Ok((stream, addr))) => Poll::Ready(Ok((
67					stream,
68					SocketAddr::Unix(addr)
69				))),
70				Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
71				Poll::Pending => Poll::Pending
72			}
73		}
74	}
75
76}
77
78pub trait ListenerExt: Listener {
79	/// Equivalent to
80	/// `async fn accept(&self) -> io::Result<(Self::Stream, SocketAddr)>`
81	fn accept<'a>(&'a self) -> Accept<'a, Self>
82	where Self: Sized;
83}
84
85impl<L> ListenerExt for L
86where L: Listener {
87	fn accept<'a>(&'a self) -> Accept<'a, Self>
88	where Self: Sized {
89		Accept { inner: self }
90	}
91}
92
93// this should be !unpin
94pub struct Accept<'a, L> {
95	inner: &'a L
96}
97
98impl<'a, L> Future for Accept<'a, L>
99where L: Listener {
100	type Output = io::Result<(L::Stream, SocketAddr)>;
101
102	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
103		self.inner.poll_accept(cx)
104	}
105}