ws-lite 0.6.1

Naive websocket server implemented using async / await
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

#[derive(Debug, Clone)]
pub enum Message {
    Text(String),
    Binary(Vec<u8>),
    Ping(Vec<u8>),
    Pong(Vec<u8>),
    Close,
}
impl Message {
    pub fn is_closed(&self) -> bool {
        matches!(self, Message::Close)
    }
}
impl AsRef<[u8]> for Message {
    fn as_ref(&self) -> &[u8] {
        match self {
            Message::Text(x) => x.as_bytes(),
            Message::Binary(x) => x.as_slice(),
            Message::Ping(x) => x.as_slice(),
            Message::Pong(x) => x.as_slice(),
            Message::Close => [136, 3, 98, 121, 101].as_slice(),
        }
    }
}
pub struct WriteMessage {
    output: Vec<u8>,
}

impl WriteMessage {
    pub fn new<D>(input: D) -> WriteMessage
    where
        D: AsRef<[u8]>,
    {
        WriteMessage {
            output: message_to_tcp_write_data(D::as_ref(&input)),
        }
    }
    pub fn get_output(&self) -> &Vec<u8> {
        &self.output
    }
}
impl AsRef<[u8]> for WriteMessage {
    fn as_ref(&self) -> &[u8] {
        self.get_output().as_slice()
    }
}
fn message_to_tcp_write_data<D>(data: D) -> Vec<u8>
where
    D: AsRef<[u8]>,
{
    let data = D::as_ref(&data);
    let mut buffer: Vec<u8> = Vec::with_capacity(data.len() + 10);
    buffer.push(129);

    match data.len() as u64 {
        size @ 0..=125 => {
            buffer.push(size as u8);
        }
        size if size <= u32::MAX as u64 => {
            let new_bytes: [u8; 2] = (size as u16).to_be_bytes();

            buffer.push(126);
            buffer.extend_from_slice(&new_bytes);
        }
        size if size > u32::MAX as u64 => {
            let new_bytes: [u8; 8] = (size as u64).to_be_bytes();

            buffer.push(127);
            buffer.extend_from_slice(&new_bytes);
        }
        _ => panic!("Don't know what to do here..."),
    };

    buffer.extend_from_slice(data);

    buffer
}

impl From<Message> for WriteMessage {
    fn from(message: Message) -> Self {
        WriteMessage::new(message.as_ref())
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_close_frame() {
        let expected_result = [136, 3, 98, 121, 101];
        let message = Message::Close;
        let result = message.as_ref();
        assert_eq!(result, expected_result);
    }
}