Repository: quartiq/minimq
Pull request #68 (2021-12-07 00:01:30)

Implement finishing partial writes
I've been running into issues with getting RST packets from brokers fairly often if there was high load. This is painful especially because starting the connection takes really long time for some reason. It turns out my TCP stack just couldn't keep up and in the current implementation the `WriteFail` error is returned if the write would either block or if it was partially written. The latter usually breaks the MQTT connection because there is currently no way to finish the packet send.

The issue is with the partial writes which usually contained the header including length of the MQTT message but only the first half was actually transferred, the rest was silently dropped. On a subsequent write the first half of the next packet would be interpreted by the broker as the second half of the preceding packet and the rest as a new MQTT packet, which usually turned out to be malformed so the broker sent a disconnect packet and after some time closed the TCP connection.

This solution is fairly simple although it eats up a bit more space (because the network manager must be able to store up to a full packet). I wasn't very successful with convincing the borrow checker that what I want to do is safe so the unfinished packet ends up being copied once after the unsuccessful send and once when we want to try to finish.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f512c54..66504f3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,21 +2,9 @@
 
 This document describes the changes to Minimq between releases.
 
-# [Unreleased]
+# Unreleased
 
-## Added
-## Removed
-## Changed
-
-# Version [0.5.0]
-Version 0.5.0 was published on 2021-12-06
-
-## Added
-* Support for the `Will` message specification.
-* [breaking] Adding `retained` flag to `publish()` to allow messages to be published in a retained
-  manner.
-
-# Version [0.4.0]
+# Version 0.4.0
 Version 0.4.0 was published on 2021-10-08
 
 * Updating to `std-embedded-nal` v0.1 (dev dependency only; now again used for integration tests)
@@ -28,7 +16,7 @@ keep-alive interval
 * Refactoring network stack management into a separate container class
 * Keep-alive settings now take a u16 integer number of seconds
 
-# Version [0.3.0]
+# Version 0.3.0
 Version 0.3.0 was published on 2021-08-06
 
 * Client ID may now be unspecified to allow the broker to automatically assign an ID.
@@ -43,7 +31,7 @@ Version 0.3.0 was published on 2021-08-06
 * Removing using of `generic-array` in favor of const-generics.
 * Correcting an issue where the client would not reconnect if the broker was restarted.
 
-# Version [0.2.0]
+# Version 0.2.0
 Version 0.2.0 was published on 2021-02-15
 
 * Updating the `MqttClient::poll()` function to take a `FnMut` closure to allow internal state
@@ -54,14 +42,7 @@ Version 0.2.0 was published on 2021-02-15
   stack errors out to the user.
 * Updating library to re-export `generic-array` publicly.
 
-## Version [0.1.0]
+## Version 0.1.0
 Version 0.1.0 was published on 2020-08-27
 
 * Initial library release and publish to crates.io
-
-[Unreleased]: https://github.com/quartiq/minimq/compare/0.5.0...HEAD
-[0.5.0]: https://github.com/quartiq/minimq/releases/tag/0.5.0
-[0.4.0]: https://github.com/quartiq/minimq/releases/tag/0.4.0
-[0.3.0]: https://github.com/quartiq/minimq/releases/tag/0.3.0
-[0.2.0]: https://github.com/quartiq/minimq/releases/tag/0.2.0
-[0.1.0]: https://github.com/quartiq/minimq/releases/tag/0.1.0
diff --git a/Cargo.toml b/Cargo.toml
index 8bbd0a5..24f3ce0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "minimq"
-version = "0.5.0"
+version = "0.4.0"
 authors = ["Ryan Summers <ryan.summers@vertigo-designs.com>", "Max Rottenkolber <max@mr.gy>"]
 edition = "2018"
 
diff --git a/src/lib.rs b/src/lib.rs
index 7a3001a..1a0476d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -8,7 +8,8 @@
 //! # Limitations
 //! This library does not currently support the following elements:
 //! * Quality-of-service `ExactlyOnce`
-//! * Quality-of-service above `AtMostOnce` for inbound messages.
+//! * Session timeouts
+//! * Server keep alive timeouts (ping)
 //! * Bulk subscriptions
 //! * Server Authentication
 //! * Encryption
