Repository: quartiq/minimq (rust)

Issue:
Add Ping and keep-alive support
Add support for utilizing the keep-alive internal and the PingRequest/PingResponse pair. This will help to resolve issues in the TCP transport layer, like those in https://github.com/quartiq/stabilizer/issues/433.

Pull request:
Adding keep-alive support
This PR refactors #36 after a restructure of `minimq` and implements ping request/response support.

This PR fixes #51 

## Testing
The integraton test was modified to never exit and I monitored network traffic on the loopback. I observed the ping request/response at the expected 30s interval on wireshark.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2ab2dc2..96d3aed 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,8 @@ This document describes the changes to Minimq between releases.
 # Unreleased
 
 * Updating to `std-embedded-nal` v0.1 (dev dependency only; now again used for integration tests)
+* Added support for PingReq/PingResp to handle broken TCP connections and configuration of the
+keep-alive interval
 
 # Version 0.3.0
 Version 0.3.0 was published on 2021-08-06
diff --git a/Cargo.toml b/Cargo.toml
index e727242..a674b34 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,6 +20,7 @@ enum-iterator = "0.6.0"
 heapless = "0.7"
 log = {version = "0.4", optional = true}
 smlang = "0.4"
+embedded-time = "0.12"
 
 [dependencies.embedded-nal]
 version = "0.6"
diff --git a/README.md b/README.md
index 74cb7ab..be72f0c 100644
--- a/README.md
+++ b/README.md
@@ -15,5 +15,4 @@ There is also an example on a standard computer in `tests/integration_test.rs`
 ## Not yet implemented features.
 
 - Support all QoS levels
-- Implement keepalive timeouts
 - Allow batch subscriptions to multiple topics
diff --git a/src/de/deserialize.rs b/src/de/deserialize.rs
index 1832c32..5e298ee 100644
--- a/src/de/deserialize.rs
+++ b/src/de/deserialize.rs
@@ -45,6 +45,7 @@ pub enum ReceivedPacket<'a> {
     ConnAck(ConnAck<'a>),
     Publish(Pub<'a>),
     SubAck(SubAck<'a>),
+    PingResp,
 }
 
 impl<'a> ReceivedPacket<'a> {
@@ -84,6 +85,14 @@ impl<'a> ReceivedPacket<'a> {
                 Ok(ReceivedPacket::SubAck(parse_suback(packet_reader)?))
             }
 
+            MessageType::PingResp => {
+                if flags != 0 || remaining_length != 0 {
+                    return Err(Error::MalformedPacket);
+                }
+
+                Ok(ReceivedPacket::PingResp)
+            }
+
             _ => Err(Error::UnsupportedPacket),
         }
     }
@@ -210,4 +219,19 @@ mod test {
             _ => panic!("Invalid message"),
         }
     }
+
+    #[test]
+    fn deserialize_good_ping_resp() {
+        let mut serialized_ping_req: [u8; 2] = [
+            0xd0, // Ping resp
+            0x00, // Remaining length (0)
+        ];
+
+        let mut reader = PacketReader::<32>::from_serialized(&mut serialized_ping_req);
+        let ping_req = ReceivedPacket::parse_message(&mut reader).unwrap();
+        match ping_req {
+            ReceivedPacket::PingResp => {}
+            _ => panic!("Invalid message"),
+        }
+    }
 }
diff --git a/src/mqtt_client.rs b/src/mqtt_client.rs
index 892c853..37a27ab 100644
--- a/src/mqtt_client.rs
+++ b/src/mqtt_client.rs
@@ -2,15 +2,22 @@ use crate::{
     de::{deserialize::ReceivedPacket, PacketReader},
     ser::serialize,
     session_state::SessionState,
-    Property, {debug, error, info},
+    Property, {debug, error, info, warn},
 };
 
 use embedded_nal::{nb, IpAddr, SocketAddr, TcpClientStack};
+use embedded_time::{
+    duration::{Extensions, Seconds},
+    Clock,
+};
 
 use heapless::String;
 
 use core::str::FromStr;
 
