diff --git a/src/ws/client.rs b/src/ws/client.rs index bc857bd8..565e7580 100644 --- a/src/ws/client.rs +++ b/src/ws/client.rs @@ -29,7 +29,7 @@ use client::{Connect, Connection, ClientConnector, ClientConnectorError}; use super::Message; use super::proto::{CloseCode, OpCode}; -use super::frame::Frame; +use super::frame::{Frame, FrameData}; pub type WsClientFuture = Future; @@ -371,7 +371,7 @@ impl Stream for WsClientReader { let _ = inner.writer.poll_completed(&mut inner.conn, false); // read - match Frame::parse(&mut inner.parser_buf) { + match Frame::parse(&mut inner.parser_buf, false) { Ok(Some(frame)) => { // trace!("WsFrame {}", frame); let (_finished, opcode, payload) = frame.unpack(); @@ -444,55 +444,49 @@ impl WsClientWriter { /// Write payload #[inline] - fn write>(&mut self, data: B) { + fn write(&mut self, data: FrameData) { if !self.as_mut().closed { - let _ = self.as_mut().writer.write(&data.into()); + match data { + FrameData::Complete(data) => { + let _ = self.as_mut().writer.write(&data); + }, + FrameData::Split(headers, payload) => { + let _ = self.as_mut().writer.write(&headers); + let _ = self.as_mut().writer.write(&payload); + } + } } else { warn!("Trying to write to disconnected response"); } } /// Send text frame + #[inline] pub fn text(&mut self, text: &str) { - let mut frame = Frame::message(Vec::from(text), OpCode::Text, true); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); - - self.write(buf); + self.write(Frame::message(Vec::from(text), OpCode::Text, true).generate(true)); } /// Send binary frame + #[inline] pub fn binary>(&mut self, data: B) { - let mut frame = Frame::message(data, OpCode::Binary, true); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); - - self.write(buf); + self.write(Frame::message(data, OpCode::Binary, true).generate(true)); } /// Send ping frame + #[inline] pub fn ping(&mut self, message: &str) { - let mut frame = Frame::message(Vec::from(message), OpCode::Ping, true); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); - - self.write(buf); + self.write(Frame::message(Vec::from(message), OpCode::Ping, true).generate(true)); } /// Send pong frame + #[inline] pub fn pong(&mut self, message: &str) { - let mut frame = Frame::message(Vec::from(message), OpCode::Pong, true); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); - - self.write(buf); + self.write(Frame::message(Vec::from(message), OpCode::Pong, true).generate(true)); } /// Send close frame + #[inline] pub fn close(&mut self, code: CloseCode, reason: &str) { - let mut frame = Frame::close(code, reason); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); - self.write(buf); + self.write(Frame::close(code, reason).generate(true)); } } diff --git a/src/ws/context.rs b/src/ws/context.rs index 55b4e67e..c74410aa 100644 --- a/src/ws/context.rs +++ b/src/ws/context.rs @@ -14,7 +14,7 @@ use error::{Error, ErrorInternalServerError}; use httprequest::HttpRequest; use context::{Frame as ContextFrame, ActorHttpContext, Drain}; -use ws::frame::Frame; +use ws::frame::{Frame, FrameData}; use ws::proto::{OpCode, CloseCode}; @@ -105,9 +105,21 @@ impl WebsocketContext where A: Actor { /// Write payload #[inline] - fn write>(&mut self, data: B) { + fn write(&mut self, data: FrameData) { if !self.disconnected { - self.add_frame(ContextFrame::Chunk(Some(data.into()))); + if self.stream.is_none() { + self.stream = Some(SmallVec::new()); + } + let stream = self.stream.as_mut().unwrap(); + + match data { + FrameData::Complete(data) => + stream.push(ContextFrame::Chunk(Some(data))), + FrameData::Split(headers, payload) => { + stream.push(ContextFrame::Chunk(Some(headers))); + stream.push(ContextFrame::Chunk(Some(payload))); + } + } } else { warn!("Trying to write to disconnected response"); } @@ -126,47 +138,33 @@ impl WebsocketContext where A: Actor { } /// Send text frame + #[inline] pub fn text(&mut self, text: &str) { - let mut frame = Frame::message(Vec::from(text), OpCode::Text, true); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); - - self.write(buf); + self.write(Frame::message(Vec::from(text), OpCode::Text, true).generate(false)); } /// Send binary frame + #[inline] pub fn binary>(&mut self, data: B) { - let mut frame = Frame::message(data, OpCode::Binary, true); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); - - self.write(buf); + self.write(Frame::message(data, OpCode::Binary, true).generate(false)); } /// Send ping frame + #[inline] pub fn ping(&mut self, message: &str) { - let mut frame = Frame::message(Vec::from(message), OpCode::Ping, true); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); - - self.write(buf); + self.write(Frame::message(Vec::from(message), OpCode::Ping, true).generate(false)); } /// Send pong frame + #[inline] pub fn pong(&mut self, message: &str) { - let mut frame = Frame::message(Vec::from(message), OpCode::Pong, true); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); - - self.write(buf); + self.write(Frame::message(Vec::from(message), OpCode::Pong, true).generate(false)); } /// Send close frame + #[inline] pub fn close(&mut self, code: CloseCode, reason: &str) { - let mut frame = Frame::close(code, reason); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); - self.write(buf); + self.write(Frame::close(code, reason).generate(false)); } /// Returns drain future @@ -183,6 +181,7 @@ impl WebsocketContext where A: Actor { !self.disconnected } + #[inline] fn add_frame(&mut self, frame: ContextFrame) { if self.stream.is_none() { self.stream = Some(SmallVec::new()); diff --git a/src/ws/frame.rs b/src/ws/frame.rs index d149f37a..e6e0f535 100644 --- a/src/ws/frame.rs +++ b/src/ws/frame.rs @@ -1,13 +1,22 @@ use std::{fmt, mem}; -use std::io::{Write, Error, ErrorKind}; +use std::io::{Error, ErrorKind}; use std::iter::FromIterator; -use bytes::BytesMut; -use byteorder::{ByteOrder, BigEndian}; +use bytes::{BytesMut, BufMut}; +use byteorder::{ByteOrder, BigEndian, NetworkEndian}; +use rand; use body::Binary; use ws::proto::{OpCode, CloseCode}; use ws::mask::apply_mask; +#[derive(Debug, PartialEq)] +pub(crate) enum FrameData { + Complete(Binary), + Split(Binary, Binary), +} + +const MAX_LEN: usize = 122; + /// A struct representing a `WebSocket` frame. #[derive(Debug)] pub(crate) struct Frame { @@ -16,7 +25,6 @@ pub(crate) struct Frame { rsv2: bool, rsv3: bool, opcode: OpCode, - mask: Option<[u8; 4]>, payload: Binary, } @@ -27,26 +35,6 @@ impl Frame { (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; - let payload_len = self.payload.len(); - if payload_len > 125 { - if payload_len <= u16::max_value() as usize { - header_length += 2; - } else { - header_length += 8; - } - } - if self.mask.is_some() { - header_length += 4; - } - - header_length + payload_len - } - /// Create a new data frame. #[inline] pub fn message>(data: B, code: OpCode, finished: bool) -> Frame { @@ -82,7 +70,7 @@ impl Frame { } /// Parse the input stream into a frame. - pub fn parse(buf: &mut BytesMut) -> Result, Error> { + pub fn parse(buf: &mut BytesMut, server: bool) -> Result, Error> { let mut idx = 2; let mut size = buf.len(); @@ -94,18 +82,27 @@ impl Frame { let second = buf[1]; let finished = first & 0x80 != 0; + // check masking + let masked = second & 0x80 != 0; + if !masked && server { + return Err(Error::new( + ErrorKind::Other, "Received an unmasked frame from client")) + } else if masked && !server { + return Err(Error::new( + ErrorKind::Other, "Received a masked frame from server")) + } + let rsv1 = first & 0x40 != 0; let rsv2 = first & 0x20 != 0; let rsv3 = first & 0x10 != 0; let opcode = OpCode::from(first & 0x0F); - let masked = second & 0x80 != 0; let len = second & 0x7F; let length = if len == 126 { if size < 2 { return Ok(None) } - let len = u64::from(BigEndian::read_u16(&buf[idx..])); + let len = NetworkEndian::read_uint(&buf[idx..], 2) as usize; size -= 2; idx += 2; len @@ -113,19 +110,19 @@ impl Frame { if size < 8 { return Ok(None) } - let len = BigEndian::read_u64(&buf[idx..]); + let len = NetworkEndian::read_uint(&buf[idx..], 8) as usize; size -= 8; idx += 8; len } else { - u64::from(len) + len as usize }; - let mask = if masked { - let mut mask_bytes = [0u8; 4]; + let mask = if server { if size < 4 { return Ok(None) } else { + let mut mask_bytes = [0u8; 4]; size -= 4; mask_bytes.copy_from_slice(&buf[idx..idx+4]); idx += 4; @@ -135,7 +132,6 @@ impl Frame { None }; - let length = length as usize; if size < length { return Ok(None) } @@ -182,13 +178,12 @@ impl Frame { rsv2: rsv2, rsv3: rsv3, opcode: opcode, - mask: mask, payload: data.into(), })) } - /// Write a frame out to a buffer - pub fn format(&mut self, w: &mut W) -> Result<(), Error> { + /// Generate binary representation + pub fn generate(self, genmask: bool) -> FrameData { let mut one = 0u8; let code: u8 = self.opcode.into(); if self.finished { @@ -205,55 +200,80 @@ impl Frame { } one |= code; - let mut two = 0u8; - - if self.mask.is_some() { - two |= 0x80; - } - - if self.payload.len() < 126 { - two |= self.payload.len() as u8; - let headers = [one, two]; - 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; - mem::transmute(short.to_be()) - }; - let headers = [one, two, length_bytes[0], length_bytes[1]]; - w.write_all(&headers)?; + let (two, mask_size) = if genmask { + (0x80, 4) } else { - two |= 127; - let length_bytes: [u8; 8] = unsafe { - let long = self.payload.len() as u64; - mem::transmute(long.to_be()) - }; - 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], - ]; - w.write_all(&headers)?; - } + (0, 0) + }; - if self.mask.is_some() { - let mask = self.mask.take().unwrap(); + let payload_len = self.payload.len(); + let mut buf = if payload_len < MAX_LEN { + if genmask { + let len = payload_len + 6; + let mask: [u8; 4] = rand::random(); + let mut buf = BytesMut::with_capacity(len); + { + let buf_mut = unsafe{buf.bytes_mut()}; + buf_mut[0] = one; + buf_mut[1] = two | payload_len as u8; + buf_mut[2..6].copy_from_slice(&mask); + buf_mut[6..payload_len+6].copy_from_slice(self.payload.as_ref()); + apply_mask(&mut buf_mut[6..], &mask); + } + unsafe{buf.advance_mut(len)}; + return FrameData::Complete(buf.into()) + } else { + let len = payload_len + 2; + let mut buf = BytesMut::with_capacity(len); + { + let buf_mut = unsafe{buf.bytes_mut()}; + buf_mut[0] = one; + buf_mut[1] = two | payload_len as u8; + buf_mut[2..payload_len+2].copy_from_slice(self.payload.as_ref()); + } + unsafe{buf.advance_mut(len)}; + return FrameData::Complete(buf.into()) + } + } else if payload_len < 126 { + let mut buf = BytesMut::with_capacity(mask_size + 2); + { + let buf_mut = unsafe{buf.bytes_mut()}; + buf_mut[0] = one; + buf_mut[1] = two | payload_len as u8; + } + unsafe{buf.advance_mut(2)}; + buf + } else if payload_len <= 65_535 { + let mut buf = BytesMut::with_capacity(mask_size + 4); + { + let buf_mut = unsafe{buf.bytes_mut()}; + buf_mut[0] = one; + buf_mut[1] = two | 126; + BigEndian::write_u16(&mut buf_mut[2..4], payload_len as u16); + } + unsafe{buf.advance_mut(4)}; + buf + } else { + let mut buf = BytesMut::with_capacity(mask_size + 10); + { + let buf_mut = unsafe{buf.bytes_mut()}; + buf_mut[0] = one; + buf_mut[1] = two | 127; + BigEndian::write_u64(&mut buf_mut[2..10], payload_len as u64); + } + unsafe{buf.advance_mut(10)}; + buf + }; + + if genmask { let mut payload = Vec::from(self.payload.as_ref()); + let mask: [u8; 4] = rand::random(); apply_mask(&mut payload, &mask); - w.write_all(&mask)?; - w.write_all(payload.as_ref())?; + buf.extend_from_slice(&mask); + FrameData::Split(buf.into(), payload.into()) } else { - w.write_all(self.payload.as_ref())?; + FrameData::Split(buf.into(), self.payload) } - Ok(()) } } @@ -265,7 +285,6 @@ impl Default for Frame { rsv2: false, rsv3: false, opcode: OpCode::Close, - mask: None, payload: Binary::from(&b""[..]), } } @@ -279,7 +298,6 @@ impl fmt::Display for Frame { final: {} reserved: {} {} {} opcode: {} - length: {} payload length: {} payload: 0x{} ", @@ -288,8 +306,6 @@ impl fmt::Display for Frame { self.rsv2, self.rsv3, self.opcode, - // self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()), - self.len(), self.payload.len(), self.payload.as_ref().iter().map( |byte| format!("{:x}", byte)).collect::()) @@ -303,9 +319,9 @@ mod tests { #[test] fn test_parse() { let mut buf = BytesMut::from(&[0b00000001u8, 0b00000001u8][..]); - assert!(Frame::parse(&mut buf).unwrap().is_none()); + assert!(Frame::parse(&mut buf, false).unwrap().is_none()); buf.extend(b"1"); - let frame = Frame::parse(&mut buf).unwrap().unwrap(); + let frame = Frame::parse(&mut buf, false).unwrap().unwrap(); println!("FRAME: {}", frame); assert!(!frame.finished); assert_eq!(frame.opcode, OpCode::Text); @@ -315,7 +331,7 @@ mod tests { #[test] fn test_parse_length0() { let mut buf = BytesMut::from(&[0b00000001u8, 0b00000000u8][..]); - let frame = Frame::parse(&mut buf).unwrap().unwrap(); + let frame = Frame::parse(&mut buf, false).unwrap().unwrap(); assert!(!frame.finished); assert_eq!(frame.opcode, OpCode::Text); assert!(frame.payload.is_empty()); @@ -324,11 +340,11 @@ mod tests { #[test] fn test_parse_length2() { let mut buf = BytesMut::from(&[0b00000001u8, 126u8][..]); - assert!(Frame::parse(&mut buf).unwrap().is_none()); + assert!(Frame::parse(&mut buf, false).unwrap().is_none()); buf.extend(&[0u8, 4u8][..]); buf.extend(b"1234"); - let frame = Frame::parse(&mut buf).unwrap().unwrap(); + let frame = Frame::parse(&mut buf, false).unwrap().unwrap(); assert!(!frame.finished); assert_eq!(frame.opcode, OpCode::Text); assert_eq!(frame.payload.as_ref(), &b"1234"[..]); @@ -337,11 +353,11 @@ mod tests { #[test] fn test_parse_length4() { let mut buf = BytesMut::from(&[0b00000001u8, 127u8][..]); - assert!(Frame::parse(&mut buf).unwrap().is_none()); + assert!(Frame::parse(&mut buf, false).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(); + let frame = Frame::parse(&mut buf, false).unwrap().unwrap(); assert!(!frame.finished); assert_eq!(frame.opcode, OpCode::Text); assert_eq!(frame.payload.as_ref(), &b"1234"[..]); @@ -353,7 +369,22 @@ mod tests { buf.extend(b"0001"); buf.extend(b"1"); - let frame = Frame::parse(&mut buf).unwrap().unwrap(); + assert!(Frame::parse(&mut buf, false).is_err()); + + let frame = Frame::parse(&mut buf, true).unwrap().unwrap(); + 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(&[0b00000001u8, 0b00000001u8][..]); + buf.extend(&[1u8]); + + assert!(Frame::parse(&mut buf, true).is_err()); + + let frame = Frame::parse(&mut buf, false).unwrap().unwrap(); assert!(!frame.finished); assert_eq!(frame.opcode, OpCode::Text); assert_eq!(frame.payload, vec![1u8].into()); @@ -361,34 +392,31 @@ mod tests { #[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 frame = Frame::message(Vec::from("data"), OpCode::Ping, true); + let res = frame.generate(false); let mut v = vec![137u8, 4u8]; v.extend(b"data"); - assert_eq!(buf, v); + assert_eq!(res, FrameData::Complete(v.into())); } #[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 frame = Frame::message(Vec::from("data"), OpCode::Pong, true); + let res = frame.generate(false); let mut v = vec![138u8, 4u8]; v.extend(b"data"); - assert_eq!(buf, v); + assert_eq!(res, FrameData::Complete(v.into())); } #[test] fn test_close_frame() { - let mut frame = Frame::close(CloseCode::Normal, "data"); - let mut buf = Vec::new(); - frame.format(&mut buf).unwrap(); + let frame = Frame::close(CloseCode::Normal, "data"); + let res = frame.generate(false); let mut v = vec![136u8, 6u8, 3u8, 232u8]; v.extend(b"data"); - assert_eq!(buf, v); + assert_eq!(res, FrameData::Complete(v.into())); } } diff --git a/src/ws/mod.rs b/src/ws/mod.rs index caecefc6..17501a7d 100644 --- a/src/ws/mod.rs +++ b/src/ws/mod.rs @@ -214,7 +214,7 @@ impl Stream for WsStream { } loop { - match Frame::parse(&mut self.buf) { + match Frame::parse(&mut self.buf, true) { Ok(Some(frame)) => { // trace!("WsFrame {}", frame); let (_finished, opcode, payload) = frame.unpack();