Repository: quartiq/minimq
Pull request #45 (2021-05-27 11:00:39)

Feature/nal update
This PR does the following:
* Update embedded-nal version to v0.6
* Remove all usage of `RefCells`
* Refactor of the API to have an `Minimq` struct and an `MqttClient`
      * this was done to resolve ownership issues with using content in the received message buffer to publish/subscribe
  * Updating to utilize const-generics.


Patch:
diff --git a/Cargo.toml b/Cargo.toml
index 4f289d0..49cb24a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,10 +17,11 @@ license = "MIT"
 [dependencies]
 bit_field = "0.10.0"
 enum-iterator = "0.6.0"
-heapless = "0.6.1"
-generic-array = ">=0.14.4"
+heapless = "0.7"
 log = {version = "0.4", optional = true}
-embedded-nal = "0.1.0"
+
+[dependencies.embedded-nal]
+version = "0.6"
 
 [features]
 default = []
diff --git a/src/de/deserialize.rs b/src/de/deserialize.rs
index aa699d4..1832c32 100644
--- a/src/de/deserialize.rs
+++ b/src/de/deserialize.rs
@@ -2,8 +2,7 @@ use crate::{
     de::PacketReader, message_types::MessageType, mqtt_client::ProtocolError as Error, Property,
 };
 use bit_field::BitField;
-use generic_array::ArrayLength;
-use heapless::{consts, Vec};
+use heapless::Vec;
 
 #[derive(Debug)]
 pub struct ConnAck<'a> {
@@ -14,7 +13,7 @@ pub struct ConnAck<'a> {
     pub reason_code: u8,
 
     /// A list of properties associated with the connection.
-    pub properties: Vec<Property<'a>, consts::U8>,
+    pub properties: Vec<Property<'a>, 8>,
 }
 
 #[derive(Debug)]
@@ -23,7 +22,10 @@ pub struct Pub<'a> {
     pub topic: &'a str,
 
     /// The properties transmitted with the publish data.
-    pub properties: Vec<Property<'a>, consts::U8>,
+    pub properties: Vec<Property<'a>, 8>,
+
+    /// The payload of the message.
+    pub payload: &'a [u8],
 }
 
 #[derive(Debug)]
@@ -35,7 +37,7 @@ pub struct SubAck<'a> {
     pub reason_code: u8,
 
     /// A list of properties associated with the subscription.
-    pub properties: Vec<Property<'a>, consts::U8>,
+    pub properties: Vec<Property<'a>, 8>,
 }
 
 #[derive(Debug)]
@@ -53,7 +55,7 @@ impl<'a> ReceivedPacket<'a> {
     ///
     /// # Returns
     /// A packet describing the received content.
-    pub(crate) fn parse_message<'reader: 'a, T: ArrayLength<u8>>(
+    pub(crate) fn parse_message<'reader: 'a, const T: usize>(
         packet_reader: &'reader PacketReader<T>,
     ) -> Result<ReceivedPacket<'a>, Error> {
         let (message_type, flags, remaining_length) = packet_reader.read_fixed_header()?;
@@ -87,7 +89,7 @@ impl<'a> ReceivedPacket<'a> {
     }
 }
 
-fn parse_connack<T: ArrayLength<u8>>(p: &PacketReader<T>) -> Result<ConnAck, Error> {
+fn parse_connack<const T: usize>(p: &PacketReader<T>) -> Result<ConnAck, Error> {
     // Read the connect acknowledgement flags.
     let flags = p.read_u8()?;
     if flags != 0 && flags != 1 {
@@ -108,7 +110,7 @@ fn parse_connack<T: ArrayLength<u8>>(p: &PacketReader<T>) -> Result<ConnAck, Err
     })
 }
 
-fn parse_publish<'a, 'reader: 'a, T: ArrayLength<u8>>(
+fn parse_publish<'a, 'reader: 'a, const T: usize>(
     p: &'reader PacketReader<T>,
 ) -> Result<Pub<'a>, Error> {
     let topic = p.read_utf8_string()?;
@@ -116,12 +118,16 @@ fn parse_publish<'a, 'reader: 'a, T: ArrayLength<u8>>(
     let properties = p.read_properties()?;
     // TODO: Validate properties associated with this message.
 
