agentlink_native/

mqtt.rs

//! Native MQTT Client Implementation
//!
//! Uses rumqttc for MQTT connections.

use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use agentlink_core::mqtt::{
    MqttClient, MqttConfig, MqttConnectionState, MqttEvent, MqttMessage, MqttQoS,
};
use agentlink_core::error::SdkResult;
use rumqttc::{AsyncClient, Event as MqttEventLoopEvent, EventLoop, Incoming, MqttOptions, QoS, Transport};
use tokio::sync::Mutex;
use tokio::task::JoinHandle;

/// Native MQTT client using rumqttc
pub struct NativeMqttClient {
    client: Arc<Mutex<Option<AsyncClient>>>,
    state: Arc<Mutex<MqttConnectionState>>,
    event_callback: Arc<Mutex<Option<Box<dyn Fn(MqttEvent) + Send + Sync>>>>,
    event_loop_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
}

impl NativeMqttClient {
    pub fn new() -> Self {
        Self {
            client: Arc::new(Mutex::new(None)),
            state: Arc::new(Mutex::new(MqttConnectionState::Disconnected)),
            event_callback: Arc::new(Mutex::new(None)),
            event_loop_handle: Arc::new(Mutex::new(None)),
        }
    }

    fn parse_broker_url(url: &str) -> SdkResult<(String, u16, Option<Transport>)> {
        let is_tls = url.starts_with("mqtts://");
        let is_tcp = url.starts_with("mqtt://");

        if !is_tls && !is_tcp {
            return Err(agentlink_core::error::SdkError::Config(
                format!("Unsupported MQTT protocol: {}", url)
            ));
        }

        let url_part = if is_tls {
            &url[8..]
        } else {
            &url[7..]
        };

        let parts: Vec<&str> = url_part.split('/').next().unwrap().split(':').collect();
        let host = parts[0];

        let default_port = if is_tls { 8883 } else { 1883 };
        let port = parts.get(1)
            .and_then(|p| p.parse::<u16>().ok())
            .unwrap_or(default_port);

        let transport = if is_tls {
            Some(Self::create_tls_transport()?)
        } else {
            None
        };

        Ok((host.to_string(), port, transport))
    }

    fn create_tls_transport() -> SdkResult<Transport> {
        use rumqttc::tokio_rustls::rustls::client::danger::{
            HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier,
        };
        use rumqttc::tokio_rustls::rustls::pki_types::{CertificateDer, ServerName, UnixTime};
        use rumqttc::tokio_rustls::rustls::{DigitallySignedStruct, Error, SignatureScheme};

        #[derive(Debug)]
        struct NoVerification;

        impl ServerCertVerifier for NoVerification {
            fn verify_server_cert(
                &self, _end_entity: &CertificateDer<'_>, _intermediates: &[CertificateDer<'_>],
                _server_name: &ServerName<'_>, _ocsp_response: &[u8], _now: UnixTime,
            ) -> Result<ServerCertVerified, Error> {
                Ok(ServerCertVerified::assertion())
            }

            fn verify_tls12_signature(
                &self, _message: &[u8], _cert: &CertificateDer<'_>, _dss: &DigitallySignedStruct,
            ) -> Result<HandshakeSignatureValid, Error> {
                Ok(HandshakeSignatureValid::assertion())
            }

            fn verify_tls13_signature(
                &self, _message: &[u8], _cert: &CertificateDer<'_>, _dss: &DigitallySignedStruct,
            ) -> Result<HandshakeSignatureValid, Error> {
                Ok(HandshakeSignatureValid::assertion())
            }

            fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
                vec![
                    SignatureScheme::RSA_PKCS1_SHA256,
                    SignatureScheme::ECDSA_NISTP256_SHA256,
                    SignatureScheme::ECDSA_NISTP384_SHA384,
                    SignatureScheme::ED25519,
                    SignatureScheme::RSA_PSS_SHA256,
                ]
            }
        }

        let config = rumqttc::tokio_rustls::rustls::ClientConfig::builder()
            .dangerous()
            .with_custom_certificate_verifier(Arc::new(NoVerification))
            .with_no_client_auth();

        Ok(Transport::Tls(rumqttc::TlsConfiguration::Rustls(Arc::new(config))))
    }

    async fn run_event_loop(
        mut eventloop: EventLoop,
        state: Arc<Mutex<MqttConnectionState>>,
        callback: Arc<Mutex<Option<Box<dyn Fn(MqttEvent) + Send + Sync>>>>,
    ) {
        loop {
            match eventloop.poll().await {
                Ok(notification) => {
                    match notification {
                        MqttEventLoopEvent::Incoming(incoming) => {
                            match incoming {
                                Incoming::Publish(packet) => {
                                    let msg = MqttMessage {
                                        topic: packet.topic,
                                        payload: packet.payload.to_vec(),
                                        qos: match packet.qos {
                                            QoS::AtMostOnce => MqttQoS::AtMostOnce,
                                            QoS::AtLeastOnce => MqttQoS::AtLeastOnce,
                                            QoS::ExactlyOnce => MqttQoS::ExactlyOnce,
                                        },
                                    };
                                    if let Some(cb) = callback.lock().await.as_ref() {
                                        cb(MqttEvent::MessageReceived(msg));
                                    }
                                }
                                Incoming::ConnAck(_) => {
                                    *state.lock().await = MqttConnectionState::Connected;
                                    if let Some(cb) = callback.lock().await.as_ref() {
                                        cb(MqttEvent::Connected);
                                    }
                                }
                                Incoming::Disconnect => {
                                    *state.lock().await = MqttConnectionState::Disconnected;
                                    if let Some(cb) = callback.lock().await.as_ref() {
                                        cb(MqttEvent::Disconnected);
                                    }
                                }
                                _ => {}
                            }
                        }
                        _ => {}
                    }
                }
                Err(e) => {
                    if let Some(cb) = callback.lock().await.as_ref() {
                        cb(MqttEvent::Error { error: e.to_string() });
                    }
                }
            }
        }
    }
}

