Skip to main content

specter/websocket/
message.rs

1use bytes::Bytes;
2
3use crate::websocket::error::{WebSocketError, WebSocketResult};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub enum Message {
7    Text(String),
8    Binary(Bytes),
9    Ping(Bytes),
10    Pong(Bytes),
11    Close(Option<CloseFrame>),
12}
13
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct CloseFrame {
16    pub code: CloseCode,
17    pub reason: String,
18}
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum CloseCode {
22    Normal,
23    Away,
24    Protocol,
25    Unsupported,
26    Status,
27    Abnormal,
28    Invalid,
29    Policy,
30    Size,
31    Extension,
32    Error,
33    Restart,
34    Again,
35    Tls,
36    Library(u16),
37    Iana(u16),
38    Private(u16),
39}
40
41impl CloseCode {
42    pub fn as_u16(self) -> u16 {
43        match self {
44            Self::Normal => 1000,
45            Self::Away => 1001,
46            Self::Protocol => 1002,
47            Self::Unsupported => 1003,
48            Self::Status => 1005,
49            Self::Abnormal => 1006,
50            Self::Invalid => 1007,
51            Self::Policy => 1008,
52            Self::Size => 1009,
53            Self::Extension => 1010,
54            Self::Error => 1011,
55            Self::Restart => 1012,
56            Self::Again => 1013,
57            Self::Tls => 1015,
58            Self::Library(code) | Self::Iana(code) | Self::Private(code) => code,
59        }
60    }
61
62    pub fn from_u16(code: u16) -> Option<Self> {
63        Some(match code {
64            1000 => Self::Normal,
65            1001 => Self::Away,
66            1002 => Self::Protocol,
67            1003 => Self::Unsupported,
68            1005 => Self::Status,
69            1006 => Self::Abnormal,
70            1007 => Self::Invalid,
71            1008 => Self::Policy,
72            1009 => Self::Size,
73            1010 => Self::Extension,
74            1011 => Self::Error,
75            1012 => Self::Restart,
76            1013 => Self::Again,
77            1015 => Self::Tls,
78            3000..=3999 => Self::Iana(code),
79            4000..=4999 => Self::Private(code),
80            _ if (1000..=2999).contains(&code) => Self::Library(code),
81            _ => return None,
82        })
83    }
84
85    pub fn is_valid_wire_code(self) -> bool {
86        is_valid_wire_close_code(self.as_u16())
87    }
88
89    pub(crate) fn from_wire(code: u16) -> Option<Self> {
90        if is_valid_wire_close_code(code) {
91            Self::from_u16(code)
92        } else {
93            None
94        }
95    }
96}
97
98impl CloseFrame {
99    pub(crate) fn encode(&self, url: &url::Url) -> WebSocketResult<Vec<u8>> {
100        self.validate_for_send(url)?;
101        let mut payload = Vec::with_capacity(2 + self.reason.len());
102        payload.extend_from_slice(&self.code.as_u16().to_be_bytes());
103        payload.extend_from_slice(self.reason.as_bytes());
104        Ok(payload)
105    }
106
107    pub(crate) fn validate_for_send(&self, url: &url::Url) -> WebSocketResult<()> {
108        let code = self.code.as_u16();
109        if !is_valid_wire_close_code(code) {
110            return Err(WebSocketError::protocol(
111                url,
112                format!("close code {code} must not be sent on the wire"),
113            ));
114        }
115
116        if self.reason.len() > 123 {
117            return Err(WebSocketError::protocol(
118                url,
119                "close reason exceeds 123 bytes",
120            ));
121        }
122
123        Ok(())
124    }
125
126    pub(crate) fn decode(url: &url::Url, payload: &[u8]) -> WebSocketResult<Option<Self>> {
127        if payload.is_empty() {
128            return Ok(None);
129        }
130        if payload.len() == 1 {
131            return Err(WebSocketError::protocol(
132                url,
133                "close frame payload must be empty or at least two bytes",
134            ));
135        }
136
137        let code = u16::from_be_bytes([payload[0], payload[1]]);
138        let code = CloseCode::from_wire(code)
139            .ok_or_else(|| WebSocketError::protocol(url, format!("invalid close code {code}")))?;
140        let reason = std::str::from_utf8(&payload[2..])
141            .map_err(|e| WebSocketError::utf8(url, e.to_string()))?
142            .to_owned();
143
144        Ok(Some(Self { code, reason }))
145    }
146}
147
148fn is_valid_wire_close_code(code: u16) -> bool {
149    matches!(code, 1000..=1003 | 1007..=1014 | 3000..=4999)
150}