Oregami
Repositories/oxedyne/fe2o3

oxedyne/fe2o3/fe2o3_net/src/ws/client.rs

12.6 KiB, 3 runs

created by r1870400018:59659, which is this file's identity for as long as the history lasts, whatever it is later renamed to

download · who wrote it · its history

1use crate::{
2 constant,
3 http::msg::HttpMessage,
4 ws::{
5 core::{
6 accept_key,
7 assemble,
8 connect_request,
9 encode_message,
10 read_frame,
11 Assembled,
12 WebSocketLimits,
13 WebSocketMessage,
14 },
15 status::WebSocketStatusCode,
16 },
17};
18
19use oxedyne_fe2o3_core::prelude::*;
20
21use std::{
22 io::Cursor,
23 pin::Pin,
24 time::Duration,
25};
26
27use tokio::{
28 io::{
29 AsyncReadExt,
30 AsyncWriteExt,
31 Chain,
32 },
33 net::TcpStream,
34 time::{
35 timeout_at,
36 Instant,
37 },
38};
39
40
41const READ_CHUNK: usize = 4_096; // payload bytes per read call
42
43#[derive(Clone, Copy, Debug, Eq, PartialEq)]
44enum State {
45 Open,
46 Closed, // close frames exchanged
47 Broken, // the stream lost its place in a frame, or the peer vanished
48}
49
50/// A websocket client that owns its TCP connection, for a caller that wants to talk to a server
51/// and has no database or handler to give [`crate::ws::core::WebSocket`].
52///
53/// It answers pings, reassembles fragmented messages (answering a ping that arrives between the
54/// fragments), masks every frame it sends, and applies [`WebSocketLimits`] to everything it
55/// reads. A timed-out [`recv`](Self::recv) loses nothing, because it only ever gives up between
56/// frames.
57pub struct WsClient {
58 stream: Chain<Cursor<Vec<u8>>, TcpStream>, // bytes read past the handshake, then the socket
59 limits: WebSocketLimits,
60 opcode: Option<u8>, // of the fragmented message in progress
61 buffer: Vec<u8>, // its payload so far
62 state: State,
63}
64
65impl WsClient {
66
67 /// Connects to `addr` (`host:port`), upgrading `path` to a websocket, and verifies the
68 /// server's `Sec-WebSocket-Accept` against the random key sent. `addr` doubles as the `Host`
69 /// field. A refusal is an error naming the status; a wrong accept key, one naming both keys.
70 ///
71 /// Nothing bounds the handshake's duration, and dropping the future abandons it cleanly, so a
72 /// caller wanting a bound wraps this in `tokio::time::timeout`.
73 pub async fn connect(
74 addr: &str,
75 path: &str,
76 origin: Option<&str>,
77 )
78 -> Outcome<Self>
79 {
80 let (request, key) = res!(connect_request(addr, path, origin));
81 let mut tcp = match TcpStream::connect(addr).await {
82 Ok(tcp) => tcp,
83 Err(e) => return Err(err!(e,
84 "While connecting to websocket server {}.", addr;
85 IO, Network)),
86 };
87 // Frames are small and latency is the point.
88 res!(tcp.set_nodelay(true), IO, Network);
89 if let Err(e) = request.write_all(&mut tcp).await {
90 return Err(err!(e,
91 "While sending the websocket upgrade request for {} to {}.", path, addr;
92 IO, Network, Wire, Write));
93 }
94 let (response, remnant) = match HttpMessage::read::<
95 { constant::HTTP_DEFAULT_HEADER_CHUNK_SIZE },
96 { constant::HTTP_DEFAULT_BODY_CHUNK_SIZE },
97 _,
98 >(Pin::new(&mut tcp), &Vec::new(), Some(false), None).await {
99 Ok((Some(response), remnant)) => (response, remnant),
100 Ok((None, _)) => return Err(err!(
101 "The websocket server {} closed the connection before answering the upgrade \
102 request for {}.", addr, path;
103 IO, Network, Wire, Read)),
104 Err(e) => return Err(err!(e,
105 "While reading the websocket server {}'s answer to the upgrade request for {}.",
106 addr, path;
107 IO, Network, Wire, Read)),
108 };
109 if let Err(e) = response.check_websocket_handshake(&accept_key(&key)) {
110 return Err(err!(e,
111 "The websocket upgrade of {} on {} failed.", path, addr;
112 IO, Network, Wire, Invalid, Input));
113 }
114 // The server may speak straight after its 101, and whatever of that the header reader
115 // took is the start of the frame stream.
116 Ok(Self {
117 stream: Cursor::new(remnant).chain(tcp),
118 limits: WebSocketLimits::default(),
119 opcode: None,
120 buffer: Vec::new(),
121 state: State::Open,
122 })
123 }
124
125 /// Replaces the bounds applied to incoming messages.
126 pub fn with_limits(mut self, limits: WebSocketLimits) -> Self {
127 self.limits = limits;
128 self
129 }
130
131 pub fn limits(&self) -> WebSocketLimits { self.limits }
132
133 /// Is the connection still usable for sending and receiving?
134 pub fn is_open(&self) -> bool { self.state == State::Open }
135
136 pub async fn send_text(&mut self, text: &str) -> Outcome<()> {
137 self.send(&WebSocketMessage::Text(text.to_string())).await
138 }
139
140 /// Sends `message` as one masked frame.
141 pub async fn send(&mut self, message: &WebSocketMessage) -> Outcome<()> {
142 res!(self.check_open());
143 self.write(message).await
144 }
145
146 /// Waits up to `timeout` for the next text, binary or close message. `Ok(None)` means the time
147 /// ran out; a message half-assembled then is kept for the next call.
148 ///
149 /// Pings are answered and pongs dropped on the way. A close from the server is echoed and
150 /// returned, after which the client is closed. The peer vanishing without a close, a read
151 /// failing, or the time running out part way through a frame each leave the stream unusable,
152 /// and are errors.
153 pub async fn recv(&mut self, timeout: Duration) -> Outcome<Option<WebSocketMessage>> {
154 res!(self.check_open());
155 let deadline = Instant::now() + timeout;
156 loop {
157 if !res!(self.frame_begun(deadline).await) {
158 return Ok(None);
159 }
160 let msg = match res!(self.next_frame(deadline).await) {
161 Assembled::Pending => continue,
162 Assembled::Message(msg) => return Ok(Some(msg)),
163 Assembled::Control(msg) => msg,
164 };
165 match msg {
166 WebSocketMessage::Ping(data) => {
167 if let Err(e) = self.write(&WebSocketMessage::Pong(data)).await {
168 return Err(err!(e,
169 "While answering a websocket ping."; IO, Network, Wire, Write));
170 }
171 }
172 WebSocketMessage::Pong(_) => (),
173 WebSocketMessage::Close(code, reason) => {
174 // RFC 6455 §5.5.1: answer a close with a close, echoing the status.
175 let echo = self.write(&WebSocketMessage::Close(code, None)).await;
176 self.state = State::Closed;
177 let _ = self.tcp().shutdown().await;
178 if let Err(e) = echo {
179 return Err(err!(e,
180 "While echoing the server's websocket close."; IO, Network, Wire, Write));
181 }
182 return Ok(Some(WebSocketMessage::Close(code, reason)));
183 }
184 WebSocketMessage::Text(_) | WebSocketMessage::Binary(_) => return Err(err!(
185 "A websocket data message was classed as a control frame."; Bug)),
186 }
187 }
188 }
189
190 /// Sends a normal close and waits up to `timeout` for the server's, discarding anything that
191 /// arrives first, then shuts the socket. Closing a client already closed does nothing.
192 pub async fn close(&mut self, timeout: Duration) -> Outcome<()> {
193 match self.state {
194 State::Closed => return Ok(()),
195 State::Broken => {
196 let _ = self.tcp().shutdown().await;
197 return Ok(());
198 }
199 State::Open => (),
200 }
201 let sent = self.write(&WebSocketMessage::Close(
202 Some(WebSocketStatusCode::NormalClosure), None)).await;
203 self.state = State::Closed;
204 if sent.is_ok() {
205 let deadline = Instant::now() + timeout;
206 // The server's close, its vanishing, or the deadline all end the wait; none of them
207 // leaves anything more to do than shut the socket.
208 while let Ok(true) = self.frame_begun(deadline).await {
209 match self.next_frame(deadline).await {
210 Ok(Assembled::Control(WebSocketMessage::Close(..))) => break,
211 Ok(_) => (),
212 Err(_) => break,
213 }
214 }
215 }
216 // A failed read above may already have shut the socket, and a shutdown refused now leaves
217 // the caller nothing to do, so only the close frame's fate is reported.
218 let _ = self.tcp().shutdown().await;
219 if let Err(e) = sent {
220 return Err(err!(e, "While sending a websocket close."; IO, Network, Wire, Write));
221 }
222 Ok(())
223 }
224
225 fn tcp(&mut self) -> &mut TcpStream { self.stream.get_mut().1 }
226
227 fn check_open(&self) -> Outcome<()> {
228 match self.state {
229 State::Open => Ok(()),
230 State::Closed => Err(err!(
231 "The websocket client is closed."; IO, Network, Invalid, Input)),
232 State::Broken => Err(err!(
233 "The websocket client's connection failed earlier and cannot be used.";
234 IO, Network, Invalid, Input)),
235 }
236 }
237
238 async fn write(&mut self, message: &WebSocketMessage) -> Outcome<()> {
239 // One frame per message: a client is required to mask (RFC 6455 §5.3), and it need not
240 // fragment.
241 let byts = res!(encode_message(message, true, usize::MAX, usize::MAX));
242 let tcp = self.tcp();
243 if let Err(e) = tcp.write_all(&byts).await {
244 self.state = State::Broken;
245 return Err(err!(e, "While writing a websocket frame."; IO, Network, Wire, Write));
246 }
247 if let Err(e) = tcp.flush().await {
248 self.state = State::Broken;
249 return Err(err!(e, "While flushing a websocket frame."; IO, Network, Wire, Write));
250 }
251 Ok(())
252 }
253
254 /// Waits until `deadline` for a frame's first byte, consuming nothing, so that giving up here
255 /// leaves the stream exactly where it was. End of stream counts as begun, for `read_frame` to
256 /// report.
257 async fn frame_begun(&mut self, deadline: Instant) -> Outcome<bool> {
258 let (remnant, tcp) = self.stream.get_mut();
259 if (remnant.position() as usize) < remnant.get_ref().len() {
260 return Ok(true);
261 }
262 let mut probe = [0u8; 1];
263 match timeout_at(deadline, tcp.peek(&mut probe)).await {
264 Err(_) => Ok(false),
265 Ok(Ok(_)) => Ok(true),
266 Ok(Err(e)) => {
267 self.state = State::Broken;
268 Err(err!(e, "While waiting for a websocket frame."; IO, Network, Wire, Read))
269 }
270 }
271 }
272
273 /// Reads the frame that has begun and folds it into the message in progress. Any failure here
274 /// leaves the stream part way through a frame, so the client is marked broken.
275 async fn next_frame(&mut self, deadline: Instant) -> Outcome<Assembled> {
276 let limits = self.limits;
277 let buffered = self.buffer.len();
278 let result = match timeout_at(
279 deadline,
280 read_frame(&mut self.stream, READ_CHUNK, limits, buffered),
281 ).await {
282 Err(_) => Err(err!(
283 "The deadline passed part way through a websocket frame, so the stream has lost \
284 its place.";
285 IO, Network, Wire, Read, Timeout)),
286 Ok(Ok(None)) => Err(err!(
287 "The websocket server closed the connection without a close frame.";
288 IO, Network, Wire, Read)),
289 Ok(Ok(Some(frame))) => assemble(&mut self.opcode, &mut self.buffer, frame),
290 Ok(Err(e)) => Err(e),
291 };
292 match result {
293 Ok(step) => Ok(step),
294 Err(e) => {
295 self.state = State::Broken;
296 // As `WebSocket::read`: an over-limit message is the one refusal the peer is told
297 // about, with status 1009 (RFC 6455 §7.4.1).
298 if e.tags().contains(&ErrTag::TooBig) {
299 let close = WebSocketMessage::Close(
300 Some(WebSocketStatusCode::MessageTooBig),
301 Some(fmt!("Message too big")),
302 );
303 if let Err(e2) = self.write(&close).await {
304 error!(err!(e2,
305 "While sending a 1009 close to a server whose message was over the \
306 limit.";
307 IO, Network, Wire, Write));
308 }
309 }
310 let _ = self.tcp().shutdown().await;
311 Err(e)
312 }
313 }
314 }
315}