Repository: radicle-dev/heartwood (rust)

Issue:
Make sure all `wire` messages fit with `64KB`
Due to the planned noise protocol framing, we need to ensure that all messages sent between peers fit within 64KB.

We should audit the code (mainly the `wire` module), find a strategy for large messages (eg. splitting them or limiting the amount of data that can go in), and if possible, add some assertions to catch any messages that are bigger than 64KB.

Pull request description:
Enforce max size for `Message` type on wire
Closes #21 

Patch:
diff --git a/radicle-node/src/service.rs b/radicle-node/src/service.rs
index a5055206..791dca6c 100644
--- a/radicle-node/src/service.rs
+++ b/radicle-node/src/service.rs
@@ -31,7 +31,7 @@ use crate::git::Url;
 use crate::identity::{Doc, Id};
 use crate::node;
 use crate::service::config::ProjectTracking;
-use crate::service::message::{Address, Announcement, AnnouncementMessage};
+use crate::service::message::{Address, Announcement, AnnouncementMessage, Ping};
 use crate::service::message::{NodeAnnouncement, RefsAnnouncement};
 use crate::service::peer::{PingState, SessionError, SessionState};
 use crate::storage;
@@ -834,11 +834,17 @@ where
                 );
                 return Err(SessionError::Misbehavior);
             }
