use crate::{Endpoint, upgrade::{InboundUpgrade, OutboundUpgrade, ProtocolName, UpgradeInfo}};
use futures::prelude::*;
use std::iter;
pub fn from_fn<P, F, C, Fut, Out, Err>(protocol_name: P, fun: F) -> FromFnUpgrade<P, F>
where
P: ProtocolName + Clone,
F: FnOnce(C, Endpoint) -> Fut,
Fut: Future<Output = Result<Out, Err>>,
{
FromFnUpgrade { protocol_name, fun }
}
#[derive(Debug, Clone)]
pub struct FromFnUpgrade<P, F> {
protocol_name: P,
fun: F,
}
impl<P, F> UpgradeInfo for FromFnUpgrade<P, F>
where
P: ProtocolName + Clone,
{
type Info = P;
type InfoIter = iter::Once<P>;
fn protocol_info(&self) -> Self::InfoIter {
iter::once(self.protocol_name.clone())
}
}
impl<C, P, F, Fut, Err, Out> InboundUpgrade<C> for FromFnUpgrade<P, F>
where
P: ProtocolName + Clone,
F: FnOnce(C, Endpoint) -> Fut,
Fut: Future<Output = Result<Out, Err>>,
{
type Output = Out;
type Error = Err;
type Future = Fut;
fn upgrade_inbound(self, sock: C, _: Self::Info) -> Self::Future {
(self.fun)(sock, Endpoint::Listener)
}
}
impl<C, P, F, Fut, Err, Out> OutboundUpgrade<C> for FromFnUpgrade<P, F>
where
P: ProtocolName + Clone,
F: FnOnce(C, Endpoint) -> Fut,
Fut: Future<Output = Result<Out, Err>>,
{
type Output = Out;
type Error = Err;
type Future = Fut;
fn upgrade_outbound(self, sock: C, _: Self::Info) -> Self::Future {
(self.fun)(sock, Endpoint::Dialer)
}
}