use nom7::bytes::streaming::take;
use nom7::combinator::cond;
use nom7::number::streaming::{be_u16, be_u32, be_u64, be_u8};
use nom7::IResult;
use suricata_derive::EnumStringU8;
#[derive(Clone, Debug, Default, EnumStringU8)]
#[repr(u8)]
pub enum WebSocketOpcode {
#[default]
Continuation = 0,
Text = 1,
Binary = 2,
Close = 8,
Ping = 9,
Pong = 10,
}
#[derive(Clone, Debug, Default)]
pub struct WebSocketPdu {
pub flags: u8,
pub fin: bool,
pub compress: bool,
pub opcode: u8,
pub mask: Option<u32>,
pub payload: Vec<u8>,
pub to_skip: u64,
}
pub fn parse_message(i: &[u8], max_pl_size: u32) -> IResult<&[u8], WebSocketPdu> {
let (i, flags_op) = be_u8(i)?;
let fin = (flags_op & 0x80) != 0;
let compress = (flags_op & 0x40) != 0;
let flags = flags_op & 0xF0;
let opcode = flags_op & 0xF;
let (i, mask_plen) = be_u8(i)?;
let mask_flag = (mask_plen & 0x80) != 0;
let (i, payload_len) = match mask_plen & 0x7F {
126 => {
let (i, val) = be_u16(i)?;
Ok((i, val.into()))
}
127 => be_u64(i),
_ => Ok((i, (mask_plen & 0x7F).into())),
}?;
let (i, xormask) = cond(mask_flag, take(4usize))(i)?;
let mask = if mask_flag {
let (_, m) = be_u32(xormask.unwrap())?;
Some(m)
} else {
None
};
let (to_skip, payload_len) = if payload_len < u64::from(max_pl_size) {
(0, payload_len as u32)
} else {
(payload_len - (max_pl_size as u64), max_pl_size)
};
let (i, payload_raw) = take(payload_len)(i)?;
let mut payload = payload_raw.to_vec();
if let Some(xorkey) = xormask {
for i in 0..payload.len() {
payload[i] ^= xorkey[i % 4];
}
}
Ok((
i,
WebSocketPdu {
flags,
fin,
compress,
opcode,
mask,
payload,
to_skip,
},
))
}