char/avril

(wip) why'd you have to go and make things so complicaaaaated?

git clone https://git.t4t.associates/char/avril

Charlotte Somrearrange libs into lib and extensions into extsf4a0f5b

main
branchesmainall refs →
5.2 KiB142 linesraw
1use fastwebsockets::{OpCode, WebSocketError};
2use tokio::io::{AsyncRead, AsyncReadExt};
3
4use super::Received;
5
6const MAX_MESSAGE_SIZE: usize = 16 * 1024 * 1024;
7
8pub(super) struct Reader<R> {
9    pub(super) stream: R,
10    message: Option<Received>,
11    validated: usize,
12    masked: bool,
13}
14
15impl<R: AsyncRead + Unpin> Reader<R> {
16    pub(super) fn new(stream: R, masked: bool) -> Self {
17        Self {
18            stream,
19            message: None,
20            validated: 0,
21            masked,
22        }
23    }
24
25    pub(super) async fn read(&mut self) -> Result<Received, WebSocketError> {
26        loop {
27            let mut header = [0; 2];
28            self.stream.read_exact(&mut header).await?;
29            if header[0] & 0x70 != 0 {
30                return Err(WebSocketError::ReservedBitsNotZero);
31            }
32            if (header[1] & 0x80 != 0) != self.masked {
33                return Err(WebSocketError::InvalidValue);
34            }
35            let fin = header[0] & 0x80 != 0;
36            let opcode = OpCode::try_from(header[0] & 0x0f)?;
37            let control = matches!(opcode, OpCode::Close | OpCode::Ping | OpCode::Pong);
38            let length_code = header[1] & 0x7f;
39            if control {
40                if !fin {
41                    return Err(WebSocketError::ControlFrameFragmented);
42                }
43                if length_code > 125 {
44                    return Err(WebSocketError::InvalidValue);
45                }
46            }
47            let length = match length_code {
48                126 => {
49                    let length = u64::from(self.stream.read_u16().await?);
50                    if length < 126 {
51                        return Err(WebSocketError::InvalidValue);
52                    }
53                    length
54                }
55                127 => {
56                    let length = self.stream.read_u64().await?;
57                    if !(65536..=i64::MAX as u64).contains(&length) {
58                        return Err(WebSocketError::InvalidValue);
59                    }
60                    length
61                }
62                length => u64::from(length),
63            };
64            let mut mask = [0; 4];
65            if self.masked {
66                self.stream.read_exact(&mut mask).await?;
67            }
68            if control {
69                let mut payload = vec![0; length as usize];
70                self.stream.read_exact(&mut payload).await?;
71                if self.masked {
72                    for (index, byte) in payload.iter_mut().enumerate() {
73                        *byte ^= mask[index % 4];
74                    }
75                }
76                return Ok(Received { opcode, payload });
77            }
78            match opcode {
79                OpCode::Text | OpCode::Binary => {
80                    if self.message.is_some() {
81                        return Err(WebSocketError::InvalidFragment);
82                    }
83                    self.message = Some(Received { opcode, payload: Vec::new() });
84                    self.validated = 0;
85                }
86                OpCode::Continuation if self.message.is_none() => {
87                    return Err(WebSocketError::InvalidContinuationFrame);
88                }
89                _ => {}
90            }
91            let message = self.message.as_mut().unwrap();
92            let start = message.payload.len();
93            // Check the aggregate before allocating or waiting for the payload.
94            if length > (MAX_MESSAGE_SIZE - start) as u64 {
95                return Err(WebSocketError::FrameTooLarge);
96            }
97            let read = (&mut self.stream)
98                .take(length)
99                .read_to_end(&mut message.payload)
100                .await?;
101            if read as u64 != length {
102                return Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof).into());
103            }
104            if self.masked {
105                for (index, byte) in message.payload[start..].iter_mut().enumerate() {
106                    *byte ^= mask[index % 4];
107                }
108            }
109            if message.opcode == OpCode::Text {
110                match std::str::from_utf8(&message.payload[self.validated..]) {
111                    Ok(_) => self.validated = message.payload.len(),
112                    Err(error) if !fin && error.error_len().is_none() => {
113                        // Recheck only the incomplete suffix when the next fragment arrives.
114                        self.validated += error.valid_up_to();
115                    }
116                    Err(_) => return Err(WebSocketError::InvalidUTF8),
117                }
118            }
119            if fin {
120                return Ok(self.message.take().unwrap());
121            }
122        }
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[tokio::test]
131    async fn truncated_large_frame_is_unexpected_eof() {
132        let mut frame = vec![0x82, 0xff];
133        frame.extend_from_slice(&(MAX_MESSAGE_SIZE as u64).to_be_bytes());
134        frame.extend_from_slice(&[1, 2, 3, 4]);
135        frame.extend_from_slice(&[5, 6, 7]);
136        let mut reader = Reader::new(frame.as_slice(), true);
137        assert!(matches!(
138            reader.read().await,
139            Err(WebSocketError::IoError(error)) if error.kind() == std::io::ErrorKind::UnexpectedEof
140        ));
141    }
142}