use futures::prelude::*;
use multiaddr::Multiaddr;
use std::io::Error as IoError;
use transport::Transport;
#[derive(Debug, Copy, Clone)]
pub struct MapErr<T, F> {
transport: T,
map: F,
}
impl<T, F> MapErr<T, F> {
#[inline]
pub(crate) fn new(transport: T, map: F) -> MapErr<T, F> {
MapErr { transport, map }
}
}
impl<T, F> Transport for MapErr<T, F>
where
T: Transport,
F: FnOnce(IoError) -> IoError + Clone,
{
type Output = T::Output;
type Listener = MapErrListener<T, F>;
type ListenerUpgrade = MapErrListenerUpgrade<T, F>;
type Dial = MapErrDial<T, F>;
fn listen_on(self, addr: Multiaddr) -> Result<(Self::Listener, Multiaddr), (Self, Multiaddr)> {
let map = self.map;
match self.transport.listen_on(addr) {
Ok((stream, listen_addr)) => {
let stream = MapErrListener { inner: stream, map };
Ok((stream, listen_addr))
}
Err((transport, addr)) => Err((MapErr { transport, map }, addr)),
}
}
fn dial(self, addr: Multiaddr) -> Result<Self::Dial, (Self, Multiaddr)> {
let map = self.map;
match self.transport.dial(addr) {
Ok(future) => Ok(MapErrDial { inner: future, map: Some(map) }),
Err((transport, addr)) => Err((MapErr { transport, map }, addr)),
}
}
#[inline]
fn nat_traversal(&self, server: &Multiaddr, observed: &Multiaddr) -> Option<Multiaddr> {
self.transport.nat_traversal(server, observed)
}
}
pub struct MapErrListener<T, F>
where T: Transport {
inner: T::Listener,
map: F,
}
impl<T, F> Stream for MapErrListener<T, F>
where T: Transport,
F: FnOnce(IoError) -> IoError + Clone,
{
type Item = (MapErrListenerUpgrade<T, F>, Multiaddr);
type Error = IoError;
#[inline]
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
match try_ready!(self.inner.poll()) {
Some((value, addr)) => Ok(Async::Ready(
Some((MapErrListenerUpgrade { inner: value, map: Some(self.map.clone()) }, addr)))),
None => Ok(Async::Ready(None))
}
}
}
pub struct MapErrListenerUpgrade<T, F>
where T: Transport {
inner: T::ListenerUpgrade,
map: Option<F>,
}
impl<T, F> Future for MapErrListenerUpgrade<T, F>
where T: Transport,
F: FnOnce(IoError) -> IoError,
{
type Item = T::Output;
type Error = IoError;
#[inline]
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
match self.inner.poll() {
Ok(Async::Ready(value)) => {
Ok(Async::Ready(value))
},
Ok(Async::NotReady) => Ok(Async::NotReady),
Err(err) => {
let map = self.map.take().expect("poll() called again after error");
Err(map(err))
}
}
}
}
pub struct MapErrDial<T, F>
where T: Transport,
F: FnOnce(IoError) -> IoError,
{
inner: T::Dial,
map: Option<F>,
}
impl<T, F> Future for MapErrDial<T, F>
where T: Transport,
F: FnOnce(IoError) -> IoError,
{
type Item = T::Output;
type Error = IoError;
#[inline]
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
match self.inner.poll() {
Ok(Async::Ready(value)) => {
Ok(Async::Ready(value))
},
Ok(Async::NotReady) => Ok(Async::NotReady),
Err(err) => {
let map = self.map.take().expect("poll() called again after error");
Err(map(err))
}
}
}
}