use crate::RequestId;
use crate::codec::RequestResponseCodec;
use futures::{channel::oneshot, future::BoxFuture, prelude::*};
use tetsy_libp2p_core::upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo};
use tetsy_libp2p_swarm::NegotiatedSubstream;
use smallvec::SmallVec;
use std::io;
#[derive(Debug, Clone)]
pub enum ProtocolSupport {
Inbound,
Outbound,
Full
}
impl ProtocolSupport {
pub fn inbound(&self) -> bool {
match self {
ProtocolSupport::Inbound | ProtocolSupport::Full => true,
ProtocolSupport::Outbound => false,
}
}
pub fn outbound(&self) -> bool {
match self {
ProtocolSupport::Outbound | ProtocolSupport::Full => true,
ProtocolSupport::Inbound => false,
}
}
}
#[derive(Debug)]
pub struct ResponseProtocol<TCodec>
where
TCodec: RequestResponseCodec
{
pub(crate) codec: TCodec,
pub(crate) protocols: SmallVec<[TCodec::Protocol; 2]>,
pub(crate) request_sender: oneshot::Sender<(RequestId, TCodec::Request)>,
pub(crate) response_receiver: oneshot::Receiver<TCodec::Response>,
pub(crate) request_id: RequestId
}
impl<TCodec> UpgradeInfo for ResponseProtocol<TCodec>
where
TCodec: RequestResponseCodec
{
type Info = TCodec::Protocol;
type InfoIter = smallvec::IntoIter<[Self::Info; 2]>;
fn protocol_info(&self) -> Self::InfoIter {
self.protocols.clone().into_iter()
}
}
impl<TCodec> InboundUpgrade<NegotiatedSubstream> for ResponseProtocol<TCodec>
where
TCodec: RequestResponseCodec + Send + 'static,
{
type Output = bool;
type Error = io::Error;
type Future = BoxFuture<'static, Result<Self::Output, Self::Error>>;
fn upgrade_inbound(mut self, mut io: NegotiatedSubstream, protocol: Self::Info) -> Self::Future {
async move {
let read = self.codec.read_request(&protocol, &mut io);
let request = read.await?;
if let Ok(()) = self.request_sender.send((self.request_id, request)) {
if let Ok(response) = self.response_receiver.await {
let write = self.codec.write_response(&protocol, &mut io, response);
write.await?;
} else {
io.close().await?;
return Ok(false)
}
}
io.close().await?;
Ok(true)
}.boxed()
}
}
#[derive(Debug)]
pub struct RequestProtocol<TCodec>
where
TCodec: RequestResponseCodec
{
pub(crate) codec: TCodec,
pub(crate) protocols: SmallVec<[TCodec::Protocol; 2]>,
pub(crate) request_id: RequestId,
pub(crate) request: TCodec::Request,
}
impl<TCodec> UpgradeInfo for RequestProtocol<TCodec>
where
TCodec: RequestResponseCodec
{
type Info = TCodec::Protocol;
type InfoIter = smallvec::IntoIter<[Self::Info; 2]>;
fn protocol_info(&self) -> Self::InfoIter {
self.protocols.clone().into_iter()
}
}
impl<TCodec> OutboundUpgrade<NegotiatedSubstream> for RequestProtocol<TCodec>
where
TCodec: RequestResponseCodec + Send + 'static,
{
type Output = TCodec::Response;
type Error = io::Error;
type Future = BoxFuture<'static, Result<Self::Output, Self::Error>>;
fn upgrade_outbound(mut self, mut io: NegotiatedSubstream, protocol: Self::Info) -> Self::Future {
async move {
let write = self.codec.write_request(&protocol, &mut io, self.request);
write.await?;
io.close().await?;
let read = self.codec.read_response(&protocol, &mut io);
let response = read.await?;
Ok(response)
}.boxed()
}
}