-    // Note that we intentionally don't read the payload from the data reader so that it is
-    // available later to be borrowed directly to the handler for the payload data.
-    Ok(Pub { topic, properties })
+    let payload = p.payload()?;
+
+    Ok(Pub {
+        topic,
+        properties,
+        payload,
+    })
 }
 
-fn parse_suback<T: ArrayLength<u8>>(p: &PacketReader<T>) -> Result<SubAck, Error> {
+fn parse_suback<const T: usize>(p: &PacketReader<T>) -> Result<SubAck, Error> {
     // Read the variable length header.
     let id = p.read_u16()?;
 
@@ -142,7 +148,6 @@ fn parse_suback<T: ArrayLength<u8>>(p: &PacketReader<T>) -> Result<SubAck, Error
 #[cfg(test)]
 mod test {
     use super::{PacketReader, ReceivedPacket};
-    use generic_array::typenum;
 
     #[test]
     fn deserialize_good_connack() {
@@ -154,7 +159,7 @@ mod test {
                   // No payload.
         ];
 
-        let mut reader = PacketReader::<typenum::U32>::from_serialized(&mut serialized_connack);
+        let mut reader = PacketReader::<32>::from_serialized(&mut serialized_connack);
         let connack = ReceivedPacket::parse_message(&mut reader).unwrap();
         match connack {
             ReceivedPacket::ConnAck(conn_ack) => {
@@ -175,7 +180,7 @@ mod test {
             0x05, // Payload
         ];
 
-        let mut reader = PacketReader::<typenum::U32>::from_serialized(&mut serialized_publish);
+        let mut reader = PacketReader::<32>::from_serialized(&mut serialized_publish);
         let publish = ReceivedPacket::parse_message(&mut reader).unwrap();
         match publish {
             ReceivedPacket::Publish(pub_info) => {
@@ -195,7 +200,7 @@ mod test {
             0x02, // Response Code
         ];
 
-        let mut reader = PacketReader::<typenum::U32>::from_serialized(&mut serialized_suback);
+        let mut reader = PacketReader::<32>::from_serialized(&mut serialized_suback);
         let suback = ReceivedPacket::parse_message(&mut reader).unwrap();
         match suback {
             ReceivedPacket::SubAck(sub_ack) => {
diff --git a/src/de/packet_reader.rs b/src/de/packet_reader.rs
index b3e0a3e..85095e5 100644
--- a/src/de/packet_reader.rs
+++ b/src/de/packet_reader.rs
@@ -4,42 +4,34 @@ use crate::{
     Property, {debug, warn},
 };
 use bit_field::BitField;
-use generic_array::{ArrayLength, GenericArray};
-use heapless::{consts, Vec};
+use heapless::Vec;
 
 // The maximum size of the fixed header. This is calculated as the type+flag byte and the maximum
 // variable length integer size (4).
 const FIXED_HEADER_MAX: usize = 5;
 
-pub(crate) struct PacketReader<T: ArrayLength<u8>> {
-    pub buffer: GenericArray<u8, T>,
+pub(crate) struct PacketReader<const T: usize> {
+    pub buffer: [u8; T],
     read_bytes: usize,
     packet_length: Option<usize>,
     index: core::cell::RefCell<usize>,
 }
 
-impl<T> PacketReader<T>
-where
-    T: ArrayLength<u8>,
-{
+impl<const T: usize> PacketReader<T> {
     pub fn new() -> PacketReader<T> {
         PacketReader {
-            buffer: GenericArray::default(),
+            buffer: [0; T],
             read_bytes: 0,
             packet_length: None,
             index: core::cell::RefCell::new(0),
         }
     }
 
-    pub fn maximum_packet_length(&self) -> usize {
-        self.buffer.len()
-    }
-
     #[cfg(test)]
     pub fn from_serialized<'a>(buffer: &'a mut [u8]) -> PacketReader<T> {
         let len = buffer.len();
         let mut reader = PacketReader {
-            buffer: GenericArray::default(),
+            buffer: [0; T],
             read_bytes: len,
             packet_length: None,
             index: core::cell::RefCell::new(0),
@@ -148,8 +140,8 @@ where
         Ok(byte[0])
     }
 
-    pub fn read_properties<'a, 'me: 'a>(&'me self) -> Result<Vec<Property<'a>, consts::U8>, Error> {
-        let mut properties: Vec<Property, consts::U8> = Vec::new();
+    pub fn read_properties<'a, 'me: 'a>(&'me self) -> Result<Vec<Property<'a>, 8>, Error> {
+        let mut properties: Vec<Property, 8> = Vec::new();
 
         let properties_size = self.read_variable_length_integer()?;
         let mut property_bytes_processed = 0;
diff --git a/src/lib.rs b/src/lib.rs
index 3f14a19..2975111 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -70,9 +70,7 @@ use message_types::MessageType;
 pub use properties::Property;
 
 pub use embedded_nal;
-pub use generic_array;
-pub use generic_array::typenum::consts;
-pub use mqtt_client::{Error, MqttClient, ProtocolError, QoS};
+pub use mqtt_client::{Error, Minimq, ProtocolError, QoS};
 
 #[cfg(feature = "logging")]
 pub(crate) use log::{debug, error, info, warn};
diff --git a/src/mqtt_client.rs b/src/mqtt_client.rs
index 944cd62..e7159ee 100644
--- a/src/mqtt_client.rs
+++ b/src/mqtt_client.rs
@@ -5,27 +5,11 @@ use crate::{
     Property, {debug, error, info},
 };
 
-use core::{cell::RefCell, str::FromStr};
+use embedded_nal::{nb, IpAddr, SocketAddr, TcpClientStack};
 
-use embedded_nal::{nb, IpAddr, Mode, SocketAddr};
-
-use generic_array::{ArrayLength, GenericArray};
 use heapless::String;
 
-/// A client for interacting with an MQTT Broker.
-pub struct MqttClient<T, N>
-where
-    T: ArrayLength<u8>,
-    N: embedded_nal::TcpStack,
-{
-    /// The network stack originally provided to the client.
-    pub network_stack: N,
-
-    socket: RefCell<Option<N::TcpSocket>>,
-    session_state: SessionState,
-    packet_reader: PacketReader<T>,
-    connect_sent: bool,
-}
+use core::str::FromStr;
 
 /// The quality-of-service for an MQTT message.
 #[derive(Debug, Copy, Clone, PartialEq)]
@@ -75,74 +59,53 @@ impl<E> From<E> for Error<E> {
     }
 }
 
-impl<T, N> MqttClient<T, N>
+/// The general structure for managing MQTT via Minimq.
+pub struct Minimq<N, const T: usize>
 where
-    T: ArrayLength<u8>,
-    N: embedded_nal::TcpStack,
+    N: TcpClientStack,
 {
-    /// Construct a new MQTT client.
-    ///
-    /// # Args
-    /// * `broker` - The IP address of the broker to connect to.
-    /// * `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.
-    ///
-    /// # Returns
-    /// An `MqttClient` that can be used for publishing messages and subscribing to topics.
-    pub fn new(broker: IpAddr, client_id: &str, network_stack: N) -> Result<Self, Error<N::Error>> {
-        let session_state = SessionState::new(
-            broker,
-            String::from_str(client_id).or(Err(Error::ProvidedClientIdTooLong))?,
-        );
-
-        let mut client = MqttClient {
-            network_stack: network_stack,
-            socket: RefCell::new(None),
-            session_state,
-            packet_reader: PacketReader::new(),
-            connect_sent: false,
-        };
-
-        client.reset()?;
-
-        Ok(client)
-    }
-
-    fn read(&self, mut buf: &mut [u8]) -> Result<usize, Error<N::Error>> {
-        let mut socket_ref = self.socket.borrow_mut();
-        let mut socket = socket_ref.take().unwrap();
-        let read = match self.network_stack.read(&mut socket, &mut buf) {
-            Ok(count) => count,
-            Err(nb::Error::WouldBlock) => 0,
-            Err(nb::Error::Other(error)) => return Err(Error::Network(error)),
-        };
+    pub client: MqttClient<N, T>,
+    packet_reader: PacketReader<T>,
+}
 
-        // Put the socket back into the option.
-        socket_ref.replace(socket);
+/// A client for interacting with an MQTT Broker.
+pub struct MqttClient<N: TcpClientStack, const T: usize> {
+    network_stack: N,
+    socket: Option<N::TcpSocket>,
+    session_state: SessionState,
+    connect_sent: bool,
+}
 
-        Ok(read)
+impl<N, const T: usize> MqttClient<N, T>
+where
+    N: TcpClientStack,
+{
+    fn read(&mut self, mut buf: &mut [u8]) -> Result<usize, Error<N::Error>> {
+        // Atomically access the socket.
+        let socket = self.socket.as_mut().unwrap();
+        let result = self.network_stack.receive(socket, &mut buf);
+
+        result.or_else(|err| match err {
+            nb::Error::WouldBlock => Ok(0),
+            nb::Error::Other(err) => Err(Error::Network(err)),
+        })
     }
 
-    fn write(&self, buf: &[u8]) -> Result<(), Error<N::Error>> {
-        let mut socket_ref = self.socket.borrow_mut();
-        let mut socket = socket_ref.take().unwrap();
-        let written = self
-            .network_stack
-            .write(&mut socket, &buf)
+    fn write(&mut self, buf: &[u8]) -> Result<(), Error<N::Error>> {
+        let socket = self.socket.as_mut().unwrap();
+        self.network_stack
+            .send(socket, &buf)
             .map_err(|err| match err {
                 nb::Error::WouldBlock => Error::WriteFail,
                 nb::Error::Other(err) => Error::Network(err),
-            })?;
-
-        // Put the socket back into the option.
-        socket_ref.replace(socket);
-
-        if written != buf.len() {
-            Err(Error::WriteFail)
-        } else {
-            Ok(())
-        }
+            })
+            .and_then(|written| {
+                if written != buf.len() {
+                    Err(Error::WriteFail)
+                } else {
+                    Ok(())
+                }
+            })
     }
 
     /// Subscribe to a topic.
@@ -165,7 +128,7 @@ where
 
         let packet_id = self.session_state.get_packet_identifier();
 
-        let mut buffer: GenericArray<u8, T> = GenericArray::default();
+        let mut buffer: [u8; T] = [0; T];
         let packet = serialize::subscribe_message(&mut buffer, topic, packet_id, properties)
             .map_err(|e| Error::Protocol(e))?;
 
@@ -192,7 +155,7 @@ where
     ///
     /// # Returns
     /// True if the client is connected to the broker.
-    pub fn is_connected(&self) -> Result<bool, Error<N::Error>> {
+    pub fn is_connected(&mut self) -> Result<bool, Error<N::Error>> {
         Ok(self.socket_is_connected()? && self.session_state.connected)
     }
 
@@ -211,7 +174,7 @@ where
     /// * `properties` - A list of properties to associate with the message being published. May be
     ///   empty.
     pub fn publish<'a, 'b>(
-        &self,
+        &mut self,
         topic: &'a str,
         data: &[u8],
         qos: QoS,
@@ -230,46 +193,37 @@ where
             topic, data, properties
         );
 
-        let mut buffer: GenericArray<u8, T> = GenericArray::default();
+        let mut buffer: [u8; T] = [0; T];
         let packet = serialize::publish_message(&mut buffer, topic, data, properties)
             .map_err(|e| Error::Protocol(e))?;
         self.write(packet)
     }
 
-    fn socket_is_connected(&self) -> Result<bool, N::Error> {
-        let mut socket_ref = self.socket.borrow_mut();
-
-        if socket_ref.is_none() {
-            // Attempt to open a socket from the network stack.
-            socket_ref.replace(self.network_stack.open(Mode::NonBlocking)?);
+    fn socket_is_connected(&mut self) -> Result<bool, N::Error> {
+        if self.socket.is_none() {
+            return Ok(false);
         }
 
-        let socket = socket_ref.take().unwrap();
-        let result = self.network_stack.is_connected(&socket);
-        socket_ref.replace(socket);
-
-        result
+        let socket = self.socket.as_ref().unwrap();
+        self.network_stack.is_connected(socket)
     }
 
-    fn reset(&mut self) -> Result<(), Error<N::Error>> {
-        self.packet_reader.reset();
+    fn reset(&mut self) {
         self.connect_sent = false;
         self.session_state.connected = false;
-
-        Ok(())
     }
 
     fn connect_to_broker(&mut self) -> Result<(), Error<N::Error>> {
-        self.reset()?;
+        self.reset();
 
         let properties = [
             // Tell the broker our maximum packet size.
-            Property::MaximumPacketSize(self.packet_reader.maximum_packet_length() as u32),
+            Property::MaximumPacketSize(T as u32),
             // The session does not expire.
             Property::SessionExpiryInterval(u32::MAX),
         ];
 
-        let mut buffer: GenericArray<u8, T> = GenericArray::default();
+        let mut buffer: [u8; T] = [0; T];
         let packet = serialize::connect_message(
             &mut buffer,
             self.session_state.client_id.as_str().as_bytes(),
@@ -288,15 +242,32 @@ where
         Ok(())
     }
 
-    fn handle_packet<F>(&mut self, f: &mut F) -> Result<(), Error<N::Error>>
-    where
-        for<'a> F: FnMut(&Self, &'a str, &[u8], &[Property<'a>]),
-    {
-        let packet =
-            ReceivedPacket::parse_message(&self.packet_reader).map_err(|e| Error::Protocol(e))?;
+    fn connect_socket(&mut self) -> Result<(), Error<N::Error>> {
+        if self.socket.is_none() {
+            let socket = self.network_stack.socket()?;
+            self.socket.replace(socket);
+        }
 
-        info!("Received {:?}", packet);
+        let socket = self.socket.as_mut().unwrap();
+
+        // Connect to the broker's TCP port with a new socket.
+        // TODO: Limit the time between connect attempts to prevent network spam.
+        self.network_stack
+            .connect(socket, SocketAddr::new(self.session_state.broker, 1883))
+            .map_err(|err| match err {
+                nb::Error::WouldBlock => Error::WriteFail,
+                nb::Error::Other(err) => Error::Network(err),
+            })
+    }
 
+    fn handle_packet<'a, F>(
+        &mut self,
+        packet: ReceivedPacket<'a>,
+        f: &mut F,
+    ) -> Result<(), Error<N::Error>>
+    where
+        F: FnMut(&mut MqttClient<N, T>, &'a str, &[u8], &[Property<'a>]),
+    {
         if !self.session_state.connected {
             if let ReceivedPacket::ConnAck(acknowledge) = packet {
                 let mut result = Ok(());
@@ -350,11 +321,7 @@ where
         match packet {
             ReceivedPacket::Publish(info) => {
                 // Call a handler function to deal with the received data.
-                let payload = self
-                    .packet_reader
-                    .payload()
-                    .map_err(|e| Error::Protocol(e))?;
-                f(self, info.topic, payload, &info.properties);
+                f(self, info.topic, info.payload, &info.properties);
 
                 Ok(())
             }
@@ -383,21 +350,37 @@ where
             _ => Err(Error::Unsupported),
         }
     }
+}
 
-    fn connect_socket(&mut self) -> Result<(), Error<N::Error>> {
-        let mut socket_ref = self.socket.borrow_mut();
-        let socket = socket_ref.take().unwrap();
-
-        // Connect to the broker's TCP port with a new socket.
-        // TODO: Limit the time between connect attempts to prevent network spam.
-        let socket = self
-            .network_stack
-            .connect(socket, SocketAddr::new(self.session_state.broker, 1883))?;
+impl<N: TcpClientStack, const T: usize> Minimq<N, T> {
+    /// Construct a new MQTT interface.
+    ///
+    /// # Args
+    /// * `broker` - The IP address of the broker to connect to.
+    /// * `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.
+    ///
+    /// # 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>> {
+        let session_state = SessionState::new(
+            broker,
+            String::from_str(client_id).or(Err(Error::ProvidedClientIdTooLong))?,
+        );
 
-        // Store the new socket for future use.
-        socket_ref.replace(socket);
+        let minimq = Minimq {
+            client: MqttClient {
+                network_stack,
+                socket: None,
+                session_state,
+                connect_sent: false,
+            },
+            packet_reader: PacketReader::new(),
+        };
 
-        Ok(())
+        Ok(minimq)
     }
 
     /// Check the MQTT interface for available messages.
@@ -407,23 +390,24 @@ where
     /// 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(&Self, &'a str, &[u8], &[Property<'a>]),
+        for<'a> F: FnMut(&mut MqttClient<N, T>, &'a str, &[u8], &[Property<'a>]),
     {
         // If the socket is not connected, we can't do anything.
-        if self.socket_is_connected()? == false {
-            self.reset()?;
-            self.connect_socket()?;
+        if self.client.socket_is_connected()? == false {
+            self.client.reset();
+            self.packet_reader.reset();
+            self.client.connect_socket()?;
             return Ok(());
         }
 
         // Connect to the MQTT broker.
-        if !self.connect_sent {
-            self.connect_to_broker()?;
+        if !self.client.connect_sent {
+            self.client.connect_to_broker()?;
             return Ok(());
         }
 
         let mut buf: [u8; 1024] = [0; 1024];
-        let received = self.read(&mut buf)?;
+        let received = self.client.read(&mut buf)?;
         if received > 0 {
             debug!("Received {} bytes", received);
         }
@@ -437,14 +421,20 @@ where
                 }
 
                 Err(e) => {
-                    self.reset()?;
+                    self.client.reset();
+                    self.packet_reader.reset();
                     return Err(Error::Protocol(e));
                 }
             }
 
             // Handle any received packets.
             while self.packet_reader.packet_available() {
-                let result = self.handle_packet(&mut f);
+                let packet = ReceivedPacket::parse_message(&self.packet_reader)
+                    .map_err(|e| Error::Protocol(e))?;
+
+                info!("Received {:?}", packet);
+
+                let result = self.client.handle_packet(packet, &mut f);
 
                 self.packet_reader
                     .pop_packet()
diff --git a/src/properties.rs b/src/properties.rs
index e377d95..82a7ecb 100644
--- a/src/properties.rs
+++ b/src/properties.rs
@@ -5,7 +5,6 @@ use crate::{
 };
 
 use enum_iterator::IntoEnumIterator;
-use generic_array::ArrayLength;
 
 #[derive(Copy, Clone, IntoEnumIterator)]
 pub(crate) enum PropertyIdentifier {
@@ -132,7 +131,7 @@ impl<'a> Property<'a> {
         }
     }
 
-    pub(crate) fn parse<'reader: 'a, T: ArrayLength<u8>>(
+    pub(crate) fn parse<'reader: 'a, const T: usize>(
         packet: &'reader PacketReader<T>,
     ) -> Result<Property<'a>, Error> {
         let identifier: PropertyIdentifier = packet.read_variable_length_integer()?.into();
diff --git a/src/session_state.rs b/src/session_state.rs
index f3534fe..e8b1366 100644
--- a/src/session_state.rs
+++ b/src/session_state.rs
@@ -2,7 +2,7 @@
 ///
 ///
 use embedded_nal::IpAddr;
-use heapless::{consts, String, Vec};
+use heapless::{String, Vec};
 
 pub struct SessionState {
     // Indicates that we are connected to a broker.
@@ -10,14 +10,14 @@ pub struct SessionState {
     pub keep_alive_interval: u16,
     pub broker: IpAddr,
     pub maximum_packet_size: Option<u32>,
-    pub client_id: String<consts::U64>,
-    pub pending_subscriptions: Vec<u16, consts::U32>,
+    pub client_id: String<64>,
+    pub pending_subscriptions: Vec<u16, 32>,
     packet_id: u16,
     active: bool,
 }
 
 impl SessionState {
-    pub fn new<'a>(broker: IpAddr, id: String<consts::U64>) -> SessionState {
+    pub fn new<'a>(broker: IpAddr, id: String<64>) -> SessionState {
         SessionState {
             connected: false,
             active: false,


Test patch:
diff --git a/tests/integration_test.rs b/tests/integration_test.rs
index ee3a154..9c8ceb6 100644
--- a/tests/integration_test.rs
+++ b/tests/integration_test.rs
@@ -1,50 +1,123 @@
-use minimq::{consts, MqttClient, Property, QoS};
+use minimq::{Minimq, Property, QoS};
+use std::io::{Read, Write};
 
-use embedded_nal::{self, IpAddr, Ipv4Addr};
+use embedded_nal::{self, nb, IpAddr, Ipv4Addr, SocketAddr};
+
+#[derive(Copy, Clone, Debug)]
+struct STACK;
+
+pub struct TcpSocket {
+    state: Option<std::net::TcpStream>,
+}
+
+impl TcpSocket {
+    pub fn new() -> Self {
+        Self { state: None }
+    }
+}
+
+fn to_nb(e: std::io::Error) -> nb::Error<std::io::Error> {
+    use std::io::ErrorKind::{TimedOut, WouldBlock};
+
+    match e.kind() {
+        WouldBlock | TimedOut => nb::Error::WouldBlock,
+        _ => e.into(),
+    }
+}
+
+impl embedded_nal::TcpClientStack for STACK {
+    type Error = std::io::Error;
+    type TcpSocket = TcpSocket;
+
+    fn socket(&mut self) -> Result<TcpSocket, Self::Error> {
+        Ok(TcpSocket::new())
+    }
+
+    fn connect(
+        &mut self,
+        socket: &mut TcpSocket,
+        remote: SocketAddr,
+    ) -> nb::Result<(), Self::Error> {
+        let octets: [u8; 4] = match remote.ip() {
+            IpAddr::V4(addr) => addr.octets(),
+            _ => panic!("Unsupported IP version"),
+        };
+
+        let soc = std::net::TcpStream::connect(std::net::SocketAddr::new(
+            std::net::IpAddr::V4(std::net::Ipv4Addr::from(octets)),
+            remote.port(),
+        ))?;
+        soc.set_nonblocking(true)?;
+        socket.state.replace(soc);
+
+        Ok(())
+    }
+
+    fn is_connected(&mut self, socket: &TcpSocket) -> Result<bool, Self::Error> {
+        Ok(socket.state.is_some())
+    }
+
+    fn send(&mut self, socket: &mut TcpSocket, buffer: &[u8]) -> nb::Result<usize, Self::Error> {
+        let socket = socket.state.as_mut().unwrap();
+        socket.write(buffer).map_err(to_nb)
+    }
+
+    fn receive(
+        &mut self,
+        socket: &mut TcpSocket,
+        buffer: &mut [u8],
+    ) -> nb::Result<usize, Self::Error> {
+        let socket = socket.state.as_mut().unwrap();
+        socket.read(buffer).map_err(to_nb)
+    }
+
+    fn close(&mut self, _socket: TcpSocket) -> Result<(), Self::Error> {
+        Ok(())
+    }
+}
 
 #[test]
 fn main() -> std::io::Result<()> {
     env_logger::init();
 
-    let stack = std_embedded_nal::STACK.clone();
+    let stack = STACK.clone();
     let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
-    let mut client = MqttClient::<consts::U256, _>::new(localhost, "", stack).unwrap();
+    let mut mqtt = Minimq::<_, 256>::new(localhost, "", stack).unwrap();
 
     let mut published = false;
     let mut subscribed = false;
 
     loop {
-        client
-            .poll(|client, topic, payload, properties| {
-                println!("{} < {}", topic, core::str::from_utf8(payload).unwrap());
-
-                for property in properties {
-                    match property {
-                        Property::ResponseTopic(topic) => client
-                            .publish(topic, "Pong".as_bytes(), QoS::AtMostOnce, &[])
-                            .unwrap(),
-                        _ => {}
-                    };
-                }
+        mqtt.poll(|client, topic, payload, properties| {
+            println!("{} < {}", topic, core::str::from_utf8(payload).unwrap());
 
-                if topic == "response" {
-                    std::process::exit(0);
-                }
-            })
-            .unwrap();
+            for property in properties {
+                match property {
+                    Property::ResponseTopic(topic) => client
+                        .publish(topic, "Pong".as_bytes(), QoS::AtMostOnce, &[])
+                        .unwrap(),
+                    _ => {}
+                };
+            }
+
+            if topic == "response" {
+                std::process::exit(0);
+            }
+        })
+        .unwrap();
 
         if !subscribed {
-            if client.is_connected().unwrap() {
-                client.subscribe("response", &[]).unwrap();
-                client.subscribe("request", &[]).unwrap();
+            if mqtt.client.is_connected().unwrap() {
+                mqtt.client.subscribe("response", &[]).unwrap();
+                mqtt.client.subscribe("request", &[]).unwrap();
                 subscribed = true;
             }
         } else {
-            if client.subscriptions_pending() == false {
+            if mqtt.client.subscriptions_pending() == false {
                 if !published {
                     println!("PUBLISH request");
                     let properties = [Property::ResponseTopic("response")];
-                    client
+                    mqtt.client
                         .publish("request", "Ping".as_bytes(), QoS::AtMostOnce, &properties)
                         .unwrap();