#[async_trait]
impl MqttClient for NativeMqttClient {
    async fn connect(&self, config: MqttConfig) -> SdkResult<()> {
        let (host, port, transport) = Self::parse_broker_url(&config.broker_url)?;

        let mut mqtt_options = MqttOptions::new(&config.client_id, &host, port);
        mqtt_options.set_keep_alive(Duration::from_secs(config.keep_alive_secs));
        mqtt_options.set_clean_session(config.clean_session);

        if let Some(transport) = transport {
            mqtt_options.set_transport(transport);
        }

        if let Some(username) = config.username {
            let password = config.password.unwrap_or_default();
            mqtt_options.set_credentials(username, password);
        }

        let (client, eventloop) = AsyncClient::new(mqtt_options, 10);

        *self.client.lock().await = Some(client);
        *self.state.lock().await = MqttConnectionState::Connecting;

        // Start event loop
        let state = self.state.clone();
        let callback = self.event_callback.clone();
        let handle = tokio::spawn(Self::run_event_loop(eventloop, state, callback));
        *self.event_loop_handle.lock().await = Some(handle);

        Ok(())
    }

    async fn disconnect(&self) -> SdkResult<()> {
        if let Some(handle) = self.event_loop_handle.lock().await.take() {
            handle.abort();
        }

        if let Some(client) = self.client.lock().await.take() {
            client.disconnect().await.map_err(|e| {
                agentlink_core::error::SdkError::Mqtt(e.to_string())
            })?;
        }

        *self.state.lock().await = MqttConnectionState::Disconnected;

        Ok(())
    }

    async fn subscribe(&self, topic: &str, qos: MqttQoS) -> SdkResult<()> {
        let client = self.client.lock().await;
        if let Some(ref c) = *client {
            let rumqtt_qos = match qos {
                MqttQoS::AtMostOnce => QoS::AtMostOnce,
                MqttQoS::AtLeastOnce => QoS::AtLeastOnce,
                MqttQoS::ExactlyOnce => QoS::ExactlyOnce,
            };
            c.subscribe(topic, rumqtt_qos).await.map_err(|e| {
                agentlink_core::error::SdkError::Mqtt(e.to_string())
            })?;
            Ok(())
        } else {
            Err(agentlink_core::error::SdkError::NotConnected)
        }
    }

    async fn unsubscribe(&self, topic: &str) -> SdkResult<()> {
        let client = self.client.lock().await;
        if let Some(ref c) = *client {
            c.unsubscribe(topic).await.map_err(|e| {
                agentlink_core::error::SdkError::Mqtt(e.to_string())
            })?;
            Ok(())
        } else {
            Err(agentlink_core::error::SdkError::NotConnected)
        }
    }

    async fn publish(&self, message: MqttMessage) -> SdkResult<()> {
        let client = self.client.lock().await;
        if let Some(ref c) = *client {
            let qos = match message.qos {
                MqttQoS::AtMostOnce => QoS::AtMostOnce,
                MqttQoS::AtLeastOnce => QoS::AtLeastOnce,
                MqttQoS::ExactlyOnce => QoS::ExactlyOnce,
            };
            c.publish(&message.topic, qos, false, message.payload).await.map_err(|e| {
                agentlink_core::error::SdkError::Mqtt(e.to_string())
            })?;
            Ok(())
        } else {
            Err(agentlink_core::error::SdkError::NotConnected)
        }
    }

    fn connection_state(&self) -> MqttConnectionState {
        // This is synchronous, so we need to use try_lock or block
        // For simplicity, we'll return Disconnected if we can't get the lock
        // In practice, this should be called from an async context
        MqttConnectionState::Disconnected
    }
}

/// Extension trait for MQTT clients with event support
pub trait MqttClientExt: MqttClient {
    fn set_event_callback<F>(&self, callback: F)
    where
        F: Fn(MqttEvent) + Send + Sync + 'static;
}

impl MqttClientExt for NativeMqttClient {
    fn set_event_callback<F>(&self, callback: F)
    where
        F: Fn(MqttEvent) + Send + Sync + 'static,
    {
        // This is a bit hacky since we can't easily store the callback
        // In a real implementation, we'd need to use a channel or similar
        // For now, we'll just ignore this in the synchronous context
        let _ = callback;
    }
}