fire_stream/util/
listener.rs1use 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
11pub 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 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
93pub 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}