1 files changed, 485 insertions, 0 deletions

diff --git a/src/mqtt/mod.rs b/src/mqtt/mod.rs
new file mode 100644
index 0000000..30e3a33
--- /dev/null
+++ b/src/mqtt/mod.rs
@@ -0,0 +1,485 @@
+mod connect_code;
+mod field;
+mod packet_id;
+mod protocol;
+mod qos;
+mod rx;
+mod transport;
+mod tx;
+mod varint;
+
+pub use connect_code::ConnectCode;
+use embedded_io_async::ReadExactError;
+pub use packet_id::PacketId;
+pub use qos::Qos;
+pub use transport::Transport;
+
+use self::{field::FieldBuffer, transport::TransportExt as _};
+
+const DEFAULT_CLIENT_RX_BUFFER_SIZE: usize = 512;
+const DEFAULT_CLIENT_TX_BUFFER_SIZE: usize = 512;
+
+pub enum Error<T: Transport> {
+    Transport(T::Error),
+    TransportEOF,
+    InsufficientBufferSpace,
+    ProtocolError(&'static str),
+    ConnectFailed(ConnectCode),
+}
+
+impl<T: Transport> core::fmt::Debug for Error<T> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            Error::Transport(err) => f.debug_tuple("Transport").field(err).finish(),
+            Error::TransportEOF => f.write_str("TransportEOF"),
+            Error::InsufficientBufferSpace => f.write_str("InsufficientBufferSpace"),
+            Error::ProtocolError(msg) => f.debug_tuple("ProtocolError").field(msg).finish(),
+            Error::ConnectFailed(code) => f.debug_tuple("ConnectFailed").field(code).finish(),
+        }
+    }
+}
+
+impl<T: Transport> core::fmt::Display for Error<T> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            Error::Transport(err) => write!(f, "transport error: {:?}", err),
+            Error::TransportEOF => write!(f, "unexpected end of transport stream"),
+            Error::InsufficientBufferSpace => {
+                write!(f, "insufficient buffer space to receive packet")
+            }
+            Error::ProtocolError(msg) => write!(f, "MQTT protocol error: {}", msg),
+            Error::ConnectFailed(code) => write!(f, "connection failed: {}", code),
+        }
+    }
+}
+
+impl<T: Transport> core::error::Error for Error<T>
+where
+    T::Error: core::error::Error + 'static,
+{
+    fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
+        match self {
+            Error::Transport(err) => Some(err),
+            _ => None,
+        }
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct ConnectParams<'a> {
+    pub will_topic: Option<&'a str>,
+    pub will_payload: Option<&'a [u8]>,
+    pub will_retain: bool,
+    pub username: Option<&'a str>,
+    pub password: Option<&'a [u8]>,
+    pub keepalive: Option<u16>,
+}
+
+#[derive(Debug, Default)]
+pub struct PublishParams {
+    pub qos: Qos,
+    pub retain: bool,
+}
+
+#[derive(Debug)]
+pub enum PublishData<'a> {
+    Inline(&'a [u8]),
+    Deferred(usize),
+}
+
+#[derive(Debug)]
+pub struct Publish<'a> {
+    pub topic: &'a str,
+    pub packet_id: Option<PacketId>,
+    pub qos: Qos,
+    pub retain: bool,
+    pub data_len: usize,
+}
+
+#[derive(Debug)]
+pub struct PublishAck {
+    pub packet_id: PacketId,
+}
+
+#[derive(Debug)]
+pub struct SubscribeAck {
+    pub packet_id: PacketId,
+    pub success: bool,
+}
+
+#[derive(Debug)]
+pub struct UnsubscribeAck {
+    pub packet_id: PacketId,
+}
+
+#[derive(Debug)]
+pub enum Packet<'a> {
+    Publish(Publish<'a>),
+    PublishAck(PublishAck),
+    SubscribeAck(SubscribeAck),
+    UnsubscribeAck(UnsubscribeAck),
+}
+
+pub struct ClientResources<
+    const RX: usize = DEFAULT_CLIENT_RX_BUFFER_SIZE,
+    const TX: usize = DEFAULT_CLIENT_TX_BUFFER_SIZE,
+> {
+    rx_buffer: [u8; RX],
+    tx_buffer: [u8; TX],
+}
+
+impl<const RX: usize, const TX: usize> Default for ClientResources<RX, TX> {
+    fn default() -> Self {
+        Self {
+            rx_buffer: [0u8; RX],
+            tx_buffer: [0u8; TX],
+        }
+    }
+}
+
+pub struct Client<'a, T> {
+    transport: T,
+    rx_buffer: &'a mut [u8],
+    rx_buffer_len: usize,
+    rx_buffer_skip: usize,
+    rx_buffer_data: usize,
+    tx_buffer: &'a mut [u8],
+    next_packet_id: u16,
+}
+
+impl<'a, T> Client<'a, T> {
+    pub fn new<const RX: usize, const TX: usize>(
+        resources: &'a mut ClientResources<RX, TX>,
+        transport: T,
+    ) -> Self {
+        Self {
+            transport,
+            rx_buffer: &mut resources.rx_buffer,
+            rx_buffer_len: 0,
+            rx_buffer_skip: 0,
+            rx_buffer_data: 0,
+            tx_buffer: &mut resources.tx_buffer,
+            next_packet_id: 1,
+        }
+    }
+}
+
+impl<'a, T> Client<'a, T>
+where
+    T: Transport,
+{
+    fn allocate_packet_id(&mut self) -> PacketId {
+        let packet_id = self.next_packet_id;
+        self.next_packet_id = self.next_packet_id.wrapping_add(1);
+        if self.next_packet_id == 0 {
+            self.next_packet_id = 1;
+        }
+        PacketId::from(packet_id)
+    }
+
+    pub async fn connect(&mut self, client_id: &str) -> Result<(), Error<T>> {
+        self.connect_with(client_id, Default::default()).await
+    }
+
+    pub async fn connect_with(
+        &mut self,
+        client_id: &str,
+        params: ConnectParams<'_>,
+    ) -> Result<(), Error<T>> {
+        let mut buffer = FieldBuffer::default();
+        tx::connect(
+            &mut buffer,
+            tx::Connect {
+                client_id,
+                clean_session: true,
+                username: params.username,
+                password: params.password,
+                will_topic: params.will_topic,
+                will_payload: params.will_payload,
+                will_retain: params.will_retain,
+                keepalive: None,
+            },
+        );
+        self.transport
+            .write_fields(&buffer)
+            .await
+            .map_err(Error::Transport)?;
+        self.transport.flush().await.map_err(Error::Transport)?;
+
+        // Wait for CONNACK response
+        match self.receive_inner().await? {
+            rx::Packet::ConnAck {
+                session_present,
+                code,
+            } => {
+                if code == ConnectCode::ConnectionAccepted {
+                    Ok(())
+                } else {
+                    Err(Error::ConnectFailed(code))
+                }
+            }
+            _ => Err(Error::ProtocolError(
+                "expected CONNACK packet after CONNECT",
+            )),
+        }
+    }
+
+    pub async fn publish(&mut self, topic: &str, data: &[u8]) -> Result<PacketId, Error<T>> {
+        self.publish_with(topic, data, Default::default()).await
+    }
+
+    pub async fn publish_with(
+        &mut self,
+        topic: &str,
+        data: &[u8],
+        params: PublishParams,
+    ) -> Result<PacketId, Error<T>> {
+        let packet_id = if params.qos.to_u8() > 0 {
+            Some(self.allocate_packet_id())
+        } else {
+            None
+        };
+
+        let mut buffer = FieldBuffer::default();
+        tx::publish(
+            &mut buffer,
+            tx::Publish {
+                topic,
+                payload: data,
+                qos: params.qos,
+                retain: params.retain,
+                dup: false,
+                packet_id,
+            },
+        );
+
+        self.transport
+            .write_fields(&buffer)
+            .await
+            .map_err(Error::Transport)?;
+        self.transport.flush().await.map_err(Error::Transport)?;
+
+        Ok(packet_id.unwrap_or(PacketId::from(0)))
+    }
+
+    pub async fn publish_ack(&mut self, packet_id: PacketId, qos: Qos) -> Result<(), Error<T>> {
+        let mut buffer = FieldBuffer::default();
+
+        match qos {
+            Qos::AtMostOnce => {
+                // QoS 0: No acknowledgment needed
+                return Ok(());
+            }
+            Qos::AtLeastOnce => {
+                // QoS 1: Send PUBACK
+                tx::puback(&mut buffer, packet_id);
+            }
+            Qos::ExactlyOnce => todo!("not implemented"),
+        }
+
+        self.transport
+            .write_fields(&buffer)
+            .await
+            .map_err(Error::Transport)?;
+        self.transport.flush().await.map_err(Error::Transport)?;
+
+        Ok(())
+    }
+
+    pub async fn subscribe(&mut self, topic: &str) -> Result<PacketId, Error<T>> {
+        self.subscribe_with(topic, Qos::AtMostOnce).await
+    }
+
+    pub async fn subscribe_with(&mut self, topic: &str, qos: Qos) -> Result<PacketId, Error<T>> {
+        let packet_id = self.allocate_packet_id();
+
+        let mut buffer = FieldBuffer::default();
+        tx::subscribe(
+            &mut buffer,
+            tx::Subscribe {
+                topic,
+                qos,
+                packet_id,
+            },
+        );
+
+        self.transport
+            .write_fields(&buffer)
+            .await
+            .map_err(Error::Transport)?;
+        self.transport.flush().await.map_err(Error::Transport)?;
+
+        Ok(packet_id)
+    }
+
+    pub async fn unsubscribe(&mut self, topic: &str) -> Result<PacketId, Error<T>> {
+        let packet_id = self.allocate_packet_id();
+
+        let mut buffer = FieldBuffer::default();
+        tx::unsubscribe(&mut buffer, tx::Unsubscribe { topic, packet_id });
+
+        self.transport
+            .write_fields(&buffer)
+            .await
+            .map_err(Error::Transport)?;
+        self.transport.flush().await.map_err(Error::Transport)?;
+
+        Ok(packet_id)
+    }
+
+    async fn receive_inner<'s>(&'s mut self) -> Result<rx::Packet<'s>, Error<T>> {
+        self.skip_if_required();
+        self.discard_data().await?;
+
+        loop {
+            let buf = &self.rx_buffer[..self.rx_buffer_len];
+            match rx::decode(buf) {
+                Ok(_) => {
+                    // NOTE: stupid workaround for borrow checker, should not
+                    // need to decode twice
+                    let buf = &self.rx_buffer[..self.rx_buffer_len];
+                    let (packet, n) = rx::decode(buf).unwrap();
+                    self.rx_buffer_skip = n;
+                    if let rx::Packet::Publish { data_len, .. } = &packet {
+                        self.rx_buffer_data = *data_len;
+                    }
+                    return Ok(packet);
+                }
+                Err(err) => match err {
+                    rx::Error::NeedMoreData => {
+                        if self.rx_buffer.len() == self.rx_buffer_len {
+                            return Err(Error::InsufficientBufferSpace);
+                        }
+                    }
+                    rx::Error::InvalidPacket(msg) => return Err(Error::ProtocolError(msg)),
+                    rx::Error::UnsupportedPacket { packet_type, .. } => {
+                        return Err(Error::ProtocolError("unsupported packet type"));
+                    }
+                    rx::Error::UnknownPacket { packet_type, .. } => {
+                        return Err(Error::ProtocolError("unknown packet type"));
+                    }
+                },
+            }
+
+            self.fill_rx_buffer().await?;
+        }
+    }
+
+    pub async fn receive<'s>(&'s mut self) -> Result<Packet<'s>, Error<T>> {
+        match self.receive_inner().await? {
+            rx::Packet::ConnAck { .. } => {
+                return Err(Error::ProtocolError("unexpected CONNACK packet"));
+            }
+            rx::Packet::Publish {
+                topic,
+                packet_id,
+                qos,
+                retain,
+                dup: _dup,
+                data_len,
+            } => {
+                return Ok(Packet::Publish(Publish {
+                    topic,
+                    packet_id,
+                    qos,
+                    retain,
+                    data_len,
+                }));
+            }
+            rx::Packet::PubAck { packet_id } => {
+                return Ok(Packet::PublishAck(PublishAck { packet_id }));
+            }
+            rx::Packet::SubscribeAck { packet_id, success } => {
+                return Ok(Packet::SubscribeAck(SubscribeAck { packet_id, success }));
+            }
+            rx::Packet::UnsubscribeAck { packet_id } => {
+                return Ok(Packet::UnsubscribeAck(UnsubscribeAck { packet_id }));
+            }
+        }
+    }
+
+    pub async fn receive_data(&mut self, buf: &mut [u8]) -> Result<(), Error<T>> {
+        self.skip_if_required();
+        if buf.len() != self.rx_buffer_data {
+            return Err(Error::InsufficientBufferSpace);
+        }
+
+        assert_eq!(self.rx_buffer_skip, 0);
+        let from_buffer = self.rx_buffer_data.min(self.rx_buffer_len);
+        let from_transport = self.rx_buffer_data.strict_sub(from_buffer);
+
+        buf[..from_buffer].copy_from_slice(&self.rx_buffer[..from_buffer]);
+        self.rx_buffer_len -= from_buffer;
+
+        if from_transport > 0 {
+            assert_eq!(self.rx_buffer_len, 0);
+            self.transport
+                .read_exact(&mut buf[from_buffer..])
+                .await
+                .map_err(|err| match err {
+                    ReadExactError::UnexpectedEof => Error::<T>::TransportEOF,
+                    ReadExactError::Other(e) => Error::Transport(e),
+                })?;
+        }
+        self.rx_buffer_data = 0;
+
+        Ok(())
+    }
+
+    pub async fn disconnect(&mut self) -> Result<(), Error<T>> {
+        let mut buffer = FieldBuffer::default();
+        tx::disconnect(&mut buffer);
+
+        self.transport
+            .write_fields(&buffer)
+            .await
+            .map_err(Error::Transport)?;
+        self.transport.flush().await.map_err(Error::Transport)?;
+
+        Ok(())
+    }
+
+    async fn fill_rx_buffer(&mut self) -> Result<(), Error<T>> {
+        let n = self
+            .transport
+            .read(&mut self.rx_buffer[self.rx_buffer_len..])
+            .await
+            .map_err(Error::Transport)?;
+        if n == 0 {
+            return Err(Error::TransportEOF);
+        }
+        self.rx_buffer_len += n;
+
+        Ok(())
+    }
+
+    fn skip_if_required(&mut self) {
+        assert!(self.rx_buffer_len >= self.rx_buffer_skip);
+        if self.rx_buffer_skip != 0 {
+            self.rx_buffer.copy_within(self.rx_buffer_skip.., 0);
+            self.rx_buffer_len = self.rx_buffer_len.strict_sub(self.rx_buffer_skip);
+            self.rx_buffer_skip = 0;
+        }
+    }
+
+    async fn discard_data(&mut self) -> Result<(), Error<T>> {
+        if self.rx_buffer_data == 0 {
+            return Ok(());
+        }
+
+        assert_eq!(self.rx_buffer_skip, 0);
+        while self.rx_buffer_data > 0 {
+            if self.rx_buffer_len <= self.rx_buffer_data {
+                self.rx_buffer_data -= self.rx_buffer_len;
+                self.rx_buffer_len = 0;
+            } else {
+                self.rx_buffer.copy_within(self.rx_buffer_data.., 0);
+                self.rx_buffer_len -= self.rx_buffer_data;
+                self.rx_buffer_data = 0;
+            }
+            self.fill_rx_buffer().await?;
+        }
+
+        Ok(())
+    }
+}