use thiserror::Error;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u16)]
pub enum CloseCode {
Normal = 1000,
GoingAway = 1001,
ProtocolError = 1002,
UnsupportedData = 1003,
InvalidPayload = 1007,
PolicyViolation = 1008,
MessageTooBig = 1009,
MandatoryExt = 1010,
InternalError = 1011,
}
impl CloseCode {
#[inline]
#[must_use]
pub const fn as_u16(self) -> u16 {
self as u16
}
}
#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
pub enum CloseError {
#[error("close payload has 1 byte (must be 0 or ≥2)")]
OneByte,
#[error("close payload reason is not valid UTF-8")]
InvalidUtf8,
#[error("close code {0} is reserved or invalid")]
InvalidCode(u16),
}
#[must_use]
pub const fn is_valid_endpoint_sent(code: u16) -> bool {
matches!(
code,
1000..=1003 | 1007..=1011 | 3000..=4999
)
}
#[must_use]
pub const fn is_valid_received(code: u16) -> bool {
matches!(
code,
1000..=1003 | 1007..=1011 | 3000..=4999
)
}
pub fn parse_close_payload(payload: &[u8]) -> Result<Option<(u16, &str)>, CloseError> {
match payload.len() {
0 => Ok(None),
1 => Err(CloseError::OneByte),
_ => {
let code = u16::from_be_bytes([payload[0], payload[1]]);
if !is_valid_received(code) {
return Err(CloseError::InvalidCode(code));
}
let reason_bytes = &payload[2..];
let reason = std::str::from_utf8(reason_bytes).map_err(|_| CloseError::InvalidUtf8)?;
Ok(Some((code, reason)))
}
}
}
pub fn encode_close_payload(dst: &mut [u8], code: u16, reason: &str) -> usize {
debug_assert!(dst.len() >= 2 + reason.len());
debug_assert!(reason.len() <= 123);
debug_assert!(std::str::from_utf8(reason.as_bytes()).is_ok());
let code_bytes = code.to_be_bytes();
dst[0] = code_bytes[0];
dst[1] = code_bytes[1];
let reason_len = reason.len();
dst[2..2 + reason_len].copy_from_slice(reason.as_bytes());
2 + reason_len
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::indexing_slicing)]
mod tests {
use super::*;
#[test]
fn empty_payload_means_no_code() {
assert_eq!(parse_close_payload(&[]).unwrap(), None);
}
#[test]
fn one_byte_is_protocol_error() {
assert_eq!(
parse_close_payload(&[0x03]).unwrap_err(),
CloseError::OneByte
);
}
#[test]
fn code_only_normal() {
let payload = [0x03, 0xE8]; let (code, reason) = parse_close_payload(&payload).unwrap().unwrap();
assert_eq!(code, 1000);
assert_eq!(reason, "");
}
#[test]
fn code_plus_reason() {
let mut payload: Vec<u8> = vec![0x03, 0xE9]; payload.extend_from_slice(b"bye");
let (code, reason) = parse_close_payload(&payload).unwrap().unwrap();
assert_eq!(code, 1001);
assert_eq!(reason, "bye");
}
#[test]
fn invalid_code_rejected() {
let payload = [0x03, 0xEE];
let err = parse_close_payload(&payload).unwrap_err();
assert!(matches!(err, CloseError::InvalidCode(1006)));
let payload = [0x0B, 0xB7];
let err = parse_close_payload(&payload).unwrap_err();
assert!(matches!(err, CloseError::InvalidCode(2999)));
}
#[test]
fn application_code_4xxx_ok() {
let payload = [0x0F, 0xA0]; let (code, _) = parse_close_payload(&payload).unwrap().unwrap();
assert_eq!(code, 4000);
}
#[test]
fn invalid_utf8_reason_rejected() {
let payload = [0x03, 0xE8, 0xFF, 0xFE];
let err = parse_close_payload(&payload).unwrap_err();
assert_eq!(err, CloseError::InvalidUtf8);
}
#[test]
fn encode_then_parse_roundtrip() {
let mut buf = [0_u8; 125];
let n = encode_close_payload(&mut buf, 1000, "bye");
assert_eq!(n, 5);
let (code, reason) = parse_close_payload(&buf[..n]).unwrap().unwrap();
assert_eq!(code, 1000);
assert_eq!(reason, "bye");
}
#[test]
fn endpoint_sendable_validation() {
assert!(is_valid_endpoint_sent(1000));
assert!(is_valid_endpoint_sent(1011));
assert!(is_valid_endpoint_sent(4000));
assert!(!is_valid_endpoint_sent(1004));
assert!(!is_valid_endpoint_sent(1005));
assert!(!is_valid_endpoint_sent(1006));
assert!(!is_valid_endpoint_sent(1015));
assert!(!is_valid_endpoint_sent(2999));
assert!(!is_valid_endpoint_sent(5000));
}
}