use crate::rpc_proto;
use crate::topic::Topic;
use tet_libp2p_core::{InboundUpgrade, OutboundUpgrade, UpgradeInfo, PeerId, upgrade};
use prost::Message;
use std::{error, fmt, io, iter, pin::Pin};
use futures::{Future, io::{AsyncRead, AsyncWrite}};
#[derive(Debug, Clone, Default)]
pub struct FloodsubProtocol {}
impl FloodsubProtocol {
pub fn new() -> FloodsubProtocol {
FloodsubProtocol {}
}
}
impl UpgradeInfo for FloodsubProtocol {
type Info = &'static [u8];
type InfoIter = iter::Once<Self::Info>;
fn protocol_info(&self) -> Self::InfoIter {
iter::once(b"/floodsub/1.0.0")
}
}
impl<TSocket> InboundUpgrade<TSocket> for FloodsubProtocol
where
TSocket: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
type Output = FloodsubRpc;
type Error = FloodsubDecodeError;
type Future = Pin<Box<dyn Future<Output = Result<Self::Output, Self::Error>> + Send>>;
fn upgrade_inbound(self, mut socket: TSocket, _: Self::Info) -> Self::Future {
Box::pin(async move {
let packet = upgrade::read_one(&mut socket, 2048).await?;
let rpc = rpc_proto::Rpc::decode(&packet[..])?;
let mut messages = Vec::with_capacity(rpc.publish.len());
for publish in rpc.publish.into_iter() {
messages.push(FloodsubMessage {
source: PeerId::from_bytes(&publish.from.unwrap_or_default()).map_err(|_| {
FloodsubDecodeError::InvalidPeerId
})?,
data: publish.data.unwrap_or_default(),
sequence_number: publish.seqno.unwrap_or_default(),
topics: publish.topic_ids
.into_iter()
.map(Topic::new)
.collect(),
});
}
Ok(FloodsubRpc {
messages,
subscriptions: rpc.subscriptions
.into_iter()
.map(|sub| FloodsubSubscription {
action: if Some(true) == sub.subscribe {
FloodsubSubscriptionAction::Subscribe
} else {
FloodsubSubscriptionAction::Unsubscribe
},
topic: Topic::new(sub.topic_id.unwrap_or_default()),
})
.collect(),
})
})
}
}
#[derive(Debug)]
pub enum FloodsubDecodeError {
ReadError(upgrade::ReadOneError),
ProtobufError(prost::DecodeError),
InvalidPeerId,
}
impl From<upgrade::ReadOneError> for FloodsubDecodeError {
fn from(err: upgrade::ReadOneError) -> Self {
FloodsubDecodeError::ReadError(err)
}
}
impl From<prost::DecodeError> for FloodsubDecodeError {
fn from(err: prost::DecodeError) -> Self {
FloodsubDecodeError::ProtobufError(err)
}
}
impl fmt::Display for FloodsubDecodeError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match *self {
FloodsubDecodeError::ReadError(ref err) =>
write!(f, "Error while reading from socket: {}", err),
FloodsubDecodeError::ProtobufError(ref err) =>
write!(f, "Error while decoding protobuf: {}", err),
FloodsubDecodeError::InvalidPeerId =>
write!(f, "Error while decoding PeerId from message"),
}
}
}
impl error::Error for FloodsubDecodeError {
fn source(&self) -> Option<&(dyn error::Error + 'static)> {
match *self {
FloodsubDecodeError::ReadError(ref err) => Some(err),
FloodsubDecodeError::ProtobufError(ref err) => Some(err),
FloodsubDecodeError::InvalidPeerId => None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FloodsubRpc {
pub messages: Vec<FloodsubMessage>,
pub subscriptions: Vec<FloodsubSubscription>,
}
impl UpgradeInfo for FloodsubRpc {
type Info = &'static [u8];
type InfoIter = iter::Once<Self::Info>;
fn protocol_info(&self) -> Self::InfoIter {
iter::once(b"/floodsub/1.0.0")
}
}
impl<TSocket> OutboundUpgrade<TSocket> for FloodsubRpc
where
TSocket: AsyncWrite + AsyncRead + Send + Unpin + 'static,
{
type Output = ();
type Error = io::Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Output, Self::Error>> + Send>>;
fn upgrade_outbound(self, mut socket: TSocket, _: Self::Info) -> Self::Future {
Box::pin(async move {
let bytes = self.into_bytes();
upgrade::write_one(&mut socket, bytes).await?;
Ok(())
})
}
}
impl FloodsubRpc {
fn into_bytes(self) -> Vec<u8> {
let rpc = rpc_proto::Rpc {
publish: self.messages.into_iter()
.map(|msg| {
rpc_proto::Message {
from: Some(msg.source.to_bytes()),
data: Some(msg.data),
seqno: Some(msg.sequence_number),
topic_ids: msg.topics
.into_iter()
.map(|topic| topic.into())
.collect()
}
})
.collect(),
subscriptions: self.subscriptions.into_iter()
.map(|topic| {
rpc_proto::rpc::SubOpts {
subscribe: Some(topic.action == FloodsubSubscriptionAction::Subscribe),
topic_id: Some(topic.topic.into())
}
})
.collect()
};
let mut buf = Vec::with_capacity(rpc.encoded_len());
rpc.encode(&mut buf).expect("Vec<u8> provides capacity as needed");
buf
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FloodsubMessage {
pub source: PeerId,
pub data: Vec<u8>,
pub sequence_number: Vec<u8>,
pub topics: Vec<Topic>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FloodsubSubscription {
pub action: FloodsubSubscriptionAction,
pub topic: Topic,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum FloodsubSubscriptionAction {
Subscribe,
Unsubscribe,
}