@@ -28,7 +29,7 @@
 //! temperature sensor installed.
 //!
 //! ```no_run
-//! use minimq::{Minimq, QoS, Retain};
+//! use minimq::{Minimq, QoS};
 //!
 //! // Construct an MQTT client with a maximum packet size of 256 bytes.
 //! // Connect to a broker at localhost - Use a client ID of "test".
@@ -50,7 +51,7 @@
 //!         match topic {
 //!             "topic" => {
 //!                println!("{:?}", message);
-//!                client.publish("echo", message, QoS::AtMostOnce, Retain::NotRetained, &[]).unwrap();
+//!                client.publish("echo", message, QoS::AtMostOnce, &[]).unwrap();
 //!             },
 //!             topic => println!("Unknown topic: {}", topic),
 //!         };
@@ -66,41 +67,17 @@ mod mqtt_client;
 mod network_manager;
 mod properties;
 mod session_state;
-mod will;
 
 use message_types::MessageType;
 pub use properties::Property;
 
 pub use embedded_nal;
 pub use embedded_time;
-pub use mqtt_client::Minimq;
+pub use mqtt_client::{Minimq, QoS};
 
 #[cfg(feature = "logging")]
 pub(crate) use log::{debug, error, info, warn};
 
-/// The quality-of-service for an MQTT message.
-#[derive(Debug, Copy, Clone, PartialEq)]
-pub enum QoS {
-    /// A packet will be delivered at most once, but may not be delivered at all.
-    AtMostOnce = 0,
-
-    /// A packet will be delivered at least one time, but possibly more than once.
-    AtLeastOnce = 1,
-
-    /// A packet will be delivered exactly one time.
-    ExactlyOnce = 2,
-}
-
-/// The retained status for an MQTT message.
-#[derive(Debug, Copy, Clone, PartialEq)]
-pub enum Retain {
-    /// The message shall not be retained by the broker.
-    NotRetained = 0,
-
-    /// The message shall be marked for retention by the broker.
-    Retained = 1,
-}
-
 /// Errors that are specific to the MQTT protocol implementation.
 #[derive(Debug, Copy, Clone, PartialEq)]
 pub enum ProtocolError {
@@ -122,6 +99,7 @@ pub enum ProtocolError {
 pub enum Error<E> {
     Network(E),
     WriteFail,
+    PartialWrite,
     NotReady,
     Unsupported,
     ProvidedClientIdTooLong,
diff --git a/src/mqtt_client.rs b/src/mqtt_client.rs
index a5f652e..5cd394f 100644
--- a/src/mqtt_client.rs
+++ b/src/mqtt_client.rs
@@ -6,8 +6,7 @@ use crate::{
     network_manager::InterfaceHolder,
     ser::serialize,
     session_state::SessionState,
-    will::Will,
-    Error, Property, ProtocolError, QoS, Retain, {debug, error, info},
+    Error, Property, ProtocolError, {debug, error, info},
 };
 
 use embedded_nal::{IpAddr, SocketAddr, TcpClientStack};
@@ -16,6 +15,19 @@ use heapless::String;
 
 use core::str::FromStr;
 
+/// The quality-of-service for an MQTT message.
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub enum QoS {
+    /// A packet will be delivered at most once, but may not be delivered at all.
+    AtMostOnce = 0,
+
+    /// A packet will be delivered at least one time, but possibly more than once.
+    AtLeastOnce = 1,
+
+    /// A packet will be delivered exactly one time.
+    ExactlyOnce = 2,
+}
+
 mod sm {
 
     use smlang::statemachine;
@@ -56,11 +68,10 @@ pub struct MqttClient<
     const MSG_SIZE: usize,
     const MSG_COUNT: usize,
 > {
-    pub(crate) network: InterfaceHolder<TcpStack>,
+    pub(crate) network: InterfaceHolder<TcpStack, MSG_SIZE>,
     clock: Clock,
     session_state: SessionState<Clock, MSG_SIZE, MSG_COUNT>,
     connection_state: StateMachine<Context>,
-    will: Option<Will<MSG_SIZE>>,
 }
 
 impl<TcpStack, Clock, const MSG_SIZE: usize, const MSG_COUNT: usize>
@@ -110,7 +121,6 @@ where
                     &properties,
                     // Only perform a clean start if we do not have any session state.
                     !self.session_state.is_present(),
-                    self.will.as_ref(),
                 )?;
 
                 info!("Sending CONNECT");
@@ -131,30 +141,6 @@ where
         Ok(())
     }
 
-    /// Specify the Will message to be sent if the client disconnects.
-    ///
-    /// # Args
-    /// * `topic` - The topic to send the message on
-    /// * `data` - The message to transmit
-    /// * `qos` - The quality of service at which to send the message.
-    /// * `retained` - Specifies whether the will message should be retained by the broker.
-    /// * `properties` - Any properties to send with the will message.
-    pub fn set_will(
-        &mut self,
-        topic: &str,
-        data: &[u8],
-        qos: QoS,
-        retained: Retain,
-        properties: &[Property],
-    ) -> Result<(), Error<TcpStack::Error>> {
-        let mut will = Will::new(topic, data, properties)?;
-        will.retained(retained);
-        will.qos(qos);
-
-        self.will.replace(will);
-        Ok(())
-    }
-
     fn reset(&mut self) {
         self.connection_state.process_event(Events::Disconnect).ok();
     }
@@ -268,7 +254,6 @@ where
         topic: &str,
         data: &[u8],
         qos: QoS,
-        retain: Retain,
         properties: &[Property],
     ) -> Result<(), Error<TcpStack::Error>> {
         // TODO: QoS 2 support.
@@ -288,8 +273,7 @@ where
         let id = self.session_state.get_packet_identifier();
 
         let mut buffer: [u8; MSG_SIZE] = [0; MSG_SIZE];
-        let packet =
-            serialize::publish_message(&mut buffer, topic, data, qos, retain, id, properties)?;
+        let packet = serialize::publish_message(&mut buffer, topic, data, qos, id, properties)?;
 
         self.network.write(packet)?;
         self.session_state.increment_packet_identifier();
@@ -301,6 +285,16 @@ where
         Ok(())
     }
 
+    /// Finish sending last packet if the previous operation returned Error::PartialWrite
+    ///
+    /// # Note
+    /// If this function is not called after partial write until the packet is
+    /// fully sent, the packets will most likely seem malformed to the receiving
+    /// party and they will close the connection.
+    pub fn finish_write(&mut self) -> Result<(), Error<TcpStack::Error>> {
+        self.network.finish_write()
+    }
+
     fn handle_connection_acknowledge(
         &mut self,
         acknowledge: ConnAck,
@@ -500,7 +494,6 @@ impl<
                 clock,
                 session_state,
                 connection_state: StateMachine::new(Context),
-                will: None,
             },
             packet_reader: PacketReader::new(),
         };
