Skip to main content

sz_orm_mqtt/
broker.rs

1use crate::error::MqttError;
2use crate::topics::{topic_matches, TopicFilter};
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokio::sync::RwLock;
6
7pub use crate::qos::QoS;
8
9#[derive(Debug, Clone)]
10pub struct MqttTopic {
11    pub name: String,
12    pub qos: QoS,
13}
14
15impl MqttTopic {
16    pub fn new(name: impl Into<String>) -> Self {
17        Self {
18            name: name.into(),
19            qos: QoS::default(),
20        }
21    }
22
23    pub fn with_qos(mut self, qos: QoS) -> Self {
24        self.qos = qos;
25        self
26    }
27
28    pub fn wildcard(&self) -> bool {
29        self.name.contains('#') || self.name.contains('+')
30    }
31
32    pub fn levels(&self) -> Vec<&str> {
33        self.name.split('/').collect()
34    }
35
36    pub fn matches(&self, topic: &str) -> bool {
37        topic_matches(topic, &self.name)
38    }
39}
40
41impl From<&str> for MqttTopic {
42    fn from(s: &str) -> Self {
43        MqttTopic::new(s)
44    }
45}
46
47impl From<String> for MqttTopic {
48    fn from(s: String) -> Self {
49        MqttTopic::new(s)
50    }
51}
52
53#[derive(Debug, Clone)]
54pub struct MqttMessage {
55    pub topic: String,
56    pub payload: Vec<u8>,
57    pub qos: QoS,
58    pub retain: bool,
59    pub client_id: Option<String>,
60    pub timestamp: i64,
61}
62
63impl MqttMessage {
64    pub fn new(topic: impl Into<String>, payload: Vec<u8>) -> Self {
65        Self {
66            topic: topic.into(),
67            payload,
68            qos: QoS::default(),
69            retain: false,
70            client_id: None,
71            timestamp: current_timestamp(),
72        }
73    }
74
75    pub fn with_qos(mut self, qos: QoS) -> Self {
76        self.qos = qos;
77        self
78    }
79
80    pub fn retain(mut self) -> Self {
81        self.retain = true;
82        self
83    }
84
85    pub fn with_client(mut self, client_id: impl Into<String>) -> Self {
86        self.client_id = Some(client_id.into());
87        self
88    }
89
90    pub fn text(&self) -> Option<&str> {
91        std::str::from_utf8(&self.payload).ok()
92    }
93
94    pub fn json<T: serde::de::DeserializeOwned>(&self) -> Option<T> {
95        serde_json::from_slice(&self.payload).ok()
96    }
97}
98
99impl MqttMessage {
100    pub fn text_message(topic: impl Into<String>, text: impl Into<String>) -> Self {
101        Self::new(topic, text.into().into_bytes())
102    }
103
104    pub fn json_message<T: serde::Serialize>(
105        topic: impl Into<String>,
106        data: &T,
107    ) -> Result<Self, MqttError> {
108        let payload = serde_json::to_vec(data)?;
109        Ok(Self::new(topic, payload))
110    }
111}
112
113fn current_timestamp() -> i64 {
114    use std::time::{SystemTime, UNIX_EPOCH};
115    SystemTime::now()
116        .duration_since(UNIX_EPOCH)
117        .unwrap_or_default()
118        .as_millis() as i64
119}
120
121pub struct MqttConfig {
122    pub broker_url: String,
123    pub client_id: Option<String>,
124    pub username: Option<String>,
125    pub password: Option<String>,
126    pub keep_alive: u16,
127    pub clean_session: bool,
128    pub topics: Vec<MqttTopic>,
129}
130
131impl Default for MqttConfig {
132    fn default() -> Self {
133        Self {
134            broker_url: "tcp://localhost:1883".to_string(),
135            client_id: None,
136            username: None,
137            password: None,
138            keep_alive: 60,
139            clean_session: true,
140            topics: Vec::new(),
141        }
142    }
143}
144
145impl MqttConfig {
146    pub fn new(broker_url: impl Into<String>) -> Self {
147        Self {
148            broker_url: broker_url.into(),
149            ..Default::default()
150        }
151    }
152
153    pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {
154        self.client_id = Some(client_id.into());
155        self
156    }
157
158    pub fn with_auth(mut self, username: impl Into<String>, password: impl Into<String>) -> Self {
159        self.username = Some(username.into());
160        self.password = Some(password.into());
161        self
162    }
163
164    pub fn with_keep_alive(mut self, seconds: u16) -> Self {
165        self.keep_alive = seconds;
166        self
167    }
168
169    pub fn with_topics(mut self, topics: Vec<MqttTopic>) -> Self {
170        self.topics = topics;
171        self
172    }
173
174    pub fn add_topic(mut self, topic: MqttTopic) -> Self {
175        self.topics.push(topic);
176        self
177    }
178}
179
180#[derive(Debug, Clone)]
181struct Subscription {
182    filter: String,
183    qos: QoS,
184}
185
186pub struct MqttPlugin {
187    config: MqttConfig,
188    connected: bool,
189    messages: Arc<RwLock<Vec<MqttMessage>>>,
190    retained: Arc<RwLock<HashMap<String, MqttMessage>>>,
191    subscriptions: Arc<RwLock<Vec<Subscription>>>,
192}
193
194impl MqttPlugin {
195    pub fn new(config: MqttConfig) -> Self {
196        Self {
197            config,
198            connected: false,
199            messages: Arc::new(RwLock::new(Vec::new())),
200            retained: Arc::new(RwLock::new(HashMap::new())),
201            subscriptions: Arc::new(RwLock::new(Vec::new())),
202        }
203    }
204
205    pub async fn connect(&mut self) -> Result<(), MqttError> {
206        self.connected = true;
207        Ok(())
208    }
209
210    pub async fn disconnect(&mut self) -> Result<(), MqttError> {
211        self.connected = false;
212        Ok(())
213    }
214
215    pub fn is_connected(&self) -> bool {
216        self.connected
217    }
218
219    pub async fn publish(&self, topic: &str, payload: Vec<u8>, qos: QoS) -> Result<(), MqttError> {
220        if !self.connected {
221            return Err(MqttError::Connection("Not connected".to_string()));
222        }
223
224        let client_id = self.config.client_id.clone();
225        let msg = MqttMessage::new(topic, payload)
226            .with_qos(qos)
227            .with_client(client_id.unwrap_or_else(|| "default".to_string()));
228
229        if msg.retain {
230            let mut retained = self.retained.write().await;
231            retained.insert(topic.to_string(), msg.clone());
232        }
233
234        let mut messages = self.messages.write().await;
235        messages.push(msg);
236        Ok(())
237    }
238
239    pub async fn publish_retain(
240        &self,
241        topic: &str,
242        payload: Vec<u8>,
243        qos: QoS,
244    ) -> Result<(), MqttError> {
245        if !self.connected {
246            return Err(MqttError::Connection("Not connected".to_string()));
247        }
248
249        let client_id = self.config.client_id.clone();
250        let msg = MqttMessage::new(topic, payload)
251            .with_qos(qos)
252            .with_client(client_id.unwrap_or_else(|| "default".to_string()))
253            .retain();
254
255        {
256            let mut retained = self.retained.write().await;
257            retained.insert(topic.to_string(), msg.clone());
258        }
259
260        let mut messages = self.messages.write().await;
261        messages.push(msg);
262        Ok(())
263    }
264
265    pub async fn subscribe(&self, topic: &str, qos: QoS) -> Result<(), MqttError> {
266        if !self.connected {
267            return Err(MqttError::Connection("Not connected".to_string()));
268        }
269
270        let _filter = TopicFilter::new(topic).map_err(|e| MqttError::Subscribe(e.to_string()))?;
271
272        let mut subscriptions = self.subscriptions.write().await;
273        if let Some(existing) = subscriptions.iter_mut().find(|s| s.filter == topic) {
274            existing.qos = qos;
275        } else {
276            subscriptions.push(Subscription {
277                filter: topic.to_string(),
278                qos,
279            });
280        }
281        Ok(())
282    }
283
284    pub async fn unsubscribe(&self, topic: &str) -> Result<(), MqttError> {
285        if !self.connected {
286            return Err(MqttError::Connection("Not connected".to_string()));
287        }
288
289        let mut subscriptions = self.subscriptions.write().await;
290        subscriptions.retain(|s| s.filter != topic);
291        Ok(())
292    }
293
294    pub fn topic_matches(&self, topic: &str, filter: &str) -> bool {
295        topic_matches(topic, filter)
296    }
297
298    pub async fn message_count(&self) -> usize {
299        let messages = self.messages.read().await;
300        messages.len()
301    }
302
303    pub async fn messages_for(&self, topic: &str) -> Vec<MqttMessage> {
304        let messages = self.messages.read().await;
305        messages
306            .iter()
307            .filter(|m| m.topic == topic)
308            .cloned()
309            .collect()
310    }
311
312    pub async fn messages_matching(&self, filter: &str) -> Vec<MqttMessage> {
313        let messages = self.messages.read().await;
314        messages
315            .iter()
316            .filter(|m| topic_matches(&m.topic, filter))
317            .cloned()
318            .collect()
319    }
320
321    pub async fn subscription_count(&self) -> usize {
322        let subscriptions = self.subscriptions.read().await;
323        subscriptions.len()
324    }
325
326    pub async fn is_subscribed(&self, topic: &str) -> bool {
327        let subscriptions = self.subscriptions.read().await;
328        subscriptions.iter().any(|s| s.filter == topic)
329    }
330
331    pub async fn subscription_qos(&self, topic: &str) -> Option<QoS> {
332        let subscriptions = self.subscriptions.read().await;
333        subscriptions
334            .iter()
335            .find(|s| s.filter == topic)
336            .map(|s| s.qos)
337    }
338
339    pub async fn retained_count(&self) -> usize {
340        let retained = self.retained.read().await;
341        retained.len()
342    }
343
344    pub async fn retained_get(&self, topic: &str) -> Option<MqttMessage> {
345        let retained = self.retained.read().await;
346        retained.get(topic).cloned()
347    }
348
349    pub async fn clear_messages(&self) {
350        let mut messages = self.messages.write().await;
351        messages.clear();
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358
359    #[tokio::test]
360    async fn test_publish_stores_message() {
361        let config = MqttConfig::default();
362        let mut plugin = MqttPlugin::new(config);
363        plugin.connect().await.unwrap();
364
365        plugin
366            .publish("sensor/temp", b"23.5".to_vec(), QoS::AtLeastOnce)
367            .await
368            .unwrap();
369        assert_eq!(plugin.message_count().await, 1);
370
371        let messages = plugin.messages_for("sensor/temp").await;
372        assert_eq!(messages.len(), 1);
373        assert_eq!(messages[0].payload, b"23.5");
374        assert_eq!(messages[0].qos, QoS::AtLeastOnce);
375    }
376
377    #[tokio::test]
378    async fn test_publish_multiple_messages() {
379        let config = MqttConfig::default();
380        let mut plugin = MqttPlugin::new(config);
381        plugin.connect().await.unwrap();
382
383        plugin
384            .publish("topic/a", b"1".to_vec(), QoS::AtMostOnce)
385            .await
386            .unwrap();
387        plugin
388            .publish("topic/b", b"2".to_vec(), QoS::AtMostOnce)
389            .await
390            .unwrap();
391        plugin
392            .publish("topic/a", b"3".to_vec(), QoS::AtMostOnce)
393            .await
394            .unwrap();
395
396        assert_eq!(plugin.message_count().await, 3);
397        assert_eq!(plugin.messages_for("topic/a").await.len(), 2);
398        assert_eq!(plugin.messages_for("topic/b").await.len(), 1);
399    }
400
401    #[tokio::test]
402    async fn test_publish_not_connected_fails() {
403        let config = MqttConfig::default();
404        let plugin = MqttPlugin::new(config);
405        let result = plugin.publish("topic", vec![], QoS::AtMostOnce).await;
406        assert!(result.is_err());
407    }
408
409    #[tokio::test]
410    async fn test_subscribe_registers_subscription() {
411        let config = MqttConfig::default();
412        let mut plugin = MqttPlugin::new(config);
413        plugin.connect().await.unwrap();
414
415        plugin.subscribe("home/#", QoS::ExactlyOnce).await.unwrap();
416        assert_eq!(plugin.subscription_count().await, 1);
417        assert!(plugin.is_subscribed("home/#").await);
418        assert_eq!(
419            plugin.subscription_qos("home/#").await,
420            Some(QoS::ExactlyOnce)
421        );
422    }
423
424    #[tokio::test]
425    async fn test_subscribe_updates_qos_for_existing() {
426        let config = MqttConfig::default();
427        let mut plugin = MqttPlugin::new(config);
428        plugin.connect().await.unwrap();
429
430        plugin.subscribe("home/#", QoS::AtMostOnce).await.unwrap();
431        plugin.subscribe("home/#", QoS::ExactlyOnce).await.unwrap();
432        assert_eq!(plugin.subscription_count().await, 1);
433        assert_eq!(
434            plugin.subscription_qos("home/#").await,
435            Some(QoS::ExactlyOnce)
436        );
437    }
438
439    #[tokio::test]
440    async fn test_unsubscribe_removes_subscription() {
441        let config = MqttConfig::default();
442        let mut plugin = MqttPlugin::new(config);
443        plugin.connect().await.unwrap();
444
445        plugin.subscribe("home/#", QoS::AtMostOnce).await.unwrap();
446        assert!(plugin.is_subscribed("home/#").await);
447
448        plugin.unsubscribe("home/#").await.unwrap();
449        assert!(!plugin.is_subscribed("home/#").await);
450        assert_eq!(plugin.subscription_count().await, 0);
451    }
452
453    #[tokio::test]
454    async fn test_subscribe_not_connected_fails() {
455        let config = MqttConfig::default();
456        let plugin = MqttPlugin::new(config);
457        let result = plugin.subscribe("home/#", QoS::AtMostOnce).await;
458        assert!(result.is_err());
459    }
460
461    #[tokio::test]
462    async fn test_messages_matching_wildcard() {
463        let config = MqttConfig::default();
464        let mut plugin = MqttPlugin::new(config);
465        plugin.connect().await.unwrap();
466
467        plugin
468            .publish("home/living/temp", b"23".to_vec(), QoS::AtMostOnce)
469            .await
470            .unwrap();
471        plugin
472            .publish("home/kitchen/temp", b"20".to_vec(), QoS::AtMostOnce)
473            .await
474            .unwrap();
475        plugin
476            .publish("office/temp", b"25".to_vec(), QoS::AtMostOnce)
477            .await
478            .unwrap();
479
480        let matched = plugin.messages_matching("home/+/temp").await;
481        assert_eq!(matched.len(), 2);
482
483        let matched_all = plugin.messages_matching("home/#").await;
484        assert_eq!(matched_all.len(), 2);
485    }
486
487    #[tokio::test]
488    async fn test_retained_messages() {
489        let config = MqttConfig::default();
490        let mut plugin = MqttPlugin::new(config);
491        plugin.connect().await.unwrap();
492
493        plugin
494            .publish_retain("config/version", b"1.0".to_vec(), QoS::AtLeastOnce)
495            .await
496            .unwrap();
497
498        assert_eq!(plugin.retained_count().await, 1);
499        let retained = plugin
500            .retained_get("config/version")
501            .await
502            .expect("retained message should exist");
503        assert_eq!(retained.payload, b"1.0");
504        assert!(retained.retain);
505    }
506
507    #[tokio::test]
508    async fn test_disconnect_clears_connection_state() {
509        let config = MqttConfig::default();
510        let mut plugin = MqttPlugin::new(config);
511        plugin.connect().await.unwrap();
512        assert!(plugin.is_connected());
513
514        plugin.disconnect().await.unwrap();
515        assert!(!plugin.is_connected());
516
517        let result = plugin.publish("topic", vec![], QoS::AtMostOnce).await;
518        assert!(result.is_err());
519    }
520
521    #[tokio::test]
522    async fn test_clear_messages() {
523        let config = MqttConfig::default();
524        let mut plugin = MqttPlugin::new(config);
525        plugin.connect().await.unwrap();
526
527        plugin
528            .publish("topic", b"data".to_vec(), QoS::AtMostOnce)
529            .await
530            .unwrap();
531        assert_eq!(plugin.message_count().await, 1);
532
533        plugin.clear_messages().await;
534        assert_eq!(plugin.message_count().await, 0);
535    }
536
537    #[test]
538    fn test_mqtt_topic_matches_method() {
539        let topic = MqttTopic::new("home/+/temp");
540        assert!(topic.matches("home/living/temp"));
541        assert!(!topic.matches("home/living/humidity"));
542    }
543}