1
0
mirror of https://github.com/actix/actix-extras.git synced 2024-11-28 09:42:40 +01:00
actix-extras/src/wsframe.rs

429 lines
12 KiB
Rust
Raw Normal View History

2017-10-09 02:47:41 +02:00
use std::{fmt, mem};
2017-10-08 06:48:00 +02:00
use std::io::{Write, Error, ErrorKind};
use std::iter::FromIterator;
use bytes::BytesMut;
use wsproto::{OpCode, CloseCode};
fn apply_mask(buf: &mut [u8], mask: &[u8; 4]) {
let iter = buf.iter_mut().zip(mask.iter().cycle());
for (byte, &key) in iter {
*byte ^= key
}
}
/// A struct representing a `WebSocket` frame.
#[derive(Debug, Clone)]
2017-10-08 07:41:02 +02:00
pub(crate) struct Frame {
2017-10-08 06:48:00 +02:00
finished: bool,
rsv1: bool,
rsv2: bool,
rsv3: bool,
opcode: OpCode,
mask: Option<[u8; 4]>,
payload: Vec<u8>,
}
impl Frame {
/// Desctructe frame
pub fn unpack(self) -> (bool, OpCode, Vec<u8>) {
(self.finished, self.opcode, self.payload)
}
/// Get the length of the frame.
/// This is the length of the header + the length of the payload.
#[inline]
pub fn len(&self) -> usize {
let mut header_length = 2;
2017-10-09 02:47:41 +02:00
let payload_len = self.payload.len();
2017-10-08 06:48:00 +02:00
if payload_len > 125 {
if payload_len <= u16::max_value() as usize {
header_length += 2;
} else {
header_length += 8;
}
}
2017-10-09 02:47:41 +02:00
if self.mask.is_some() {
2017-10-08 06:48:00 +02:00
header_length += 4;
}
header_length + payload_len
}
/// Create a new data frame.
#[inline]
pub fn message(data: Vec<u8>, code: OpCode, finished: bool) -> Frame {
Frame {
finished: finished,
opcode: code,
payload: data,
.. Frame::default()
}
}
/// Create a new Close control frame.
#[inline]
pub fn close(code: CloseCode, reason: &str) -> Frame {
let raw: [u8; 2] = unsafe {
let u: u16 = code.into();
2017-10-09 02:47:41 +02:00
mem::transmute(u.to_be())
2017-10-08 06:48:00 +02:00
};
let payload = if let CloseCode::Empty = code {
Vec::new()
} else {
Vec::from_iter(
raw[..].iter()
.chain(reason.as_bytes().iter())
.cloned())
};
Frame {
payload: payload,
.. Frame::default()
}
}
/// Parse the input stream into a frame.
pub fn parse(buf: &mut BytesMut) -> Result<Option<Frame>, Error> {
let mut idx = 2;
let (frame, length) = {
let mut size = buf.len();
if size < 2 {
return Ok(None)
}
let mut head = [0u8; 2];
size -= 2;
head.copy_from_slice(&buf[..2]);
trace!("Parsed headers {:?}", head);
let first = head[0];
let second = head[1];
trace!("First: {:b}", first);
trace!("Second: {:b}", second);
let finished = first & 0x80 != 0;
let rsv1 = first & 0x40 != 0;
let rsv2 = first & 0x20 != 0;
let rsv3 = first & 0x10 != 0;
let opcode = OpCode::from(first & 0x0F);
trace!("Opcode: {:?}", opcode);
let masked = second & 0x80 != 0;
trace!("Masked: {:?}", masked);
let mut header_length = 2;
let mut length = u64::from(second & 0x7F);
if length == 126 {
if size < 2 {
return Ok(None)
}
2017-10-23 22:31:22 +02:00
let mut length_bytes = [0u8; 2];
2017-10-08 06:48:00 +02:00
length_bytes.copy_from_slice(&buf[idx..idx+2]);
size -= 2;
idx += 2;
length = u64::from(unsafe{
2017-10-09 02:47:41 +02:00
let mut wide: u16 = mem::transmute(length_bytes);
2017-10-08 06:48:00 +02:00
wide = u16::from_be(wide);
wide});
header_length += 2;
} else if length == 127 {
if size < 8 {
return Ok(None)
}
2017-10-23 22:31:22 +02:00
let mut length_bytes = [0u8; 8];
2017-10-08 06:48:00 +02:00
length_bytes.copy_from_slice(&buf[idx..idx+8]);
size -= 8;
2017-10-23 22:31:22 +02:00
idx += 8;
2017-10-08 06:48:00 +02:00
2017-10-09 02:47:41 +02:00
unsafe { length = mem::transmute(length_bytes); }
2017-10-08 06:48:00 +02:00
length = u64::from_be(length);
header_length += 8;
}
trace!("Payload length: {}", length);
let mask = if masked {
let mut mask_bytes = [0u8; 4];
if size < 4 {
return Ok(None)
} else {
header_length += 4;
size -= 4;
mask_bytes.copy_from_slice(&buf[idx..idx+4]);
idx += 4;
Some(mask_bytes)
}
} else {
None
};
let length = length as usize;
if size < length {
return Ok(None)
}
let mut data = Vec::with_capacity(length);
if length > 0 {
data.extend_from_slice(&buf[idx..idx+length]);
}
// Disallow bad opcode
if let OpCode::Bad = opcode {
return Err(
Error::new(
ErrorKind::Other,
format!("Encountered invalid opcode: {}", first & 0x0F)))
}
// control frames must have length <= 125
match opcode {
OpCode::Ping | OpCode::Pong if length > 125 => {
return Err(
Error::new(
ErrorKind::Other,
format!("Rejected WebSocket handshake.Received control frame with length: {}.", length)))
}
OpCode::Close if length > 125 => {
debug!("Received close frame with payload length exceeding 125. Morphing to protocol close frame.");
return Ok(Some(Frame::close(CloseCode::Protocol, "Received close frame with payload length exceeding 125.")))
}
_ => ()
}
// unmask
if let Some(ref mask) = mask {
apply_mask(&mut data, mask);
}
let frame = Frame {
finished: finished,
rsv1: rsv1,
rsv2: rsv2,
rsv3: rsv3,
opcode: opcode,
mask: mask,
payload: data,
};
(frame, header_length + length)
};
buf.split_to(length);
Ok(Some(frame))
}
/// Write a frame out to a buffer
pub fn format<W>(&mut self, w: &mut W) -> Result<(), Error>
where W: Write
{
let mut one = 0u8;
let code: u8 = self.opcode.into();
2017-10-09 02:47:41 +02:00
if self.finished {
2017-10-08 06:48:00 +02:00
one |= 0x80;
}
2017-10-09 02:47:41 +02:00
if self.rsv1 {
2017-10-08 06:48:00 +02:00
one |= 0x40;
}
2017-10-09 02:47:41 +02:00
if self.rsv2 {
2017-10-08 06:48:00 +02:00
one |= 0x20;
}
2017-10-09 02:47:41 +02:00
if self.rsv3 {
2017-10-08 06:48:00 +02:00
one |= 0x10;
}
one |= code;
let mut two = 0u8;
2017-10-09 02:47:41 +02:00
if self.mask.is_some() {
2017-10-08 06:48:00 +02:00
two |= 0x80;
}
if self.payload.len() < 126 {
two |= self.payload.len() as u8;
let headers = [one, two];
try!(w.write_all(&headers));
} else if self.payload.len() <= 65_535 {
two |= 126;
let length_bytes: [u8; 2] = unsafe {
let short = self.payload.len() as u16;
2017-10-09 02:47:41 +02:00
mem::transmute(short.to_be())
2017-10-08 06:48:00 +02:00
};
let headers = [one, two, length_bytes[0], length_bytes[1]];
try!(w.write_all(&headers));
} else {
two |= 127;
let length_bytes: [u8; 8] = unsafe {
let long = self.payload.len() as u64;
2017-10-09 02:47:41 +02:00
mem::transmute(long.to_be())
2017-10-08 06:48:00 +02:00
};
let headers = [
one,
two,
length_bytes[0],
length_bytes[1],
length_bytes[2],
length_bytes[3],
length_bytes[4],
length_bytes[5],
length_bytes[6],
length_bytes[7],
];
try!(w.write_all(&headers));
}
2017-10-09 02:47:41 +02:00
if self.mask.is_some() {
2017-10-08 06:48:00 +02:00
let mask = self.mask.take().unwrap();
apply_mask(&mut self.payload, &mask);
try!(w.write_all(&mask));
}
try!(w.write_all(&self.payload));
Ok(())
}
}
impl Default for Frame {
fn default() -> Frame {
Frame {
finished: true,
rsv1: false,
rsv2: false,
rsv3: false,
opcode: OpCode::Close,
mask: None,
payload: Vec::new(),
}
}
}
impl fmt::Display for Frame {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f,
"
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
length: {}
payload length: {}
payload: 0x{}
</FRAME>",
self.finished,
self.rsv1,
self.rsv2,
self.rsv3,
self.opcode,
// self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
self.len(),
self.payload.len(),
self.payload.iter().map(|byte| format!("{:x}", byte)).collect::<String>())
}
}
2017-10-23 22:31:22 +02:00
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse() {
let mut buf = BytesMut::from(&[0b00000001u8, 0b00000001u8][..]);
assert!(Frame::parse(&mut buf).unwrap().is_none());
buf.extend(b"1");
let frame = Frame::parse(&mut buf).unwrap().unwrap();
2017-10-23 23:08:11 +02:00
println!("FRAME: {}", frame);
2017-10-23 22:31:22 +02:00
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert_eq!(frame.payload, &b"1"[..]);
}
#[test]
fn test_parse_length0() {
let mut buf = BytesMut::from(&[0b00000001u8, 0b00000000u8][..]);
let frame = Frame::parse(&mut buf).unwrap().unwrap();
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert!(frame.payload.is_empty());
}
#[test]
fn test_parse_length2() {
let mut buf = BytesMut::from(&[0b00000001u8, 126u8][..]);
assert!(Frame::parse(&mut buf).unwrap().is_none());
buf.extend(&[0u8, 4u8][..]);
buf.extend(b"1234");
let frame = Frame::parse(&mut buf).unwrap().unwrap();
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert_eq!(frame.payload, &b"1234"[..]);
}
#[test]
fn test_parse_length4() {
let mut buf = BytesMut::from(&[0b00000001u8, 127u8][..]);
assert!(Frame::parse(&mut buf).unwrap().is_none());
buf.extend(&[0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 4u8][..]);
buf.extend(b"1234");
let frame = Frame::parse(&mut buf).unwrap().unwrap();
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert_eq!(frame.payload, &b"1234"[..]);
}
#[test]
fn test_parse_frame_mask() {
let mut buf = BytesMut::from(&[0b00000001u8, 0b10000001u8][..]);
buf.extend(b"0001");
buf.extend(b"1");
let frame = Frame::parse(&mut buf).unwrap().unwrap();
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert_eq!(frame.payload, vec![1u8]);
}
#[test]
fn test_ping_frame() {
let mut frame = Frame::message(Vec::from("data"), OpCode::Ping, true);
let mut buf = Vec::new();
frame.format(&mut buf).unwrap();
let mut v = vec![137u8, 4u8];
v.extend(b"data");
assert_eq!(buf, v);
}
#[test]
fn test_pong_frame() {
let mut frame = Frame::message(Vec::from("data"), OpCode::Pong, true);
let mut buf = Vec::new();
frame.format(&mut buf).unwrap();
let mut v = vec![138u8, 4u8];
v.extend(b"data");
assert_eq!(buf, v);
}
#[test]
fn test_close_frame() {
let mut frame = Frame::close(CloseCode::Normal, "data");
let mut buf = Vec::new();
frame.format(&mut buf).unwrap();
let mut v = vec![136u8, 6u8, 3u8, 232u8];
v.extend(b"data");
assert_eq!(buf, v);
}
}