diff --git a/src/network_manager.rs b/src/network_manager.rs
index ddefb95..f56fb61 100644
--- a/src/network_manager.rs
+++ b/src/network_manager.rs
@@ -5,17 +5,21 @@
 //! simple ownership semantics of reading and writing to the network stack. This allows the network
 //! stack to be used to transmit buffers that may be stored internally in other structs without
 //! violating Rust's borrow rules.
+use core::mem;
+
 use embedded_nal::{nb, SocketAddr, TcpClientStack};
+use heapless::Vec;
 
 use crate::Error;
 
 /// Simple structure for maintaining state of the network connection.
-pub(crate) struct InterfaceHolder<TcpStack: TcpClientStack> {
+pub(crate) struct InterfaceHolder<TcpStack: TcpClientStack, const MSG_SIZE: usize> {
     socket: Option<TcpStack::TcpSocket>,
     network_stack: TcpStack,
+    unfinished_packet: Option<Vec<u8, MSG_SIZE>>,
 }
 
-impl<TcpStack> InterfaceHolder<TcpStack>
+impl<TcpStack, const MSG_SIZE: usize> InterfaceHolder<TcpStack, MSG_SIZE>
 where
     TcpStack: TcpClientStack,
 {
@@ -24,6 +28,7 @@ where
         Self {
             socket: None,
             network_stack: stack,
+            unfinished_packet: None,
         }
     }
 
@@ -88,13 +93,24 @@ where
             })
             .and_then(|written| {
                 if written != packet.len() {
-                    Err(Error::WriteFail)
+                    self.unfinished_packet = Some(Vec::from_slice(&packet[written..]).unwrap());
+                    Err(Error::PartialWrite)
                 } else {
                     Ok(())
                 }
             })
     }
 
