add maximum message size

eigerco · Feb 7, 2024 · 01f4913 · 01f4913
1 parent 04c6a0d
commit 01f4913
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/src/incoming_stream.rs b/src/incoming_stream.rs
@@ -87,7 +87,7 @@ impl<const S: usize> futures::Stream for IncomingStream<S> {
             let msg = match self.stream.poll_next_unpin(cx) {
                 Poll::Ready(Some(Ok(msg))) => msg,
                 Poll::Ready(Some(Err(e))) => {
-                    error!("{e}");
+                    error!("Message decoding failed: {e}");
                     return Poll::Ready(None);
                 }
                 Poll::Ready(None) => return Poll::Ready(None),

diff --git a/src/message.rs b/src/message.rs
@@ -8,6 +8,9 @@ use quick_protobuf::{BytesReader, BytesWriter, MessageWrite, Writer, WriterBacke
 use crate::proto::message::mod_Message::mod_Wantlist::{Entry, WantType};
 use crate::proto::message::Message;
 
+/// Bitswap spec defines maximum `Message` size to 4MiB.
+pub(crate) const MAX_MESSAGE_SIZE: usize = 4 * 1024 * 1024;
+
 pub(crate) struct Codec;
 
 #[derive(Debug, thiserror::Error)]
@@ -45,18 +48,19 @@ impl Decoder for Codec {
 
         let varint_len = src.len() - rest.len();
 
-        // TODO: if len > MAX return error
+        if varint_len > MAX_MESSAGE_SIZE {
+            return Err(io::Error::new(io::ErrorKind::Other, "Message too large"));
+        }
 
         if rest.len() < len {
             return Ok(None);
         }
 
         let mut reader = BytesReader::from_bytes(rest);
 
-        let msg = reader.read_message_by_len(rest, len).map_err(|e| {
-            let s = format!("Message decoding failed: {e}");
-            io::Error::new(io::ErrorKind::Other, s)
-        })?;
+        let msg = reader
+            .read_message_by_len(rest, len)
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
 
         src.advance(varint_len + len);