wrpc_transport/frame/conn/
accept.rs

1use core::future::Future;
2use core::ops::{Deref, DerefMut};
3
4use tokio::io::{AsyncRead, AsyncWrite};
5
6/// Accepts connections on a transport
7pub trait Accept {
8    /// Transport-specific invocation context
9    type Context: Send + Sync + 'static;
10
11    /// Outgoing byte stream
12    type Outgoing: AsyncWrite + Send + Sync + Unpin + 'static;
13
14    /// Incoming byte stream
15    type Incoming: AsyncRead + Send + Sync + Unpin + 'static;
16
17    /// Accept a connection returning a pair of streams and connection context
18    fn accept(
19        &self,
20    ) -> impl Future<Output = std::io::Result<(Self::Context, Self::Outgoing, Self::Incoming)>>;
21}
22
23/// Wrapper returned by [`AcceptExt::map_context`]
24pub struct AcceptMapContext<T, F> {
25    inner: T,
26    f: F,
27}
28
29impl<T, F> Deref for AcceptMapContext<T, F> {
30    type Target = T;
31
32    fn deref(&self) -> &Self::Target {
33        &self.inner
34    }
35}
36
37impl<T, F> DerefMut for AcceptMapContext<T, F> {
38    fn deref_mut(&mut self) -> &mut Self::Target {
39        &mut self.inner
40    }
41}
42
43/// Extension trait for [Accept]
44pub trait AcceptExt: Accept + Sized {
45    /// Maps [`Self::Context`](Accept::Context) to a type `T` using `F`
46    fn map_context<T, F: Fn(Self::Context) -> T>(self, f: F) -> AcceptMapContext<Self, F> {
47        AcceptMapContext { inner: self, f }
48    }
49}
50
51impl<T: Accept> AcceptExt for T {}
52
53impl<T, U, F> Accept for AcceptMapContext<T, F>
54where
55    T: Accept,
56    U: Send + Sync + 'static,
57    F: Fn(T::Context) -> U,
58{
59    type Context = U;
60    type Outgoing = T::Outgoing;
61    type Incoming = T::Incoming;
62
63    async fn accept(&self) -> std::io::Result<(Self::Context, Self::Outgoing, Self::Incoming)> {
64        (&self).accept().await
65    }
66}
67
68impl<T, U, F> Accept for &AcceptMapContext<T, F>
69where
70    T: Accept,
71    U: Send + Sync + 'static,
72    F: Fn(T::Context) -> U,
73{
74    type Context = U;
75    type Outgoing = T::Outgoing;
76    type Incoming = T::Incoming;
77
78    async fn accept(&self) -> std::io::Result<(Self::Context, Self::Outgoing, Self::Incoming)> {
79        let (cx, tx, rx) = self.inner.accept().await?;
80        Ok(((self.f)(cx), tx, rx))
81    }
82}