protocols/dcutr: Implement Encoder and Decoder for DCUtR prost messages

mxinden · mxinden · commit 354900fd1eaf · 2022-02-02T17:53:13.000+01:00
diff --git a/protocols/dcutr/src/protocol.rs b/protocols/dcutr/src/protocol.rs
@@ -18,9 +18,8 @@
 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
+mod codec;
 pub mod inbound;
 pub mod outbound;
 
 const PROTOCOL_NAME: &[u8; 13] = b"/libp2p/dcutr";
-
-const MAX_MESSAGE_SIZE_BYTES: usize = 4096;
diff --git a/protocols/dcutr/src/protocol/codec.rs b/protocols/dcutr/src/protocol/codec.rs
@@ -0,0 +1,88 @@
+// 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,
+    ),
+}
diff --git a/protocols/dcutr/src/protocol/inbound.rs b/protocols/dcutr/src/protocol/inbound.rs
@@ -20,16 +20,12 @@
 
 use crate::message_proto::{hole_punch, HolePunch};
 use asynchronous_codec::Framed;
-use bytes::BytesMut;
 use futures::{future::BoxFuture, prelude::*};
 use libp2p_core::{multiaddr::Protocol, upgrade, Multiaddr};
 use libp2p_swarm::NegotiatedSubstream;
-use prost::Message;
 use std::convert::TryFrom;
-use std::io::Cursor;
 use std::iter;
 use thiserror::Error;
-use unsigned_varint::codec::UviBytes;
 
 pub struct Upgrade {}
 
@@ -48,17 +44,17 @@ 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 codec = UviBytes::default();
-        codec.set_max_len(super::MAX_MESSAGE_SIZE_BYTES);
-        let mut substream = Framed::new(substream, codec);
+        let mut substream = Framed::new(substream, super::codec::Codec::new());
 
         async move {
-            let msg: bytes::BytesMut = substream
-                .next()
-                .await
-                .ok_or(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, ""))??;
-
-            let HolePunch { r#type, obs_addrs } = HolePunch::decode(Cursor::new(msg))?;
+            let HolePunch { r#type, obs_addrs } =
+                substream
+                    .next()
+                    .await
+                    .ok_or(super::codec::Error::Io(std::io::Error::new(
+                        std::io::ErrorKind::UnexpectedEof,
+                        "",
+                    )))??;
 
             let obs_addrs = if obs_addrs.is_empty() {
                 return Err(UpgradeError::NoAddresses);
@@ -92,7 +88,7 @@ impl upgrade::InboundUpgrade<NegotiatedSubstream> for Upgrade {
 }
 
 pub struct PendingConnect {
-    substream: Framed<NegotiatedSubstream, UviBytes>,
+    substream: Framed<NegotiatedSubstream, super::codec::Codec>,
     remote_obs_addrs: Vec<Multiaddr>,
 }
 
@@ -106,18 +102,15 @@ impl PendingConnect {
             obs_addrs: local_obs_addrs.into_iter().map(|a| a.to_vec()).collect(),
         };
 
-        let mut encoded_msg = BytesMut::new();
-        msg.encode(&mut encoded_msg)
-            .expect("BytesMut to have sufficient capacity.");
-
-        self.substream.send(encoded_msg.freeze()).await?;
-        let msg: bytes::BytesMut = self
-            .substream
-            .next()
-            .await
-            .ok_or(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, ""))??;
-
-        let HolePunch { r#type, .. } = HolePunch::decode(Cursor::new(msg))?;
+        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 r#type = hole_punch::Type::from_i32(r#type).ok_or(UpgradeError::ParseTypeField)?;
         match r#type {
@@ -131,17 +124,11 @@ impl PendingConnect {
 
 #[derive(Debug, Error)]
 pub enum UpgradeError {
-    #[error("Failed to decode response: {0}.")]
-    Decode(
-        #[from]
-        #[source]
-        prost::DecodeError,
-    ),
-    #[error("Io error {0}")]
-    Io(
+    #[error("Failed to encode or decode: {0}")]
+    Codec(
         #[from]
         #[source]
-        std::io::Error,
+        super::codec::Error,
     ),
     #[error("Expected at least one address in reservation.")]
     NoAddresses,
diff --git a/protocols/dcutr/src/protocol/outbound.rs b/protocols/dcutr/src/protocol/outbound.rs
@@ -20,18 +20,14 @@
 
 use crate::message_proto::{hole_punch, HolePunch};
 use asynchronous_codec::Framed;
-use bytes::BytesMut;
 use futures::{future::BoxFuture, prelude::*};
 use futures_timer::Delay;
 use libp2p_core::{multiaddr::Protocol, upgrade, Multiaddr};
 use libp2p_swarm::NegotiatedSubstream;
-use prost::Message;
 use std::convert::TryFrom;
-use std::io::Cursor;
 use std::iter;
 use std::time::Instant;
 use thiserror::Error;
-use unsigned_varint::codec::UviBytes;
 
 pub struct Upgrade {
     obs_addrs: Vec<Multiaddr>,
@@ -58,33 +54,29 @@ 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 msg = HolePunch {
             r#type: hole_punch::Type::Connect.into(),
             obs_addrs: self.obs_addrs.into_iter().map(|a| a.to_vec()).collect(),
         };
 
-        let mut encoded_msg = BytesMut::new();
-        msg.encode(&mut encoded_msg)
-            .expect("BytesMut to have sufficient capacity.");
-
-        let mut codec = UviBytes::default();
-        codec.set_max_len(super::MAX_MESSAGE_SIZE_BYTES);
-        let mut substream = Framed::new(substream, codec);
-
         async move {
-            substream.send(encoded_msg.freeze()).await?;
+            substream.send(msg).await?;
 
             let sent_time = Instant::now();
 
-            let msg: bytes::BytesMut = substream
-                .next()
-                .await
-                .ok_or(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, ""))??;
+            let HolePunch { r#type, obs_addrs } =
+                substream
+                    .next()
+                    .await
+                    .ok_or(super::codec::Error::Io(std::io::Error::new(
+                        std::io::ErrorKind::UnexpectedEof,
+                        "",
+                    )))??;
 
             let rtt = sent_time.elapsed();
 
-            let HolePunch { r#type, obs_addrs } = HolePunch::decode(Cursor::new(msg))?;
-
             let r#type = hole_punch::Type::from_i32(r#type).ok_or(UpgradeError::ParseTypeField)?;
             match r#type {
                 hole_punch::Type::Connect => {}
@@ -111,11 +103,7 @@ impl upgrade::OutboundUpgrade<NegotiatedSubstream> for Upgrade {
                 obs_addrs: vec![],
             };
 
-            let mut encoded_msg = BytesMut::new();
-            msg.encode(&mut encoded_msg)
-                .expect("BytesMut to have sufficient capacity.");
-
-            substream.send(encoded_msg.freeze()).await?;
+            substream.send(msg).await?;
 
             Delay::new(rtt / 2).await;
 
@@ -131,17 +119,11 @@ pub struct Connect {
 
 #[derive(Debug, Error)]
 pub enum UpgradeError {
-    #[error("Failed to decode response: {0}.")]
-    Decode(
-        #[from]
-        #[source]
-        prost::DecodeError,
-    ),
-    #[error("Io error {0}")]
-    Io(
+    #[error("Failed to encode or decode: {0}")]
+    Codec(
         #[from]
         #[source]
-        std::io::Error,
+        super::codec::Error,
     ),
     #[error("Expected 'status' field to be set.")]
     MissingStatusField,
