use byteorder::{ByteOrder, LittleEndian, NetworkEndian};
use bytes::{BufMut, Bytes, BytesMut};
use futures::{Async, Poll, Stream};
use rand;
use std::fmt;
use body::Binary;
use error::PayloadError;
use payload::PayloadBuffer;
use ws::mask::apply_mask;
use ws::proto::{CloseCode, CloseReason, OpCode};
use ws::ProtocolError;
#[derive(Debug)]
pub struct Frame {
finished: bool,
opcode: OpCode,
payload: Binary,
}
impl Frame {
pub fn unpack(self) -> (bool, OpCode, Binary) {
(self.finished, self.opcode, self.payload)
}
#[inline]
pub fn close(reason: Option<CloseReason>, genmask: bool) -> FramedMessage {
let payload = match reason {
None => Vec::new(),
Some(reason) => {
let mut code_bytes = [0; 2];
NetworkEndian::write_u16(&mut code_bytes, reason.code.into());
let mut payload = Vec::from(&code_bytes[..]);
if let Some(description) = reason.description {
payload.extend(description.as_bytes());
}
payload
}
};
Frame::message(payload, OpCode::Close, true, genmask)
}
#[cfg_attr(feature = "cargo-clippy", allow(type_complexity))]
fn read_copy_md<S>(
pl: &mut PayloadBuffer<S>, server: bool, max_size: usize,
) -> Poll<Option<(usize, bool, OpCode, usize, Option<u32>)>, ProtocolError>
where
S: Stream<Item = Bytes, Error = PayloadError>,
{
let mut idx = 2;
let buf = match pl.copy(2)? {
Async::Ready(Some(buf)) => buf,
Async::Ready(None) => return Ok(Async::Ready(None)),
Async::NotReady => return Ok(Async::NotReady),
};
let first = buf[0];
let second = buf[1];
let finished = first & 0x80 != 0;
let masked = second & 0x80 != 0;
if !masked && server {
return Err(ProtocolError::UnmaskedFrame);
} else if masked && !server {
return Err(ProtocolError::MaskedFrame);
}
let opcode = OpCode::from(first & 0x0F);
if let OpCode::Bad = opcode {
return Err(ProtocolError::InvalidOpcode(first & 0x0F));
}
let len = second & 0x7F;
let length = if len == 126 {
let buf = match pl.copy(4)? {
Async::Ready(Some(buf)) => buf,
Async::Ready(None) => return Ok(Async::Ready(None)),
Async::NotReady => return Ok(Async::NotReady),
};
let len = NetworkEndian::read_uint(&buf[idx..], 2) as usize;
idx += 2;
len
} else if len == 127 {
let buf = match pl.copy(10)? {
Async::Ready(Some(buf)) => buf,
Async::Ready(None) => return Ok(Async::Ready(None)),
Async::NotReady => return Ok(Async::NotReady),
};
let len = NetworkEndian::read_uint(&buf[idx..], 8);
if len > max_size as u64 {
return Err(ProtocolError::Overflow);
}
idx += 8;
len as usize
} else {
len as usize
};
if length > max_size {
return Err(ProtocolError::Overflow);
}
let mask = if server {
let buf = match pl.copy(idx + 4)? {
Async::Ready(Some(buf)) => buf,
Async::Ready(None) => return Ok(Async::Ready(None)),
Async::NotReady => return Ok(Async::NotReady),
};
let mask: &[u8] = &buf[idx..idx + 4];
let mask_u32 = LittleEndian::read_u32(mask);
idx += 4;
Some(mask_u32)
} else {
None
};
Ok(Async::Ready(Some((idx, finished, opcode, length, mask))))
}
fn read_chunk_md(
chunk: &[u8], server: bool, max_size: usize,
) -> Poll<(usize, bool, OpCode, usize, Option<u32>), ProtocolError> {
let chunk_len = chunk.len();
let mut idx = 2;
if chunk_len < 2 {
return Ok(Async::NotReady);
}
let first = chunk[0];
let second = chunk[1];
let finished = first & 0x80 != 0;
let masked = second & 0x80 != 0;
if !masked && server {
return Err(ProtocolError::UnmaskedFrame);
} else if masked && !server {
return Err(ProtocolError::MaskedFrame);
}
let opcode = OpCode::from(first & 0x0F);
if let OpCode::Bad = opcode {
return Err(ProtocolError::InvalidOpcode(first & 0x0F));
}
let len = second & 0x7F;
let length = if len == 126 {
if chunk_len < 4 {
return Ok(Async::NotReady);
}
let len = NetworkEndian::read_uint(&chunk[idx..], 2) as usize;
idx += 2;
len
} else if len == 127 {
if chunk_len < 10 {
return Ok(Async::NotReady);
}
let len = NetworkEndian::read_uint(&chunk[idx..], 8);
if len > max_size as u64 {
return Err(ProtocolError::Overflow);
}
idx += 8;
len as usize
} else {
len as usize
};
if length > max_size {
return Err(ProtocolError::Overflow);
}
let mask = if server {
if chunk_len < idx + 4 {
return Ok(Async::NotReady);
}
let mask: &[u8] = &chunk[idx..idx + 4];
let mask_u32 = LittleEndian::read_u32(mask);
idx += 4;
Some(mask_u32)
} else {
None
};
Ok(Async::Ready((idx, finished, opcode, length, mask)))
}
pub fn parse<S>(
pl: &mut PayloadBuffer<S>, server: bool, max_size: usize,
) -> Poll<Option<Frame>, ProtocolError>
where
S: Stream<Item = Bytes, Error = PayloadError>,
{
let result = match pl.get_chunk()? {
Async::NotReady => return Ok(Async::NotReady),
Async::Ready(None) => return Ok(Async::Ready(None)),
Async::Ready(Some(chunk)) => Frame::read_chunk_md(chunk, server, max_size)?,
};
let (idx, finished, opcode, length, mask) = match result {
Async::NotReady => match Frame::read_copy_md(pl, server, max_size)? {
Async::Ready(Some(item)) => item,
Async::NotReady => return Ok(Async::NotReady),
Async::Ready(None) => return Ok(Async::Ready(None)),
},
Async::Ready(item) => item,
};
match pl.can_read(idx + length)? {
Async::Ready(Some(true)) => (),
Async::Ready(None) => return Ok(Async::Ready(None)),
Async::Ready(Some(false)) | Async::NotReady => return Ok(Async::NotReady),
}
pl.drop_bytes(idx);
if length == 0 {
return Ok(Async::Ready(Some(Frame {
finished,
opcode,
payload: Binary::from(""),
})));
}
let data = match pl.read_exact(length)? {
Async::Ready(Some(buf)) => buf,
Async::Ready(None) => return Ok(Async::Ready(None)),
Async::NotReady => panic!(),
};
match opcode {
OpCode::Ping | OpCode::Pong if length > 125 => {
return Err(ProtocolError::InvalidLength(length))
}
OpCode::Close if length > 125 => {
debug!("Received close frame with payload length exceeding 125. Morphing to protocol close frame.");
return Ok(Async::Ready(Some(Frame::default())));
}
_ => (),
}
let data = if let Some(mask) = mask {
let mut buf = BytesMut::new();
buf.extend_from_slice(&data);
apply_mask(&mut buf, mask);
buf.freeze()
} else {
data
};
Ok(Async::Ready(Some(Frame {
finished,
opcode,
payload: data.into(),
})))
}
pub fn parse_close_payload(payload: &Binary) -> Option<CloseReason> {
if payload.len() >= 2 {
let raw_code = NetworkEndian::read_u16(payload.as_ref());
let code = CloseCode::from(raw_code);
let description = if payload.len() > 2 {
Some(String::from_utf8_lossy(&payload.as_ref()[2..]).into())
} else {
None
};
Some(CloseReason { code, description })
} else {
None
}
}
pub fn message<B: Into<Binary>>(
data: B, code: OpCode, finished: bool, genmask: bool,
) -> FramedMessage {
let payload = data.into();
let one: u8 = if finished {
0x80 | Into::<u8>::into(code)
} else {
code.into()
};
let payload_len = payload.len();
let (two, p_len) = if genmask {
(0x80, payload_len + 4)
} else {
(0, payload_len)
};
let mut buf = if payload_len < 126 {
let mut buf = BytesMut::with_capacity(p_len + 2);
buf.put_slice(&[one, two | payload_len as u8]);
buf
} else if payload_len <= 65_535 {
let mut buf = BytesMut::with_capacity(p_len + 4);
buf.put_slice(&[one, two | 126]);
buf.put_u16_be(payload_len as u16);
buf
} else {
let mut buf = BytesMut::with_capacity(p_len + 10);
buf.put_slice(&[one, two | 127]);
buf.put_u64_be(payload_len as u64);
buf
};
let binary = if genmask {
let mask = rand::random::<u32>();
buf.put_u32_le(mask);
buf.extend_from_slice(payload.as_ref());
let pos = buf.len() - payload_len;
apply_mask(&mut buf[pos..], mask);
buf.into()
} else {
buf.put_slice(payload.as_ref());
buf.into()
};
FramedMessage(binary)
}
}
impl Default for Frame {
fn default() -> Frame {
Frame {
finished: true,
opcode: OpCode::Close,
payload: Binary::from(&b""[..]),
}
}
}
impl fmt::Display for Frame {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"
<FRAME>
final: {}
opcode: {}
payload length: {}
payload: 0x{}
</FRAME>",
self.finished,
self.opcode,
self.payload.len(),
self.payload
.as_ref()
.iter()
.map(|byte| format!("{:x}", byte))
.collect::<String>()
)
}
}
#[derive(Debug)]
pub struct FramedMessage(pub(crate) Binary);
#[cfg(test)]
mod tests {
use super::*;
use futures::stream::once;
fn is_none(frm: &Poll<Option<Frame>, ProtocolError>) -> bool {
match *frm {
Ok(Async::Ready(None)) => true,
_ => false,
}
}
fn extract(frm: Poll<Option<Frame>, ProtocolError>) -> Frame {
match frm {
Ok(Async::Ready(Some(frame))) => frame,
_ => unreachable!("error"),
}
}
#[test]
fn test_parse() {
let mut buf = PayloadBuffer::new(once(Ok(BytesMut::from(
&[0b0000_0001u8, 0b0000_0001u8][..],
).freeze())));
assert!(is_none(&Frame::parse(&mut buf, false, 1024)));
let mut buf = BytesMut::from(&[0b0000_0001u8, 0b0000_0001u8][..]);
buf.extend(b"1");
let mut buf = PayloadBuffer::new(once(Ok(buf.freeze())));
let frame = extract(Frame::parse(&mut buf, false, 1024));
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert_eq!(frame.payload.as_ref(), &b"1"[..]);
}
#[test]
fn test_parse_length0() {
let buf = BytesMut::from(&[0b0000_0001u8, 0b0000_0000u8][..]);
let mut buf = PayloadBuffer::new(once(Ok(buf.freeze())));
let frame = extract(Frame::parse(&mut buf, false, 1024));
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert!(frame.payload.is_empty());
}
#[test]
fn test_parse_length2() {
let buf = BytesMut::from(&[0b0000_0001u8, 126u8][..]);
let mut buf = PayloadBuffer::new(once(Ok(buf.freeze())));
assert!(is_none(&Frame::parse(&mut buf, false, 1024)));
let mut buf = BytesMut::from(&[0b0000_0001u8, 126u8][..]);
buf.extend(&[0u8, 4u8][..]);
buf.extend(b"1234");
let mut buf = PayloadBuffer::new(once(Ok(buf.freeze())));
let frame = extract(Frame::parse(&mut buf, false, 1024));
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert_eq!(frame.payload.as_ref(), &b"1234"[..]);
}
#[test]
fn test_parse_length4() {
let buf = BytesMut::from(&[0b0000_0001u8, 127u8][..]);
let mut buf = PayloadBuffer::new(once(Ok(buf.freeze())));
assert!(is_none(&Frame::parse(&mut buf, false, 1024)));
let mut buf = BytesMut::from(&[0b0000_0001u8, 127u8][..]);
buf.extend(&[0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 4u8][..]);
buf.extend(b"1234");
let mut buf = PayloadBuffer::new(once(Ok(buf.freeze())));
let frame = extract(Frame::parse(&mut buf, false, 1024));
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert_eq!(frame.payload.as_ref(), &b"1234"[..]);
}
#[test]
fn test_parse_frame_mask() {
let mut buf = BytesMut::from(&[0b0000_0001u8, 0b1000_0001u8][..]);
buf.extend(b"0001");
buf.extend(b"1");
let mut buf = PayloadBuffer::new(once(Ok(buf.freeze())));
assert!(Frame::parse(&mut buf, false, 1024).is_err());
let frame = extract(Frame::parse(&mut buf, true, 1024));
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert_eq!(frame.payload, vec![1u8].into());
}
#[test]
fn test_parse_frame_no_mask() {
let mut buf = BytesMut::from(&[0b0000_0001u8, 0b0000_0001u8][..]);
buf.extend(&[1u8]);
let mut buf = PayloadBuffer::new(once(Ok(buf.freeze())));
assert!(Frame::parse(&mut buf, true, 1024).is_err());
let frame = extract(Frame::parse(&mut buf, false, 1024));
assert!(!frame.finished);
assert_eq!(frame.opcode, OpCode::Text);
assert_eq!(frame.payload, vec![1u8].into());
}
#[test]
fn test_parse_frame_max_size() {
let mut buf = BytesMut::from(&[0b0000_0001u8, 0b0000_0010u8][..]);
buf.extend(&[1u8, 1u8]);
let mut buf = PayloadBuffer::new(once(Ok(buf.freeze())));
assert!(Frame::parse(&mut buf, true, 1).is_err());
if let Err(ProtocolError::Overflow) = Frame::parse(&mut buf, false, 0) {
} else {
unreachable!("error");
}
}
#[test]
fn test_ping_frame() {
let frame = Frame::message(Vec::from("data"), OpCode::Ping, true, false);
let mut v = vec![137u8, 4u8];
v.extend(b"data");
assert_eq!(frame.0, v.into());
}
#[test]
fn test_pong_frame() {
let frame = Frame::message(Vec::from("data"), OpCode::Pong, true, false);
let mut v = vec![138u8, 4u8];
v.extend(b"data");
assert_eq!(frame.0, v.into());
}
#[test]
fn test_close_frame() {
let reason = (CloseCode::Normal, "data");
let frame = Frame::close(Some(reason.into()), false);
let mut v = vec![136u8, 6u8, 3u8, 232u8];
v.extend(b"data");
assert_eq!(frame.0, v.into());
}
#[test]
fn test_empty_close_frame() {
let frame = Frame::close(None, false);
assert_eq!(frame.0, vec![0x88, 0x00].into());
}
}