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/4] 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, + + /// 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,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::::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/4] 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. From 56a8e528758110fe3432c47f9b7c54e37ff8b575 Mon Sep 17 00:00:00 2001 From: Ximon Eighteen <3304436+ximon18@users.noreply.github.com> Date: Tue, 9 Jul 2024 15:25:24 +0200 Subject: [PATCH 3/4] Add `PushError::LimitExceeded`. Note: This is a breaking change. --- src/base/message_builder.rs | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/base/message_builder.rs b/src/base/message_builder.rs index fa10ecd40..0e114b82f 100644 --- a/src/base/message_builder.rs +++ b/src/base/message_builder.rs @@ -421,7 +421,7 @@ impl MessageBuilder { /// /// This message is atop the octets slices derived from the builder, so /// it can be created cheaply. - pub fn as_message(&self) -> Message<&[u8]> + pub fn as_message(&self) -> Message<&[u8]>git where Target: AsRef<[u8]>, { @@ -448,7 +448,7 @@ impl MessageBuilder { let new_pos = self.target.as_ref().len(); if new_pos >= self.limit { self.target.truncate(pos); - return Err(PushError::ShortBuf); + return Err(PushError::LimitExceeded); } if inc(self.counts_mut()).is_err() { @@ -2332,6 +2332,7 @@ impl Truncate for TreeCompressor { #[derive(Clone, Copy, Debug)] pub enum PushError { CountOverflow, + LimitedExceeded, ShortBuf, } @@ -2345,6 +2346,7 @@ impl fmt::Display for PushError { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match *self { PushError::CountOverflow => f.write_str("counter overflow"), + PushError::LimitExceeded => f.write_str("limit exceeded"), PushError::ShortBuf => ShortBuf.fmt(f), } } @@ -2451,7 +2453,10 @@ mod test { msg.set_push_limit(25); // Verify that push fails. - assert!(msg.push(|t| t.append_slice(&[0u8; 1]), |_| Ok(())).is_err()); + assert!(matches!( + msg.push(|t| t.append_slice(&[0u8; 1]), |_| Ok(())), + Err(PushError::LimitExceeded) + )); assert_eq!(msg.as_slice().len(), hdr_len + 50); // Remove the limit. @@ -2464,7 +2469,10 @@ mod test { 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!(matches!( + msg.push(|t| t.append_slice(&[0u8; 1]), |_| Ok(())), + Err(PushError::ShortBuf) + )); assert_eq!(msg.as_slice().len(), 100); } From 37f6d31c0cb78d7ffb8d06bb237047a6616aea33 Mon Sep 17 00:00:00 2001 From: Ximon Eighteen <3304436+ximon18@users.noreply.github.com> Date: Tue, 9 Jul 2024 15:29:20 +0200 Subject: [PATCH 4/4] Compilation fixes. --- src/base/message_builder.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/base/message_builder.rs b/src/base/message_builder.rs index 0e114b82f..134fe6141 100644 --- a/src/base/message_builder.rs +++ b/src/base/message_builder.rs @@ -421,7 +421,7 @@ impl MessageBuilder { /// /// This message is atop the octets slices derived from the builder, so /// it can be created cheaply. - pub fn as_message(&self) -> Message<&[u8]>git + pub fn as_message(&self) -> Message<&[u8]> where Target: AsRef<[u8]>, { @@ -2332,7 +2332,7 @@ impl Truncate for TreeCompressor { #[derive(Clone, Copy, Debug)] pub enum PushError { CountOverflow, - LimitedExceeded, + LimitExceeded, ShortBuf, }