misc/prost-codec: Introduce codec for varint prefixed Protobuf messages (#2630)

Extracts the Protobuf en-/decoding pattern into its separate crate
and applies it to `libp2p-identify`.
This commit is contained in:
Max Inden
2022-05-05 18:28:47 +02:00
committed by GitHub
parent 3cfbf89a3a
commit bbd2f8f009
14 changed files with 234 additions and 192 deletions

View File

@ -1,88 +0,0 @@
// Copyright 2022 Protocol Labs.
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.
use crate::message_proto;
use bytes::BytesMut;
use prost::Message;
use std::io::Cursor;
use thiserror::Error;
use unsigned_varint::codec::UviBytes;
const MAX_MESSAGE_SIZE_BYTES: usize = 4096;
pub struct Codec(UviBytes);
impl Codec {
pub fn new() -> Self {
let mut codec = UviBytes::default();
codec.set_max_len(MAX_MESSAGE_SIZE_BYTES);
Self(codec)
}
}
impl asynchronous_codec::Encoder for Codec {
type Item = message_proto::HolePunch;
type Error = Error;
fn encode(
&mut self,
item: Self::Item,
dst: &mut asynchronous_codec::BytesMut,
) -> Result<(), Self::Error> {
let mut encoded_msg = BytesMut::new();
item.encode(&mut encoded_msg)
.expect("BytesMut to have sufficient capacity.");
self.0
.encode(encoded_msg.freeze(), dst)
.map_err(|e| e.into())
}
}
impl asynchronous_codec::Decoder for Codec {
type Item = message_proto::HolePunch;
type Error = Error;
fn decode(
&mut self,
src: &mut asynchronous_codec::BytesMut,
) -> Result<Option<Self::Item>, Self::Error> {
Ok(self
.0
.decode(src)?
.map(|msg| message_proto::HolePunch::decode(Cursor::new(msg)))
.transpose()?)
}
}
#[derive(Debug, Error)]
pub enum Error {
#[error("Failed to decode response: {0}.")]
Decode(
#[from]
#[source]
prost::DecodeError,
),
#[error("Io error {0}")]
Io(
#[from]
#[source]
std::io::Error,
),
}

View File

@ -44,17 +44,14 @@ impl upgrade::InboundUpgrade<NegotiatedSubstream> for Upgrade {
type Future = BoxFuture<'static, Result<Self::Output, Self::Error>>;
fn upgrade_inbound(self, substream: NegotiatedSubstream, _: Self::Info) -> Self::Future {
let mut substream = Framed::new(substream, super::codec::Codec::new());
let mut substream = Framed::new(
substream,
prost_codec::Codec::new(super::MAX_MESSAGE_SIZE_BYTES),
);
async move {
let HolePunch { r#type, obs_addrs } =
substream
.next()
.await
.ok_or(super::codec::Error::Io(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"",
)))??;
substream.next().await.ok_or(UpgradeError::StreamClosed)??;
let obs_addrs = if obs_addrs.is_empty() {
return Err(UpgradeError::NoAddresses);
@ -88,7 +85,7 @@ impl upgrade::InboundUpgrade<NegotiatedSubstream> for Upgrade {
}
pub struct PendingConnect {
substream: Framed<NegotiatedSubstream, super::codec::Codec>,
substream: Framed<NegotiatedSubstream, prost_codec::Codec<HolePunch>>,
remote_obs_addrs: Vec<Multiaddr>,
}
@ -103,14 +100,11 @@ impl PendingConnect {
};
self.substream.send(msg).await?;
let HolePunch { r#type, .. } =
self.substream
.next()
.await
.ok_or(super::codec::Error::Io(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"",
)))??;
let HolePunch { r#type, .. } = self
.substream
.next()
.await
.ok_or(UpgradeError::StreamClosed)??;
let r#type = hole_punch::Type::from_i32(r#type).ok_or(UpgradeError::ParseTypeField)?;
match r#type {
@ -124,12 +118,14 @@ impl PendingConnect {
#[derive(Debug, Error)]
pub enum UpgradeError {
#[error("Failed to encode or decode: {0}")]
#[error("Failed to encode or decode")]
Codec(
#[from]
#[source]
super::codec::Error,
prost_codec::Error,
),
#[error("Stream closed")]
StreamClosed,
#[error("Expected at least one address in reservation.")]
NoAddresses,
#[error("Invalid addresses.")]

View File

@ -54,7 +54,10 @@ impl upgrade::OutboundUpgrade<NegotiatedSubstream> for Upgrade {
type Future = BoxFuture<'static, Result<Self::Output, Self::Error>>;
fn upgrade_outbound(self, substream: NegotiatedSubstream, _: Self::Info) -> Self::Future {
let mut substream = Framed::new(substream, super::codec::Codec::new());
let mut substream = Framed::new(
substream,
prost_codec::Codec::new(super::MAX_MESSAGE_SIZE_BYTES),
);
let msg = HolePunch {
r#type: hole_punch::Type::Connect.into(),
@ -67,13 +70,7 @@ impl upgrade::OutboundUpgrade<NegotiatedSubstream> for Upgrade {
let sent_time = Instant::now();
let HolePunch { r#type, obs_addrs } =
substream
.next()
.await
.ok_or(super::codec::Error::Io(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"",
)))??;
substream.next().await.ok_or(UpgradeError::StreamClosed)??;
let rtt = sent_time.elapsed();
@ -123,8 +120,10 @@ pub enum UpgradeError {
Codec(
#[from]
#[source]
super::codec::Error,
prost_codec::Error,
),
#[error("Stream closed")]
StreamClosed,
#[error("Expected 'status' field to be set.")]
MissingStatusField,
#[error("Expected 'reservation' field to be set.")]