+    /// Finish writing an MQTT control packet to the interface if the last write returned Error::PartialWrite
+    ///
+    pub fn finish_write(&mut self) -> Result<(), Error<TcpStack::Error>> {
+        let packet = mem::replace(&mut self.unfinished_packet, None);
+        if let Some(packet) = packet {
+            self.write(packet.as_slice())?;
+        }
+        Ok(())
+    }
+
     /// Read data from the TCP interface.
     ///
     /// # Args
diff --git a/src/ser/packet_writer.rs b/src/ser/packet_writer.rs
index a66db1d..0037f57 100644
--- a/src/ser/packet_writer.rs
+++ b/src/ser/packet_writer.rs
@@ -173,6 +173,7 @@ impl<'a> ReversedPacketWriter<'a> {
         }
     }
 
+    #[cfg(test)]
     pub fn finish(self) -> &'a [u8] {
         &self.buffer[self.index..]
     }
diff --git a/src/ser/serialize.rs b/src/ser/serialize.rs
index 797b338..dd8b41b 100644
--- a/src/ser/serialize.rs
+++ b/src/ser/serialize.rs
@@ -1,6 +1,7 @@
+use crate::QoS;
 use crate::{
     message_types::MessageType, properties::PropertyIdentifier, ser::ReversedPacketWriter,
-    will::Will, Property, ProtocolError as Error, QoS, Retain,
+    Property, ProtocolError as Error,
 };
 
 use bit_field::BitField;
@@ -19,14 +20,13 @@ pub fn integer_size(value: usize) -> usize {
     }
 }
 
