Skip to main content

volans_core/transport/
map_err.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use crate::{Listener, ListenerEvent, Multiaddr, Transport, TransportError};
7
8#[derive(Debug, Copy, Clone)]
9pub struct MapErr<T, F> {
10    transport: T,
11    map: F,
12}
13
14impl<T, F> MapErr<T, F> {
15    pub(crate) fn new(transport: T, map: F) -> Self {
16        Self { transport, map }
17    }
18}
19
20impl<T, F, TErr> Transport for MapErr<T, F>
21where
22    T: Transport,
23    F: FnOnce(T::Error) -> TErr + Clone,
24    TErr: std::error::Error,
25{
26    type Output = T::Output;
27    type Error = TErr;
28    type Dial = MapErrDial<T, F>;
29    type Incoming = MapErrUpgrade<T, F>;
30    type Listener = MapErrListener<T, F>;
31
32    fn dial(&self, addr: Multiaddr) -> Result<Self::Dial, TransportError<Self::Error>> {
33        let map = self.map.clone();
34
35        match self.transport.dial(addr) {
36            Ok(dial) => Ok(MapErrDial {
37                inner: dial,
38                map: Some(map),
39            }),
40            Err(err) => Err(err.map(map)),
41        }
42    }
43
44    fn listen(&self, addr: Multiaddr) -> Result<Self::Listener, TransportError<Self::Error>> {
45        let map = self.map.clone();
46        match self.transport.listen(addr) {
47            Ok(listener) => Ok(MapErrListener {
48                inner: listener,
49                map,
50            }),
51            Err(err) => Err(err.map(map)),
52        }
53    }
54}
55
56#[pin_project::pin_project]
57pub struct MapErrDial<T: Transport, F> {
58    #[pin]
59    inner: T::Dial,
60    map: Option<F>,
61}
62
63impl<T, F, TErr> Future for MapErrDial<T, F>
64where
65    T: Transport,
66    F: FnOnce(T::Error) -> TErr + Clone,
67    TErr: std::error::Error,
68{
69    type Output = Result<T::Output, TErr>;
70
71    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
72        let this = self.project();
73        match this.inner.poll(cx) {
74            Poll::Ready(Ok(output)) => Poll::Ready(Ok(output)),
75            Poll::Ready(Err(err)) => {
76                let map = this.map.take().expect("MapErrDial can only be polled once");
77                Poll::Ready(Err(map(err)))
78            }
79            Poll::Pending => Poll::Pending,
80        }
81    }
82}
83
84#[pin_project::pin_project]
85pub struct MapErrUpgrade<T: Transport, F> {
86    #[pin]
87    inner: T::Incoming,
88    map: Option<F>,
89}
90
91impl<T, F, TErr> Future for MapErrUpgrade<T, F>
92where
93    T: Transport,
94    F: FnOnce(T::Error) -> TErr + Clone,
95    TErr: std::error::Error,
96{
97    type Output = Result<T::Output, TErr>;
98
99    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
100        let this = self.project();
101        match this.inner.poll(cx) {
102            Poll::Ready(Ok(output)) => Poll::Ready(Ok(output)),
103            Poll::Ready(Err(err)) => {
104                let map = this
105                    .map
106                    .take()
107                    .expect("MapErrUpgrade can only be polled once");
108                Poll::Ready(Err(map(err)))
109            }
110            Poll::Pending => Poll::Pending,
111        }
112    }
113}
114
115#[pin_project::pin_project]
116pub struct MapErrListener<T: Transport, F> {
117    #[pin]
118    inner: T::Listener,
119    map: F,
120}
121
122impl<T, F, TErr> Listener for MapErrListener<T, F>
123where
124    T: Transport,
125    F: FnOnce(T::Error) -> TErr + Clone,
126    TErr: std::error::Error,
127{
128    type Output = T::Output;
129    type Error = TErr;
130    type Upgrade = MapErrUpgrade<T, F>;
131
132    fn poll_event(
133        self: Pin<&mut Self>,
134        cx: &mut Context<'_>,
135    ) -> Poll<ListenerEvent<Self::Upgrade, TErr>> {
136        let this = self.project();
137        let map = &*this.map;
138        this.inner.poll_event(cx).map(|ev| {
139            ev.map_upgrade(move |u| MapErrUpgrade {
140                inner: u,
141                map: Some(map.clone()),
142            })
143            .map_err(map.clone())
144        })
145    }
146
147    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), TErr>> {
148        let this = self.project();
149        match this.inner.poll_close(cx) {
150            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
151            Poll::Ready(Err(err)) => {
152                let map = this.map.clone();
153                Poll::Ready(Err(map(err)))
154            }
155            Poll::Pending => Poll::Pending,
156        }
157    }
158}