tls_listener/
accept.rs

1use pin_project_lite::pin_project;
2use std::fmt::Debug;
3use std::future::Future;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6
7use tokio::io::{AsyncRead, AsyncWrite};
8
9/// Asynchronously accept connections.
10pub trait AsyncAccept {
11    /// The type of the connection that is accepted.
12    type Connection: AsyncRead + AsyncWrite;
13    /// The type of the remote address, such as [`std::net::SocketAddr`].
14    ///
15    /// If no remote address can be determined (such as for mock connections),
16    /// `()` or a similar dummy type can be used.
17    type Address: Debug;
18    /// The type of error that may be returned.
19    type Error: std::error::Error;
20
21    /// Poll to accept the next connection.
22    ///
23    /// On success return the new connection, and the address of the peer.
24    #[allow(clippy::type_complexity)]
25    fn poll_accept(
26        self: Pin<&mut Self>,
27        cx: &mut Context<'_>,
28    ) -> Poll<Result<(Self::Connection, Self::Address), Self::Error>>;
29}
30
31pin_project! {
32    struct AcceptGenerator<F, A>
33    {
34        accept: A,
35        #[pin]
36        current: F,
37    }
38}
39
40impl<Conn, Addr, E, F, A> AsyncAccept for AcceptGenerator<F, A>
41where
42    A: FnMut() -> F,
43    Conn: AsyncRead + AsyncWrite,
44    E: std::error::Error,
45    Addr: Debug,
46    F: Future<Output = Result<(Conn, Addr), E>>,
47{
48    type Connection = Conn;
49    type Address = Addr;
50    type Error = E;
51
52    fn poll_accept(
53        self: Pin<&mut Self>,
54        cx: &mut Context<'_>,
55    ) -> Poll<Result<(Self::Connection, Self::Address), Self::Error>> {
56        let mut this = self.project();
57
58        let result = this.current.as_mut().poll(cx);
59        if result.is_ready() {
60            // Prime the future for the next poll
61            let next = (this.accept)();
62            this.current.set(next);
63        }
64        result
65    }
66}
67
68/// Create a new `AsyncAccept` from a generator function
69///
70/// This allows you to create an `AsyncAccept` implementation from any
71/// function that can generate futures of connections.
72///
73/// `accept_fn` will be called immediately, to create the initial future, then again
74/// whenever polling the `Future` returned by the previous call returns a ready status.
75///
76/// This function is experimental, and may be changed or removed in future versions
77pub fn accept_generator<Acc, Conn, Addr, F, E>(
78    mut accept_fn: Acc,
79) -> impl AsyncAccept<Connection = Conn, Error = E, Address = Addr> + Send
80where
81    Acc: (FnMut() -> F) + Send,
82    Conn: AsyncRead + AsyncWrite + 'static,
83    Addr: Debug + 'static,
84    F: Future<Output = Result<(Conn, Addr), E>> + Send,
85    E: std::error::Error + 'static,
86{
87    let first_future = (accept_fn)();
88    AcceptGenerator {
89        accept: accept_fn,
90        current: first_future,
91    }
92}
93
94///An AsyncListener that can also report its local address
95pub trait AsyncListener: AsyncAccept {
96    /// The local address of the listener, if available.
97    fn local_addr(&self) -> Result<Self::Address, Self::Error>;
98}