-            (SessionState::Negotiated { .. }, Message::Ping { ponglen, .. }) => {
-                let resp = Message::Pong {
-                    zeroes: ZeroBytes::new(ponglen),
-                };
-                self.reactor.write(peer.addr, resp);
+            (SessionState::Negotiated { .. }, Message::Ping(Ping { ponglen, .. })) => {
+                // Ignore pings which ask for too much data.
+                if ponglen > Ping::MAX_PONG_ZEROES {
+                    return Ok(());
+                }
+                self.reactor.write(
+                    peer.addr,
+                    Message::Pong {
+                        zeroes: ZeroBytes::new(ponglen),
+                    },
+                );
             }
             (SessionState::Negotiated { ping, .. }, Message::Pong { zeroes }) => {
                 if let PingState::AwaitingResponse(ponglen) = *ping {
diff --git a/radicle-node/src/service/message.rs b/radicle-node/src/service/message.rs
index 4f90b755..6031c3d1 100644
--- a/radicle-node/src/service/message.rs
+++ b/radicle-node/src/service/message.rs
@@ -1,5 +1,5 @@
 use std::str::FromStr;
-use std::{fmt, io, net};
+use std::{fmt, io, mem, net};
 
 use thiserror::Error;
 
@@ -318,14 +318,9 @@ pub enum Message {
 
     /// Ask a connected peer for a Pong.
     ///
-    /// Use to check if the remote peer is responsive or a side-effect free way to keep a
+    /// Used to check if the remote peer is responsive, or a side-effect free way to keep a
     /// connection alive.
-    Ping {
-        /// The desired response length
-        ponglen: u16,
-        /// The ping payload.
-        zeroes: ZeroBytes,
-    },
+    Ping(Ping),
 
     /// Response to `Ping` message.
     Pong {
@@ -374,6 +369,36 @@ impl Message {
     }
 }
 
+/// A ping message.
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct Ping {
+    /// The requested length of the pong message.
+    pub ponglen: wire::Size,
+    /// Zero bytes (ignored).
+    pub zeroes: ZeroBytes,
+}
+
+impl Ping {
+    /// Maximum number of zero bytes in a ping message.
+    pub const MAX_PING_ZEROES: wire::Size = Message::MAX_SIZE // Message size without the type.
+        - mem::size_of::<wire::Size>() as wire::Size // Account for pong length.
+        - mem::size_of::<wire::Size>() as wire::Size; // Account for zeroes length prefix.
+
+    /// Maximum number of zero bytes in a pong message.
+    pub const MAX_PONG_ZEROES: wire::Size =
+        Message::MAX_SIZE - mem::size_of::<wire::Size>() as wire::Size; // Account for zeroes length
+                                                                        // prefix.
+
+    pub fn new(rng: &mut fastrand::Rng) -> Self {
+        let ponglen = rng.u16(0..Self::MAX_PONG_ZEROES);
+
+        Ping {
+            ponglen,
+            zeroes: ZeroBytes::new(rng.u16(0..Self::MAX_PING_ZEROES)),
+        }
+    }
+}
+
 impl From<Announcement> for Message {
     fn from(ann: Announcement) -> Self {
         Self::Announcement(ann)
@@ -390,18 +415,19 @@ impl fmt::Debug for Message {
             Self::Announcement(Announcement { node, message, .. }) => {
                 write!(f, "Announcement({}, {:?})", node, message)
             }
-            Self::Ping { ponglen, zeroes } => write!(f, "Ping({ponglen}, {:?})", zeroes),
+            Self::Ping(Ping { ponglen, zeroes }) => write!(f, "Ping({ponglen}, {:?})", zeroes),
             Self::Pong { zeroes } => write!(f, "Pong({:?})", zeroes),
         }
     }
 }
 
+/// Represents a vector of zeroes of a certain length.
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub struct ZeroBytes(u16);
+pub struct ZeroBytes(wire::Size);
 
 impl ZeroBytes {
-    pub fn new(arg: u16) -> Self {
-        ZeroBytes(arg)
+    pub fn new(size: wire::Size) -> Self {
+        ZeroBytes(size)
     }
 
     pub fn is_empty(&self) -> bool {
diff --git a/radicle-node/src/service/peer.rs b/radicle-node/src/service/peer.rs
index 1b2c152e..cab2f370 100644
--- a/radicle-node/src/service/peer.rs
+++ b/radicle-node/src/service/peer.rs
@@ -1,5 +1,3 @@
-use std::mem::size_of;
-
 use crate::service::message::*;
 use crate::service::*;
 
@@ -117,17 +115,10 @@ impl Session {
 
     pub fn ping(&mut self, reactor: &mut Reactor) -> Result<(), SessionError> {
         if let SessionState::Negotiated { ping, .. } = &mut self.state {
-            let ponglen = self.rng.u16(0..Message::MAX_SIZE);
-            let msg = Message::Ping {
-                ponglen,
-                zeroes: message::ZeroBytes::new(
-                    self.rng
-                        .u16(0..(Message::MAX_SIZE - size_of::<u16>() as u16)),
-                ),
-            };
-            reactor.write(self.addr, msg);
+            let msg = message::Ping::new(&mut self.rng);
+            *ping = PingState::AwaitingResponse(msg.ponglen);
 
-            *ping = PingState::AwaitingResponse(ponglen);
+            reactor.write(self.addr, Message::Ping(msg));
         }
         Ok(())
     }
diff --git a/radicle-node/src/wire/message.rs b/radicle-node/src/wire/message.rs
index 9984a27b..d8271c2e 100644
--- a/radicle-node/src/wire/message.rs
+++ b/radicle-node/src/wire/message.rs
@@ -198,7 +198,7 @@ impl wire::Encode for Message {
                 n += message.encode(writer)?;
                 n += signature.encode(writer)?;
             }
-            Self::Ping { ponglen, zeroes } => {
+            Self::Ping(Ping { ponglen, zeroes }) => {
                 n += ponglen.encode(writer)?;
                 n += zeroes.encode(writer)?;
             }
@@ -206,6 +206,13 @@ impl wire::Encode for Message {
                 n += zeroes.encode(writer)?;
             }
         }
+
+        if n > wire::Size::MAX as usize {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                "Message exceeds maximum size",
+            ));
+        }
         Ok(n)
     }
 }
@@ -278,7 +285,7 @@ impl wire::Decode for Message {
             Ok(MessageType::Ping) => {
                 let ponglen = u16::decode(reader)?;
                 let zeroes = ZeroBytes::decode(reader)?;
-                Ok(Self::Ping { ponglen, zeroes })
+                Ok(Self::Ping(Ping { ponglen, zeroes }))
             }
             Ok(MessageType::Pong) => {
                 let zeroes = ZeroBytes::decode(reader)?;


Test patch:
diff --git a/radicle-node/src/test/arbitrary.rs b/radicle-node/src/test/arbitrary.rs
index 4df3f03e..0b4e5d5a 100644
--- a/radicle-node/src/test/arbitrary.rs
+++ b/radicle-node/src/test/arbitrary.rs
@@ -7,7 +7,7 @@ use crate::crypto;
 use crate::prelude::{Id, NodeId, Refs, Timestamp};
 use crate::service::filter::{Filter, FILTER_SIZE_L, FILTER_SIZE_M, FILTER_SIZE_S};
 use crate::service::message::{
-    Address, Announcement, Envelope, InventoryAnnouncement, Message, NodeAnnouncement,
+    Address, Announcement, Envelope, InventoryAnnouncement, Message, NodeAnnouncement, Ping,
     RefsAnnouncement, Subscribe, ZeroBytes,
 };
 use crate::wire::message::MessageType;
@@ -95,12 +95,13 @@ impl Arbitrary for Message {
                 since: Timestamp::arbitrary(g),
                 until: Timestamp::arbitrary(g),
             }),
-            MessageType::Ping => Self::Ping {
-                ponglen: u16::arbitrary(g),
-                zeroes: ZeroBytes::arbitrary(g),
-            },
+            MessageType::Ping => {
+                let mut rng = fastrand::Rng::with_seed(u64::arbitrary(g));
+
+                Self::Ping(Ping::new(&mut rng))
+            }
             MessageType::Pong => Self::Pong {
-                zeroes: ZeroBytes::arbitrary(g),
+                zeroes: ZeroBytes::new(u16::arbitrary(g).min(Ping::MAX_PONG_ZEROES)),
             },
             _ => unreachable!(),
         }
diff --git a/radicle-node/src/test/tests.rs b/radicle-node/src/test/tests.rs
index 7f00eb14..42f2edd2 100644
--- a/radicle-node/src/test/tests.rs
+++ b/radicle-node/src/test/tests.rs
@@ -44,10 +44,10 @@ fn test_ping_response() {
     alice.connect_to(&bob);
     alice.receive(
         &bob.addr(),
-        Message::Ping {
+        Message::Ping(Ping {
             ponglen: 21,
             zeroes: ZeroBytes::new(42),
-        },
+        }),
     );
     assert_matches!(
         alice.messages(&bob.addr()).next(),