moirai_http/
websocket.rs

1//! Bounded RFC 6455 message framing over a Moirai async byte stream.
2
3use std::io;
4
5mod close;
6pub use close::WebSocketPeerClosed;
7use close::{parse_close_payload, protocol_error};
8
9use moirai_async::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
10use moirai_async::timer::timeout;
11
12use crate::upgrade::WebSocketConfig;
13
14const FIN: u8 = 0x80;
15const RSV_MASK: u8 = 0x70;
16const OPCODE_MASK: u8 = 0x0f;
17const MASK: u8 = 0x80;
18const CLOSE: u8 = 0x8;
19const PING: u8 = 0x9;
20const PONG: u8 = 0xa;
21const BINARY: u8 = 0x2;
22
23/// A bounded, message-oriented RFC 6455 stream.
24pub struct WebSocketStream<S> {
25    stream: S,
26    config: WebSocketConfig,
27    prefix: Vec<u8>,
28    prefix_position: usize,
29    closed: bool,
30}
31
32impl<S> WebSocketStream<S> {
33    pub(crate) fn new(stream: S, config: WebSocketConfig, prefix: Vec<u8>) -> Self {
34        Self {
35            stream,
36            config,
37            prefix,
38            prefix_position: 0,
39            closed: false,
40        }
41    }
42}
43
44impl<S: AsyncRead + AsyncWrite + Unpin> WebSocketStream<S> {
45    /// Receive the next complete binary message.
46    ///
47    /// Ping frames are answered with pong frames and are not returned. Text,
48    /// continuation, fragmented, reserved-bit and unmasked client frames are
49    /// rejected, and non-minimal payload lengths are invalid. A close frame or
50    /// any frame I/O error transitions the stream to a terminal state; a
51    /// partially consumed or written frame is never retried on the same stream.
52    ///
53    /// # Errors
54    /// Returns malformed-frame, timeout, connection, or write failures.
55    pub async fn recv_message(&mut self) -> io::Result<Vec<u8>> {
56        if self.closed {
57            return Err(closed_error());
58        }
59        let timeout_duration = self.config.frame_timeout;
60        match timeout(timeout_duration, self.recv_message_inner()).await {
61            Ok(Ok(message)) => Ok(message),
62            Ok(Err(error)) => {
63                self.closed = true;
64                Err(error)
65            }
66            Err(_) => {
67                self.closed = true;
68                Err(timed_out("WebSocket frame receive"))
69            }
70        }
71    }
72
73    /// Send one unfragmented binary message to the peer.
74    ///
75    /// A frame I/O error or timeout transitions the stream to a terminal state
76    /// because a partial write cannot be retried without duplicating bytes.
77    ///
78    /// # Errors
79    /// Returns a size, timeout, connection, or write failure.
80    pub async fn send_binary(&mut self, payload: &[u8]) -> io::Result<()> {
81        if self.closed {
82            return Err(closed_error());
83        }
84        if payload.len() > self.config.max_message_bytes {
85            return Err(io::Error::new(
86                io::ErrorKind::InvalidData,
87                "WebSocket message exceeds configured byte bound",
88            ));
89        }
90        let timeout_duration = self.config.frame_timeout;
91        match timeout(timeout_duration, self.send_frame(BINARY, payload)).await {
92            Ok(Ok(())) => Ok(()),
93            Ok(Err(error)) => {
94                self.closed = true;
95                Err(error)
96            }
97            Err(_) => {
98                self.closed = true;
99                Err(timed_out("WebSocket binary send"))
100            }
101        }
102    }
103
104    /// Send a close frame and transition the stream to a terminal state.
105    ///
106    /// # Errors
107    /// Returns invalid-input, timeout, connection, or write failures.
108    pub async fn close(&mut self, code: u16, reason: &[u8]) -> io::Result<()> {
109        if self.closed {
110            return Ok(());
111        }
112        if reason.len() > 123 || std::str::from_utf8(reason).is_err() {
113            return Err(io::Error::new(
114                io::ErrorKind::InvalidInput,
115                "WebSocket close reason must be valid UTF-8 and at most 123 bytes",
116            ));
117        }
118        let capacity = reason.len().checked_add(2).ok_or_else(|| {
119            io::Error::new(io::ErrorKind::InvalidInput, "close reason is too large")
120        })?;
121        let mut payload = Vec::with_capacity(capacity);
122        payload.extend_from_slice(&code.to_be_bytes());
123        payload.extend_from_slice(reason);
124        parse_close_payload(&payload)?;
125        let timeout_duration = self.config.frame_timeout;
126        let result = match timeout(timeout_duration, self.send_frame(CLOSE, &payload)).await {
127            Ok(result) => result,
128            Err(_) => Err(timed_out("WebSocket close send")),
129        };
130        self.closed = true;
131        result
132    }
133
134    async fn recv_message_inner(&mut self) -> io::Result<Vec<u8>> {
135        loop {
136            let mut header = [0u8; 2];
137            self.read_exact(&mut header).await?;
138            let first = header[0];
139            let second = header[1];
140            if first & RSV_MASK != 0 {
141                return Err(protocol_error("WebSocket reserved bits are not supported"));
142            }
143            let opcode = first & OPCODE_MASK;
144            let final_frame = first & FIN != 0;
145            let masked = second & MASK != 0;
146            let length_code = second & 0x7f;
147            let payload_length = self.read_length(length_code).await?;
148            if (length_code == 126 && payload_length < 126)
149                || (length_code == 127 && payload_length <= 65_535)
150            {
151                return Err(protocol_error(
152                    "WebSocket payload length is not minimally encoded",
153                ));
154            }
155            let control = opcode >= CLOSE;
156            if control {
157                if !final_frame || payload_length > 125 {
158                    return Err(protocol_error("WebSocket control frame is invalid"));
159                }
160            } else if !final_frame {
161                return Err(protocol_error(
162                    "Fragmented WebSocket messages are not supported",
163                ));
164            }
165            if (!control && opcode != BINARY) || (control && !matches!(opcode, CLOSE | PING | PONG))
166            {
167                return Err(protocol_error(
168                    "WebSocket frame opcode is not a supported message or control frame",
169                ));
170            }
171            if !masked {
172                return Err(protocol_error("Client WebSocket frames must be masked"));
173            }
174            let mask = self.read_mask().await?;
175            if opcode == BINARY {
176                if payload_length > self.config.max_message_bytes {
177                    return Err(io::Error::new(
178                        io::ErrorKind::InvalidData,
179                        "WebSocket message exceeds configured byte bound",
180                    ));
181                }
182                return self.read_payload(payload_length, mask).await;
183            }
184            if opcode == PING {
185                let payload = self.read_payload(payload_length, mask).await?;
186                self.send_frame(PONG, &payload).await?;
187                continue;
188            }
189            if opcode == PONG {
190                let _ = self.read_payload(payload_length, mask).await?;
191                continue;
192            }
193            if opcode == CLOSE {
194                let payload = self.read_payload(payload_length, mask).await?;
195                let peer_close = WebSocketPeerClosed::from(parse_close_payload(&payload)?);
196                self.closed = true;
197                self.send_frame(CLOSE, &payload).await?;
198                return Err(io::Error::new(io::ErrorKind::UnexpectedEof, peer_close));
199            }
200            return Err(protocol_error("WebSocket frame opcode is not supported"));
201        }
202    }
203
204    async fn read_length(&mut self, length_code: u8) -> io::Result<usize> {
205        match length_code {
206            0..=125 => Ok(usize::from(length_code)),
207            126 => {
208                let mut bytes = [0u8; 2];
209                self.read_exact(&mut bytes).await?;
210                Ok(usize::from(u16::from_be_bytes(bytes)))
211            }
212            127 => {
213                let mut bytes = [0u8; 8];
214                self.read_exact(&mut bytes).await?;
215                let length = u64::from_be_bytes(bytes);
216                if length & (1u64 << 63) != 0 {
217                    return Err(protocol_error(
218                        "WebSocket payload length has its high bit set",
219                    ));
220                }
221                usize::try_from(length).map_err(|_| {
222                    io::Error::new(
223                        io::ErrorKind::InvalidData,
224                        "WebSocket payload length cannot be represented",
225                    )
226                })
227            }
228            _ => Err(protocol_error("WebSocket payload length code is invalid")),
229        }
230    }
231
232    async fn read_mask(&mut self) -> io::Result<[u8; 4]> {
233        let mut mask = [0u8; 4];
234        self.read_exact(&mut mask).await?;
235        Ok(mask)
236    }
237
238    async fn read_payload(&mut self, length: usize, mask: [u8; 4]) -> io::Result<Vec<u8>> {
239        let mut payload = vec![0u8; length];
240        self.read_exact(&mut payload).await?;
241        for (byte, mask_byte) in payload.iter_mut().zip(mask.iter().cycle()) {
242            *byte ^= *mask_byte;
243        }
244        Ok(payload)
245    }
246
247    async fn read_exact(&mut self, output: &mut [u8]) -> io::Result<()> {
248        let prefix_available = self.prefix.len().saturating_sub(self.prefix_position);
249        let from_prefix = output.len().min(prefix_available);
250        if from_prefix != 0 {
251            let start = self.prefix_position;
252            let end = start.checked_add(from_prefix).ok_or_else(|| {
253                io::Error::other("WebSocket prefix position arithmetic overflowed")
254            })?;
255            let source = self
256                .prefix
257                .get(start..end)
258                .ok_or_else(|| io::Error::other("WebSocket prefix bounds are inconsistent"))?;
259            let destination = output
260                .get_mut(..from_prefix)
261                .ok_or_else(|| io::Error::other("WebSocket output bounds are inconsistent"))?;
262            destination.copy_from_slice(source);
263            self.prefix_position = end;
264        }
265        let mut filled = from_prefix;
266        while filled < output.len() {
267            let destination = output
268                .get_mut(filled..)
269                .ok_or_else(|| io::Error::other("WebSocket output bounds are inconsistent"))?;
270            let count = self.stream.read(destination).await?;
271            if count == 0 {
272                return Err(io::Error::new(
273                    io::ErrorKind::UnexpectedEof,
274                    "connection closed inside a WebSocket frame",
275                ));
276            }
277            filled = filled
278                .checked_add(count)
279                .ok_or_else(|| io::Error::other("WebSocket read length overflow"))?;
280        }
281        Ok(())
282    }
283
284    async fn send_frame(&mut self, opcode: u8, payload: &[u8]) -> io::Result<()> {
285        let capacity = payload
286            .len()
287            .checked_add(10)
288            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "message is too large"))?;
289        let mut frame = Vec::with_capacity(capacity);
290        frame.push(FIN | opcode);
291        match payload.len() {
292            length @ 0..=125 => frame.push(u8::try_from(length).expect("invariant: length <= 125")),
293            length @ 126..=65_535 => {
294                frame.push(126);
295                let length =
296                    u16::try_from(length).expect("invariant: extended WebSocket length fits u16");
297                frame.extend_from_slice(&length.to_be_bytes());
298            }
299            length => {
300                frame.push(127);
301                let length = u64::try_from(length).map_err(|_| {
302                    io::Error::new(io::ErrorKind::InvalidInput, "message is too large")
303                })?;
304                frame.extend_from_slice(&length.to_be_bytes());
305            }
306        }
307        frame.extend_from_slice(payload);
308        self.stream.write_all(&frame).await?;
309        self.stream.flush().await
310    }
311
312    #[cfg(test)]
313    pub(crate) fn initial_bytes(&self) -> &[u8] {
314        self.prefix
315            .get(self.prefix_position..)
316            .expect("invariant: WebSocket prefix position stays within the prefix")
317    }
318
319    #[cfg(test)]
320    pub(crate) fn output_bytes(&self) -> &[u8]
321    where
322        S: OutputBytes,
323    {
324        self.stream.output_bytes()
325    }
326}
327
328fn timed_out(operation: &str) -> io::Error {
329    io::Error::new(io::ErrorKind::TimedOut, operation)
330}
331
332fn closed_error() -> io::Error {
333    io::Error::new(io::ErrorKind::BrokenPipe, "WebSocket is closed")
334}
335
336#[cfg(test)]
337pub(crate) trait OutputBytes {
338    fn output_bytes(&self) -> &[u8];
339}
340
341#[cfg(test)]
342#[path = "websocket_tests.rs"]
343mod tests;
moirai_http/websocket.rs

moirai_http/
websocket.rs
