From 538080a6f12e727ab7688d6265e246e02de05acd Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Fri, 6 Aug 2021 16:21:00 +0100
Subject: [PATCH] Added MAX_BODY_SIZE
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fortune for rsh's current commit: Small blessing − 小吉
---
 src/message.rs        | 16 +++++++++++++++-
 src/message/binary.rs |  9 +++++++--
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/message.rs b/src/message.rs
index 4cf3f0c..3c74b2f 100644
--- a/src/message.rs
+++ b/src/message.rs
@@ -36,6 +36,16 @@ pub const RSA_BLOCK_SIZE: usize = 512;
 /// Max size to pre-allocate when reading a message buffer.
 pub const MAX_ALLOC_SIZE: usize = 4096; // 4kb
 
+/// Max size to allow reading for a message buffer component.
+///
+/// Not including the message body, see `MAX_BODY_SIZE` for body.
+pub const MAX_READ_SIZE: usize = 2048 * 1024; // 2MB.
+
+/// Max allowed size of a single message body.
+///
+/// Set to 0 for unlimited.
+pub const MAX_BODY_SIZE: usize = (1024 * 1024) * 1024; // 1GB
+
 /// A message that can send a value into bytes.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct Message<V: ?Sized + MessageValue>
@@ -158,7 +168,11 @@ impl<V: ?Sized + MessageValue> Message<V>
     {
 	let send_with: &S = send_with.borrow();
 	let data = serde_cbor::to_vec(&self.value)?;
-	
+	if MAX_BODY_SIZE > 0 && data.len() > MAX_BODY_SIZE {
+	    return Err(eyre!("Encoded body is too large"))
+		.with_section(|| data.len().header("Body size was"))
+		.with_section(|| MAX_BODY_SIZE.header("Max size is"));
+	}
 	let sig = if self.sign {
 	    Some(send_with.sign_data(&data[..]).expect("Message expected signing, sender did not support it"))
 	} else {
diff --git a/src/message/binary.rs b/src/message/binary.rs
index 4f1c549..71be494 100644
--- a/src/message/binary.rs
+++ b/src/message/binary.rs
@@ -127,9 +127,9 @@ impl<V: ?Sized + MessageValue> SerializedMessage<V>
 	    (: $ser:ty) => {
 		{
 		    let len = try_from!(ref usize, bytes.get_u64())?;
-		    if len > MAX_ALLOC_SIZE {
+		    if len > MAX_READ_SIZE {
 			return Err(eyre!("Invalid length read: {}", len)
-				   .with_section(|| format!("Max length read: {}", MAX_ALLOC_SIZE)))
+				   .with_section(|| MAX_READ_SIZE.header("Max length read")));
 		    }
 		    alloc_local_bytes(len, |de| {
 			read!(&mut de[..]);
@@ -156,6 +156,11 @@ impl<V: ?Sized + MessageValue> SerializedMessage<V>
 	
 	let header = read!(: SerHeader);
 	let data_len = try_from!(ref usize, bytes.get_u64())?;
+	if MAX_BODY_SIZE > 0 && data_len > MAX_BODY_SIZE {
+	    return Err(eyre!("Body size too large"))
+		.with_section(|| data_len.header("Encoded size was"))
+		.with_section(|| MAX_BODY_SIZE.header("Max size is"));
+	}
 	let mut data = Vec::with_capacity(std::cmp::min(data_len, MAX_ALLOC_SIZE)); //XXX: Redesign so we don't allocate OR try to read massive buffers by accident on corrupted/malformed messages
 	read!(&mut data, data_len);
 	if data.len()!=data_len {