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