simple_pub_sub/
message.rs

1use log::{error, trace};
2use tokio::sync::broadcast::Sender;
3
4/// Packet Type Publish
5pub const PUBLISH: u8 = 0x02;
6/// Packet Type Subscribe
7pub const SUBSCRIBE: u8 = 0x03;
8/// Packet Type Unsubscribe
9pub const UNSUBSCRIBE: u8 = 0x04;
10/// Packet Type Query
11pub const QUERY: u8 = 0x05;
12/// Packet Type Publish Acknowledgement
13pub const PUBLISHACK: u8 = 0x0B;
14/// Packet Type Subscribe Acknowledgement
15pub const SUBSCRIBEACK: u8 = 0x0C;
16/// Packet Type Unsubscribe Acknowledgement
17pub const UNSUBSCRIBEACK: u8 = 0x0D;
18/// Packet Type Query Response
19pub const QUERYRESP: u8 = 0x0E;
20
21/// Packet type
22#[derive(Debug, Clone, PartialEq)]
23pub enum PktType {
24    /// publish
25    PUBLISH,
26    /// subscribe
27    SUBSCRIBE,
28    /// unsubscribe
29    UNSUBSCRIBE,
30    /// query the topics
31    QUERY,
32    /// acknoledgement to publish
33    PUBLISHACK,
34    /// acknoledgement to subscribe
35    SUBSCRIBEACK,
36    /// acknoledgement to unsubscribe
37    UNSUBSCRIBEACK,
38    /// response to the query packet
39    QUERYRESP,
40}
41impl PktType {
42    /// returns the byte value for a given packet type.
43    pub fn to_byte(&self) -> u8 {
44        match self {
45            PktType::PUBLISH => PUBLISH,
46            PktType::SUBSCRIBE => SUBSCRIBE,
47            PktType::UNSUBSCRIBE => UNSUBSCRIBE,
48            PktType::QUERY => QUERY,
49            PktType::PUBLISHACK => PUBLISHACK,
50            PktType::SUBSCRIBEACK => SUBSCRIBEACK,
51            PktType::UNSUBSCRIBEACK => UNSUBSCRIBEACK,
52            PktType::QUERYRESP => QUERYRESP,
53        }
54    }
55}
56impl ToString for PktType {
57    /// returns the string representation for the given packet type.
58    fn to_string(&self) -> String {
59        match self {
60            PktType::PUBLISH => "PUBLISH".to_string(),
61            PktType::SUBSCRIBE => "SUBSCRIBE".to_string(),
62            PktType::UNSUBSCRIBE => "UNSUBSCRIBE".to_string(),
63            PktType::QUERY => "QUERY".to_string(),
64            PktType::PUBLISHACK => "PUBLISH_ACK".to_string(),
65            PktType::SUBSCRIBEACK => "SUBSCRIBE_ACK".to_string(),
66            PktType::UNSUBSCRIBEACK => "UNSUBSCRIBE_ACK".to_string(),
67            PktType::QUERYRESP => "QUERY_RESP".to_string(),
68        }
69    }
70}
71
72/// supported versions for the pub-sub header format/protocol.
73pub const SUPPORTED_VERSIONS: [[u8; 2]; 1] = [[0x00, 0x01]];
74
75/// default version for the pub-sub header format/protocol.
76pub const DEFAULT_VERSION: [u8; 2] = [0x00, 0x01];
77
78/// error types for the `Header`.
79#[derive(Debug, Clone)]
80pub enum HeaderError {
81    /// invalid header buffer length
82    InvalidHeaderBufferLength,
83    /// invalid header value or padding
84    InvalidHeadOrTail,
85    /// unsupported version of the packet
86    UnsupportedVersion,
87    /// invalid message type
88    InvalidMessageType,
89    /// invalid topic length
90    InvalidTopicLength,
91    /// invalid message length
92    InvalidMessageLength,
93    /// invalid request/response type
94    InvalidResuestResponseType,
95}
96
97/// Header for the pub/sub packet
98/// total length 8 bytes.
99#[derive(Debug, Clone)]
100pub struct Header {
101    /// start byte of the packet, default value: 0x0F
102    pub header: u8,
103    /// pub-sub version: two bytes.
104    pub version: [u8; 2],
105    /// packet type: `PktType`
106    pub pkt_type: PktType,
107    /// topic length for publishing/subscribing/querying.
108    pub topic_length: u8,
109    /// message length: Max length 16 MB.
110    pub message_length: u16,
111    /// padding/endo of the header: 0x00
112    pub padding: u8,
113}
114
115/// structure containing the complete information about a message.
116#[derive(Debug, Clone)]
117pub struct Msg {
118    /// `Header`: the header of the message.
119    pub header: Header,
120    /// The topic for the message.
121    pub topic: String,
122    /// the actual message, bytes.
123    pub message: Vec<u8>,
124    /// `tokio::broadcast::sync::Sender` the channel for passing the messages across.
125    pub channel: Option<Sender<Msg>>,
126    /// client_id: to identify each socket connection/client.
127    pub client_id: Option<String>,
128}
129
130impl Msg {
131    /// Creates a new `Msg` with the given data.
132    pub fn new(pkt_type: PktType, topic: String, message: Option<Vec<u8>>) -> Msg {
133        let msg: Vec<u8> = match message {
134            Some(m) => m,
135            None => vec![],
136        };
137
138        Msg {
139            header: Header::new(pkt_type, topic.len() as u8, msg.len() as u16),
140            topic,
141            message: msg,
142            channel: None,
143            client_id: None,
144        }
145    }
146
147    /// adds the given channel to the message.
148    pub fn channel(&mut self, chan: Sender<Msg>) {
149        self.channel = Some(chan);
150    }
151
152    // returns the client id for the message.
153    pub fn client_id(&mut self, client_id: String) {
154        self.client_id = Some(client_id);
155    }
156
157    /// generates the response `Msg` with the given data.
158    pub fn response_msg(&self, _message: Vec<u8>) -> Result<Msg, String> {
159        let mut header: Header = match self.header.response_header() {
160            Ok(h) => h,
161            Err(e) => {
162                error!("unable to generate the response header: {:?}", e);
163                return Err("unable to generate response header".to_string());
164            }
165        };
166        header.message_length = _message.len() as u16;
167        Ok(Msg {
168            header,
169            topic: self.topic.clone(),
170            message: _message,
171            channel: None,
172            client_id: None,
173        })
174    }
175
176    /// returns bytes for the `Msg` that can be sent to the stream.
177    pub fn bytes(&self) -> Vec<u8> {
178        let mut buffer: Vec<u8> = self.header.bytes();
179        buffer.extend(self.topic.as_bytes().to_vec());
180        buffer.extend(self.message.clone());
181        trace!("the generated buffer is: {:?}", buffer);
182        buffer
183    }
184}
185
186impl Header {
187    /// creates a new `Header` with the given data.
188    pub fn new(pkt_type: PktType, topic_len: u8, message_len: u16) -> Header {
189        Header {
190            header: 0x0F,
191            version: DEFAULT_VERSION,
192            pkt_type,
193            topic_length: topic_len,
194            message_length: message_len,
195            padding: 0x00,
196        }
197    }
198
199    /// returns a `Header` for the response `Msg`.
200    pub fn response_header(&self) -> Result<Header, HeaderError> {
201        let resp_type: PktType = match self.pkt_type {
202            PktType::SUBSCRIBE => PktType::SUBSCRIBEACK,
203            PktType::PUBLISH => PktType::PUBLISHACK,
204            PktType::UNSUBSCRIBE => PktType::UNSUBSCRIBEACK,
205            PktType::QUERY => PktType::QUERYRESP,
206            _ => {
207                error!("invalid request/response type");
208                return Err(HeaderError::InvalidResuestResponseType);
209            }
210        };
211        Ok(Header {
212            header: 0x0F,
213            version: self.version,
214            pkt_type: resp_type,
215            topic_length: self.topic_length,
216            message_length: self.message_length,
217            padding: 0x00,
218        })
219    }
220
221    /// returns the bytes for `Header`.
222    pub fn bytes(&self) -> Vec<u8> {
223        let bytes_ = self.message_length.to_be_bytes();
224        vec![
225            self.header,
226            self.version[0],
227            self.version[1],
228            self.pkt_type.to_byte(),
229            self.topic_length,
230            bytes_[0],
231            bytes_[1],
232            self.padding,
233        ]
234    }
235
236    /// returns a `Header` from bytes.
237    pub fn from_vec(bytes: Vec<u8>) -> Result<Header, HeaderError> {
238        if !bytes.len() == 8 {
239            error!("invalid header buffer length, aborting");
240            return Err(HeaderError::InvalidHeaderBufferLength);
241        }
242
243        if !(bytes[0] == 0x0F && bytes[7] == 0x00) {
244            error!("invalid header value, aborting");
245            return Err(HeaderError::InvalidHeadOrTail);
246        }
247
248        if !SUPPORTED_VERSIONS.contains(&[bytes[1], bytes[2]]) {
249            error!("unsupported packet version, aborting");
250            return Err(HeaderError::UnsupportedVersion);
251        }
252
253        let pkt_type: PktType = match bytes[3] {
254            PUBLISH => PktType::PUBLISH,
255            SUBSCRIBE => PktType::SUBSCRIBE,
256            UNSUBSCRIBE => PktType::UNSUBSCRIBE,
257            QUERY => PktType::QUERY,
258            PUBLISHACK => PktType::PUBLISHACK,
259            SUBSCRIBEACK => PktType::SUBSCRIBEACK,
260            QUERYRESP => PktType::QUERYRESP,
261            _ => {
262                error!("invalid message type, aborting");
263                return Err(HeaderError::InvalidMessageType);
264            }
265        };
266
267        if bytes[4] == 0 {
268            match pkt_type {
269                PktType::PUBLISH => {
270                    error!("invalid topic length, aborting");
271                    return Err(HeaderError::InvalidTopicLength);
272                }
273                PktType::SUBSCRIBE => {
274                    error!("invalid topic length, aborting");
275                    return Err(HeaderError::InvalidTopicLength);
276                }
277                PktType::UNSUBSCRIBE => {
278                    error!("invalid topic length, aborting");
279                    return Err(HeaderError::InvalidTopicLength);
280                }
281                PktType::QUERY => {
282                    error!("invalid topic length, aborting");
283                    return Err(HeaderError::InvalidTopicLength);
284                }
285                _ => {}
286            };
287        }
288
289        let message_length = ((bytes[5] as u16) << 8) | bytes[6] as u16;
290
291        if message_length == 0 {
292            match pkt_type {
293                PktType::PUBLISH => {
294                    error!("invalid message length, aborting");
295                    return Err(HeaderError::InvalidMessageLength);
296                }
297                PktType::QUERY => {
298                    error!("invalid message length, aborting");
299                    return Err(HeaderError::InvalidMessageLength);
300                }
301                _ => {}
302            };
303        }
304        Ok(Header {
305            header: 0x0F,
306            version: [bytes[1], bytes[2]],
307            pkt_type,
308            topic_length: bytes[4],
309            message_length,
310            padding: 0x00,
311        })
312    }
313}
314
315/// returns a response `Msg`.
316pub fn get_msg_response(msg: Msg) -> Result<Vec<u8>, String> {
317    let mut resp: Vec<u8> = match msg.response_msg(msg.message.clone()) {
318        Ok(m) => m.header.bytes(),
319        Err(e) => {
320            error!(
321                "error occurred while generating the response message: {}",
322                e
323            );
324            return Err(e);
325        }
326    };
327    resp.extend(msg.topic.bytes());
328    Ok(resp)
329}