use crate::core::upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo};
use bytes::Bytes;
use futures::prelude::*;
use std::{iter, sync::Arc};
#[derive(Debug)]
pub struct SimpleProtocol<F> {
info: Bytes,
upgrade: Arc<F>,
}
impl<F> SimpleProtocol<F> {
pub fn new<N>(info: N, upgrade: F) -> SimpleProtocol<F>
where
N: Into<Bytes>,
{
SimpleProtocol {
info: info.into(),
upgrade: Arc::new(upgrade),
}
}
}
impl<F> Clone for SimpleProtocol<F> {
fn clone(&self) -> Self {
SimpleProtocol {
info: self.info.clone(),
upgrade: self.upgrade.clone(),
}
}
}
impl<F> UpgradeInfo for SimpleProtocol<F> {
type Info = Bytes;
type InfoIter = iter::Once<Self::Info>;
fn protocol_info(&self) -> Self::InfoIter {
iter::once(self.info.clone())
}
}
impl<C, F, O, A, E> InboundUpgrade<C> for SimpleProtocol<F>
where
C: AsyncRead + AsyncWrite,
F: Fn(C) -> O,
O: Future<Output = Result<A, E>>,
{
type Output = A;
type Error = E;
type Future = O;
fn upgrade_inbound(self, socket: C, _: Self::Info) -> Self::Future {
let upgrade = &self.upgrade;
upgrade(socket)
}
}
impl<C, F, O, A, E> OutboundUpgrade<C> for SimpleProtocol<F>
where
C: AsyncRead + AsyncWrite,
F: Fn(C) -> O,
O: Future<Output = Result<A, E>>,
{
type Output = A;
type Error = E;
type Future = O;
fn upgrade_outbound(self, socket: C, _: Self::Info) -> Self::Future {
let upgrade = &self.upgrade;
upgrade(socket)
}
}