-pub fn connect_message<'a, const S: usize>(
-    dest: &'a mut [u8],
+pub fn connect_message<'a, 'b>(
+    dest: &'b mut [u8],
     client_id: &[u8],
     keep_alive: u16,
-    properties: &[Property],
+    properties: &[Property<'a>],
     clean_start: bool,
-    will: Option<&Will<S>>,
-) -> Result<&'a [u8], Error> {
+) -> Result<&'b [u8], Error> {
     // Validate the properties for this packet.
     for property in properties {
         match property.id() {
@@ -49,18 +49,7 @@ pub fn connect_message<'a, const S: usize>(
 
     let mut packet = ReversedPacketWriter::new(dest);
 
-    if let Some(will) = will {
-        // Serialize the will data into the packet
-        packet.write(&will.payload)?;
-
-        // Update the flags for the will parameters. Indicate that the will is present, the QoS of
-        // the will message, and whether or not the will message should be retained.
-        flags.set_bit(2, true);
-        flags.set_bits(3..=4, will.qos as u8);
-        flags.set_bit(5, will.retain == Retain::Retained);
-    }
-
-    // Write the the client ID payload.
+    // Write the payload, which is the client ID.
     packet.write_binary_data(client_id)?;
 
     // Write the variable header.
@@ -83,7 +72,6 @@ pub fn publish_message<'a, 'b, 'c>(
     topic: &'a str,
     payload: &[u8],
     qos: QoS,
-    retain: Retain,
     id: u16,
     properties: &[Property<'c>],
 ) -> Result<&'b [u8], Error> {
@@ -118,9 +106,7 @@ pub fn publish_message<'a, 'b, 'c>(
     packet.write_utf8_string(topic)?;
 
     // Set qos to fixed header bits 1 and 2 in binary
-    let flags = *0u8
-        .set_bits(1..=2, qos as u8)
-        .set_bit(0, retain == Retain::Retained);
+    let flags = *0u8.set_bits(1..=2, qos as u8);
 
     // Write the fixed length header.
     packet.finalize(MessageType::Publish, flags)
@@ -168,16 +154,7 @@ pub fn serialize_publish() {
 
     let mut buffer: [u8; 900] = [0; 900];
     let payload: [u8; 2] = [0xAB, 0xCD];
-    let message = publish_message(
-        &mut buffer,
-        "ABC",
-        &payload,
-        QoS::AtMostOnce,
-        Retain::NotRetained,
-        0,
-        &[],
-    )
-    .unwrap();
+    let message = publish_message(&mut buffer, "ABC", &payload, QoS::AtMostOnce, 0, &[]).unwrap();
 
     assert_eq!(message, good_publish);
 }
@@ -195,16 +172,8 @@ pub fn serialize_publish_qos1() {
 
     let mut buffer: [u8; 900] = [0; 900];
     let payload: [u8; 2] = [0xAB, 0xCD];
-    let message = publish_message(
-        &mut buffer,
-        "ABC",
-        &payload,
-        QoS::AtLeastOnce,
-        Retain::NotRetained,
-        0xbeef,
-        &[],
-    )
-    .unwrap();
+    let message =
+        publish_message(&mut buffer, "ABC", &payload, QoS::AtLeastOnce, 0xbeef, &[]).unwrap();
 
     assert_eq!(message, good_publish);
 }
@@ -244,7 +213,6 @@ pub fn serialize_publish_with_properties() {
         "ABC",
         &payload,
         QoS::AtMostOnce,
-        Retain::NotRetained,
         0,
         &[Property::ResponseTopic("A")],
     )
@@ -267,46 +235,7 @@ fn serialize_connect() {
 
     let mut buffer: [u8; 900] = [0; 900];
     let client_id = "ABC".as_bytes();
-    let message = connect_message::<100>(&mut buffer, client_id, 10, &[], true, None).unwrap();
-
-    assert_eq!(message, good_serialized_connect)
-}
-
-#[test]
-fn serialize_connect_with_will() {
-    #[rustfmt::skip]
-    let good_serialized_connect: [u8; 28] = [
-        0x10, // Connect
-        26, // Remaining length
-
-        // Header: "MQTT5"
-        0x00, 0x04, 0x4d, 0x51, 0x54, 0x54, 0x05,
-
-        // Flags: Clean start, will present, will not retained, will QoS = 0,
-        0b0000_0110,
-
-        // Keep-alive: 10 seconds
-        0x00, 0x0a,
-
-        // Connected Properties: None
-        0x00,
-        // Client ID: "ABC"
-        0x00, 0x03, 0x41, 0x42, 0x43,
-        // Will properties: None
-        0x00,
-        // Will topic: "EFG"
-        0x00, 0x03, 0x45, 0x46, 0x47,
-        // Will payload: [0xAB, 0xCD]
-        0x00, 0x02, 0xAB, 0xCD,
-    ];
-
-    let mut buffer: [u8; 900] = [0; 900];
-    let client_id = "ABC".as_bytes();
-    let mut will = Will::<100>::new("EFG", &[0xAB, 0xCD], &[]).unwrap();
-    will.qos(QoS::AtMostOnce);
-    will.retained(Retain::NotRetained);
-
-    let message = connect_message(&mut buffer, client_id, 10, &[], true, Some(&will)).unwrap();
+    let message = connect_message(&mut buffer, client_id, 10, &[], true).unwrap();
 
     assert_eq!(message, good_serialized_connect)
 }
diff --git a/src/will.rs b/src/will.rs
deleted file mode 100644
index bc86afe..0000000
--- a/src/will.rs
+++ /dev/null
@@ -1,69 +0,0 @@
-use crate::{
-    properties::{Property, PropertyIdentifier},
-    ser::ReversedPacketWriter,
-    ProtocolError, QoS, Retain,
-};
-
-use heapless::Vec;
-
-pub struct Will<const MSG_SIZE: usize> {
-    pub payload: Vec<u8, MSG_SIZE>,
-    pub qos: QoS,
-    pub retain: Retain,
-}
-
-impl<const MSG_SIZE: usize> Will<MSG_SIZE> {
-    /// Construct a new will message.
-    ///
-    /// # Args
-    /// * `topic` - The topic to send the message on
-    /// * `data` - The message to transmit
-    /// * `properties` - Any properties to send with the will message.
-    pub fn new(topic: &str, data: &[u8], properties: &[Property]) -> Result<Self, ProtocolError> {
-        // Check that the input properties are valid for a will.
-        for property in properties {
-            match property.id() {
-                PropertyIdentifier::WillDelayInterval
-                | PropertyIdentifier::PayloadFormatIndicator
-                | PropertyIdentifier::MessageExpiryInterval
-                | PropertyIdentifier::ContentType
-                | PropertyIdentifier::ResponseTopic
-                | PropertyIdentifier::CorrelationData
-                | PropertyIdentifier::UserProperty => {}
-                _ => return Err(ProtocolError::InvalidProperty),
-            }
-        }
-
-        let mut buffer: [u8; MSG_SIZE] = [0; MSG_SIZE];
-        let mut packet = ReversedPacketWriter::new(&mut buffer);
-
-        // Serialize the will payload
-        packet.write_binary_data(data)?;
-        packet.write_utf8_string(topic)?;
-        packet.write_properties(properties)?;
-
-        Ok(Self {
-            qos: QoS::AtMostOnce,
-            retain: Retain::NotRetained,
-            // Note(unwrap): The vectro is declared as identical size to the vector, so it will
-            // always fit.
-            payload: Vec::from_slice(packet.finish()).unwrap(),
-        })
-    }
-
-    /// Set the retained status of the will.
-    ///
-    /// # Args
-    /// * `retained` - Specifies the retained state of the will.
-    pub fn retained(&mut self, retained: Retain) {
-        self.retain = retained;
-    }
-
-    /// Set the quality of service at which the will message is sent.
-    ///
-    /// # Args
-    /// * `qos` - The desired quality-of-service level to send the message at.
-    pub fn qos(&mut self, qos: QoS) {
-        self.qos = qos;
-    }
-}


Test patch:
diff --git a/tests/integration_test.rs b/tests/integration_test.rs
index 606224b..ad3758a 100644
--- a/tests/integration_test.rs
+++ b/tests/integration_test.rs
@@ -1,4 +1,4 @@
-use minimq::{Minimq, Property, QoS, Retain};
+use minimq::{Minimq, Property, QoS};
 
 use embedded_nal::{self, IpAddr, Ipv4Addr};
 use std_embedded_time::StandardClock;
@@ -16,16 +16,6 @@ fn main() -> std::io::Result<()> {
     let mut subscribed = false;
     let mut responses = 0;
 
-    mqtt.client
-        .set_will(
-            "exit",
-            "Test complete".as_bytes(),
-            QoS::AtMostOnce,
-            Retain::NotRetained,
-            &[],
-        )
-        .unwrap();
-
     loop {
         mqtt.poll(|client, topic, payload, properties| {
             println!("{} < {}", topic, core::str::from_utf8(payload).unwrap());
@@ -33,13 +23,7 @@ fn main() -> std::io::Result<()> {
             for property in properties {
                 if let Property::ResponseTopic(topic) = property {
                     client
-                        .publish(
-                            topic,
-                            "Pong".as_bytes(),
-                            QoS::AtMostOnce,
-                            Retain::NotRetained,
-                            &[],
-                        )
+                        .publish(topic, "Pong".as_bytes(), QoS::AtMostOnce, &[])
                         .unwrap();
                 }
             }
@@ -66,23 +50,11 @@ fn main() -> std::io::Result<()> {
                     println!("PUBLISH request");
                     let properties = [Property::ResponseTopic("response")];
                     mqtt.client
-                        .publish(
-                            "request",
-                            "Ping".as_bytes(),
-                            QoS::AtMostOnce,
-                            Retain::NotRetained,
-                            &properties,
-                        )
+                        .publish("request", "Ping".as_bytes(), QoS::AtMostOnce, &properties)
                         .unwrap();
 
                     mqtt.client
-                        .publish(
-                            "request",
-                            "Ping".as_bytes(),
-                            QoS::AtLeastOnce,
-                            Retain::NotRetained,
-                            &properties,
-                        )
+                        .publish("request", "Ping".as_bytes(), QoS::AtLeastOnce, &properties)
                         .unwrap();
 
                     // The message cannot be ack'd until the next poll call