Skip to main content

rustlavel_ws/
base64.rs

1//! Base64, needed by the opening handshake.
2//!
3//! `Sec-WebSocket-Key` arrives base64-encoded and `Sec-WebSocket-Accept` goes
4//! back base64-encoded, so the crate needs both directions. Fifty lines is
5//! cheaper than a dependency, and a WebSocket crate that dragged in a base64
6//! crate would be the only thing in the framework that did.
7
8const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
9
10pub fn encode(input: &[u8]) -> String {
11    let mut out = String::with_capacity(input.len().div_ceil(3) * 4);
12
13    for chunk in input.chunks(3) {
14        let b = [chunk[0], *chunk.get(1).unwrap_or(&0), *chunk.get(2).unwrap_or(&0)];
15        let bits = (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]);
16
17        out.push(ALPHABET[(bits >> 18 & 0x3f) as usize] as char);
18        out.push(ALPHABET[(bits >> 12 & 0x3f) as usize] as char);
19        out.push(if chunk.len() > 1 { ALPHABET[(bits >> 6 & 0x3f) as usize] as char } else { '=' });
20        out.push(if chunk.len() > 2 { ALPHABET[(bits & 0x3f) as usize] as char } else { '=' });
21    }
22
23    out
24}
25
26pub fn decode(input: &str) -> Option<Vec<u8>> {
27    let mut out = Vec::with_capacity(input.len() / 4 * 3);
28    let mut buffer = 0u32;
29    let mut bits = 0u32;
30
31    for byte in input.bytes() {
32        let value = match byte {
33            b'A'..=b'Z' => byte - b'A',
34            b'a'..=b'z' => byte - b'a' + 26,
35            b'0'..=b'9' => byte - b'0' + 52,
36            b'+' => 62,
37            b'/' => 63,
38            b'=' | b'\n' | b'\r' => continue,
39            _ => return None,
40        };
41
42        buffer = (buffer << 6) | u32::from(value);
43        bits += 6;
44        if bits >= 8 {
45            bits -= 8;
46            out.push((buffer >> bits) as u8);
47        }
48    }
49
50    Some(out)
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56
57    #[test]
58    fn matches_the_rfc_4648_vectors() {
59        for (plain, encoded) in [
60            ("", ""),
61            ("f", "Zg=="),
62            ("fo", "Zm8="),
63            ("foo", "Zm9v"),
64            ("foob", "Zm9vYg=="),
65            ("fooba", "Zm9vYmE="),
66            ("foobar", "Zm9vYmFy"),
67        ] {
68            assert_eq!(encode(plain.as_bytes()), encoded, "encoding {plain:?}");
69            assert_eq!(decode(encoded).unwrap(), plain.as_bytes(), "decoding {encoded:?}");
70        }
71    }
72
73    #[test]
74    fn round_trips_the_sixteen_random_bytes_a_handshake_key_carries() {
75        let key: Vec<u8> = (0..16).map(|i| i * 17).collect();
76        assert_eq!(decode(&encode(&key)).unwrap(), key);
77    }
78
79    #[test]
80    fn rejects_characters_outside_the_alphabet() {
81        assert!(decode("not base64!").is_none());
82    }
83}