fundamentum_sdk_mqtt/

client.rs

//! Fundamentum Client
//!

use std::{
    collections::HashMap,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};

use backon::{ExponentialBuilder, Retryable};
use displaydoc::Display;
use futures::Stream;
#[cfg(any(feature = "use-native-tls", feature = "use-rustls"))]
use rumqttc::TlsConfiguration;
use rumqttc::{
    Transport,
    v5::{
        AsyncClient, ClientError, ConnectionError, Event, EventLoop, MqttOptions,
        mqttbytes::{
            QoS,
            v5::{ConnectReturnCode, LastWill, PublishProperties},
        },
    },
};
use tokio::{
    sync::{
        Mutex,
        broadcast::{self, Receiver, Sender},
        watch,
    },
    time::Duration,
};
use tokio_stream::wrappers::WatchStream;
use tracing::{error, info};

use crate::{
    Device, Error, Message, PublishOptions, Publishable, Publisher, Security, error,
    publish_options::PublishOptionsResolved,
};

/// Max Packet Size Override
pub struct MQTTMaxPacketSize(u32);

/// Option Size Override
#[derive(Default)]
pub struct MQTTOptionsOverrides {
    /// Override clean session
    pub clean_session: Option<bool>,
    /// Override session expiry interval
    pub session_expiry_interval: Option<Duration>,
    /// Override keep alive
    pub keep_alive: Option<Duration>,
    /// Override max packet size
    pub max_packet_size: Option<MQTTMaxPacketSize>,
    /// Override request channel capacity
    pub request_channel_capacity: Option<usize>,
    /// Override pending throttle
    pub pending_throttle: Option<Duration>,
    /// Override inflight
    pub inflight: Option<u16>,
    /// Override last will
    pub last_will: Option<LastWill>,
    /// Override transport
    pub transport: Option<Transport>,
}

/// Fundamentum `IoT` Settings
pub struct ClientSettings {
    /// Security strategy
    security: Security,
    /// Device's representation
    device: Device,
    /// Fundamentum endpoint
    endpoint: String,
    /// MQTT options overrides
    mqtt_options_overrides: Option<MQTTOptionsOverrides>,
}

impl ClientSettings {
    /// Create a new `FundamentumIoTSettings`
    ///
    /// # Params
    ///
    /// * `security`: Security generator
    /// * `device`: Device's definition
    /// * `iot_endpoint`: Uri endpoint MQTT
    /// * `mqtt_options_overrides`: MQTT options overrides
    ///
    #[must_use]
    pub const fn new(
        security: Security,
        device: Device,
        endpoint: String,
        mqtt_options_overrides: Option<MQTTOptionsOverrides>,
    ) -> Self {
        Self {
            security,
            device,
            endpoint,
            mqtt_options_overrides,
        }
    }

    /// Converts the `ClientSettings` struct into `MqttOptions`.
    ///
    /// This method transforms the current `ClientSettings` instance into a set of
    /// options that can be used to configure rumqtt client.
    ///
    /// # Returns
    ///
    /// A struct containing the MQTT options derived from the current `ClientSettings`.
    ///
    /// # Errors
    ///
    /// Returns an `error::Error` if an error occurs while generating a new signed token.
    ///
    pub async fn to_mqtt_options(&self) -> Result<MqttOptions, error::Error> {
        let mut mqtt_options =
            MqttOptions::new(self.device.client_id(), self.endpoint.clone(), 8883);

        if let Some(ref overrides) = self.mqtt_options_overrides {
            if let Some(clean_session) = overrides.clean_session {
                mqtt_options.set_clean_start(clean_session);
            }
            if let Some(session_expiry_interval) = overrides.session_expiry_interval {
                let mut connect_properties = mqtt_options.connect_properties().unwrap_or_default();
                connect_properties.session_expiry_interval =
                    Some(session_expiry_interval.as_secs().try_into()?);
                mqtt_options.set_connect_properties(connect_properties);
            }
            if let Some(transport) = overrides.transport.clone() {
                mqtt_options.set_transport(transport);
            } else {
                let transport = get_default_transport();
                mqtt_options.set_transport(transport);
            }
            if let Some(keep_alive) = overrides.keep_alive {
                mqtt_options.set_keep_alive(keep_alive);
            }
            if let Some(ref packet_size) = overrides.max_packet_size {
                mqtt_options.set_max_packet_size(Some(packet_size.0));
            }
            if let Some(request_channel_capacity) = overrides.request_channel_capacity {
                mqtt_options.set_request_channel_capacity(request_channel_capacity);
            }
            if let Some(pending_throttle) = overrides.pending_throttle {
                mqtt_options.set_pending_throttle(pending_throttle);
            }
            if let Some(inflight) = overrides.inflight {
                mqtt_options.set_outgoing_inflight_upper_limit(inflight);
            }
            if let Some(last_will) = overrides.last_will.clone() {
                mqtt_options.set_last_will(last_will);
            }
        }

        // Generate a new signed token
        let token = self.security.generate_token().await?;
        mqtt_options.set_credentials("unused", token);

        Ok(mqtt_options)
    }
}

