use crate::core::upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo, Negotiated};
use bytes::Bytes;
use futures::{future::FromErr, prelude::*};
use std::{iter, io::Error as IoError, sync::Arc};
use tokio_io::{AsyncRead, AsyncWrite};
#[derive(Debug)]
pub struct SimpleProtocol<F> {
info: Bytes,
upgrade: Arc<F>,
}
impl<F> SimpleProtocol<F> {
#[inline]
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> {
#[inline]
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>;
#[inline]
fn protocol_info(&self) -> Self::InfoIter {
iter::once(self.info.clone())
}
}
impl<C, F, O> InboundUpgrade<C> for SimpleProtocol<F>
where
C: AsyncRead + AsyncWrite,
F: Fn(Negotiated<C>) -> O,
O: IntoFuture<Error = IoError>
{
type Output = O::Item;
type Error = IoError;
type Future = FromErr<O::Future, IoError>;
#[inline]
fn upgrade_inbound(self, socket: Negotiated<C>, _: Self::Info) -> Self::Future {
let upgrade = &self.upgrade;
upgrade(socket).into_future().from_err()
}
}
impl<C, F, O> OutboundUpgrade<C> for SimpleProtocol<F>
where
C: AsyncRead + AsyncWrite,
F: Fn(Negotiated<C>) -> O,
O: IntoFuture<Error = IoError>
{
type Output = O::Item;
type Error = IoError;
type Future = FromErr<O::Future, IoError>;
#[inline]
fn upgrade_outbound(self, socket: Negotiated<C>, _: Self::Info) -> Self::Future {
let upgrade = &self.upgrade;
upgrade(socket).into_future().from_err()
}
}