+/// The default duration to wait for a ping response from the broker.
+const PING_TIMEOUT: embedded_time::duration::Seconds = embedded_time::duration::Seconds(5);
+
 /// The quality-of-service for an MQTT message.
 #[derive(Debug, Copy, Clone, PartialEq)]
 pub enum QoS {
@@ -25,7 +32,7 @@ pub enum QoS {
 }
 
 /// Possible errors encountered during an MQTT connection.
-#[derive(Debug, Copy, Clone, PartialEq)]
+#[derive(Debug, PartialEq)]
 pub enum Error<E> {
     Network(E),
     WriteFail,
@@ -35,6 +42,7 @@ pub enum Error<E> {
     Failed(u8),
     Protocol(ProtocolError),
     SessionReset,
+    Clock(embedded_time::clock::Error),
 }
 
 /// Errors that are specific to the MQTT protocol implementation.
@@ -53,9 +61,15 @@ pub enum ProtocolError {
     InvalidProperty,
 }
 
-impl<E> From<E> for Error<E> {
-    fn from(e: E) -> Error<E> {
-        Error::Network(e)
+impl<E> From<embedded_time::clock::Error> for Error<E> {
+    fn from(clock: embedded_time::clock::Error) -> Self {
+        Error::Clock(clock)
+    }
+}
+
+impl<E> From<ProtocolError> for Error<E> {
+    fn from(error: ProtocolError) -> Self {
+        Error::Protocol(error)
     }
 }
 
@@ -81,29 +95,35 @@ mod sm {
 use sm::{Context, Events, StateMachine, States};
 
 /// The general structure for managing MQTT via Minimq.
-pub struct Minimq<N, const T: usize>
+pub struct Minimq<N, C, const T: usize>
 where
     N: TcpClientStack,
+    C: Clock,
 {
-    pub client: MqttClient<N, T>,
+    pub client: MqttClient<N, C, T>,
     packet_reader: PacketReader<T>,
 }
 
 /// A client for interacting with an MQTT Broker.
-pub struct MqttClient<N: TcpClientStack, const T: usize> {
+pub struct MqttClient<N: TcpClientStack, C: Clock, const T: usize> {
     network_stack: N,
+    clock: C,
     socket: Option<N::TcpSocket>,
-    session_state: SessionState,
+    session_state: SessionState<C>,
     connection_state: StateMachine<Context>,
 }
 
-impl<N, const T: usize> MqttClient<N, T>
+impl<N, C, const T: usize> MqttClient<N, C, T>
 where
     N: TcpClientStack,
+    C: Clock,
 {
     fn process(&mut self) -> Result<(), Error<N::Error>> {
         // Potentially update the state machine depending on the current socket connection status.
-        if !self.socket_is_connected()? {
+        if !self
+            .socket_is_connected()
+            .map_err(|err| Error::Network(err))?
+        {
             self.connection_state.process_event(Events::Disconnect).ok();
         } else {
             self.connection_state.process_event(Events::Connect).ok();
@@ -112,12 +132,18 @@ where
         match self.connection_state.state() {
             // In the RESTART state, we need to reopen the TCP socket.
             &States::Restart => {
-                if self.socket.is_some() {
-                    self.network_stack.close(self.socket.take().unwrap())?;
+                if let Some(socket) = self.socket.take() {
+                    self.network_stack
+                        .close(socket)
+                        .map_err(|err| Error::Network(err))?;
                 }
 
                 // Allocate a new socket to use and begin connecting it.
-                self.socket.replace(self.network_stack.socket()?);
+                self.socket.replace(
+                    self.network_stack
+                        .socket()
+                        .map_err(|err| Error::Network(err))?,
+                );
                 self.connection_state
                     .process_event(Events::GotSocket)
                     .unwrap();
@@ -151,12 +177,14 @@ where
                 let packet = serialize::connect_message(
                     &mut buffer,
                     self.session_state.client_id.as_str().as_bytes(),
-                    self.session_state.keep_alive_interval,
+                    self.session_state
+                        .keep_alive_interval
+                        .unwrap_or(0.seconds())
+                        .0 as u16,
                     &properties,
                     // Only perform a clean start if we do not have any session state.
                     !self.session_state.is_present(),
-                )
-                .map_err(|e| Error::Protocol(e))?;
+                )?;
 
                 info!("Sending CONNECT");
                 self.write(packet)?;
@@ -169,6 +197,8 @@ where
             _ => {}
         }
 
+        self.handle_timers()?;
+
         Ok(())
     }
 
@@ -195,11 +225,44 @@ where
                 if written != buf.len() {
                     Err(Error::WriteFail)
                 } else {
+                    self.session_state
+                        .register_transmission(self.clock.try_now()?);
                     Ok(())
                 }
             })
     }
 
+    /// Configure the MQTT keep-alive interval.
+    ///
+    /// # Note
+    /// This must be completed before connecting to a broker.
+    ///
+    /// # Note
+    /// The broker may override the requested keep-alive interval. Any value requested by the
+    /// broker will be used instead.
+    ///
+    /// # Args
+    /// * `interval` - The keep-alive interval in seconds. A ping will be transmitted if no other
+    /// messages are sent within 50% of the keep-alive interval.
+    pub fn set_keepalive_interval(
+        &mut self,
+        interval: impl Into<Seconds<u32>>,
+    ) -> Result<(), Error<N::Error>> {
+        let interval = interval.into();
+        if interval.0 > u16::MAX as u32 {
+            return Err(ProtocolError::Invalid.into());
+        }
+        match self.connection_state.state() {
+            &States::Active => Err(Error::NotReady),
+            _ => {
+                self.session_state
+                    .keep_alive_interval
+                    .replace(interval.into());
+                Ok(())
+            }
+        }
+    }
+
     /// Subscribe to a topic.
     ///
     /// # Note
@@ -221,8 +284,7 @@ where
         let packet_id = self.session_state.get_packet_identifier();
 
         let mut buffer: [u8; T] = [0; T];
-        let packet = serialize::subscribe_message(&mut buffer, topic, packet_id, properties)
-            .map_err(|e| Error::Protocol(e))?;
+        let packet = serialize::subscribe_message(&mut buffer, topic, packet_id, properties)?;
 
         self.write(packet).and_then(|_| {
             info!("Subscribing to `{}`: {}", topic, packet_id);
@@ -248,7 +310,10 @@ where
     /// # Returns
     /// True if the client is connected to the broker.
     pub fn is_connected(&mut self) -> Result<bool, Error<N::Error>> {
-        Ok(self.socket_is_connected()? && self.session_state.connected)
+        Ok(self
+            .socket_is_connected()
+            .map_err(|err| Error::Network(err))?
+            && self.session_state.connected)
     }
 
     /// Publish a message over MQTT.
@@ -286,8 +351,7 @@ where
         );
 
         let mut buffer: [u8; T] = [0; T];
-        let packet = serialize::publish_message(&mut buffer, topic, data, properties)
-            .map_err(|e| Error::Protocol(e))?;
+        let packet = serialize::publish_message(&mut buffer, topic, data, properties)?;
         self.write(packet)
     }
 
@@ -310,7 +374,7 @@ where
         f: &mut F,
     ) -> Result<(), Error<N::Error>>
     where
-        F: FnMut(&mut MqttClient<N, T>, &'a str, &[u8], &[Property<'a>]),
+        F: FnMut(&mut MqttClient<N, C, T>, &'a str, &[u8], &[Property<'a>]),
     {
         if !self.session_state.connected {
             if let ReceivedPacket::ConnAck(acknowledge) = packet {
@@ -344,7 +408,9 @@ where
                                 String::from_str(id).or(Err(Error::ProvidedClientIdTooLong))?;
                         }
                         Property::ServerKeepAlive(keep_alive) => {
-                            self.session_state.keep_alive_interval = keep_alive;
+                            self.session_state
+                                .keep_alive_interval
+                                .replace((keep_alive as u32).seconds());
                         }
                         _prop => info!("Ignoring property: {:?}", _prop),
                     };
@@ -391,12 +457,55 @@ where
                 Ok(())
             }
 
+            ReceivedPacket::PingResp => {
+                // Cancel the ping response timeout.
+                if self.session_state.ping_timeout.is_some() {
+                    self.session_state.ping_timeout = None;
+                } else {
+                    warn!("Got unexpected ping response");
+                }
+
+                Ok(())
+            }
+
             _ => Err(Error::Unsupported),
         }
     }
+
+    fn handle_timers(&mut self) -> Result<(), Error<N::Error>> {
+        // If we are not connected, there's no session state to manage.
+        if self.connection_state.state() != &States::Active {
+            return Ok(());
+        }
+
+        let now = self.clock.try_now()?;
+
+        if let Some(timeout) = self.session_state.ping_timeout {
+            if now > timeout {
+                // Reset network connection.
+                self.connection_state.process_event(Events::Disconnect).ok();
+            }
+        } else {
+            // Check if we need to ping the server.
+            if self.session_state.ping_is_due(&now) {
+                let mut buffer: [u8; T] = [0; T];
+
+                // Note: The ping timeout is set at this point so that it's running even if we fail
+                // to write the ping message. This is intentional incase the underlying transport
+                // mechanism has stalled. The ping timeout will then allow us to recover the
+                // underlying TCP connection.
+                self.session_state.ping_timeout.replace(now + PING_TIMEOUT);
+
+                let packet = serialize::ping_req_message(&mut buffer)?;
+                self.write(packet)?;
+            }
+        }
+
+        Ok(())
+    }
 }
 
-impl<N: TcpClientStack, const T: usize> Minimq<N, T> {
+impl<N: TcpClientStack, C: Clock, const T: usize> Minimq<N, C, T> {
     /// Construct a new MQTT interface.
     ///
     /// # Args
@@ -404,11 +513,17 @@ impl<N: TcpClientStack, const T: usize> Minimq<N, T> {
     /// * `client_id` The client ID to use for communicating with the broker. If empty, rely on the
     ///   broker to automatically assign a client ID.
     /// * `network_stack` - The network stack to use for communication.
+    /// * `clock` - The clock to use for managing MQTT state timing.
     ///
     /// # Returns
     /// A `Minimq` object that can be used for publishing messages, subscribing to topics, and
     /// managing the MQTT state.
-    pub fn new(broker: IpAddr, client_id: &str, network_stack: N) -> Result<Self, Error<N::Error>> {
+    pub fn new(
+        broker: IpAddr,
+        client_id: &str,
+        network_stack: N,
+        clock: C,
+    ) -> Result<Self, Error<N::Error>> {
         let session_state = SessionState::new(
             broker,
             String::from_str(client_id).or(Err(Error::ProvidedClientIdTooLong))?,
@@ -417,6 +532,7 @@ impl<N: TcpClientStack, const T: usize> Minimq<N, T> {
         let minimq = Minimq {
             client: MqttClient {
                 network_stack,
+                clock,
                 socket: None,
                 session_state,
                 connection_state: StateMachine::new(Context),
@@ -434,7 +550,7 @@ impl<N: TcpClientStack, const T: usize> Minimq<N, T> {
     /// topic, message, and list of proprties (in that order).
     pub fn poll<F>(&mut self, mut f: F) -> Result<(), Error<N::Error>>
     where
-        for<'a> F: FnMut(&mut MqttClient<N, T>, &'a str, &[u8], &[Property<'a>]),
+        for<'a> F: FnMut(&mut MqttClient<N, C, T>, &'a str, &[u8], &[Property<'a>]),
     {
         self.client.process()?;
 
@@ -468,16 +584,13 @@ impl<N: TcpClientStack, const T: usize> Minimq<N, T> {
 
             // Handle any received packets.
             while self.packet_reader.packet_available() {
-                let packet = ReceivedPacket::parse_message(&self.packet_reader)
-                    .map_err(|e| Error::Protocol(e))?;
+                let packet = ReceivedPacket::parse_message(&self.packet_reader)?;
 
                 info!("Received {:?}", packet);
 
                 let result = self.client.handle_packet(packet, &mut f);
 
-                self.packet_reader
-                    .pop_packet()
-                    .map_err(|e| Error::Protocol(e))?;
+                self.packet_reader.pop_packet()?;
 
                 // If there was an error, return it now. Note that we ensure the packet is removed
                 // from buffering after processing even in error conditions..
diff --git a/src/ser/serialize.rs b/src/ser/serialize.rs
index ca43fdf..e5af88e 100644
--- a/src/ser/serialize.rs
+++ b/src/ser/serialize.rs
@@ -62,6 +62,10 @@ pub fn connect_message<'a, 'b>(
     packet.finalize(MessageType::Connect, 0)
 }
 
+pub fn ping_req_message<'a>(dest: &'a mut [u8]) -> Result<&'a [u8], Error> {
+    ReversedPacketWriter::new(dest).finalize(MessageType::PingReq, 0x00)
+}
+
 pub fn publish_message<'a, 'b, 'c>(
     dest: &'b mut [u8],
     topic: &'a str,
@@ -205,3 +209,14 @@ fn serialize_connect() {
 
     assert_eq!(message, good_serialized_connect)
 }
+
+#[test]
+fn serialize_ping_req() {
+    let good_ping_req: [u8; 2] = [
+        0xc0, // Ping
+        0x00, // Remaining length (0)
+    ];
+
+    let mut buffer: [u8; 1024] = [0; 1024];
+    assert_eq!(ping_req_message(&mut buffer).unwrap(), good_ping_req);
+}
diff --git a/src/session_state.rs b/src/session_state.rs
index e8b1366..c75d3da 100644
--- a/src/session_state.rs
+++ b/src/session_state.rs
@@ -4,27 +4,36 @@
 use embedded_nal::IpAddr;
 use heapless::{String, Vec};
 
-pub struct SessionState {
+use embedded_time::{
+    duration::{Extensions, Seconds},
+    Clock, Instant,
+};
+
+pub struct SessionState<C: Clock> {
     // Indicates that we are connected to a broker.
     pub connected: bool,
-    pub keep_alive_interval: u16,
+    pub keep_alive_interval: Option<Seconds<u32>>,
+    pub ping_timeout: Option<Instant<C>>,
     pub broker: IpAddr,
     pub maximum_packet_size: Option<u32>,
     pub client_id: String<64>,
+    last_transmission: Option<Instant<C>>,
     pub pending_subscriptions: Vec<u16, 32>,
     packet_id: u16,
     active: bool,
 }
 
-impl SessionState {
-    pub fn new<'a>(broker: IpAddr, id: String<64>) -> SessionState {
+impl<C: Clock> SessionState<C> {
+    pub fn new<'a>(broker: IpAddr, id: String<64>) -> SessionState<C> {
         SessionState {
             connected: false,
             active: false,
+            ping_timeout: None,
             broker,
             client_id: id,
             packet_id: 1,
-            keep_alive_interval: 0,
+            keep_alive_interval: Some(59.seconds()),
+            last_transmission: None,
             pending_subscriptions: Vec::new(),
             maximum_packet_size: None,
         }
@@ -34,9 +43,10 @@ impl SessionState {
         self.active = false;
         self.connected = false;
         self.packet_id = 1;
-        self.keep_alive_interval = 0;
+        self.keep_alive_interval = Some(59.seconds());
         self.maximum_packet_size = None;
         self.pending_subscriptions.clear();
+        self.last_transmission = None;
     }
 
     /// Called whenever an active connection has been made with a broker.
@@ -63,4 +73,17 @@ impl SessionState {
             self.packet_id = result;
         }
     }
+
+    pub fn register_transmission(&mut self, now: Instant<C>) {
+        self.last_transmission = Some(now);
+    }
+
+    pub fn ping_is_due(&self, now: &Instant<C>) -> bool {
+        // Send a ping if we haven't sent a transmission in the last 50% of the keepalive internal.
+        self.keep_alive_interval
+            .zip(self.last_transmission)
+            .map_or(false, |(interval, last)| {
+                *now > last + 500.milliseconds() * interval.0
+            })
+    }
 }


Test patch:
diff --git a/tests/integration_test.rs b/tests/integration_test.rs
index e871ea8..7558a0b 100644
--- a/tests/integration_test.rs
+++ b/tests/integration_test.rs
@@ -1,6 +1,35 @@
 use minimq::{Minimq, Property, QoS};
 
 use embedded_nal::{self, IpAddr, Ipv4Addr};
+use embedded_time::{fraction::Fraction, Clock};
+use std::time::Instant;
+
+#[derive(Default)]
+struct StdClock {
+    start: core::cell::UnsafeCell<Option<Instant>>,
+}
+
+impl Clock for StdClock {
+    type T = u32;
+
+    const SCALING_FACTOR: Fraction = Fraction::new(1, 1_000);
+
+    fn try_now(&self) -> Result<embedded_time::Instant<Self>, embedded_time::clock::Error> {
+        let std_now = Instant::now();
+        let start = unsafe {
+            if (*self.start.get()).is_none() {
+                (*self.start.get()).replace(std_now);
+                std_now
+            } else {
+                (*self.start.get()).unwrap()
+            }
+        };
+
+        let elapsed = std_now - start;
+
+        Ok(embedded_time::Instant::new(elapsed.as_millis() as u32))
+    }
+}
 
 #[test]
 fn main() -> std::io::Result<()> {
@@ -8,7 +37,7 @@ fn main() -> std::io::Result<()> {
 
     let stack = std_embedded_nal::Stack::default();
     let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
-    let mut mqtt = Minimq::<_, 256>::new(localhost, "", stack).unwrap();
+    let mut mqtt = Minimq::<_, _, 256>::new(localhost, "", stack, StdClock::default()).unwrap();
 
     let mut published = false;
     let mut subscribed = false;