Add an optional push limit to MessageBuilder. (#348)

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

@@ -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,42 @@ 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.
+        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();