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}