1use crate::parse::has_token;
11use crate::sha1::sha1;
12use crate::write::HeadWriter;
13use crate::{base64_decode, base64_encode, Error, Framing, Request, Response, Version};
14
15pub const GUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
17
18pub const VERSION_REQUIRED: &[u8] =
20 b"HTTP/1.1 426 Upgrade Required\r\nSec-WebSocket-Version: 13\r\nContent-Length: 0\r\n\r\n";
21
22pub fn accept_key(key: &[u8]) -> [u8; 28] {
24 let mut out = Vec::with_capacity(28);
25 base64_encode(&sha1(&[key, GUID]), &mut out);
26 let mut a = [0u8; 28];
27 a.copy_from_slice(&out);
28 a
29}
30
31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
33pub enum HandshakeError {
34 NotGet,
36 NotHttp11,
38 NotUpgrade,
40 BadVersion,
42 BadKey,
44 HasBody,
46}
47
48impl HandshakeError {
49 pub fn status(self) -> u16 {
51 if self == HandshakeError::BadVersion {
52 426
53 } else {
54 400
55 }
56 }
57}
58
59#[derive(Clone, Copy, Debug, PartialEq, Eq)]
61pub struct Handshake<'b> {
62 pub key: &'b [u8],
64 pub protocols: Option<&'b [u8]>,
66 pub extensions: Option<&'b [u8]>,
68 pub accept: [u8; 28],
70}
71
72fn single<'b>(req: &Request<'b, '_>, name: &[u8]) -> Option<&'b [u8]> {
73 let mut it = req.headers.iter().filter(|h| h.name.eq_ignore_ascii_case(name));
74 match (it.next(), it.next()) {
75 (Some(h), None) => Some(h.value),
76 _ => None,
77 }
78}
79
80pub fn validate_handshake<'b>(req: &Request<'b, '_>) -> Result<Handshake<'b>, HandshakeError> {
82 if req.method != b"GET" {
83 return Err(HandshakeError::NotGet);
84 }
85 if req.version != Version::Http11 {
86 return Err(HandshakeError::NotHttp11);
87 }
88 if !req.upgrade.is_some_and(|u| has_token(u, b"websocket")) {
89 return Err(HandshakeError::NotUpgrade);
90 }
91 if single(req, b"sec-websocket-version") != Some(b"13") {
92 return Err(HandshakeError::BadVersion);
93 }
94 let key = single(req, b"sec-websocket-key").ok_or(HandshakeError::BadKey)?;
95 let mut raw = Vec::with_capacity(18);
96 if key.len() != 24 || !base64_decode(key, &mut raw) || raw.len() != 16 {
97 return Err(HandshakeError::BadKey);
98 }
99 if req.framing != Framing::None {
100 return Err(HandshakeError::HasBody);
101 }
102 Ok(Handshake {
103 key,
104 protocols: req.header(b"sec-websocket-protocol"),
105 extensions: req.header(b"sec-websocket-extensions"),
106 accept: accept_key(key),
107 })
108}
109
110pub fn write_accept(out: &mut Vec<u8>, hs: &Handshake<'_>, protocol: Option<&[u8]>) -> Result<(), Error> {
112 write_accept_with(out, hs, protocol, None)
113}
114
115pub fn write_accept_with(out: &mut Vec<u8>, hs: &Handshake<'_>, protocol: Option<&[u8]>, extensions: Option<&[u8]>) -> Result<(), Error> {
117 let mut w = HeadWriter::response(out, 101);
118 w.header(b"Upgrade", b"websocket")?.header(b"Connection", b"Upgrade")?.header(b"Sec-WebSocket-Accept", &hs.accept)?;
119 if let Some(p) = protocol {
120 w.header(b"Sec-WebSocket-Protocol", p)?;
121 }
122 if let Some(e) = extensions {
123 w.header(b"Sec-WebSocket-Extensions", e)?;
124 }
125 w.finish();
126 Ok(())
127}
128
129pub fn validate_response(resp: &Response<'_, '_>, key: &[u8]) -> bool {
131 let has = |name: &[u8], token: &[u8]| {
132 resp.headers.iter().any(|h| h.name.eq_ignore_ascii_case(name) && has_token(h.value, token))
133 };
134 resp.status == 101
135 && has(b"upgrade", b"websocket")
136 && has(b"connection", b"upgrade")
137 && resp.header(b"sec-websocket-accept") == Some(&accept_key(key)[..])
138}
139
140#[derive(Clone, Copy, Debug, PartialEq, Eq)]
142pub enum Opcode {
143 Continuation = 0,
145 Text = 1,
147 Binary = 2,
149 Close = 8,
151 Ping = 9,
153 Pong = 10,
155}
156
157#[derive(Clone, Copy, Debug, PartialEq, Eq)]
159pub enum Role {
160 Server,
162 Client,
164}
165
166#[derive(Clone, Copy, Debug, PartialEq, Eq)]
168pub struct WsLimits {
169 pub max_frame: u64,
171 pub max_message: u64,
173}
174
175impl Default for WsLimits {
176 fn default() -> Self {
177 WsLimits { max_frame: 16 << 20, max_message: 16 << 20 }
178 }
179}
180
181#[derive(Clone, Copy, Debug, PartialEq, Eq)]
183pub enum WsError {
184 ReservedBits,
186 BadOpcode,
188 Unmasked,
190 Masked,
192 FragmentedControl,
194 ControlTooLong,
196 BadLength,
198 FrameTooLarge,
200 MessageTooLarge,
202 UnexpectedContinuation,
204 ExpectedContinuation,
206 InvalidUtf8,
208 InvalidClose,
210 AfterClose,
212 Compression,
214}
215
216impl WsError {
217 pub fn close_code(self) -> u16 {
219 match self {
220 WsError::InvalidUtf8 | WsError::Compression => 1007,
221 WsError::FrameTooLarge | WsError::MessageTooLarge => 1009,
222 _ => 1002,
223 }
224 }
225}
226
227#[derive(Debug, PartialEq, Eq)]
229pub enum WsEvent<'a> {
230 Data {
234 text: bool,
236 payload: &'a [u8],
238 fin: bool,
240 },
241 Ping(&'a [u8]),
243 Pong(&'a [u8]),
245 Close {
247 code: Option<u16>,
249 reason: &'a str,
251 },
252 NeedMore,
254}
255
256#[derive(Clone, Copy, Debug, Default)]
258struct Utf8 {
259 need: u8,
260 lo: u8,
261 hi: u8,
262}
263
264impl Utf8 {
265 fn feed(&mut self, bytes: &[u8]) -> bool {
266 let mut i = 0;
267 while i < bytes.len() {
268 if self.need == 0 {
269 if i + 8 <= bytes.len() {
270 let w = u64::from_le_bytes(bytes[i..i + 8].try_into().unwrap_or([0; 8]));
271 if w & 0x8080_8080_8080_8080 == 0 {
272 i += 8;
273 continue;
274 }
275 }
276 let b = bytes[i];
277 let (need, lo, hi) = match b {
278 0x00..=0x7f => (0, 0x80, 0xbf),
279 0xc2..=0xdf => (1, 0x80, 0xbf),
280 0xe0 => (2, 0xa0, 0xbf),
281 0xe1..=0xec | 0xee | 0xef => (2, 0x80, 0xbf),
282 0xed => (2, 0x80, 0x9f),
283 0xf0 => (3, 0x90, 0xbf),
284 0xf1..=0xf3 => (3, 0x80, 0xbf),
285 0xf4 => (3, 0x80, 0x8f),
286 _ => return false,
287 };
288 *self = Utf8 { need, lo, hi };
289 } else {
290 let b = bytes[i];
291 if b < self.lo || b > self.hi {
292 return false;
293 }
294 *self = Utf8 { need: self.need - 1, lo: 0x80, hi: 0xbf };
295 }
296 i += 1;
297 }
298 true
299 }
300}
301
302pub fn apply_mask(data: &mut [u8], key: [u8; 4], offset: usize) {
304 let k = [key[offset & 3], key[(offset + 1) & 3], key[(offset + 2) & 3], key[(offset + 3) & 3]];
305 let k8 = u64::from_ne_bytes([k[0], k[1], k[2], k[3], k[0], k[1], k[2], k[3]]);
306 let mut chunks = data.chunks_exact_mut(8);
307 for c in &mut chunks {
308 let v = u64::from_ne_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) ^ k8;
309 c.copy_from_slice(&v.to_ne_bytes());
310 }
311 for (i, b) in chunks.into_remainder().iter_mut().enumerate() {
312 *b ^= k[i & 3];
313 }
314}
315
316fn close_code_ok(code: u16) -> bool {
317 matches!(code, 1000..=1003 | 1007..=1014 | 3000..=4999)
318}
319
320#[derive(Clone, Debug)]
322pub struct FrameDecoder {
323 role: Role,
324 limits: WsLimits,
325 remaining: u64,
327 frame_fin: bool,
328 mask: Option<[u8; 4]>,
329 mask_pos: u8,
330 in_message: bool,
331 text: bool,
332 message_len: u64,
333 utf8: Utf8,
334 closed: bool,
335 rsv1_ok: bool,
337 compressed: bool,
339}
340
341impl FrameDecoder {
342 pub fn set_permessage_deflate(&mut self, on: bool) {
344 self.rsv1_ok = on;
345 }
346
347 pub fn message_compressed(&self) -> bool {
350 self.compressed
351 }
352
353 pub fn new(role: Role, limits: WsLimits) -> Self {
355 FrameDecoder {
356 role,
357 limits,
358 remaining: 0,
359 frame_fin: false,
360 mask: None,
361 mask_pos: 0,
362 in_message: false,
363 text: false,
364 message_len: 0,
365 utf8: Utf8::default(),
366 closed: false,
367 rsv1_ok: false,
368 compressed: false,
369 }
370 }
371
372 pub fn decode<'a>(&mut self, buf: &'a mut [u8]) -> Result<(usize, WsEvent<'a>), WsError> {
375 if self.closed {
376 return if buf.is_empty() { Ok((0, WsEvent::NeedMore)) } else { Err(WsError::AfterClose) };
377 }
378 if self.remaining > 0 {
379 return self.payload(buf, 0);
380 }
381 let [b0, b1, ..] = *buf else { return Ok((0, WsEvent::NeedMore)) };
382 let fin = b0 & 0x80 != 0;
383 let rsv1 = b0 & 0x40 != 0;
384 if b0 & 0x30 != 0 || (rsv1 && (!self.rsv1_ok || !matches!(b0 & 0x0f, 1 | 2))) {
385 return Err(WsError::ReservedBits);
386 }
387 let opcode = match b0 & 0x0f {
388 0 => Opcode::Continuation,
389 1 => Opcode::Text,
390 2 => Opcode::Binary,
391 8 => Opcode::Close,
392 9 => Opcode::Ping,
393 10 => Opcode::Pong,
394 _ => return Err(WsError::BadOpcode),
395 };
396 let masked = b1 & 0x80 != 0;
397 match (self.role, masked) {
398 (Role::Server, false) => return Err(WsError::Unmasked),
399 (Role::Client, true) => return Err(WsError::Masked),
400 _ => {}
401 }
402 let len7 = b1 & 0x7f;
403 let control = (opcode as u8) & 0x08 != 0;
404 if control {
405 if !fin {
406 return Err(WsError::FragmentedControl);
407 }
408 if len7 > 125 {
409 return Err(WsError::ControlTooLong);
410 }
411 }
412 let ext = match len7 {
413 126 => 2,
414 127 => 8,
415 _ => 0,
416 };
417 let hlen = 2 + ext + if masked { 4 } else { 0 };
418 if buf.len() < hlen {
419 return Ok((0, WsEvent::NeedMore));
420 }
421 let len = match ext {
422 2 => {
423 let n = u64::from(u16::from_be_bytes([buf[2], buf[3]]));
424 if n < 126 {
425 return Err(WsError::BadLength);
426 }
427 n
428 }
429 8 => {
430 let n = u64::from_be_bytes([buf[2], buf[3], buf[4], buf[5], buf[6], buf[7], buf[8], buf[9]]);
431 if n <= 0xffff || n >> 63 != 0 {
432 return Err(WsError::BadLength);
433 }
434 n
435 }
436 _ => u64::from(len7),
437 };
438 let mask = masked.then(|| [buf[hlen - 4], buf[hlen - 3], buf[hlen - 2], buf[hlen - 1]]);
439 if control {
440 let end = hlen + len as usize;
441 if buf.len() < end {
442 return Ok((0, WsEvent::NeedMore));
443 }
444 let payload = &mut buf[hlen..end];
445 if let Some(k) = mask {
446 apply_mask(payload, k, 0);
447 }
448 let payload = &*payload;
449 let ev = match opcode {
450 Opcode::Ping => WsEvent::Ping(payload),
451 Opcode::Pong => WsEvent::Pong(payload),
452 _ => {
453 self.closed = true;
454 match payload {
455 [] => WsEvent::Close { code: None, reason: "" },
456 [_] => return Err(WsError::InvalidClose),
457 [hi, lo, reason @ ..] => {
458 let code = u16::from_be_bytes([*hi, *lo]);
459 if !close_code_ok(code) {
460 return Err(WsError::InvalidClose);
461 }
462 let reason = std::str::from_utf8(reason).map_err(|_| WsError::InvalidUtf8)?;
463 WsEvent::Close { code: Some(code), reason }
464 }
465 }
466 }
467 };
468 return Ok((end, ev));
469 }
470 if len > self.limits.max_frame {
471 return Err(WsError::FrameTooLarge);
472 }
473 match (opcode, self.in_message) {
474 (Opcode::Continuation, false) => return Err(WsError::UnexpectedContinuation),
475 (Opcode::Text | Opcode::Binary, true) => return Err(WsError::ExpectedContinuation),
476 (Opcode::Continuation, true) => {}
477 _ => {
478 self.in_message = true;
479 self.text = opcode == Opcode::Text;
480 self.message_len = 0;
481 self.utf8 = Utf8::default();
482 self.compressed = rsv1;
483 }
484 }
485 self.message_len += len;
486 if self.message_len > self.limits.max_message {
487 return Err(WsError::MessageTooLarge);
488 }
489 self.remaining = len;
490 self.frame_fin = fin;
491 self.mask = mask;
492 self.mask_pos = 0;
493 if len == 0 {
494 return self.finish_piece(hlen, &buf[..0]);
495 }
496 self.payload(buf, hlen)
497 }
498
499 fn payload<'a>(&mut self, buf: &'a mut [u8], at: usize) -> Result<(usize, WsEvent<'a>), WsError> {
500 let avail = buf.len() - at;
501 if avail == 0 {
502 return Ok((at, WsEvent::NeedMore));
503 }
504 let n = self.remaining.min(avail as u64) as usize;
505 let piece = &mut buf[at..at + n];
506 if let Some(k) = self.mask {
507 apply_mask(piece, k, usize::from(self.mask_pos));
508 self.mask_pos = ((usize::from(self.mask_pos) + n) & 3) as u8;
509 }
510 if self.text && !self.compressed && !self.utf8.feed(piece) {
511 return Err(WsError::InvalidUtf8);
512 }
513 self.remaining -= n as u64;
514 self.finish_piece(at + n, piece)
515 }
516
517 fn finish_piece<'a>(&mut self, used: usize, piece: &'a [u8]) -> Result<(usize, WsEvent<'a>), WsError> {
518 let fin = self.remaining == 0 && self.frame_fin;
519 if fin {
520 if self.text && !self.compressed && self.utf8.need != 0 {
521 return Err(WsError::InvalidUtf8);
522 }
523 self.in_message = false;
524 }
525 Ok((used, WsEvent::Data { text: self.text, payload: piece, fin }))
526 }
527}
528
529pub fn encode_header(out: &mut [u8; 14], fin: bool, opcode: Opcode, len: u64, mask: Option<[u8; 4]>) -> usize {
532 out[0] = (u8::from(fin) << 7) | opcode as u8;
533 let m = if mask.is_some() { 0x80 } else { 0 };
534 let mut n = if len < 126 {
535 out[1] = m | len as u8;
536 2
537 } else if len <= 0xffff {
538 out[1] = m | 126;
539 out[2..4].copy_from_slice(&(len as u16).to_be_bytes());
540 4
541 } else {
542 out[1] = m | 127;
543 out[2..10].copy_from_slice(&len.to_be_bytes());
544 10
545 };
546 if let Some(k) = mask {
547 out[n..n + 4].copy_from_slice(&k);
548 n += 4;
549 }
550 n
551}
552
553pub fn write_frame(out: &mut Vec<u8>, fin: bool, opcode: Opcode, payload: &[u8], mask: Option<[u8; 4]>) {
555 let mut h = [0u8; 14];
556 let n = encode_header(&mut h, fin, opcode, payload.len() as u64, mask);
557 out.extend_from_slice(&h[..n]);
558 let start = out.len();
559 out.extend_from_slice(payload);
560 if let Some(k) = mask {
561 apply_mask(&mut out[start..], k, 0);
562 }
563}
564
565pub fn write_close(out: &mut Vec<u8>, code: u16, reason: &str, mask: Option<[u8; 4]>) -> bool {
567 if reason.len() > 123 || !close_code_ok(code) {
568 return false;
569 }
570 let mut p = [0u8; 125];
571 p[..2].copy_from_slice(&code.to_be_bytes());
572 p[2..2 + reason.len()].copy_from_slice(reason.as_bytes());
573 write_frame(out, true, Opcode::Close, &p[..2 + reason.len()], mask);
574 true
575}
576
577pub fn write_frame_compressed(out: &mut Vec<u8>, fin: bool, opcode: Opcode, payload: &[u8], mask: Option<[u8; 4]>) {
580 let at = out.len();
581 write_frame(out, fin, opcode, payload, mask);
582 out[at] |= 0x40;
583}
584
585#[cfg(feature = "permessage-deflate")]
586pub mod deflate;