Repository: quartiq/minimq
Pull request #145 (2023-08-25 11:46:38)

Updating will management
This PR refactors the `will` temporary storage such that the `will` is serialized at configuration time into a user-provided buffer. This means that the `Will` message itself does not need to live very long.

Pros and Cons:
* User has to provide yet another buffer at config time
* Libraries can leverage will memory provided by the user to construct customized wills.
* Will serialization only ever happens once


Patch:
diff --git a/src/config.rs b/src/config.rs
index b6abd5e..acc8c2d 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -1,18 +1,23 @@
-use crate::{ProtocolError, Will};
+use crate::{will::SerializedWill, ProtocolError, Will};
 use core::convert::TryFrom;
 use embedded_time::duration::{Extensions, Milliseconds};
 use heapless::String;
 
+pub(crate) enum WillState<'a> {
+    BufferAvailable(&'a mut [u8]),
+    Serialized(SerializedWill<'a>),
+}
+
 /// Configuration specifying the operational state of the MQTT client.
 pub struct Config<'a, Broker: crate::Broker> {
     pub(crate) broker: Broker,
     pub(crate) rx_buffer: &'a mut [u8],
     pub(crate) tx_buffer: &'a mut [u8],
     pub(crate) state_buffer: &'a mut [u8],
+    pub(crate) will: Option<WillState<'a>>,
     pub(crate) client_id: String<64>,
     pub(crate) keepalive_interval: Milliseconds<u32>,
     pub(crate) downgrade_qos: bool,
-    pub(crate) will: Option<Will<'a>>,
 }
 
 impl<'a, Broker: crate::Broker> Config<'a, Broker> {
@@ -42,6 +47,12 @@ impl<'a, Broker: crate::Broker> Config<'a, Broker> {
         self
     }
 
+    /// Provide additional buffer space for an optional connection [Will]
+    pub fn will_buffer(mut self, buffer: &'a mut [u8]) -> Self {
+        self.will = Some(WillState::BufferAvailable(buffer));
+        self
+    }
+
     /// Specify a known client ID to use. If not assigned, the broker will auto assign an ID.
     pub fn client_id(mut self, id: &str) -> Result<Self, ProtocolError> {
         self.client_id =
@@ -70,12 +81,27 @@ impl<'a, Broker: crate::Broker> Config<'a, Broker> {
         self
     }
 
+    /// Construct a will in the provided buffer.
+    pub fn will_with_buffer(
+        mut self,
+        buf: &'a mut [u8],
+        will: Will<'_>,
+    ) -> Result<Self, crate::ser::Error> {
+        self = self.will_buffer(buf);
+        self.will(will)
+    }
+
     /// Specify the Will message to be sent if the client disconnects.
     ///
     /// # Args
     /// * `will` - The will to use.
-    pub fn will(mut self, will: Will<'a>) -> Self {
-        self.will.replace(will);
-        self
+    pub fn will(mut self, will: Will<'_>) -> Result<Self, crate::ser::Error> {
+        let Some(WillState::BufferAvailable(buf)) = self.will else {
+            return Err(crate::ser::Error::InsufficientMemory);
+        };
+
+        self.will = Some(WillState::Serialized(will.serialize(buf)?));
+
+        Ok(self)
     }
 }
diff --git a/src/mqtt_client.rs b/src/mqtt_client.rs
index 90f3b24..4ac1389 100644
--- a/src/mqtt_client.rs
+++ b/src/mqtt_client.rs
@@ -5,7 +5,7 @@ use crate::{
     reason_codes::ReasonCode,
     session_state::SessionState,
     types::{Properties, TopicFilter, Utf8String},
-    will::Will,
+    will::SerializedWill,
     Config, Error, MinimqError, Property, ProtocolError, QoS, {debug, error, info, warn},
 };
 
@@ -268,7 +268,7 @@ pub struct MqttClient<
 > {
     sm: sm::StateMachine<ClientContext<'buf, Clock>>,
     network: InterfaceHolder<'buf, TcpStack>,
-    will: Option<Will<'buf>>,
+    will: Option<SerializedWill<'buf>>,
     downgrade_qos: bool,
     broker: Broker,
     max_packet_size: usize,
@@ -439,7 +439,7 @@ impl<'buf, TcpStack: TcpClientStack, Clock: embedded_time::Clock, Broker: crate:
             keep_alive: self.sm.context().keepalive_interval(),
             properties: Properties::Slice(&properties),
             client_id: Utf8String(self.sm.context().session_state.client_id.as_str()),
-            will: self.will.as_ref(),
+            will: self.will,
             clean_start: !self.sm.context().session_state.is_present(),
         })?;
 
@@ -697,6 +697,12 @@ impl<'buf, TcpStack: TcpClientStack, Clock: embedded_time::Clock, Broker: crate:
             config.tx_buffer.len(),
         );
 
