protocols/gossipsub: Make use of prost-codec (#3070)

This commit is contained in:
Nick Loadholtes
2022-11-04 06:12:08 -04:00
committed by GitHub
parent b528d336cd
commit 1ba9e4579d
4 changed files with 31 additions and 50 deletions

View File

@ -29,7 +29,6 @@ use crate::types::{
};
use asynchronous_codec::{Decoder, Encoder, Framed};
use byteorder::{BigEndian, ByteOrder};
use bytes::Bytes;
use bytes::BytesMut;
use futures::future;
use futures::prelude::*;
@ -184,17 +183,18 @@ where
/* Gossip codec for the framing */
pub struct GossipsubCodec {
/// Codec to encode/decode the Unsigned varint length prefix of the frames.
length_codec: codec::UviBytes,
/// Determines the level of validation performed on incoming messages.
validation_mode: ValidationMode,
/// The codec to handle common encoding/decoding of protobuf messages
codec: prost_codec::Codec<rpc_proto::Rpc>,
}
impl GossipsubCodec {
pub fn new(length_codec: codec::UviBytes, validation_mode: ValidationMode) -> Self {
pub fn new(length_codec: codec::UviBytes, validation_mode: ValidationMode) -> GossipsubCodec {
let codec = prost_codec::Codec::new(length_codec.max_len());
GossipsubCodec {
length_codec,
validation_mode,
codec,
}
}
@ -267,16 +267,12 @@ impl Encoder for GossipsubCodec {
type Item = rpc_proto::Rpc;
type Error = GossipsubHandlerError;
fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
let mut buf = Vec::with_capacity(item.encoded_len());
item.encode(&mut buf)
.expect("Buffer has sufficient capacity");
// length prefix the protobuf message, ensuring the max limit is not hit
self.length_codec
.encode(Bytes::from(buf), dst)
.map_err(|_| GossipsubHandlerError::MaxTransmissionSize)
fn encode(
&mut self,
item: Self::Item,
dst: &mut BytesMut,
) -> Result<(), GossipsubHandlerError> {
Ok(self.codec.encode(item, dst)?)
}
}
@ -284,20 +280,12 @@ impl Decoder for GossipsubCodec {
type Item = HandlerEvent;
type Error = GossipsubHandlerError;
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
let packet = match self.length_codec.decode(src).map_err(|e| {
if let std::io::ErrorKind::PermissionDenied = e.kind() {
GossipsubHandlerError::MaxTransmissionSize
} else {
GossipsubHandlerError::Io(e)
}
})? {
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, GossipsubHandlerError> {
let rpc = match self.codec.decode(src)? {
Some(p) => p,
None => return Ok(None),
};
let rpc = rpc_proto::Rpc::decode(&packet[..]).map_err(std::io::Error::from)?;
// Store valid messages.
let mut messages = Vec::with_capacity(rpc.publish.len());
// Store any invalid messages.