char/avril

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

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

Charlotte Somsnapshot: embed shared main and worker startup in the CLI688e96e

main
branchesmainall refs →
8.8 KiB267 linesraw
1use std::io::{self, Read, Write};
2use std::net::{TcpListener, TcpStream};
3use std::time::{Duration, Instant};
4
5pub const CONTINUATION: u8 = 0;
6pub const TEXT: u8 = 1;
7pub const BINARY: u8 = 2;
8pub const CLOSE: u8 = 8;
9pub const PING: u8 = 9;
10pub const PONG: u8 = 10;
11
12pub fn accept(key: &str) -> String {
13    use base64::Engine;
14    use sha1::{Digest, Sha1};
15    let mut hasher = Sha1::new();
16    hasher.update(key.as_bytes());
17    hasher.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
18    base64::engine::general_purpose::STANDARD.encode(hasher.finalize())
19}
20
21pub struct Frame {
22    pub opcode: u8,
23    pub payload: Vec<u8>,
24}
25
26impl Frame {
27    pub fn close(code: u16, reason: &str) -> Frame {
28        let mut payload = code.to_be_bytes().to_vec();
29        payload.extend_from_slice(reason.as_bytes());
30        Frame { opcode: CLOSE, payload }
31    }
32
33    pub fn code(&self) -> u16 {
34        u16::from_be_bytes([self.payload[0], self.payload[1]])
35    }
36
37    pub fn reason(&self) -> String {
38        String::from_utf8_lossy(&self.payload[2..]).into_owned()
39    }
40}
41
42pub fn read_exact(stream: &mut TcpStream, count: usize) -> io::Result<Vec<u8>> {
43    let mut buffer = vec![0; count];
44    stream.read_exact(&mut buffer)?;
45    Ok(buffer)
46}
47
48pub fn read_frame(stream: &mut TcpStream) -> io::Result<Frame> {
49    let head = read_exact(stream, 2)?;
50    let opcode = head[0] & 0x0f;
51    let masked = head[1] & 0x80 != 0;
52    assert!(masked, "client frames must be masked");
53    assert_ne!(head[0] & 0x80, 0, "expected a complete client frame");
54    assert_eq!(head[0] & 0x70, 0, "client must not set reserved bits");
55    let mut length = usize::from(head[1] & 0x7f);
56    match length {
57        126 => {
58            let extra = read_exact(stream, 2)?;
59            length = u16::from_be_bytes([extra[0], extra[1]]) as usize;
60        }
61        127 => {
62            let extra = read_exact(stream, 8)?;
63            length = u64::from_be_bytes(extra.try_into().unwrap()) as usize;
64        }
65        _ => {}
66    }
67    let mask = read_exact(stream, 4)?;
68    let mut payload = read_exact(stream, length)?;
69    for (index, byte) in payload.iter_mut().enumerate() {
70        *byte ^= mask[index % 4];
71    }
72    Ok(Frame { opcode, payload })
73}
74
75pub fn write_frame(
76    stream: &mut TcpStream,
77    fin: bool,
78    opcode: u8,
79    payload: &[u8],
80) -> io::Result<()> {
81    let mut head = vec![(if fin { 0x80 } else { 0 }) | opcode];
82    let length = payload.len();
83    if length < 126 {
84        head.push(length as u8);
85    } else if length <= u16::MAX as usize {
86        head.push(126);
87        head.extend_from_slice(&(length as u16).to_be_bytes());
88    } else {
89        head.push(127);
90        head.extend_from_slice(&(length as u64).to_be_bytes());
91    }
92    stream.write_all(&head)?;
93    stream.write_all(payload)
94}
95
96/// A handshook test server connection.
97pub struct Session {
98    pub stream: TcpStream,
99    pub path: String,
100    pub headers: String,
101    pub protocol: String,
102}
103
104impl Session {
105    pub fn read(&mut self) -> io::Result<Frame> {
106        read_frame(&mut self.stream)
107    }
108
109    pub fn write(&mut self, opcode: u8, payload: &[u8]) -> io::Result<()> {
110        write_frame(&mut self.stream, true, opcode, payload)
111    }
112
113    /// Assert the handshake request matches RFC 6455 and the runtime conventions.
114    pub fn assert_handshake(&self) {
115        assert!(
116            self.path.starts_with('/'),
117            "request target must be origin-form"
118        );
119        let has = |name: &str, value: &str| {
120            self.headers.lines().any(|line| {
121                line.split_once(':').is_some_and(|(field, found)| {
122                    field == name && found.trim().eq_ignore_ascii_case(value)
123                })
124            })
125        };
126        assert!(has("upgrade", "websocket"));
127        assert!(has("connection", "upgrade"));
128        assert!(has("sec-websocket-version", "13"));
129        assert!(
130            self.headers
131                .lines()
132                .any(|line| line.starts_with("sec-websocket-key: ") && line.len() > 16)
133        );
134        assert!(has(
135            "user-agent",
136            concat!("avril/", env!("CARGO_PKG_VERSION"))
137        ));
138    }
139
140    /// Echo data messages back until the peer closes, then echo the close frame.
141    pub fn echo(mut self) -> io::Result<()> {
142        self.assert_handshake();
143        self.accept(None)?;
144        loop {
145            let frame = self.read()?;
146            match frame.opcode {
147                CLOSE => {
148                    self.write(CLOSE, &frame.payload)?;
149                    return Ok(());
150                }
151                TEXT | BINARY => self.write(frame.opcode, &frame.payload)?,
152                _ => {}
153            }
154        }
155    }
156
157    pub fn accept(&mut self, protocol: Option<&str>) -> io::Result<()> {
158        let mut response = format!(
159            "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {}\r\n",
160            accept(&self.key())
161        );
162        if let Some(protocol) = protocol {
163            response.push_str(&format!("Sec-WebSocket-Protocol: {protocol}\r\n"));
164        }
165        response.push_str("\r\n");
166        self.stream.write_all(response.as_bytes())
167    }
168
169    pub fn reject(&mut self) -> io::Result<()> {
170        self.stream
171            .write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
172    }
173
174    pub fn key(&self) -> String {
175        self.headers
176            .lines()
177            .find(|line| line.starts_with("sec-websocket-key:"))
178            .map(|line| line.split_once(':').unwrap().1.trim().to_owned())
179            .expect("Sec-WebSocket-Key header")
180    }
181}
182
183/// Accept `count` connections and run `handler` for each.
184pub fn server(
185    count: usize,
186    handler: impl Fn(Session) -> io::Result<()> + Send + 'static,
187) -> (String, std::thread::JoinHandle<()>) {
188    let (url, thread) = server_on("127.0.0.1", count, handler);
189    (url, thread)
190}
191
192/// Accept one connection, panicking instead of blocking forever when the runtime never
193/// connects: tests join the server thread, so a missed connection would otherwise hang.
194pub fn incoming(listener: &TcpListener) -> TcpStream {
195    listener.set_nonblocking(true).unwrap();
196    let deadline = Instant::now() + Duration::from_secs(10);
197    let stream = loop {
198        match listener.accept() {
199            Ok((stream, _)) => break stream,
200            Err(error)
201                if error.kind() == io::ErrorKind::WouldBlock && Instant::now() < deadline =>
202            {
203                std::thread::sleep(Duration::from_millis(5));
204            }
205            Err(error) => panic!("the runtime never connected: {error}"),
206        }
207    };
208    stream.set_nonblocking(false).unwrap();
209    stream
210        .set_read_timeout(Some(Duration::from_secs(10)))
211        .unwrap();
212    stream
213        .set_write_timeout(Some(Duration::from_secs(10)))
214        .unwrap();
215    stream
216}
217
218pub fn server_on(
219    host: &str,
220    count: usize,
221    handler: impl Fn(Session) -> io::Result<()> + Send + 'static,
222) -> (String, std::thread::JoinHandle<()>) {
223    let listener = TcpListener::bind((host, 0)).unwrap();
224    // SocketAddr serializes IPv6 literals with brackets.
225    let address = listener.local_addr().unwrap().to_string();
226    let thread = std::thread::spawn(move || {
227        for _ in 0..count {
228            let mut stream = incoming(&listener);
229            let mut request = Vec::new();
230            loop {
231                let mut byte = [0; 1];
232                match stream.read_exact(&mut byte) {
233                    Ok(()) => request.push(byte[0]),
234                    Err(_) => break,
235                }
236                if request.ends_with(b"\r\n\r\n") {
237                    break;
238                }
239            }
240            let request = String::from_utf8_lossy(&request);
241            // Header names are matched case-insensitively; values are not.
242            let request = request
243                .lines()
244                .map(|line| match line.split_once(':') {
245                    Some((name, value)) => format!("{}:{value}", name.to_ascii_lowercase()),
246                    None => line.to_owned(),
247                })
248                .collect::<Vec<_>>()
249                .join("\n");
250            let first = request.lines().next().unwrap_or_default().to_owned();
251            let path = first.split(' ').nth(1).unwrap_or_default().to_owned();
252            let headers = request.lines().skip(1).collect::<Vec<_>>().join("\n");
253            let protocol = headers
254                .lines()
255                .find(|line| line.starts_with("sec-websocket-protocol:"))
256                .map(|line| line.split_once(':').unwrap().1.trim().to_owned())
257                .unwrap_or_default();
258            let session = Session { stream, path, headers, protocol };
259            handler(session).unwrap();
260        }
261    });
262    (format!("ws://{address}"), thread)
263}
264
265pub fn connect(url: &str, body: &str) -> String {
266    format!("const ws = new WebSocket('{url}');{body}")
267}