simple_pub_sub/client/
mod.rs

1use crate::message;
2use crate::stream;
3use log::{error, info, trace};
4use std::fs::File;
5use std::io::ErrorKind;
6use std::io::Read;
7use tokio::{
8    io::{AsyncReadExt, AsyncWriteExt},
9    net::TcpStream,
10    net::UnixStream,
11};
12use tokio_native_tls::native_tls::{Certificate, TlsConnector};
13use tokio_native_tls::TlsStream;
14
15/// Simple pub sub Client for Tcp connection
16#[derive(Debug, Clone)]
17pub struct PubSubTcpClient {
18    pub server: String,
19    pub port: u16,
20    pub cert: Option<String>,
21    pub cert_password: Option<String>,
22}
23
24/// Simple pub sub Client for Unix connection
25#[derive(Debug, Clone)]
26pub struct PubSubUnixClient {
27    pub path: String,
28}
29
30/// Simple pub sub Client
31#[derive(Debug, Clone)]
32pub enum PubSubClient {
33    Tcp(PubSubTcpClient),
34    Unix(PubSubUnixClient),
35}
36
37/// Stream for Tcp and Unix connection
38#[derive(Debug)]
39pub enum StreamType {
40    Tcp(TcpStream),
41    Tls(TlsStream<TcpStream>),
42    Unix(UnixStream),
43}
44
45/// on_message callback function
46type Callback = fn(String, Vec<u8>);
47
48/// Simple pub sub Client
49#[derive(Debug)]
50pub struct Client {
51    pub client_type: PubSubClient,
52    stream: Option<StreamType>,
53    callback: Option<Callback>,
54}
55
56/// default implementation for callback function
57pub fn on_message(topic: String, message: Vec<u8>) {
58    match String::from_utf8(message.clone()) {
59        Ok(msg_str) => {
60            info!("topic: {} message: {}", topic, msg_str);
61        }
62        Err(_) => {
63            info!("topic: {} message: {:?}", topic, message);
64        }
65    };
66}
67
68impl Client {
69    /// Creates a new instance of `Client`
70    pub fn new(client_type: PubSubClient) -> Client {
71        Client {
72            client_type,
73            stream: None,
74            callback: None,
75        }
76    }
77
78    /// Sets the on_message callback function
79    pub fn on_message(&mut self, callback: Callback) {
80        self.callback = Some(callback)
81    }
82
83    async fn connect_tls(&mut self, url: String, cert: String) -> Result<(), tokio::io::Error> {
84        // Load CA certificate
85        let mut file = match File::open(cert) {
86            Ok(file) => file,
87            Err(e) => {
88                error!("unable to open the certificate file: {}", e);
89                return Err(tokio::io::Error::new(tokio::io::ErrorKind::Other, e));
90            }
91        };
92        let mut ca_cert = vec![];
93        match file.read_to_end(&mut ca_cert) {
94            Ok(size) => size,
95            Err(e) => {
96                error!("unable to read the CA file: {}", e);
97                return Err(tokio::io::Error::new(tokio::io::ErrorKind::Other, e));
98            }
99        };
100        let ca_cert = match Certificate::from_pem(&ca_cert) {
101            Ok(cert) => cert,
102            Err(e) => {
103                error!("unable to parse the CA certificate: {}", e);
104                return Err(tokio::io::Error::new(tokio::io::ErrorKind::Other, e));
105            }
106        };
107
108        // Configure TLS
109        let connector = match TlsConnector::builder()
110            .add_root_certificate(ca_cert)
111            .build()
112        {
113            Ok(connector) => connector,
114            Err(e) => {
115                error!("cannot create TLS connector: {}", e);
116                return Err(tokio::io::Error::new(tokio::io::ErrorKind::Other, e));
117            }
118        };
119
120        let connector = tokio_native_tls::TlsConnector::from(connector);
121
122        // Connect to the server
123        let stream = match TcpStream::connect(url).await {
124            Ok(stream) => stream,
125            Err(e) => {
126                error!("failed to connect to server: {}", e);
127                return Err(tokio::io::Error::new(tokio::io::ErrorKind::Other, e));
128            }
129        };
130
131        self.stream = match connector.connect("localhost", stream).await {
132            Ok(stream) => Some(StreamType::Tls(stream)),
133            Err(e) => {
134                error!("could not establish the tls connection: {}", e);
135                return Err(tokio::io::Error::new(tokio::io::ErrorKind::Other, e));
136            }
137        };
138        Ok(())
139    }
140
141    /// Connects to the server
142    pub async fn connect(&mut self) -> Result<(), tokio::io::Error> {
143        match self.client_type.clone() {
144            PubSubClient::Tcp(tcp_client) => {
145                let server_url: String = format!("{}:{}", tcp_client.server, tcp_client.port);
146                if let Some(cert) = tcp_client.cert {
147                    self.connect_tls(server_url, cert).await?;
148                } else {
149                    let stream = TcpStream::connect(server_url).await?;
150                    self.stream = Some(StreamType::Tcp(stream));
151                }
152            }
153            PubSubClient::Unix(unix_stream) => {
154                let path = unix_stream.path;
155                let stream = UnixStream::connect(path).await?;
156                self.stream = Some(StreamType::Unix(stream));
157            }
158        }
159        Ok(())
160    }
161
162    /// Sends the message to the given server and returns the ack
163    /// the server could be either a tcp or unix server
164    pub async fn post(&mut self, msg: message::Msg) -> Result<Vec<u8>, tokio::io::Error> {
165        self.write(msg.bytes()).await?;
166        let mut buf: Vec<u8>;
167        buf = vec![0; 8];
168        match self.read(&mut buf).await {
169            Ok(()) => {
170                trace!("buf: {:?}", buf);
171                match message::Header::from_vec(buf.clone()) {
172                    Ok(resp) => {
173                        trace!("resp: {:?}", resp);
174                        // read the remaining message
175                        let mut buf_buf = Vec::with_capacity(resp.message_length as usize);
176                        trace!("reading remaining bytes");
177                        match self.read_buf(&mut buf_buf).await {
178                            Ok(()) => {
179                                buf.extend(buf_buf);
180                            }
181                            Err(e) => {
182                                return Err(e);
183                            }
184                        }
185                    }
186                    Err(e) => {
187                        return Err(tokio::io::Error::new(ErrorKind::Other, format!("{:?}", e)));
188                    }
189                };
190
191                Ok(buf)
192            }
193            Err(e) => Err(e),
194        }
195    }
196
197    /// Publishes the message to the given topic
198    pub async fn publish(
199        &mut self,
200        topic: String,
201        message: Vec<u8>,
202    ) -> Result<(), tokio::io::Error> {
203        let msg: message::Msg = message::Msg::new(message::PktType::PUBLISH, topic, Some(message));
204        trace!("msg: {:?}", msg);
205
206        let buf = match self.post(msg).await {
207            Ok(resp) => resp,
208            Err(e) => return Err(e),
209        };
210
211        trace!("the raw buffer is: {:?}", buf);
212        let resp_: message::Header = match message::Header::from_vec(buf) {
213            Ok(resp) => resp,
214            Err(e) => {
215                return Err(tokio::io::Error::new(ErrorKind::Other, format!("{:?}", e)));
216            }
217        };
218        trace!("{:?}", resp_);
219
220        Ok(())
221    }
222
223    /// Sends the query message to the server
224    pub async fn query(&mut self, topic: String) -> Result<String, tokio::io::Error> {
225        let msg: message::Msg = message::Msg::new(
226            message::PktType::QUERY,
227            topic,
228            Some(" ".to_string().as_bytes().to_vec()),
229        );
230        trace!("msg: {:?}", msg);
231
232        self.write(msg.bytes()).await?;
233        let msg = self.read_message().await;
234        match msg {
235            Ok(m) => match String::from_utf8(m.message) {
236                Ok(s) => Ok(s),
237                Err(e) => Err(tokio::io::Error::new(ErrorKind::Other, e.to_string())),
238            },
239            Err(e) => Err(e),
240        }
241    }
242
243    /// subscribes to the given topic
244    pub async fn subscribe(&mut self, topic: String) -> Result<(), tokio::io::Error> {
245        let msg: message::Msg = message::Msg::new(message::PktType::SUBSCRIBE, topic, None);
246        trace!("msg: {:?}", msg);
247        self.write(msg.bytes()).await?;
248        if let Some(callback) = self.callback {
249            loop {
250                match self.read_message().await {
251                    Ok(m) => {
252                        callback(m.topic, m.message);
253                    }
254                    Err(e) => {
255                        return Err(e);
256                    }
257                }
258            }
259        } else {
260            Err(tokio::io::Error::new(
261                tokio::io::ErrorKind::Other,
262                "no callback function is set".to_string(),
263            ))
264        }
265    }
266    pub async fn write(&mut self, message: Vec<u8>) -> Result<(), tokio::io::Error> {
267        if let Some(stream) = &mut self.stream {
268            match stream {
269                StreamType::Tls(tls_stream) => match tls_stream.write_all(&message).await {
270                    Ok(size) => {
271                        trace!("{:?} bytes written", size);
272                        Ok(())
273                    }
274                    Err(e) => Err(tokio::io::Error::new(
275                        tokio::io::ErrorKind::Other,
276                        e.to_string(),
277                    )),
278                },
279                StreamType::Tcp(ref mut tcp_stream) => match tcp_stream.write_all(&message).await {
280                    Ok(size) => {
281                        trace!("{:?} bytes written", size);
282                        Ok(())
283                    }
284                    Err(e) => Err(e),
285                },
286
287                StreamType::Unix(ref mut unix_stream) => {
288                    match unix_stream.write_all(&message).await {
289                        Ok(size) => {
290                            trace!("{:?} bytes written", size);
291                            Ok(())
292                        }
293
294                        Err(e) => Err(e),
295                    }
296                }
297            }
298        } else {
299            Err(tokio::io::Error::new(
300                tokio::io::ErrorKind::Other,
301                "client not connected yet".to_string(),
302            ))
303        }
304    }
305
306    pub async fn read(&mut self, message: &mut [u8]) -> Result<(), tokio::io::Error> {
307        if let Some(stream) = &mut self.stream {
308            match stream {
309                StreamType::Tls(ref mut tls_stream) => match tls_stream.read(message).await {
310                    Ok(_) => Ok(()),
311                    Err(e) => Err(e),
312                },
313                StreamType::Tcp(ref mut tcp_stream) => match tcp_stream.read(message).await {
314                    Ok(_) => Ok(()),
315                    Err(e) => Err(e),
316                },
317                StreamType::Unix(ref mut unix_stream) => match unix_stream.read(message).await {
318                    Ok(_) => Ok(()),
319                    Err(e) => Err(e),
320                },
321            }
322        } else {
323            Err(tokio::io::Error::new(
324                tokio::io::ErrorKind::Other,
325                "client not connected yet.".to_string(),
326            ))
327        }
328    }
329    pub async fn read_buf(&mut self, message: &mut Vec<u8>) -> Result<(), tokio::io::Error> {
330        if let Some(stream) = &mut self.stream {
331            match stream {
332                StreamType::Tls(ref mut tls_stream) => match tls_stream.read_buf(message).await {
333                    Ok(_) => Ok(()),
334                    Err(e) => Err(e),
335                },
336                StreamType::Tcp(ref mut tcp_stream) => match tcp_stream.read_buf(message).await {
337                    Ok(_) => Ok(()),
338                    Err(e) => Err(e),
339                },
340                StreamType::Unix(ref mut unix_stream) => {
341                    match unix_stream.read_buf(message).await {
342                        Ok(_) => Ok(()),
343                        Err(e) => Err(e),
344                    }
345                }
346            }
347        } else {
348            Err(tokio::io::Error::new(
349                tokio::io::ErrorKind::Other,
350                "client not connected yet.".to_string(),
351            ))
352        }
353    }
354
355    pub async fn read_message(&mut self) -> Result<message::Msg, tokio::io::Error> {
356        if let Some(stream) = &mut self.stream {
357            match stream {
358                StreamType::Tcp(stream) => match stream::read_message(stream).await {
359                    Ok(msg) => Ok(msg),
360                    Err(e) => {
361                        error!("could not read the message from the tcp stream: {}", e);
362                        Err(e)
363                    }
364                },
365                StreamType::Tls(stream) => match stream::read_message(stream).await {
366                    Ok(msg) => Ok(msg),
367                    Err(e) => {
368                        error!("could not read the message from the tcp stream: {}", e);
369                        Err(e)
370                    }
371                },
372                StreamType::Unix(stream) => match stream::read_message(stream).await {
373                    Ok(msg) => Ok(msg),
374                    Err(e) => {
375                        error!("could not read the message from the tcp stream: {}", e);
376                        Err(e)
377                    }
378                },
379            }
380        } else {
381            Err(tokio::io::Error::new(
382                tokio::io::ErrorKind::Other,
383                "client not connected yet.".to_string(),
384            ))
385        }
386    }
387}
simple_pub_sub/client/mod.rs

simple_pub_sub/client/
mod.rs
