1use crate::base64;
27use rustlavel_http::{IntoResponse, Method, Request, Response, Status};
28use sha1::{Digest, Sha1};
29
30pub const ACCEPT_GUID: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
32
33pub const VERSION: &str = "13";
36
37pub fn accept_key(key: &str) -> String {
39 let mut hasher = Sha1::new();
40 hasher.update(key.as_bytes());
41 hasher.update(ACCEPT_GUID.as_bytes());
42 base64::encode(&hasher.finalize())
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
47pub struct Handshake {
48 pub key: String,
49 pub accept: String,
50 pub protocol: Option<String>,
52 pub offered: Vec<String>,
54}
55
56#[derive(Debug, Clone, PartialEq, Eq)]
61pub enum HandshakeError {
62 NotAGet(Method),
63 MissingUpgrade,
64 NotWebSocket(String),
65 MissingConnection,
66 ConnectionNotUpgrade(String),
67 MissingVersion,
68 UnsupportedVersion(String),
69 MissingKey,
70 MalformedKey(String),
71}
72
73impl HandshakeError {
74 pub fn message(&self) -> String {
75 match self {
76 HandshakeError::NotAGet(method) => {
77 format!("a WebSocket handshake must be a GET, not a {method}")
78 }
79 HandshakeError::MissingUpgrade => {
80 "missing the `Upgrade: websocket` header".to_string()
81 }
82 HandshakeError::NotWebSocket(value) => {
83 format!("`Upgrade` was `{value}`, expected `websocket`")
84 }
85 HandshakeError::MissingConnection => {
86 "missing the `Connection: Upgrade` header".to_string()
87 }
88 HandshakeError::ConnectionNotUpgrade(value) => {
89 format!("`Connection` was `{value}` and does not list `Upgrade`")
90 }
91 HandshakeError::MissingVersion => {
92 format!("missing the `Sec-WebSocket-Version: {VERSION}` header")
93 }
94 HandshakeError::UnsupportedVersion(value) => format!(
95 "`Sec-WebSocket-Version` was `{value}`; this server speaks version {VERSION}"
96 ),
97 HandshakeError::MissingKey => "missing the `Sec-WebSocket-Key` header".to_string(),
98 HandshakeError::MalformedKey(value) => format!(
99 "`Sec-WebSocket-Key` was `{value}`; it must be 16 random bytes, base64-encoded"
100 ),
101 }
102 }
103}
104
105impl std::fmt::Display for HandshakeError {
106 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
107 f.write_str(&self.message())
108 }
109}
110
111impl std::error::Error for HandshakeError {}
112
113impl IntoResponse for HandshakeError {
117 fn into_response(self) -> Response {
118 Response::new(Status::BAD_REQUEST)
119 .with_header("sec-websocket-version", VERSION)
120 .with_text(format!("WebSocket handshake failed: {}", self.message()))
121 }
122}
123
124pub fn negotiate(request: &Request, supported: &[String]) -> Result<Handshake, HandshakeError> {
129 if request.method() != Method::Get {
130 return Err(HandshakeError::NotAGet(request.method()));
131 }
132
133 match request.header("upgrade") {
136 None => return Err(HandshakeError::MissingUpgrade),
137 Some(value) if !value.eq_ignore_ascii_case("websocket") => {
138 return Err(HandshakeError::NotWebSocket(value.to_string()));
139 }
140 Some(_) => {}
141 }
142
143 match request.header("connection") {
144 None => return Err(HandshakeError::MissingConnection),
145 Some(value) if !lists_token(value, "upgrade") => {
148 return Err(HandshakeError::ConnectionNotUpgrade(value.to_string()));
149 }
150 Some(_) => {}
151 }
152
153 match request.header("sec-websocket-version") {
154 None => return Err(HandshakeError::MissingVersion),
155 Some(value) if value.trim() != VERSION => {
156 return Err(HandshakeError::UnsupportedVersion(value.to_string()));
157 }
158 Some(_) => {}
159 }
160
161 let key = request.header("sec-websocket-key").ok_or(HandshakeError::MissingKey)?.trim();
162 match base64::decode(key) {
166 Some(bytes) if bytes.len() == 16 => {}
167 _ => return Err(HandshakeError::MalformedKey(key.to_string())),
168 }
169
170 let offered = offered_protocols(request);
171 let protocol = supported.iter().find(|wanted| offered.iter().any(|o| o == *wanted)).cloned();
172
173 Ok(Handshake { key: key.to_string(), accept: accept_key(key), protocol, offered })
174}
175
176pub fn response(handshake: &Handshake) -> Response {
178 let mut response = Response::new(Status(101))
179 .with_header("upgrade", "websocket")
180 .with_header("connection", "Upgrade")
181 .with_header("sec-websocket-accept", handshake.accept.clone());
182
183 if let Some(protocol) = &handshake.protocol {
186 response = response.with_header("sec-websocket-protocol", protocol.clone());
187 }
188 response
189}
190
191fn offered_protocols(request: &Request) -> Vec<String> {
193 request
194 .headers()
195 .get_all("sec-websocket-protocol")
196 .iter()
197 .flat_map(|value| value.split(','))
198 .map(|token| token.trim().to_string())
199 .filter(|token| !token.is_empty())
200 .collect()
201}
202
203fn lists_token(value: &str, token: &str) -> bool {
205 value.split(',').any(|part| part.trim().eq_ignore_ascii_case(token))
206}
207
208#[cfg(test)]
209mod tests {
210 use super::*;
211
212 fn request() -> Request {
214 Request::new(Method::Get, "/ws")
215 .with_header("upgrade", "websocket")
216 .with_header("connection", "Upgrade")
217 .with_header("sec-websocket-version", "13")
218 .with_header("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ==")
219 }
220
221 #[test]
222 fn computes_the_accept_key_from_the_rfc_6455_example() {
223 assert_eq!(accept_key("dGhlIHNhbXBsZSBub25jZQ=="), "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=");
225 }
226
227 #[test]
228 fn accepts_a_well_formed_handshake() {
229 let handshake = negotiate(&request(), &[]).unwrap();
230
231 assert_eq!(handshake.key, "dGhlIHNhbXBsZSBub25jZQ==");
232 assert_eq!(handshake.accept, "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=");
233 assert_eq!(handshake.protocol, None);
234
235 let response = response(&handshake);
236 assert_eq!(response.status, Status(101));
237 assert_eq!(response.headers.get("upgrade"), Some("websocket"));
238 assert_eq!(response.headers.get("connection"), Some("Upgrade"));
239 assert_eq!(
240 response.headers.get("sec-websocket-accept"),
241 Some("s3pPLMBiTxaQ9kYGzzhZRbK+xOo=")
242 );
243 assert_eq!(response.headers.get("sec-websocket-protocol"), None);
244 }
245
246 #[test]
247 fn header_names_and_values_are_matched_case_insensitively() {
248 let request = Request::new(Method::Get, "/ws")
249 .with_header("Upgrade", "WebSocket")
250 .with_header("Connection", "keep-alive, Upgrade")
251 .with_header("Sec-WebSocket-Version", "13")
252 .with_header("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==");
253
254 assert!(negotiate(&request, &[]).is_ok());
255 }
256
257 #[test]
258 fn a_missing_upgrade_header_is_refused_by_name() {
259 let mut request = request();
260 request.headers_mut().remove("upgrade");
261
262 let error = negotiate(&request, &[]).unwrap_err();
263 assert_eq!(error, HandshakeError::MissingUpgrade);
264 assert!(error.message().contains("Upgrade: websocket"), "{error}");
265 }
266
267 #[test]
268 fn an_upgrade_to_something_else_is_refused() {
269 let request = request().with_header("upgrade", "h2c");
270
271 let error = negotiate(&request, &[]).unwrap_err();
272 assert_eq!(error, HandshakeError::NotWebSocket("h2c".to_string()));
273 }
274
275 #[test]
276 fn a_connection_header_that_does_not_list_upgrade_is_refused() {
277 let mut missing = request();
278 missing.headers_mut().remove("connection");
279 assert_eq!(negotiate(&missing, &[]).unwrap_err(), HandshakeError::MissingConnection);
280
281 let wrong = request().with_header("connection", "keep-alive");
282 assert_eq!(
283 negotiate(&wrong, &[]).unwrap_err(),
284 HandshakeError::ConnectionNotUpgrade("keep-alive".to_string())
285 );
286 }
287
288 #[test]
289 fn only_version_thirteen_is_supported() {
290 let mut missing = request();
291 missing.headers_mut().remove("sec-websocket-version");
292 assert_eq!(negotiate(&missing, &[]).unwrap_err(), HandshakeError::MissingVersion);
293
294 let old = request().with_header("sec-websocket-version", "8");
295 let error = negotiate(&old, &[]).unwrap_err();
296 assert_eq!(error, HandshakeError::UnsupportedVersion("8".to_string()));
297 assert!(error.message().contains("version 13"), "{error}");
298 }
299
300 #[test]
301 fn a_key_that_is_not_sixteen_base64_bytes_is_refused() {
302 let mut missing = request();
303 missing.headers_mut().remove("sec-websocket-key");
304 assert_eq!(negotiate(&missing, &[]).unwrap_err(), HandshakeError::MissingKey);
305
306 let short = request().with_header("sec-websocket-key", "Zm9v");
308 assert!(matches!(negotiate(&short, &[]), Err(HandshakeError::MalformedKey(_))));
309
310 let junk = request().with_header("sec-websocket-key", "not a key!");
312 assert!(matches!(negotiate(&junk, &[]), Err(HandshakeError::MalformedKey(_))));
313 }
314
315 #[test]
316 fn a_post_cannot_open_a_socket() {
317 let request = Request::new(Method::Post, "/ws")
318 .with_header("upgrade", "websocket")
319 .with_header("connection", "Upgrade")
320 .with_header("sec-websocket-version", "13")
321 .with_header("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ==");
322
323 assert_eq!(negotiate(&request, &[]).unwrap_err(), HandshakeError::NotAGet(Method::Post));
324 }
325
326 #[test]
327 fn a_refused_handshake_answers_400_and_says_what_was_wrong() {
328 let response = HandshakeError::MissingKey.into_response();
329
330 assert_eq!(response.status, Status::BAD_REQUEST);
331 assert!(response.body_string().contains("Sec-WebSocket-Key"), "{}", response.body_string());
332 assert_eq!(response.headers.get("sec-websocket-version"), Some("13"));
333 }
334
335 #[test]
336 fn subprotocol_negotiation_follows_the_servers_preference() {
337 let request = request().with_header("sec-websocket-protocol", "soap, chat, superchat");
338 let supported = vec!["superchat".to_string(), "chat".to_string()];
339
340 let handshake = negotiate(&request, &supported).unwrap();
341
342 assert_eq!(handshake.offered, ["soap", "chat", "superchat"]);
343 assert_eq!(handshake.protocol.as_deref(), Some("superchat"));
344 assert_eq!(
345 response(&handshake).headers.get("sec-websocket-protocol"),
346 Some("superchat")
347 );
348 }
349
350 #[test]
351 fn no_shared_subprotocol_still_opens_a_plain_socket() {
352 let request = request().with_header("sec-websocket-protocol", "soap");
353
354 let handshake = negotiate(&request, &["chat".to_string()]).unwrap();
355
356 assert_eq!(handshake.protocol, None);
360 assert_eq!(response(&handshake).headers.get("sec-websocket-protocol"), None);
361 }
362}