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}