rws_server/test_client/
websocket.rs

1//! WebSocket routes in tests: [`TestClient::websocket`](super::TestClient::websocket)
2//! sends the opening handshake through the application, and when the route
3//! answers with [`WebSocket::upgrade`]'s `101`, [`TestWebSocketRequest::connect`]
4//! runs its handler on a thread of its own over an in-memory duplex — the
5//! same handler the server would run, given the same reader and writer
6//! types it expects, with no socket in between. The test side is a
7//! [`TestWebSocket`]: it masks what it sends, as a browser does (the
8//! handler's `recv` refuses unmasked frames), answers pings, and joins
9//! fragmented messages.
10
11use std::io::{Read, Write};
12use std::sync::mpsc::{self, Receiver, RecvTimeoutError, Sender};
13use std::time::Duration;
14
15use super::{TestRequest, TestResponse};
16use crate::application::Application;
17use crate::websocket::{close_payload, encode_frame, read_raw_frame, take_upgrade, Frame, Message, WebSocket};
18
19/// How long a [`TestWebSocket`] waits for the handler by default: 5 s.
20pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
21
22/// The `Sec-WebSocket-Key` the test client sends (RFC 6455's sample nonce).
23const KEY: &str = "dGhlIHNhbXBsZSBub25jZQ==";
24
25/// A WebSocket upgrade request being built; see
26/// [`TestClient::websocket`](super::TestClient::websocket).
27pub struct TestWebSocketRequest<'a, A: Application> {
28    request: TestRequest<'a, A>,
29    timeout: Duration,
30}
31
32impl<'a, A: Application> TestWebSocketRequest<'a, A> {
33    pub(super) fn new(request: TestRequest<'a, A>) -> Self {
34        TestWebSocketRequest { request, timeout: DEFAULT_TIMEOUT }
35    }
36
37    /// Add a request header (an `Authorization` or `Cookie` for a guarded
38    /// route, a `Sec-WebSocket-Protocol`, …).
39    pub fn header(mut self, name: &str, value: &str) -> Self {
40        self.request = self.request.header(name, value);
41        self
42    }
43
44    /// How long [`TestWebSocket::recv`] and [`close`](TestWebSocket::close)
45    /// wait for the handler (default [`DEFAULT_TIMEOUT`]).
46    pub fn timeout(mut self, timeout: Duration) -> Self {
47        self.timeout = timeout;
48        self
49    }
50
51    /// Send the handshake. A `101` from [`WebSocket::upgrade`] starts the
52    /// handler and returns the connection; any other answer — `426` for a
53    /// route that isn't a WebSocket, `401` from an auth layer, a `101` built
54    /// by hand without a handler — comes back as the `Err` response.
55    pub fn connect(self) -> Result<TestWebSocket, TestResponse> {
56        let mut request = self.request;
57        for (name, value) in [
58            ("Host", "localhost"),
59            ("Upgrade", "websocket"),
60            ("Connection", "Upgrade"),
61            ("Sec-WebSocket-Key", KEY),
62            ("Sec-WebSocket-Version", "13"),
63        ] {
64            if !request.has_header(name) {
65                request = request.header(name, value);
66            }
67        }
68        let mut response = request.dispatch();
69        let upgrade = take_upgrade(&mut response);
70        let accepted = response
71            .headers
72            .iter()
73            .find(|h| h.name.eq_ignore_ascii_case("Sec-WebSocket-Accept"))
74            .is_some_and(|h| *h.value == *WebSocket::accept_key(KEY));
75        let handshake = TestResponse::from_response(response);
76        let upgrade = match upgrade {
77            Some(upgrade) if handshake.status() == 101 && accepted => upgrade,
78            _ => return Err(handshake),
79        };
80
81        let (to_server, server_in) = mpsc::channel();
82        let (server_out, from_server) = mpsc::channel();
83        let (done_tx, done) = mpsc::channel::<()>();
84        #[cfg(feature = "http2")]
85        let runtime = tokio::runtime::Handle::try_current().ok();
86        let spawned = std::thread::Builder::new().name("rws-test-websocket".to_string()).spawn(move || {
87            // A handler that reaches for the test's tokio runtime finds it,
88            // as it would on the server's blocking pool.
89            #[cfg(feature = "http2")]
90            let _entered = runtime.as_ref().map(|r| r.enter());
91            let reader = PipeReader::new(server_in, None);
92            let writer = PipeWriter(server_out);
93            upgrade.run(Box::new(reader), Box::new(writer));
94            let _ = done_tx.send(());
95        });
96        if spawned.is_err() {
97            return Err(handshake);
98        }
99        Ok(TestWebSocket {
100            handshake,
101            to_server: Some(PipeWriter(to_server)),
102            from_server: PipeReader::new(from_server, Some(self.timeout)),
103            done,
104            timeout: self.timeout,
105            close: None,
106            mask: 0x9E37_79B9,
107        })
108    }
109}
110
111/// The test side of an upgraded connection; see the module docs.
112///
113/// Dropping it ends the connection (the handler's `recv` returns an error),
114/// without waiting for the handler.
115pub struct TestWebSocket {
116    handshake: TestResponse,
117    to_server: Option<PipeWriter>,
118    from_server: PipeReader,
119    done: Receiver<()>,
120    timeout: Duration,
121    close: Option<(Option<u16>, String)>,
122    mask: u32,
123}
124
125impl TestWebSocket {
126    /// The `101` response, for its headers (`Sec-WebSocket-Protocol`, …).
127    pub fn handshake(&self) -> &TestResponse {
128        &self.handshake
129    }
130
131    /// Send a text message. Panics if the handler's side is gone.
132    #[track_caller]
133    pub fn send_text(&mut self, text: &str) {
134        self.send(Frame::Text(text.to_string())).expect("send a WebSocket text message");
135    }
136
137    /// Send a binary message. Panics if the handler's side is gone.
138    #[track_caller]
139    pub fn send_binary(&mut self, data: &[u8]) {
140        self.send(Frame::Binary(data.to_vec())).expect("send a WebSocket binary message");
141    }
142
143    /// Send a ping; the handler's pong is returned by
144    /// [`recv_frame`](Self::recv_frame) (and skipped by [`recv`](Self::recv)).
145    pub fn ping(&mut self, payload: &[u8]) -> Result<(), String> {
146        self.send(Frame::Ping(payload.to_vec()))
147    }
148
149    /// Send `frame`, masked as a client must.
150    pub fn send(&mut self, frame: Frame) -> Result<(), String> {
151        let masked = mask_frame(encode_frame(frame), self.next_mask());
152        self.send_raw(&masked)
153    }
154
155    /// Send bytes as they are — for testing how the handler copes with
156    /// frames a well-behaved client wouldn't send (unmasked, bad opcodes).
157    pub fn send_raw(&mut self, bytes: &[u8]) -> Result<(), String> {
158        match self.to_server.as_mut() {
159            Some(writer) => writer.write_all(bytes).map_err(|e| format!("write error: {e}")),
160            None => Err("the test side already closed the connection".to_string()),
161        }
162    }
163
164    fn next_mask(&mut self) -> [u8; 4] {
165        // Masks only need to vary, not to be secret (they're for proxies).
166        self.mask = self.mask.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
167        self.mask.to_be_bytes()
168    }
169
170    /// The next frame from the handler, waiting at most the timeout. Pings
171    /// are answered (and returned); nothing is joined.
172    pub fn recv_frame(&mut self) -> Result<Frame, String> {
173        self.recv_frame_fin().map(|(_, frame)| frame)
174    }
175
176    /// [`recv_frame`](Self::recv_frame) with the frame's FIN bit.
177    fn recv_frame_fin(&mut self) -> Result<(bool, Frame), String> {
178        let raw = read_raw_frame(&mut self.from_server, WebSocket::MAX_FRAME_PAYLOAD)?;
179        let fin = raw.fin;
180        let frame = match raw.opcode {
181            0x0 => Frame::Continuation { fin: raw.fin, data: raw.payload },
182            0x1 => Frame::Text(String::from_utf8(raw.payload).map_err(|_| "a text frame isn't UTF-8".to_string())?),
183            0x2 => Frame::Binary(raw.payload),
184            0x8 => {
185                let (code, reason) = close_payload(&raw.payload)?;
186                self.close = Some((code, reason.clone()));
187                Frame::Close(code, reason)
188            }
189            0x9 => {
190                let _ = self.send(Frame::Pong(raw.payload.clone()));
191                Frame::Ping(raw.payload)
192            }
193            0xA => Frame::Pong(raw.payload),
194            other => return Err(format!("unknown opcode 0x{other:X} from the handler")),
195        };
196        Ok((fin, frame))
197    }
198
199    /// The next message from the handler: `Ok(None)` once it sent a close
200    /// frame (answered here, see [`close_code`](Self::close_code)), `Err` on
201    /// a timeout or when the handler returned without closing. Pings and
202    /// pongs are skipped, fragments joined.
203    pub fn recv(&mut self) -> Result<Option<Message>, String> {
204        if self.close.is_some() {
205            return Ok(None);
206        }
207        let mut message: Option<(bool, Vec<u8>)> = None;
208        loop {
209            let (fin, frame) = self.recv_frame_fin()?;
210            match frame {
211                Frame::Text(text) => message = Some((true, text.into_bytes())),
212                Frame::Binary(data) => message = Some((false, data)),
213                Frame::Continuation { data, .. } => match message.as_mut() {
214                    Some((_, buf)) => buf.extend_from_slice(&data),
215                    None => return Err("a continuation frame with no message".to_string()),
216                },
217                Frame::Close(code, _) => {
218                    // Answer the handler's close, as a client would.
219                    let _ = self.send(Frame::Close(code.or(Some(1000)), String::new()));
220                    return Ok(None);
221                }
222                Frame::Ping(_) | Frame::Pong(_) => continue,
223            }
224            // A text or binary frame without FIN waits for its continuations.
225            if !fin {
226                continue;
227            }
228            return match message.take() {
229                Some((true, bytes)) => String::from_utf8(bytes).map(|t| Some(Message::Text(t))).map_err(|_| "a text message isn't UTF-8".to_string()),
230                Some((false, bytes)) => Ok(Some(Message::Binary(bytes))),
231                None => continue,
232            };
233        }
234    }
235
236    /// The next message, which must be text. Panics otherwise (binary, a
237    /// close, a timeout), with what came instead.
238    #[track_caller]
239    pub fn recv_text(&mut self) -> String {
240        match self.recv() {
241            Ok(Some(Message::Text(text))) => text,
242            other => panic!("expected a WebSocket text message, got {other:?}"),
243        }
244    }
245
246    /// The next message, which must be binary. Panics otherwise.
247    #[track_caller]
248    pub fn recv_binary(&mut self) -> Vec<u8> {
249        match self.recv() {
250            Ok(Some(Message::Binary(data))) => data,
251            other => panic!("expected a WebSocket binary message, got {other:?}"),
252        }
253    }
254
255    /// The close code the handler sent, once it closed (`Some(None)` for a
256    /// close frame without one).
257    pub fn close_code(&self) -> Option<Option<u16>> {
258        self.close.as_ref().map(|(code, _)| *code)
259    }
260
261    /// Run the closing handshake: send a close frame with `code` and
262    /// `reason`, skip messages until the handler's close (or the timeout),
263    /// then end the connection and wait (at most the timeout) for the
264    /// handler to return. Returns the handler's close code, `None` when it
265    /// didn't send one.
266    pub fn close(mut self, code: u16, reason: &str) -> Option<u16> {
267        if self.close.is_none() {
268            let _ = self.send(Frame::Close(Some(code), reason.to_string()));
269            while self.close.is_none() {
270                if self.recv_frame().is_err() {
271                    break;
272                }
273            }
274        }
275        self.to_server = None;
276        let _ = self.done.recv_timeout(self.timeout);
277        self.close.and_then(|(code, _)| code)
278    }
279
280    /// Wait at most `timeout` for the handler to return; `true` if it did.
281    pub fn handler_finished(&self, timeout: Duration) -> bool {
282        match self.done.recv_timeout(timeout) {
283            Ok(()) | Err(RecvTimeoutError::Disconnected) => true,
284            Err(RecvTimeoutError::Timeout) => false,
285        }
286    }
287}
288
289/// Set the mask bit on a server-style (unmasked) frame and mask its payload.
290fn mask_frame(mut frame: Vec<u8>, key: [u8; 4]) -> Vec<u8> {
291    let offset = match frame[1] & 0x7F {
292        126 => 4,
293        127 => 10,
294        _ => 2,
295    };
296    frame[1] |= 0x80;
297    for (i, byte) in frame[offset..].iter_mut().enumerate() {
298        *byte ^= key[i % 4];
299    }
300    frame.splice(offset..offset, key);
301    frame
302}
303
304/// One direction of the in-memory duplex: each write is a chunk.
305struct PipeWriter(Sender<Vec<u8>>);
306
307impl Write for PipeWriter {
308    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
309        self.0
310            .send(buf.to_vec())
311            .map(|_| buf.len())
312            .map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "the other side of the test WebSocket is gone"))
313    }
314
315    fn flush(&mut self) -> std::io::Result<()> {
316        Ok(())
317    }
318}
319
320/// The reading end: blocks for the next chunk (at most `timeout`, when
321/// set); the writer gone is the end of the stream.
322struct PipeReader {
323    rx: Receiver<Vec<u8>>,
324    buf: Vec<u8>,
325    pos: usize,
326    timeout: Option<Duration>,
327}
328
329impl PipeReader {
330    fn new(rx: Receiver<Vec<u8>>, timeout: Option<Duration>) -> Self {
331        PipeReader { rx, buf: Vec::new(), pos: 0, timeout }
332    }
333}
334
335impl Read for PipeReader {
336    fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
337        while self.pos >= self.buf.len() {
338            let next = match self.timeout {
339                Some(timeout) => self.rx.recv_timeout(timeout).map_err(|e| match e {
340                    RecvTimeoutError::Timeout => Some(std::io::Error::new(
341                        std::io::ErrorKind::TimedOut,
342                        format!("no WebSocket frame from the handler within {timeout:?}"),
343                    )),
344                    RecvTimeoutError::Disconnected => None,
345                }),
346                None => self.rx.recv().map_err(|_| None),
347            };
348            match next {
349                Ok(chunk) => {
350                    self.buf = chunk;
351                    self.pos = 0;
352                }
353                Err(Some(e)) => return Err(e),
354                Err(None) => return Ok(0),
355            }
356        }
357        let n = out.len().min(self.buf.len() - self.pos);
358        out[..n].copy_from_slice(&self.buf[self.pos..self.pos + n]);
359        self.pos += n;
360        Ok(n)
361    }
362}
rws_server/test_client/websocket.rs

rws_server/test_client/
websocket.rs
