diff --git a/src/base/message_builder.rs b/src/base/message_builder.rs index 2175d710f..fa10ecd40 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, + + /// An optional maximum message size. + /// + /// Defaults to usize::MAX. + limit: usize, } /// # Creating Message Builders @@ -187,7 +192,10 @@ impl MessageBuilder { ) -> Result { target.truncate(0); target.append_slice(HeaderSection::new().as_slice())?; - Ok(MessageBuilder { target }) + Ok(MessageBuilder { + target, + limit: usize::MAX, + }) } } @@ -270,6 +278,41 @@ impl MessageBuilder { } } +/// # Limiting message size +impl MessageBuilder { + /// 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 { + if self.limit == usize::MAX { + None + } else { + Some(self.limit) + } + } +} + /// # Access to the Message Header /// impl> MessageBuilder { @@ -401,6 +444,13 @@ impl MessageBuilder { 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::::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();