const fn get_default_transport() -> Transport {
    #[cfg(all(feature = "use-native-tls", not(feature = "use-rustls")))]
    let transport = Transport::Tls(TlsConfiguration::Native);
    #[cfg(all(feature = "use-rustls", not(feature = "use-native-tls")))]
    let transport = Transport::Tls(TlsConfiguration::default());
    #[cfg(all(feature = "use-native-tls", feature = "use-rustls"))]
    let transport = Transport::Tls(TlsConfiguration::Native);
    #[cfg(not(any(feature = "use-rustls", feature = "use-native-tls")))]
    let transport = Transport::Tcp;

    transport
}

#[derive(Clone)]
struct EventLoopManager {
    event_loop: Arc<Mutex<EventLoop>>,
    client_status: watch::Sender<ClientStatus>,
}

impl EventLoopManager {
    fn new(event_loop: EventLoop, client_status: watch::Sender<ClientStatus>) -> Self {
        Self {
            event_loop: Arc::new(Mutex::new(event_loop)),
            client_status,
        }
    }

    async fn poll(&self) -> Result<Event, ConnectionError> {
        let mut in_error = false;

        let polling_result = (|| async { self.event_loop.lock().await.poll().await })
            .retry(ExponentialBuilder::default().without_max_times())
            .when(is_backoff_error)
            .notify(|err, dur: Duration| {
                in_error = true;
                self.set_client_status(ClientStatus::InError(ClientStatusError::MqttConnection(
                    err.to_string(),
                )));

                let dur = dur.as_secs_f32();
                error!("Error while polling MQTT event loop: {err}\n -> Retrying in {dur:.1}s...");
            })
            .await;

        match polling_result.as_ref() {
            Ok(_) => {
                self.set_client_status(ClientStatus::Connected);
                if in_error {
                    info!("MQTT connection restored");
                }
            }
            Err(err) => self.set_client_status(ClientStatus::InError(
                ClientStatusError::MqttConnection(err.to_string()),
            )),
        }

        polling_result
    }

    fn set_client_status(&self, status: ClientStatus) {
        self.client_status.send_if_modified(|current_status| {
            let notify = current_status != &status;
            *current_status = status;
            notify
        });
    }
}

const fn is_backoff_error(err: &ConnectionError) -> bool {
    !matches!(
        err,
        &ConnectionError::ConnectionRefused(
            ConnectReturnCode::ProtocolError
                | ConnectReturnCode::UnsupportedProtocolVersion
                | ConnectReturnCode::ClientIdentifierNotValid
                | ConnectReturnCode::BadUserNamePassword
                | ConnectReturnCode::NotAuthorized
                | ConnectReturnCode::Banned
                | ConnectReturnCode::BadAuthenticationMethod
                | ConnectReturnCode::UseAnotherServer
                | ConnectReturnCode::ServerMoved
        )
    )
}

/// Represents the possible errors that can occur with the client connection status.
#[derive(Display, Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum ClientStatusError {
    /// An error occurred with the MQTT connection: {0}
    MqttConnection(String),
}

/// Represents the status of the client connection.
#[derive(Clone, PartialEq, Eq)]
pub enum ClientStatus {
    /// The client is currently connected.
    Connected,
    /// The client is currently disconnected.
    Disconnected,
    /// Represents an error state in the connection status.
    /// Contains a string describing the error.
    InError(ClientStatusError),
}