+        let will = if let Some(crate::config::WillState::Serialized(will)) = config.will {
+            Some(will)
+        } else {
+            None
+        };
+
         Minimq {
             client: MqttClient {
                 sm: StateMachine::new(ClientContext::new(
@@ -706,7 +712,7 @@ impl<'buf, TcpStack: TcpClientStack, Clock: embedded_time::Clock, Broker: crate:
                 )),
                 downgrade_qos: config.downgrade_qos,
                 broker: config.broker,
-                will: config.will,
+                will,
                 network: InterfaceHolder::new(network_stack, config.tx_buffer),
                 max_packet_size: config.rx_buffer.len(),
             },
diff --git a/src/packets.rs b/src/packets.rs
index 19a080f..1d23c7d 100644
--- a/src/packets.rs
+++ b/src/packets.rs
@@ -1,7 +1,7 @@
 use crate::{
     reason_codes::ReasonCode,
     types::{Properties, TopicFilter, Utf8String},
-    will::Will,
+    will::SerializedWill,
     QoS, Retain,
 };
 use bit_field::BitField;
@@ -23,7 +23,7 @@ pub struct Connect<'a> {
     pub client_id: Utf8String<'a>,
 
     /// An optional will message to be transmitted whenever the connection is lost.
-    pub will: Option<&'a Will<'a>>,
+    pub(crate) will: Option<SerializedWill<'a>>,
 
     /// Specified true there is no session state being taken in to the MQTT connection.
     pub clean_start: bool,
@@ -51,8 +51,7 @@ impl<'a> serde::Serialize for Connect<'a> {
         item.serialize_field("client_id", &self.client_id)?;
 
         if let Some(will) = &self.will {
-            let flattened = will.flatten();
-            item.serialize_field("will", &flattened)?;
+            item.serialize_field("will", will.contents)?;
         }
 
         item.end()
@@ -427,6 +426,7 @@ mod tests {
         ];
 
         let mut buffer: [u8; 900] = [0; 900];
+        let mut will_buff = [0; 64];
         let mut will = crate::will::Will::new("EFG", &[0xAB, 0xCD], &[]).unwrap();
         will.qos(crate::QoS::AtMostOnce);
         will.retained(crate::Retain::NotRetained);
@@ -436,7 +436,7 @@ mod tests {
             keep_alive: 10,
             properties: crate::types::Properties::Slice(&[]),
             client_id: crate::types::Utf8String("ABC"),
-            will: Some(&will),
+            will: Some(will.serialize(&mut will_buff).unwrap()),
         };
 
         let message = MqttSerializer::to_buffer(&mut buffer, &connect).unwrap();
diff --git a/src/ser/mod.rs b/src/ser/mod.rs
index 4466f3d..aaae9fe 100644
--- a/src/ser/mod.rs
+++ b/src/ser/mod.rs
@@ -86,6 +86,14 @@ impl<'a> MqttSerializer<'a> {
         }
     }
 
+    /// Immediately finish the packet and return the contents.
+    ///
+    /// # Note
+    /// This does not append any MQTT headers.
+    pub fn finish(self) -> &'a mut [u8] {
+        &mut self.buf[MAX_FIXED_HEADER_SIZE..self.index]
+    }
+
     /// Encode an MQTT control packet into a buffer.
     ///
     /// # Args
