net: remove intermediate Packet type

Instead we read Message directly to and from the stream. We also introduce some
explicit bounds checking when sending and receiving messages.

net: cleanup

src/net/channel.rs

@@ -25,11 +25,13 @@ use std::{
     time::UNIX_EPOCH,
 };
 
-use darkfi_serial::{async_trait, serialize, SerialDecodable, SerialEncodable};
-use log::{debug, error, info};
+use darkfi_serial::{
+    async_trait, AsyncDecodable, AsyncEncodable, SerialDecodable, SerialEncodable, VarInt,
+};
+use log::{debug, error, info, trace};
 use rand::{rngs::OsRng, Rng};
 use smol::{
-    io::{self, ReadHalf, WriteHalf},
+    io::{self, AsyncRead, AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf},
     lock::Mutex,
     Executor,
 };
@@ -39,7 +41,7 @@ use super::{
     dnet::{self, dnetev, DnetEvent},
     hosts::HostColor,
     message,
-    message::{Packet, VersionMessage},
+    message::{VersionMessage, MAGIC_BYTES},
     message_subscriber::{MessageSubscription, MessageSubsystem},
     p2p::P2pPtr,
     session::{Session, SessionBitFlag, SessionWeakPtr, SESSION_ALL, SESSION_REFINE},
@@ -214,25 +216,82 @@ impl Channel {
         Ok(())
     }
 
-    /// Implements send message functionality. Creates a new payload and
-    /// encodes it. Then creates a message packet (the base type of the
-    /// network) and copies the payload into it. Then we send the packet
-    /// over the network stream.
+    /// Sends an outbound Message by writing data to the given async stream.
     async fn send_message<M: message::Message>(&self, message: &M) -> Result<()> {
-        let packet = Packet { command: M::NAME.to_string(), payload: serialize(message) };
+        let command = M::NAME.to_string();
+        assert!(!command.is_empty());
+        assert!(std::mem::size_of::<usize>() <= std::mem::size_of::<u64>());
+
+        let stream = &mut *self.writer.lock().await;
+        let mut name_buffer = Vec::<u8>::new();
+        let mut msg_buffer = Vec::<u8>::new();
+        let mut written: usize = 0;
 
         dnetev!(self, SendMessage, {
             chan: self.info.clone(),
-            cmd: packet.command.clone(),
+            cmd: command,
             time: NanoTimestamp::current_time(),
         });
 
-        let stream = &mut *self.writer.lock().await;
-        let _ = message::send_packet(stream, packet).await?;
+        trace!(target: "net::channel::send_message()", "Sending magic...");
+        written += MAGIC_BYTES.encode_async(stream).await?;
+
+        trace!(target: "net::channel::send_message()", "Sent magic");
+        trace!(target: "net::channel::send_message()", "Sending command...");
+
+        // First encode the name to an intermediate buffer.
+        M::NAME.to_string().encode_async(&mut name_buffer).await?;
+
+        // Then extract the length of the intermediate buffer as a VarInt
+        // and write to the stream. This is the length of the name message.
+        // Then encode the name itself to the stream.
+        written += VarInt(name_buffer.len() as u64).encode_async(stream).await?;
+        written += M::NAME.to_string().encode_async(stream).await?;
+
+        trace!(target: "net::channel::send_message()", "Sent command: {}", M::NAME.to_string());
+        trace!(target: "net::channel::send_message()", "Sending payload...");
+
+        // Do the same proceedure for the Message.
+        message.encode_async(&mut msg_buffer).await?;
+
+        written += VarInt(msg_buffer.len() as u64).encode_async(stream).await?;
+        written += message.encode_async(stream).await?;
+
+        trace!(target: "net::channel::send_message()", "Sent payload {} bytes, total bytes {}",
+            msg_buffer.len(), written);
+
+        stream.flush().await?;
 
         Ok(())
     }
 
+    /// Returns a decoded Message command.
+    /// We start by extracting the length from the stream, then allocate
+    /// the precise buffer for this length using stream.take(). This provides
+    /// a basic DDOS protection.
+    pub async fn read_command<R: AsyncRead + Unpin + Send + Sized>(
+        &self,
+        stream: &mut R,
+    ) -> Result<String> {
+        // Messages should have a 4 byte header of magic digits.
+        // This is used for network debugging.
+        let mut magic = [0u8; 4];
+        trace!(target: "net::channel::read_command()", "Reading magic...");
+        stream.read_exact(&mut magic).await?;
+
+        trace!(target: "net::channel::read_command()", "Read magic {:?}", magic);
+        if magic != MAGIC_BYTES {
+            error!(target: "net::channel::read_command", "Error: Magic bytes mismatch");
+            return Err(Error::MalformedPacket)
+        }
+
+        let len = VarInt::decode_async(stream).await.unwrap().0;
+        let mut take = stream.take(len);
+        let command = String::decode_async(&mut take).await.unwrap();
+
+        Ok(command)
+    }
+
     /// Subscribe to a message on the message subsystem.
     pub async fn subscribe_msg<M: message::Message>(&self) -> Result<MessageSubscription<M>> {
         debug!(
@@ -278,8 +337,8 @@ impl Channel {
 
         // Run loop
         loop {
-            let packet = match message::read_packet(reader).await {
-                Ok(packet) => packet,
+            let command = match self.read_command(reader).await {
+                Ok(command) => command,
                 Err(err) => {
                     if Self::is_eof_error(&err) {
                         info!(
@@ -308,12 +367,12 @@ impl Channel {
 
             dnetev!(self, RecvMessage, {
                 chan: self.info.clone(),
-                cmd: packet.command.clone(),
+                cmd: command.clone(),
                 time: NanoTimestamp::current_time(),
             });
 
             // Send result to our subscribers
-            match self.message_subsystem.notify(&packet.command, &packet.payload).await {
+            match self.message_subsystem.notify(&command, reader).await {
                 Ok(()) => {}
                 // If we're getting messages without dispatchers, it's spam.
                 Err(Error::MissingDispatcher) => {

src/net/message.rs

@@ -17,19 +17,14 @@
  */
 
 use darkfi_serial::{
-    async_trait, AsyncDecodable, AsyncEncodable, Decodable, Encodable, SerialDecodable,
-    SerialEncodable, VarInt,
+    async_trait, AsyncDecodable, AsyncEncodable, SerialDecodable, SerialEncodable,
 };
-use log::trace;
-use smol::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 use url::Url;
 
-use crate::{Error, Result};
-
-const MAGIC_BYTES: [u8; 4] = [0xd9, 0xef, 0xb6, 0x7d];
+pub(in crate::net) const MAGIC_BYTES: [u8; 4] = [0xd9, 0xef, 0xb6, 0x7d];
 
 /// Generic message template.
-pub trait Message: 'static + Send + Sync + Encodable + Decodable {
+pub trait Message: 'static + Send + Sync + AsyncDecodable + AsyncEncodable {
     const NAME: &'static str;
 }
 
@@ -109,81 +104,3 @@ pub struct VerackMessage {
     pub app_version: semver::Version,
 }
 impl_p2p_message!(VerackMessage, "verack");
-
-/// Packets are the base type read from the network.
-/// Converted to messages and passed to event loop.
-#[derive(Debug, SerialEncodable, SerialDecodable)]
-pub struct Packet {
-    pub command: String,
-    pub payload: Vec<u8>,
-}
-
-/// Reads and decodes an inbound payload from the given async stream.
-/// Returns decoded [`Packet`].
-/// We start by extracting the packet length from the stream, then allocate
-/// the precise buffer for this length using stream.take(). This provides
-/// a basic DDOS protection.
-pub async fn read_packet<R: AsyncRead + Unpin + Send + Sized>(stream: &mut R) -> Result<Packet> {
-    // Packets should have a 4 byte header of magic digits.
-    // This is used for network debugging.
-    let mut magic = [0u8; 4];
-    trace!(target: "net::message", "Reading magic...");
-    stream.read_exact(&mut magic).await?;
-
-    trace!(target: "net::message", "Read magic {:?}", magic);
-    if magic != MAGIC_BYTES {
-        trace!(target: "net::message", "Error: Magic bytes mismatch");
-        return Err(Error::MalformedPacket)
-    }
-
-    // First deserialize the command, i.e. the type of the message.
-    let cmd_len = VarInt::decode_async(stream).await?.0;
-    let mut cmd_stream = stream.take(cmd_len);
-    let mut cmd_str = Vec::new();
-    cmd_str.try_reserve(cmd_len as usize)?;
-
-    for _ in 0..cmd_len {
-        cmd_str.push(AsyncDecodable::decode_async(&mut cmd_stream).await?);
-    }
-    let command = String::from_utf8(cmd_str)?;
-
-    // Then deserialize the message-dependent payload (see: message types)
-    let msg_len = VarInt::decode_async(stream).await?.0;
-    let mut msg_stream = stream.take(msg_len);
-    let mut payload = Vec::new();
-    payload.try_reserve(msg_len as usize)?;
-
-    for _ in 0..msg_len {
-        payload.push(AsyncDecodable::decode_async(&mut msg_stream).await?);
-    }
-
-    trace!(target: "net::message", "Read payload {} bytes", payload.len());
-
-    Ok(Packet { command, payload })
-}
-
-/// Sends an outbound packet by writing data to the given async stream.
-/// Returns the total written bytes.
-pub async fn send_packet<W: AsyncWrite + Unpin + Send + Sized>(
-    stream: &mut W,
-    packet: Packet,
-) -> Result<usize> {
-    assert!(!packet.command.is_empty());
-    assert!(std::mem::size_of::<usize>() <= std::mem::size_of::<u64>());
-
-    let mut written: usize = 0;
-
-    trace!(target: "net::message", "Sending magic...");
-    written += MAGIC_BYTES.encode_async(stream).await?;
-    trace!(target: "net::message", "Sent magic");
-
-    written += packet.command.encode_async(stream).await?;
-    trace!(target: "net::message", "Sent command: {}", packet.command);
-
-    written += packet.payload.encode_async(stream).await?;
-    trace!(target: "net::message", "Sent payload {} bytes", packet.payload.len() as u64);
-
-    stream.flush().await?;
-
-    Ok(written)
-}

src/net/message_subscriber.rs

@@ -16,16 +16,17 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use std::{any::Any, collections::HashMap, io::Cursor, sync::Arc, time::Duration};
+use std::{any::Any, collections::HashMap, sync::Arc, time::Duration};
 
 use async_trait::async_trait;
 use futures::stream::{FuturesUnordered, StreamExt};
-use log::{debug, warn};
+use log::{debug, error, warn};
 use rand::{rngs::OsRng, Rng};
-use smol::lock::Mutex;
+use smol::{io::AsyncReadExt, lock::Mutex};
 
 use super::message::Message;
-use crate::{system::timeout::timeout, Error, Result};
+use crate::{net::transport::PtStream, system::timeout::timeout, Error, Result};
+use darkfi_serial::{AsyncDecodable, VarInt};
 
 /// 64-bit identifier for message subscription.
 pub type MessageSubscriptionId = u64;
@@ -177,7 +178,7 @@ impl<M: Message> MessageSubscription<M> {
 /// Generic interface for the message dispatcher.
 #[async_trait]
 trait MessageDispatcherInterface: Send + Sync {
-    async fn trigger(&self, payload: &[u8]);
+    async fn trigger(&self, stream: &mut smol::io::ReadHalf<Box<dyn PtStream + 'static>>);
 
     async fn trigger_error(&self, err: Error);
 
@@ -188,18 +189,25 @@ trait MessageDispatcherInterface: Send + Sync {
 #[async_trait]
 impl<M: Message> MessageDispatcherInterface for MessageDispatcher<M> {
     /// Internal function to deserialize data into a message type
-    /// and dispatch it across subscriber channels.
-    async fn trigger(&self, payload: &[u8]) {
-        // Deserialize data into type, send down the pipes.
-        let cursor = Cursor::new(payload);
-        match M::decode(cursor) {
+    /// and dispatch it across subscriber channels. Reads directly
+    /// from an inbound stream.
+    ///
+    /// We extract the message length from the stream and use `take()`
+    /// to allocate an appropiately sized buffer as a basic DDOS protection.
+    async fn trigger(&self, stream: &mut smol::io::ReadHalf<Box<dyn PtStream + 'static>>) {
+        // TODO: check the message length does not exceed some bound.
+        let len = VarInt::decode_async(stream).await.unwrap().0;
+        let mut take = stream.take(len);
+
+        // Deserialize stream into type, send down the pipes.
+        match M::decode_async(&mut take).await {
             Ok(message) => {
                 let message = Ok(Arc::new(message));
                 self._trigger_all(message).await
             }
 
             Err(err) => {
-                debug!(
+                error!(
                     target: "net::message_subscriber::trigger()",
                     "Unable to decode data. Dropping...: {}",
                     err,
@@ -265,7 +273,11 @@ impl MessageSubsystem {
 
     /// Transmits a payload to a dispatcher.
     /// Returns an error if the payload fails to transmit.
-    pub async fn notify(&self, command: &str, payload: &[u8]) -> Result<()> {
+    pub async fn notify(
+        &self,
+        command: &str,
+        reader: &mut smol::io::ReadHalf<Box<dyn PtStream + 'static>>,
+    ) -> Result<()> {
         let Some(dispatcher) = self.dispatchers.lock().await.get(command).cloned() else {
             warn!(
                 target: "net::message_subscriber::notify",
@@ -275,7 +287,7 @@ impl MessageSubsystem {
             return Err(Error::MissingDispatcher)
         };
 
-        dispatcher.trigger(payload).await;
+        dispatcher.trigger(reader).await;
         Ok(())
     }
 
@@ -296,47 +308,3 @@ impl MessageSubsystem {
         while let Some(_r) = futures.next().await {}
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use darkfi_serial::{serialize, SerialDecodable, SerialEncodable};
-
-    #[test]
-    fn message_subscriber_test() {
-        #[derive(SerialEncodable, SerialDecodable)]
-        struct MyVersionMessage(pub u32);
-        crate::impl_p2p_message!(MyVersionMessage, "verver");
-
-        smol::block_on(async {
-            let subsystem = MessageSubsystem::new();
-            subsystem.add_dispatch::<MyVersionMessage>().await;
-
-            // Subscribe:
-            // 1. Get dispatcher
-            // 2. Cast to specific type
-            // 3. Do sub, return sub
-            let sub = subsystem.subscribe::<MyVersionMessage>().await.unwrap();
-
-            // Receive message and publish:
-            // 1. Based on string, lookup relevant dispatcher interface
-            // 2. Publish data there
-            let msg = MyVersionMessage(110);
-            let payload = serialize(&msg);
-            subsystem.notify("verver", &payload).await.unwrap();
-
-            // Receive:
-            // 1. Do a get easy
-            let msg2 = sub.receive().await.unwrap();
-            assert_eq!(msg.0, msg2.0);
-
-            // Trigger an error
-            subsystem.trigger_error(Error::ChannelStopped).await;
-
-            let msg2 = sub.receive().await;
-            assert!(msg2.is_err());
-
-            sub.unsubscribe().await;
-        });
-    }
-}