From 07dd2ca4c3ecfcb91292316da6c6b8df7c2a4177 Mon Sep 17 00:00:00 2001
From: Ximon Eighteen <3304436+ximon18@users.noreply.github.com>
Date: Tue, 9 Jul 2024 14:39:36 +0200
Subject: [PATCH 1/2] Add an optional push limit to `MessageBuilder`.

This is useful when wanting to "reserve" space such that pushes at one point in the code to the message will leave behind enough room for subsequent pushes later, e.g. of OPT or TSIG RRs.
---
 src/base/message_builder.rs | 87 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 86 insertions(+), 1 deletion(-)

diff --git a/src/base/message_builder.rs b/src/base/message_builder.rs
index 2175d710f..65f2b8c60 100644
--- a/src/base/message_builder.rs
+++ b/src/base/message_builder.rs
@@ -169,6 +169,11 @@ use std::vec::Vec;
 #[derive(Clone, Debug)]
 pub struct MessageBuilder<Target> {
     target: Target,
+
+    /// An optional maximum message size.
+    ///
+    /// Defaults to usize::MAX.
+    limit: usize,
 }
 
 /// # Creating Message Builders
@@ -187,7 +192,10 @@ impl<Target: OctetsBuilder + Truncate> MessageBuilder<Target> {
     ) -> Result<Self, Target::AppendError> {
         target.truncate(0);
         target.append_slice(HeaderSection::new().as_slice())?;
-        Ok(MessageBuilder { target })
+        Ok(MessageBuilder {
+            target,
+            limit: usize::MAX,
+        })
     }
 }
 
@@ -270,6 +278,41 @@ impl<Target: Composer> MessageBuilder<Target> {
     }
 }
 
+/// # Limiting message size
+impl<Target: Composer> MessageBuilder<Target> {
+    /// Limit how much of the underlying buffer may be used.
+    ///
+    /// When a limit is set, calling [`push()`] will fail if the limit is
+    /// exceeded just as if the actual end of the underlying buffer had been
+    /// reached.
+    ///
+    /// Note: Calling this function does NOT truncate the underlying buffer.
+    /// If the new limit is lees than the amount of the buffer that has
+    /// already been used, exisitng content beyond the limit will remain
+    /// untouched, the length will remain larger than the limit, and calls to
+    /// [`push()`] will fail until the buffer is truncated to a size less than
+    /// the limit.
+    pub fn set_push_limit(&mut self, limit: usize) {
+        self.limit = limit;
+    }
+
+    /// Clear the push limit, if set.
+    ///
+    /// Removes any push limit previously set via `[set_push_limit()`].
+    pub fn clear_push_limit(&mut self) {
+        self.limit = usize::MAX;
+    }
+
+    /// Returns the current push limit, if set.
+    pub fn push_limit(&self) -> Option<usize> {
+        if self.limit == usize::MAX {
+            None
+        } else {
+            Some(self.limit)
+        }
+    }
+}
+
 /// # Access to the Message Header
 ///
 impl<Target: OctetsBuilder + AsRef<[u8]>> MessageBuilder<Target> {
@@ -401,6 +444,13 @@ impl<Target: Composer> MessageBuilder<Target> {
             self.target.truncate(pos);
             return Err(From::from(err));
         }
+
+        let new_pos = self.target.as_ref().len();
+        if new_pos >= self.limit {
+            self.target.truncate(pos);
+            return Err(PushError::ShortBuf);
+        }
+
         if inc(self.counts_mut()).is_err() {
             self.target.truncate(pos);
             return Err(PushError::CountOverflow);
@@ -2382,6 +2432,41 @@ mod test {
         assert_eq!(rr.data(), &A::from_octets(192, 0, 2, 1));
     }
 
+    #[test]
+    fn exceed_limits() {
+        // Create a limited message builder.
+        let buf = heapless::Vec::<u8, 100>::new();
+
+        // Initialize it with a message header (12 bytes)
+        let mut msg = MessageBuilder::from_target(buf).unwrap();
+        let hdr_len = msg.as_slice().len();
+
+        // Add some bytes.
+        msg.push(|t| t.append_slice(&[0u8; 50]), |_| Ok(()))
+            .unwrap();
+        assert_eq!(msg.as_slice().len(), hdr_len + 50);
+
+        // Set a push limit below the current length.
+        msg.set_push_limit(25);
+
+        // Verify that push fails.
+        assert!(msg.push(|t| t.append_slice(&[0u8; 1]), |_| Ok(())).is_err());
+        assert_eq!(msg.as_slice().len(), hdr_len + 50);
+
+        // Remove the limit.
+        msg.clear_push_limit();
+
+        // Verify that push up until capacity succeeds.
+        for _ in (hdr_len + 50)..100 {
+            msg.push(|t| t.append_slice(&[0u8; 1]), |_| Ok(())).unwrap();
+        }
+        assert_eq!(msg.as_slice().len(), 100);
+
+        // Verify that exceeding the underlying capacity limit fails.
+        assert!(msg.push(|t| t.append_slice(&[0u8; 1]), |_| Ok(())).is_err());
+        assert_eq!(msg.as_slice().len(), 100);
+    }
+
     #[test]
     fn opt_builder() {
         let mut msg = MessageBuilder::new_vec().additional();

From 319d95e54f9d1cecba1992f4233d623b17a91744 Mon Sep 17 00:00:00 2001
From: Ximon Eighteen <3304436+ximon18@users.noreply.github.com>
Date: Tue, 9 Jul 2024 15:01:34 +0200
Subject: [PATCH 2/2] Review feedback: Add missing feature gate.

---
 src/base/message_builder.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/base/message_builder.rs b/src/base/message_builder.rs
index 65f2b8c60..fa10ecd40 100644
--- a/src/base/message_builder.rs
+++ b/src/base/message_builder.rs
@@ -2432,6 +2432,7 @@ mod test {
         assert_eq!(rr.data(), &A::from_octets(192, 0, 2, 1));
     }
 
+    #[cfg(feature = "heapless")]
     #[test]
     fn exceed_limits() {
         // Create a limited message builder.