mod codec;
mod pkid;
mod publish_store;
use std::{
sync::Arc,
time::{Duration, Instant},
};
use anyhow::{anyhow, bail, Result};
use fluent_uri::Uri;
use futures_util::{sink::SinkExt, stream::StreamExt, FutureExt};
use mqttbytes::{v5::*, QoS};
use parking_lot::RwLock;
use tokio::{net::TcpStream, time};
use tokio_util::codec::Framed;
use tracing::{debug, error, warn};
use self::{codec::MqttV5Codec, pkid::PacketId};
use super::{config::MqttClientConfig, Message};
pub(super) use self::publish_store::PublishStore;
use crate::client::config::OnMessageOwnedCallback;
const MAX_MQTT_PACKET_SIZE: usize = 10 * 1024;
const KEEP_ALIVE_INTERVAL: Duration = Duration::from_millis(5_000);
const CONNECTION_TIMEOUT: Duration = Duration::from_millis(5_000);
const RECONNECT_INTERVAL: Duration = Duration::from_millis(1_000);
const STORE_INTERVAL: Duration = Duration::from_millis(1_000);
pub struct MqttWorker {
mqtt_broker: Uri<String>,
config: Arc<RwLock<MqttClientConfig>>,
last_broker_response: Instant,
publish_store: Arc<PublishStore>,
on_message_owned: Option<Box<OnMessageOwnedCallback>>,
}
impl MqttWorker {
pub(super) fn new(
mqtt_broker: Uri<String>,
config: Arc<RwLock<MqttClientConfig>>,
on_message_owned: Option<Box<OnMessageOwnedCallback>>,
) -> (Self, Arc<PublishStore>) {
let publish_store = Arc::new(PublishStore::new());
let worker = Self {
mqtt_broker,
config,
last_broker_response: Instant::now(),
publish_store: publish_store.clone(),
on_message_owned,
};
(worker, publish_store)
}
pub async fn run(mut self) {
let mut interval = time::interval(RECONNECT_INTERVAL);
loop {
interval.tick().await;
self.publish_store.on_connect_cleanup();
let authority = self
.mqtt_broker
.authority()
.expect("missing broker authority");
let host = authority.host();
let port = authority
.port()
.filter(|p| !p.is_empty())
.and_then(|p| p.as_str().parse::<u16>().ok())
.unwrap_or(1883);
let broker_addr = (host, port);
debug!("connecting to broker {:?}", broker_addr);
let socket = match TcpStream::connect(broker_addr).await {
Ok(s) => s,
Err(err) => {
error!("TCP connection to MQTT broker failed: {}", err);
continue;
}
};
if let Err(err) = self.connection_handle(socket).await {
error!("MQTT connection closed: {}", err);
continue;
}
}
}
async fn connection_handle(&mut self, socket: TcpStream) -> Result<()> {
let mqtt_codec = MqttV5Codec::new(MAX_MQTT_PACKET_SIZE);
let mut socket = Framed::new(socket, mqtt_codec);
let _connack = self.mqtt_connect(&mut socket).await?;
self.on_connected();
let _suback = self.mqtt_subscribe(&mut socket).await?;
let mut ping_interval = time::interval(KEEP_ALIVE_INTERVAL);
let mut store_tick_interval = time::interval(STORE_INTERVAL);
let mut packet_id = PacketId::new();
loop {
tokio::select! {
_ = ping_interval.tick().fuse() => {
if self.last_broker_response.elapsed() < (KEEP_ALIVE_INTERVAL * 2)
{
debug!("Sending PingReq");
socket.send(&PingReq).await?;
} else {
bail!("MQTT broker has not responded in {} seconds", self.last_broker_response.elapsed().as_secs())
}
}
packet = socket.next() => {
let received = Instant::now();
debug!("MQTT incoming packet: {:?}", packet);
let packet = match packet.transpose()? {
Some(p) => p,
None => {
bail!("MQTT broker has finished")
}
};
match packet {
Packet::Connect(_connect) => {}
Packet::ConnAck(_connack) => {}
Packet::Publish(publish) => {
let pkid = publish.pkid;
let qos = publish.qos;
let message = publish.into();
self.on_message(&message);
if let Some(callback) = &mut self.on_message_owned {
callback(message, received);
}
self.last_broker_response = Instant::now();
match qos {
QoS::AtMostOnce => (),
QoS::AtLeastOnce => {
let puback = PubAck::new(pkid);
debug!("Sending PubAck for packet id = {}", pkid);
socket.send(&puback).await?;
}
QoS::ExactlyOnce => {
warn!("Received message with QoS 2, missing implementation")
},
}
}
Packet::PubAck(puback) => {
self.last_broker_response = Instant::now();
self.publish_store.remove_waiting_for_ack(puback.pkid);
if puback.reason as u8 > 0x00 {
warn!("Received PubAck reason code other stan success for pkid = {}: {:?}", puback.pkid, puback.reason);
}
}
Packet::PubRec(_pubrec) => {}
Packet::PubRel(_pubrel) => {}
Packet::PubComp(_pubcomp) => {}
Packet::Subscribe(_subscribe) => {}
Packet::SubAck(_suback) => {}
Packet::Unsubscribe(_unsubscribe) => {}
Packet::UnsubAck(_unsuback) => {}
Packet::PingReq => socket.send(&PingResp).await?,
Packet::PingResp => {
self.last_broker_response = Instant::now()
},
Packet::Disconnect(disconnect) => {
warn!("MQTT disconnected by broker: {:?}", disconnect);
return Ok(());
}
}
}
_ = store_tick_interval.tick().fuse() => {
self.publish_store.remove_expired();
},
_ = self.publish_store.notified().fuse() => {
if let Some(publish) = self.publish_store.next_publish(packet_id.into()) {
if let Some(0) = publish.properties.as_ref().and_then(|x| x.message_expiry_interval) {
debug!("Message timeouted {}", publish.topic);
continue;
};
packet_id.increment();
debug!("Message to send: {:?}", publish);
socket.send(publish.as_ref()).await?;
}
}
}
}
}
async fn mqtt_connect(
&mut self,
socket: &mut Framed<TcpStream, MqttV5Codec>,
) -> Result<ConnAck> {
let (username, password) = self
.mqtt_broker
.authority()
.and_then(|a| a.userinfo())
.map(|u| {
let s = u.as_str();
match s.split_once(':') {
Some((user, pass)) => (user, Some(pass)),
None => (s, None),
}
})
.unwrap_or(("", None));
let login = password.map(|p| Login::new(username, p));
let properties = ConnectProperties {
session_expiry_interval: None,
receive_maximum: None,
max_packet_size: Some(MAX_MQTT_PACKET_SIZE as u32),
topic_alias_max: Some(0),
request_response_info: None,
request_problem_info: None,
user_properties: Vec::new(),
authentication_method: None,
authentication_data: None,
};
let mut connect = { Connect::new(&self.config.read().client_id) };
let clean_session = { self.config.read().clean_session };
connect.keep_alive = KEEP_ALIVE_INTERVAL.as_secs() as u16;
connect.login = login;
connect.clean_session = clean_session;
connect.properties = Some(properties);
socket.send(&connect).await?;
let rx_packet = time::timeout(CONNECTION_TIMEOUT, async {
let rx_packet = match socket.next().await {
Some(packet) => packet?,
None => bail!("MQTT broker disconnected"),
};
let rx_packet = match rx_packet {
Packet::ConnAck(connack) if connack.code == ConnectReturnCode::Success => connack,
Packet::ConnAck(connack) => {
bail!("Broker rejected. Reason = {:?}", connack.code)
}
other => bail!("Expecting ConnAck. Received = {:?}", other),
};
Ok::<ConnAck, anyhow::Error>(rx_packet)
})
.await??;
debug!("MQTT client successfully connected to broker");
self.last_broker_response = Instant::now();
Ok(rx_packet)
}
async fn mqtt_subscribe(
&mut self,
socket: &mut Framed<TcpStream, MqttV5Codec>,
) -> Result<Option<SubAck>> {
let mut subscribe = {
let config = self.config.read();
if config.subscribe_filters.is_empty() {
return Ok(None);
}
Subscribe::new_many(config.subscribe_filters.iter().cloned())
};
let pkid = 1;
subscribe.pkid = pkid;
socket.send(&subscribe).await?;
let rx_packet = time::timeout(CONNECTION_TIMEOUT, async {
let rx_packet = match socket.next().await {
Some(packet) => packet?,
None => bail!("MQTT broker disconnected"),
};
let rx_packet = match rx_packet {
Packet::SubAck(suback) if pkid != suback.pkid => {
bail!("Packet Id mismatch")
}
Packet::SubAck(suback)
if !suback
.return_codes
.iter()
.any(|c| *c as u8 >= SubscribeReasonCode::Unspecified as u8) =>
{
suback
}
Packet::SubAck(suback) => {
bail!(
"Broker rejects subscriptions. Reasons = {:?}",
suback.return_codes
)
}
other => bail!("Expecting SubAck. Received = {:?}", other),
};
Ok::<SubAck, anyhow::Error>(rx_packet)
})
.await??;
debug!("MQTT client successfully subscribed topics");
self.last_broker_response = Instant::now();
Ok(Some(rx_packet))
}
fn on_connected(&self) {
if let Some(callback) = &self.config.write().on_connected_callback {
callback()
}
}
fn on_message(&self, message: &Message) {
if let Some(callback) = &mut self.config.write().on_message_callback {
callback(message)
}
}
}
fn mqttbytes_error_to_anyhow(error: mqttbytes::Error) -> anyhow::Error {
anyhow!("{}", error)
}