mqtt5_protocol/session/
topic_alias.rs

1use crate::error::{MqttError, Result};
2use crate::prelude::{HashMap, String, ToString};
3
4#[derive(Debug)]
5pub struct TopicAliasManager {
6    topic_alias_maximum: u16,
7    alias_to_topic: HashMap<u16, String>,
8    topic_to_alias: HashMap<String, u16>,
9    next_alias: u16,
10}
11
12impl TopicAliasManager {
13    #[must_use]
14    pub fn new(topic_alias_maximum: u16) -> Self {
15        Self {
16            topic_alias_maximum,
17            alias_to_topic: HashMap::new(),
18            topic_to_alias: HashMap::new(),
19            next_alias: 1,
20        }
21    }
22
23    #[must_use]
24    pub fn get_or_create_alias(&mut self, topic: &str) -> Option<u16> {
25        if let Some(&alias) = self.topic_to_alias.get(topic) {
26            return Some(alias);
27        }
28
29        if self.topic_alias_maximum == 0
30            || self.alias_to_topic.len() >= usize::from(self.topic_alias_maximum)
31        {
32            return None;
33        }
34
35        while self.alias_to_topic.contains_key(&self.next_alias)
36            && self.next_alias <= self.topic_alias_maximum
37        {
38            self.next_alias += 1;
39            if self.next_alias > self.topic_alias_maximum {
40                self.next_alias = 1;
41            }
42        }
43
44        let alias = self.next_alias;
45        self.alias_to_topic.insert(alias, topic.to_string());
46        self.topic_to_alias.insert(topic.to_string(), alias);
47
48        self.next_alias += 1;
49        if self.next_alias > self.topic_alias_maximum {
50            self.next_alias = 1;
51        }
52
53        Some(alias)
54    }
55
56    /// # Errors
57    /// Returns `TopicAliasInvalid` if alias is 0 or exceeds the maximum.
58    pub fn register_alias(&mut self, alias: u16, topic: &str) -> Result<()> {
59        if alias == 0 || alias > self.topic_alias_maximum {
60            return Err(MqttError::TopicAliasInvalid(alias));
61        }
62
63        if let Some(old_topic) = self.alias_to_topic.get(&alias) {
64            self.topic_to_alias.remove(old_topic);
65        }
66
67        self.alias_to_topic.insert(alias, topic.to_string());
68        self.topic_to_alias.insert(topic.to_string(), alias);
69
70        Ok(())
71    }
72
73    #[must_use]
74    pub fn get_topic(&self, alias: u16) -> Option<&str> {
75        self.alias_to_topic.get(&alias).map(String::as_str)
76    }
77
78    #[must_use]
79    pub fn get_alias(&self, topic: &str) -> Option<u16> {
80        self.topic_to_alias.get(topic).copied()
81    }
82
83    pub fn clear(&mut self) {
84        self.alias_to_topic.clear();
85        self.topic_to_alias.clear();
86        self.next_alias = 1;
87    }
88
89    #[must_use]
90    pub fn topic_alias_maximum(&self) -> u16 {
91        self.topic_alias_maximum
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn test_topic_alias_basic() {
101        let mut ta = TopicAliasManager::new(10);
102
103        let alias1 = ta.get_or_create_alias("topic/1").unwrap();
104        assert_eq!(alias1, 1);
105
106        let alias2 = ta.get_or_create_alias("topic/2").unwrap();
107        assert_eq!(alias2, 2);
108
109        let alias1_again = ta.get_or_create_alias("topic/1").unwrap();
110        assert_eq!(alias1_again, 1);
111
112        assert_eq!(ta.get_topic(1), Some("topic/1"));
113        assert_eq!(ta.get_alias("topic/1"), Some(1));
114    }
115
116    #[test]
117    fn test_topic_alias_register() {
118        let mut ta = TopicAliasManager::new(5);
119
120        ta.register_alias(3, "remote/topic").unwrap();
121        assert_eq!(ta.get_topic(3), Some("remote/topic"));
122
123        assert!(ta.register_alias(0, "topic").is_err());
124        assert!(ta.register_alias(6, "topic").is_err());
125
126        ta.register_alias(3, "new/topic").unwrap();
127        assert_eq!(ta.get_topic(3), Some("new/topic"));
128        assert!(ta.get_alias("remote/topic").is_none());
129    }
130
131    #[test]
132    fn test_topic_alias_limit() {
133        let mut ta = TopicAliasManager::new(2);
134
135        let alias1 = ta.get_or_create_alias("topic/1");
136        let alias2 = ta.get_or_create_alias("topic/2");
137        let alias3 = ta.get_or_create_alias("topic/3");
138
139        assert!(alias1.is_some());
140        assert!(alias2.is_some());
141        assert!(alias3.is_none());
142    }
143
144    #[test]
145    fn test_topic_alias_clear() {
146        let mut ta = TopicAliasManager::new(10);
147
148        let _ = ta.get_or_create_alias("topic/1");
149        let _ = ta.get_or_create_alias("topic/2");
150        ta.register_alias(5, "topic/5").unwrap();
151
152        ta.clear();
153
154        assert!(ta.get_topic(1).is_none());
155        assert!(ta.get_topic(5).is_none());
156        assert!(ta.get_alias("topic/1").is_none());
157    }
158}