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; /// A struct representing a `WebSocket` frame. #[derive(Debug)] pub struct Frame { finished: bool, opcode: OpCode, payload: Binary, } impl Frame { /// Destruct frame pub fn unpack(self) -> (bool, OpCode, Binary) { (self.finished, self.opcode, self.payload) } /// Create a new Close control frame. #[inline] pub fn close(reason: Option, 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(clippy::type_complexity))] fn read_copy_md( pl: &mut PayloadBuffer, server: bool, max_size: usize, ) -> Poll)>, ProtocolError> where S: Stream, { 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; // check masking let masked = second & 0x80 != 0; if !masked && server { return Err(ProtocolError::UnmaskedFrame); } else if masked && !server { return Err(ProtocolError::MaskedFrame); } // Op code 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 }; // check for max allowed size 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), 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; // check masking let masked = second & 0x80 != 0; if !masked && server { return Err(ProtocolError::UnmaskedFrame); } else if masked && !server { return Err(ProtocolError::MaskedFrame); } // Op code 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 }; // check for max allowed size 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))) } /// Parse the input stream into a frame. pub fn parse( pl: &mut PayloadBuffer, server: bool, max_size: usize, ) -> Poll, ProtocolError> where S: Stream, { // try to parse ws frame md from one chunk 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 { // we may need to join several chunks 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), } // remove prefix pl.drop_bytes(idx); // no need for body 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!(), }; // control frames must have length <= 125 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()))); } _ => (), } // unmask 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(), }))) } /// Parse the payload of a close frame. pub fn parse_close_payload(payload: &Binary) -> Option { 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 } } /// Generate binary representation pub fn message>( data: B, code: OpCode, finished: bool, genmask: bool, ) -> FramedMessage { let payload = data.into(); let one: u8 = if finished { 0x80 | Into::::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::(); 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, " final: {} opcode: {} payload length: {} payload: 0x{} ", self.finished, self.opcode, self.payload.len(), self.payload .as_ref() .iter() .map(|byte| format!("{:x}", byte)) .collect::() ) } } /// `WebSocket` message with framing. #[derive(Debug)] pub struct FramedMessage(pub(crate) Binary); #[cfg(test)] mod tests { use super::*; use futures::stream::once; fn is_none(frm: &Poll, ProtocolError>) -> bool { match *frm { Ok(Async::Ready(None)) => true, _ => false, } } fn extract(frm: Poll, 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()); } }