/// A stream for tracking the status of a client.
///
/// `ClientStatusStream` provides functionality to observe the current status
/// of a client and subscribe to status updates via a stream.
pub struct ClientStatusStream {
    receiver: watch::Receiver<ClientStatus>,
    stream: WatchStream<ClientStatus>,
}

impl ClientStatusStream {
    #[must_use]
    fn new(receiver: watch::Receiver<ClientStatus>) -> Self {
        let stream = WatchStream::new(receiver.clone());
        Self { receiver, stream }
    }

    /// Retrieves the current status of the client.
    ///
    /// # Returns
    ///
    /// The current `ClientStatus` value.
    #[must_use]
    pub fn current(&self) -> ClientStatus {
        self.receiver.borrow().clone()
    }
}

impl Stream for ClientStatusStream {
    type Item = ClientStatus;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        Pin::new(&mut this.stream).poll_next(cx)
    }
}

impl Unpin for ClientStatusStream {}

/// Fundamentum `IoT` Async Client
#[derive(Clone)]
pub struct Client {
    client: AsyncClient,
    manager: EventLoopManager,
    device: Device,
    incoming_event_sender: Sender<Message>,
    status_rx: watch::Receiver<ClientStatus>,
}

impl Client {
    /// Create new `FundamentumIoTAsyncClient`. Input argument should be the `FundamentumIoTSettings`. Returns a tuple where the first element is the
    /// `FundamentumIoTAsyncClient`, and the second element is a new tuple with the eventloop and incoming
    /// event sender. This tuple should be sent as an argument to the `async_event_loop_listener`.
    ///
    /// # Errors
    ///
    /// If the creation failed then the result returns an error `FundamentumIoTError`
    pub async fn new(settings: ClientSettings) -> Result<Self, error::Error> {
        let mqtt_options = settings.to_mqtt_options().await?;

        let (client, eventloop) = AsyncClient::new(mqtt_options, 10);
        let (request_tx, _) = broadcast::channel(50);
        let (status_tx, status_rx) = watch::channel(ClientStatus::Disconnected);
        let manager = EventLoopManager::new(eventloop, status_tx);

        Ok(Self {
            client,
            manager,
            device: settings.device.clone(),
            incoming_event_sender: request_tx,
            status_rx,
        })
    }

    /// Run the client’s background task.
    ///
    /// This task continuously polls the underlying even loop which
    /// ensures messages are exchanged between this client and the broker.
    ///
    /// Nothing will be exchanged until this task is run.
    ///
    /// # Errors
    ///
    /// Returns an `error::Error` if a fatal MQTT connection error occurs that cannot be retried or recovered from.
    /// The underlying `ConnectionError` is considered fatal if it matches one of the following MQTT connection refusal codes:
    /// - `ProtocolError`
    /// - `UnsupportedProtocolVersion`
    /// - `ClientIdentifierNotValid`
    /// - `BadUserNamePassword`
    /// - `NotAuthorized`
    /// - `Banned`
    /// - `BadAuthenticationMethod`
    /// - `UseAnotherServer`
    /// - `ServerMoved`
    ///
    /// In these cases, the client will terminate with an error. For all other connection errors (such as temporary network issues),
    /// the client will attempt to reconnect automatically (backoff).
    ///
    /// Any error returned by this function indicates that the client can no longer function and requires user intervention.
    pub async fn run(&self) -> Result<(), error::Error> {
        loop {
            let event = self.manager.poll().await?;
            if let Err(err) = self.handle_event(event) {
                error!("Error while handling MQTT event: {err}");
            }
        }
    }

    fn handle_event(&self, event: Event) -> Result<(), error::Error> {
        match event {
            Event::Incoming(packet) => {
                let message = Message::try_from_packet(packet)?;
                self.incoming_event_sender.send(message)?;
                Ok(())
            }
            Event::Outgoing(_) => Ok(()), // Silently discarded.
        }
    }

    /// Subscribe to a topic
    ///
    /// # Errors
    ///
    /// Returns `ClientError` if the request failed
    async fn subscribe<S: Into<String> + Send>(
        &self,
        topic: S,
        qos: QoS,
    ) -> Result<(), ClientError> {
        self.client.subscribe(topic, qos).await
    }

