Enforce max size for `Message` type on wire

Signed-off-by: Alexis Sellier <alexis@radicle.xyz>

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, .. }) => {
+            (SessionState::Negotiated { .. }, Message::Ping(Ping { ponglen, .. })) => {
-                let resp = Message::Pong {
+                // Ignore pings which ask for too much data.
-                    zeroes: ZeroBytes::new(ponglen),
+                if ponglen > Ping::MAX_PONG_ZEROES {
-                };
+                    return Ok(());
-                self.reactor.write(peer.addr, resp);
+                }
+                self.reactor.write(
+                    peer.addr,
+                    Message::Pong {
+                        zeroes: ZeroBytes::new(ponglen),
+                    },
+                );
             }
             (SessionState::Negotiated { ping, .. }, Message::Pong { zeroes }) => {
                 if let PingState::AwaitingResponse(ponglen) = *ping {

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 {
+    Ping(Ping),
-        /// The desired response length
-        ponglen: u16,
-        /// The ping payload.
-        zeroes: ZeroBytes,
-    },
 
     /// 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 {
+    pub fn new(size: wire::Size) -> Self {
-        ZeroBytes(arg)
+        ZeroBytes(size)
     }
 
     pub fn is_empty(&self) -> bool {

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::new(&mut self.rng);
-            let msg = Message::Ping {
+            *ping = PingState::AwaitingResponse(msg.ponglen);
-                ponglen,
-                zeroes: message::ZeroBytes::new(
-                    self.rng
-                        .u16(0..(Message::MAX_SIZE - size_of::<u16>() as u16)),
-                ),
-            };
-            reactor.write(self.addr, msg);
 
-            *ping = PingState::AwaitingResponse(ponglen);
+            reactor.write(self.addr, Message::Ping(msg));
         }
         Ok(())
     }

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 {
+            MessageType::Ping => {
-                ponglen: u16::arbitrary(g),
+                let mut rng = fastrand::Rng::with_seed(u64::arbitrary(g));
-                zeroes: ZeroBytes::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!(),
         }

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(),

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)?;