pawkit_net_websocket/native/
mod.rs

1use futures_util::{stream::FusedStream, SinkExt, StreamExt};
2use tokio::net::TcpStream;
3use tokio_tungstenite::{connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream};
4
5use crate::{WebsocketError, WebsocketMessage};
6
7pub struct Websocket {
8    raw_sock: WebSocketStream<MaybeTlsStream<TcpStream>>,
9}
10
11impl Websocket {
12    pub async fn new(url: &str) -> Option<Self> {
13        let (sock, _) = match connect_async(url).await {
14            Ok(value) => value,
15            Err(err) => {
16                pawkit_logger::error(&format!("{}", err));
17                return None;
18            }
19        };
20
21        return Some(Self { raw_sock: sock });
22    }
23
24    pub fn is_open(&self) -> bool {
25        return !self.raw_sock.is_terminated();
26    }
27
28    pub async fn close(&mut self) {
29        self.raw_sock.close(None).await.unwrap();
30    }
31
32    pub async fn recv(&mut self) -> Option<WebsocketMessage> {
33        if !self.is_open() {
34            return None;
35        }
36
37        loop {
38            let Some(Ok(message)) = self.raw_sock.next().await else {
39                return None;
40            };
41
42            match message {
43                Message::Text(text) => {
44                    return Some(WebsocketMessage::String(text.as_str().into()));
45                }
46
47                Message::Binary(bin) => {
48                    return Some(WebsocketMessage::Array(bin.to_vec()));
49                }
50
51                _ => {}
52            }
53        }
54    }
55
56    pub async fn send(&mut self, message: WebsocketMessage) -> Result<(), WebsocketError> {
57        if !self.is_open() {
58            return Err(WebsocketError::NotOpen);
59        }
60
61        match message {
62            WebsocketMessage::String(str) => {
63                if self.raw_sock.send(Message::Text(str.into())).await.is_err() {
64                    return Err(WebsocketError::InvalidState);
65                }
66
67                return Ok(());
68            }
69
70            WebsocketMessage::Array(arr) => {
71                if self
72                    .raw_sock
73                    .send(Message::Binary(arr.into()))
74                    .await
75                    .is_err()
76                {
77                    return Err(WebsocketError::InvalidState);
78                }
79
80                return Ok(());
81            }
82        }
83    }
84}