    /// Subscribe to device's config channel
    ///
    /// # Reference
    ///
    /// * `registries/{REGISTRY_ID}/devices/{DEVICE_SN}/config`
    ///
    /// # Errors
    ///
    /// Returns `ClientError` if the request failed
    pub async fn subscribe_config(&self, qos: QoS) -> Result<(), ClientError> {
        self.subscribe(
            format!(
                "registries/{}/devices/{}/config",
                self.device.registry_id(),
                self.device.serial()
            ),
            qos,
        )
        .await
    }

    /// Subscribe to device's port forward channel
    ///
    /// # Reference
    ///
    /// * `registries/{REGISTRY_ID}/devices/{DEVICE_SN}/portforward/tx`
    ///
    /// # Errors
    ///
    /// Returns `ClientError` if the request failed
    pub async fn subscribe_portforward(&self, qos: QoS) -> Result<(), ClientError> {
        self.subscribe(
            format!(
                "registries/{}/devices/{}/pfwd/tx",
                self.device.registry_id(),
                self.device.serial()
            ),
            qos,
        )
        .await
    }

    /// Subscribe to device's commands channel
    ///
    /// # Reference
    ///
    /// * `registries/{REGISTRY_ID}/devices/{DEVICE_SN}/commands`
    ///
    /// # Errors
    ///
    /// Returns `ClientError` if the request failed
    pub async fn subscribe_commands(&self, qos: QoS) -> Result<(), ClientError> {
        self.subscribe(
            format!(
                "registries/{}/devices/{}/commands",
                self.device.registry_id(),
                self.device.serial()
            ),
            qos,
        )
        .await
    }

    /// Subscribe to device's actions channel
    ///
    /// # Reference
    ///
    /// * `registries/{REGISTRY_ID}/devices/{DEVICE_SN}/actions`
    ///
    /// # Errors
    ///
    /// Returns `ClientError` if the request failed
    pub async fn subscribe_actions(&self, qos: QoS) -> Result<(), ClientError> {
        self.subscribe(
            format!(
                "registries/{}/devices/{}/actions",
                self.device.registry_id(),
                self.device.serial(),
            ),
            qos,
        )
        .await
    }

    /// Get a receiver of the incoming messages. Send this to any function that
    /// wants to read the incoming messages from `IoT` Core.
    ///
    /// Note that it is very important that your retrieve your *receiver* prior
    /// to any `subscribe_` of interest. Otherwise, you most likely won't
    /// receive messages sent to you in between this call and prior
    /// subscriptions.
    ///
    /// This is of particular importance for messages returned as a result to
    /// your subscription.
    #[must_use]
    pub fn get_receiver(&self) -> Receiver<Message> {
        self.incoming_event_sender.subscribe()
    }

    /// If you want to use the Rumqttc `AsyncClient` and `EventLoop` manually, this method can be used
    /// to get the `AsyncClient`.
    #[must_use]
    #[allow(clippy::nursery)]
    pub fn get_client(self) -> AsyncClient {
        self.client
    }

    /// Monitor our client's statuses over time.
    ///
    /// The returned [`ClientStatusStream`] is a [`futures::Stream`] of
    /// [`ClientStatus`] that can also provide the most recent status through
    /// its [`current`](ClientStatusStream::current) method.
    ///
    /// # Returns
    ///
    /// The streamable monitor instance.
    #[must_use]
    pub fn status_stream(&self) -> ClientStatusStream {
        ClientStatusStream::new(self.status_rx.clone())
    }
}

impl Publisher for Client {
    type Error = Error;

    async fn publish_with<P: Publishable + Send>(
        &self,
        publishable: P,
        options: PublishOptions,
    ) -> Result<(), Self::Error> {
        let options = PublishOptionsResolved {
            qos: QoS::ExactlyOnce,
            retain: false,
            user_properties: HashMap::new(),
            content_type: None,
        }
        .override_with(&publishable.publish_overrides())
        .override_with(&options);

        let properties = PublishProperties {
            user_properties: options.user_properties.into_iter().collect(),
            content_type: options.content_type,
            ..Default::default()
        };

        self.client
            .publish_with_properties(
                publishable.topic(&self.device),
                options.qos,
                options.retain,
                publishable.payload().map_err(Into::into)?,
                properties,
            )
            .await?;

        Ok(())
    }
}