cataclysm_ws/
web_socket_stream.rs

1use tokio::{
2    net::{TcpStream},
3    sync::{OwnedSemaphorePermit}
4};
5use crate::{Error, Message, Frame, WebSocketWriter, WebSocketReader};
6use crate::communication::{write_message, read_frame};
7
8/// Wrapper structure of a tcp stream with some websockets utilities
9pub struct WebSocketStream {
10    inner: TcpStream,
11    permit: Option<OwnedSemaphorePermit>
12}
13
14impl WebSocketStream {
15    /// Wraps a tcp stream withot checking the handshake or anything
16    pub fn from_tcp_stream_unchecked(stream: TcpStream) -> WebSocketStream {
17        WebSocketStream {
18            inner: stream,
19            permit: None
20        }
21    }
22
23    /// Auxiliar function that cataclysm uses to keep track of connections
24    pub fn set_permit(&mut self, permit: OwnedSemaphorePermit) {
25        self.permit = Some(permit);
26    }
27
28    /// Sends a message through the websockets connection
29    pub async fn send_message(&self, message: Message) -> Result<(), Error> {
30        write_message(&self, message).await
31    }
32
33    /// Blocks until a message is received
34    pub async fn try_read_frame(&self) -> Result<Frame, Error> {
35        read_frame(&self).await
36    }
37
38    /// Splits the stream into the reading and writting part
39    pub fn split(self) -> (WebSocketWriter, WebSocketReader) {
40        let (rx, tx) = self.inner.into_split();
41        let mut web_socket_reader = WebSocketReader::new_unchecked(rx);
42        if let Some(permit) = self.permit {
43            web_socket_reader.set_permit(permit);
44        }
45        (WebSocketWriter::new_unchecked(tx), web_socket_reader)
46    }
47}
48
49// Reference access to the inner structure
50impl AsRef<TcpStream> for WebSocketStream {
51    fn as_ref(&self) -> &TcpStream {
52        &self.inner
53    }
54}
55
56// Mutable reference access to the inner structure
57impl AsMut<TcpStream> for WebSocketStream {
58    fn as_mut(&mut self) -> &mut TcpStream {
59        &mut self.inner
60    }
61}
62
63// Conversion to inner type
64impl Into<TcpStream> for WebSocketStream {
65    fn into(self) -> TcpStream {
66        self.inner
67    }
68}