Skip to main content

websockets/websocket/
mod.rs

1pub mod builder;
2pub mod frame;
3mod handshake;
4mod parsed_addr;
5pub mod split;
6mod stream;
7
8use crate::error::WebSocketError;
9use builder::WebSocketBuilder;
10use frame::Frame;
11use split::{WebSocketReadHalf, WebSocketWriteHalf};
12
13#[derive(Debug)]
14enum FrameType {
15    Text,
16    Binary,
17    Control,
18}
19
20impl Default for FrameType {
21    fn default() -> Self {
22        Self::Control
23    }
24}
25
26/// Manages the WebSocket connection; used to connect, send data, and receive data.
27///
28/// Connect with [`WebSocket::connect()`]:
29///
30/// ```
31/// # use websockets::{WebSocket, WebSocketError};
32/// # #[tokio::main]
33/// # async fn main() -> Result<(), WebSocketError> {
34/// let mut ws = WebSocket::connect("wss://echo.websocket.org/").await?;
35/// # Ok(())
36/// # }
37/// ```
38///
39/// Cuustomize the handshake using a [`WebSocketBuilder`] obtained from [`WebSocket::builder()`]:
40///
41/// ```
42/// # use websockets::{WebSocket, WebSocketError};
43/// # #[tokio::main]
44/// # async fn main() -> Result<(), WebSocketError> {
45/// let mut ws = WebSocket::builder()
46///     .add_subprotocol("wamp")
47///     .connect("wss://echo.websocket.org")
48///     .await?;
49/// # Ok(())
50/// # }
51/// ```
52///
53/// Use the `WebSocket::send*` methods to send frames:
54///
55/// ```
56/// # use websockets::{WebSocket, WebSocketError};
57/// # #[tokio::main]
58/// # async fn main() -> Result<(), WebSocketError> {
59/// # let mut ws = WebSocket::connect("wss://echo.websocket.org")
60/// #     .await?;
61/// ws.send_text("foo".to_string()).await?;
62/// # Ok(())
63/// # }
64/// ```
65///
66/// Use [`WebSocket::receive()`] to receive frames:
67///
68/// ```
69/// # use websockets::{WebSocket, WebSocketError, Frame};
70/// # #[tokio::main]
71/// # async fn main() -> Result<(), WebSocketError> {
72/// # let mut ws = WebSocket::connect("wss://echo.websocket.org")
73/// #     .await?;
74/// # ws.send_text("foo".to_string()).await?;
75/// if let Frame::Text { payload: received_msg, .. } =  ws.receive().await? {
76///     // echo.websocket.org echoes text frames
77///     assert_eq!(received_msg, "foo".to_string());
78/// }
79/// # else { panic!() }
80/// # Ok(())
81/// # }
82/// ```
83///
84/// Close the connection with [`WebSocket::close()`]:
85///
86/// ```
87/// # use websockets::{WebSocket, WebSocketError, Frame};
88/// # #[tokio::main]
89/// # async fn main() -> Result<(), WebSocketError> {
90/// #     let mut ws = WebSocket::connect("wss://echo.websocket.org")
91/// #         .await?;
92/// ws.close(Some((1000, String::new()))).await?;
93/// if let Frame::Close{ payload: Some((status_code, _reason)) } = ws.receive().await? {
94///     assert_eq!(status_code, 1000);
95/// }
96/// # Ok(())
97/// # }
98/// ```
99///
100/// # Splitting
101///
102/// To facilitate simulataneous reads and writes, the `WebSocket` can be split
103/// into a [read half](WebSocketReadHalf) and a [write half](WebSocketWriteHalf).
104/// The read half allows frames to be received, while the write half
105/// allows frames to be sent.
106///
107/// If the read half receives a Ping or Close frame, it needs to send a
108/// Pong or echo the Close frame and close the WebSocket, respectively.
109/// The write half is notified of these events, but it cannot act on them
110/// unless it is flushed. Events can be explicitly [`flush`](WebSocketWriteHalf::flush())ed,
111/// but sending a frame will also flush events. If frames are not being
112/// sent frequently, consider explicitly flushing events.
113///
114/// Flushing is done automatically if you are using the the `WebSocket` type by itself.
115#[derive(Debug)]
116pub struct WebSocket {
117    read_half: WebSocketReadHalf,
118    write_half: WebSocketWriteHalf,
119    accepted_subprotocol: Option<String>,
120    handshake_response_headers: Option<Vec<(String, String)>>,
121}
122
123impl WebSocket {
124    /// Constructs a [`WebSocketBuilder`], which can be used to customize
125    /// the WebSocket handshake.
126    pub fn builder() -> WebSocketBuilder {
127        WebSocketBuilder::new()
128    }
129
130    /// Connects to a URL (and performs the WebSocket handshake).
131    pub async fn connect(url: &str) -> Result<Self, WebSocketError> {
132        WebSocketBuilder::new().connect(url).await
133    }
134
135    /// Receives a [`Frame`] over the WebSocket connection.
136    ///
137    /// If the received frame is a Ping frame, a Pong frame will be sent.
138    /// If the received frame is a Close frame, an echoed Close frame
139    /// will be sent and the WebSocket will close.
140    pub async fn receive(&mut self) -> Result<Frame, WebSocketError> {
141        let received_frame = self.read_half.receive().await?;
142        self.write_half.flush().await?;
143        Ok(received_frame)
144    }
145
146    /// Receives a [`Frame`] over the WebSocket connection **without handling incoming frames.**
147    /// For example, receiving a Ping frame will not queue a Pong frame to be sent,
148    /// and receiving a Close frame will not queue a Close frame to be sent nor close
149    /// the connection.
150    ///
151    /// To automatically handle incoming frames, use the [`receive()`](WebSocket::receive())
152    /// method instead.
153    pub async fn receive_without_handling(&mut self) -> Result<Frame, WebSocketError> {
154        self.read_half.receive_without_handling().await
155    }
156
157    /// Sends an already constructed [`Frame`] over the WebSocket connection.
158    pub async fn send(&mut self, frame: Frame) -> Result<(), WebSocketError> {
159        self.write_half.send(frame).await
160    }
161
162    /// Sends a Text frame over the WebSocket connection, constructed
163    /// from passed arguments. `continuation` will be `false` and `fin` will be `true`.
164    /// To use a custom `continuation` or `fin`, construct a [`Frame`] and use
165    /// [`WebSocket::send()`].
166    pub async fn send_text(&mut self, payload: String) -> Result<(), WebSocketError> {
167        self.write_half.send_text(payload).await
168    }
169
170    /// Sends a Binary frame over the WebSocket connection, constructed
171    /// from passed arguments. `continuation` will be `false` and `fin` will be `true`.
172    /// To use a custom `continuation` or `fin`, construct a [`Frame`] and use
173    /// [`WebSocket::send()`].
174    pub async fn send_binary(&mut self, payload: Vec<u8>) -> Result<(), WebSocketError> {
175        self.write_half.send_binary(payload).await
176    }
177
178    /// Sends a Close frame over the WebSocket connection, constructed
179    /// from passed arguments, and closes the WebSocket connection.
180    /// This method will attempt to wait for an echoed Close frame,
181    /// which is returned.
182    pub async fn close(&mut self, payload: Option<(u16, String)>) -> Result<(), WebSocketError> {
183        self.write_half.close(payload).await
184    }
185
186    /// Sends a Ping frame over the WebSocket connection, constructed
187    /// from passed arguments.
188    pub async fn send_ping(&mut self, payload: Option<Vec<u8>>) -> Result<(), WebSocketError> {
189        self.write_half.send_ping(payload).await
190    }
191
192    /// Sends a Pong frame over the WebSocket connection, constructed
193    /// from passed arguments.
194    pub async fn send_pong(&mut self, payload: Option<Vec<u8>>) -> Result<(), WebSocketError> {
195        self.write_half.send_pong(payload).await
196    }
197
198    /// Shuts down the WebSocket connection **without sending a Close frame**.
199    /// It is recommended to use the [`close()`](WebSocket::close()) method instead.
200    pub async fn shutdown(&mut self) -> Result<(), WebSocketError> {
201        self.write_half.shutdown().await
202    }
203
204    /// Splits the WebSocket into a read half and a write half, which can be used separately.
205    /// [Accepted subprotocol](WebSocket::accepted_subprotocol())
206    /// and [handshake response headers](WebSocket::handshake_response_headers()) data
207    /// will be lost.
208    pub fn split(self) -> (WebSocketReadHalf, WebSocketWriteHalf) {
209        (self.read_half, self.write_half)
210    }
211
212    /// Joins together a split read half and write half to reconstruct a WebSocket.
213    pub fn join(read_half: WebSocketReadHalf, write_half: WebSocketWriteHalf) -> Self {
214        Self {
215            read_half,
216            write_half,
217            accepted_subprotocol: None,
218            handshake_response_headers: None,
219        }
220    }
221
222    /// Returns the subprotocol that was accepted by the server during the handshake,
223    /// if any. This data will be lost if the WebSocket is [`split`](WebSocket::split()).
224    pub fn accepted_subprotocol(&self) -> &Option<String> {
225        // https://tools.ietf.org/html/rfc6455#section-1.9
226        &self.accepted_subprotocol
227    }
228
229    /// Returns the headers that were returned by the server during the handshake.
230    /// This data will be lost if the WebSocket is [`split`](WebSocket::split()).
231    pub fn handshake_response_headers(&self) -> &Option<Vec<(String, String)>> {
232        // https://tools.ietf.org/html/rfc6455#section-4.2.2
233        &self.handshake_response_headers
234    }
235}