diff --git a/src/will.rs b/src/will.rs
index ee0b623..b5a4650 100644
--- a/src/will.rs
+++ b/src/will.rs
@@ -10,13 +10,13 @@ use serde::Serialize;
 pub struct Will<'a> {
     topic: &'a str,
     data: &'a [u8],
-    pub(crate) qos: QoS,
-    pub(crate) retained: Retain,
+    qos: QoS,
+    retained: Retain,
     properties: &'a [Property<'a>],
 }
 
 #[derive(Serialize)]
-pub(crate) struct WillMessage<'a> {
+struct WillMessage<'a> {
     properties: Properties<'a>,
     topic: Utf8String<'a>,
     data: BinaryData<'a>,
@@ -57,12 +57,24 @@ impl<'a> Will<'a> {
         })
     }
 
-    pub(crate) fn flatten(&self) -> WillMessage<'a> {
-        WillMessage {
+    /// Serialize the will contents into a flattened, borrowed buffer.
+    pub(crate) fn serialize<'b>(
+        &self,
+        buf: &'b mut [u8],
+    ) -> Result<SerializedWill<'b>, crate::ser::Error> {
+        let message = WillMessage {
             topic: Utf8String(self.topic),
             properties: Properties::Slice(self.properties),
             data: BinaryData(self.data),
-        }
+        };
+
+        let mut serializer = crate::ser::MqttSerializer::new(buf);
+        message.serialize(&mut serializer)?;
+        Ok(SerializedWill {
+            qos: self.qos,
+            retained: self.retained,
+            contents: serializer.finish(),
+        })
     }
 
     /// Set the retained status of the will.
@@ -81,3 +93,11 @@ impl<'a> Will<'a> {
         self.qos = qos;
     }
 }
+
+/// A will where the topic, properties, and contents have already been serialized.
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub(crate) struct SerializedWill<'a> {
+    pub(crate) qos: QoS,
+    pub(crate) retained: Retain,
+    pub(crate) contents: &'a [u8],
+}


Test patch:
diff --git a/tests/integration_test.rs b/tests/integration_test.rs
index 26157b9..237520b 100644
--- a/tests/integration_test.rs
+++ b/tests/integration_test.rs
@@ -12,6 +12,7 @@ fn main() -> std::io::Result<()> {
     let mut rx_buffer = [0u8; 256];
     let mut tx_buffer = [0u8; 256];
     let mut session = [0u8; 256];
+    let mut will_buffer = [0; 64];
     let stack = std_embedded_nal::Stack::default();
     let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
     let mut mqtt: Minimq<'_, _, _, minimq::broker::IpBroker> = Minimq::new(
@@ -19,7 +20,8 @@ fn main() -> std::io::Result<()> {
         StandardClock::default(),
         minimq::Config::new(localhost.into(), &mut rx_buffer, &mut tx_buffer)
             .session_state(&mut session)
-            .will(will)
+            .will_with_buffer(&mut will_buffer, will)
+            .unwrap()
             .keepalive_interval(60),
     );
 
diff --git a/tests/stack/mod.rs b/tests/stack/mod.rs
index aa1e4de..645ee01 100644
--- a/tests/stack/mod.rs
+++ b/tests/stack/mod.rs
@@ -78,7 +78,7 @@ impl TcpSocket {
 
     fn connect(&mut self, remote: embedded_nal::SocketAddr) -> Result<(), Error> {
         let embedded_nal::IpAddr::V4(addr) = remote.ip() else {
-            return Err(Error::new(ErrorKind::Other, "Only IPv4 supported"))
+            return Err(Error::new(ErrorKind::Other, "Only IPv4 supported"));
         };
         let remote = SocketAddr::new(addr.octets().into(), remote.port());
         let soc = TcpStream::connect(remote)?;