optimize ws frame generation

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<Item=(WsClientReader, WsClientWriter), Error=WsClientError>;
@@ -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<B: Into<Binary>>(&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);
+        self.write(Frame::message(Vec::from(text), OpCode::Text, true).generate(true));
-        let mut buf = Vec::new();
-        frame.format(&mut buf).unwrap();
-
-        self.write(buf);
     }
 
     /// Send binary frame
+    #[inline]
     pub fn binary<B: Into<Binary>>(&mut self, data: B) {
-        let mut frame = Frame::message(data, OpCode::Binary, true);
+        self.write(Frame::message(data, OpCode::Binary, true).generate(true));
-        let mut buf = Vec::new();
-        frame.format(&mut buf).unwrap();
-
-        self.write(buf);
     }
 
     /// Send ping frame
+    #[inline]
     pub fn ping(&mut self, message: &str) {
-        let mut frame = Frame::message(Vec::from(message), OpCode::Ping, true);
+        self.write(Frame::message(Vec::from(message), OpCode::Ping, true).generate(true));
-        let mut buf = Vec::new();
-        frame.format(&mut buf).unwrap();
-
-        self.write(buf);
     }
 
     /// Send pong frame
+    #[inline]
     pub fn pong(&mut self, message: &str) {
-        let mut frame = Frame::message(Vec::from(message), OpCode::Pong, true);
+        self.write(Frame::message(Vec::from(message), OpCode::Pong, true).generate(true));
-        let mut buf = Vec::new();
-        frame.format(&mut buf).unwrap();
-
-        self.write(buf);
     }
 
     /// Send close frame
+    #[inline]
     pub fn close(&mut self, code: CloseCode, reason: &str) {
-        let mut frame = Frame::close(code, reason);
+        self.write(Frame::close(code, reason).generate(true));
-        let mut buf = Vec::new();
-        frame.format(&mut buf).unwrap();
-        self.write(buf);
     }
 }

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<A, S> WebsocketContext<A, S> where A: Actor<Context=Self> {
 
     /// Write payload
     #[inline]
-    fn write<B: Into<Binary>>(&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<A, S> WebsocketContext<A, S> where A: Actor<Context=Self> {
     }
 
     /// Send text frame
+    #[inline]
     pub fn text(&mut self, text: &str) {
-        let mut frame = Frame::message(Vec::from(text), OpCode::Text, true);
+        self.write(Frame::message(Vec::from(text), OpCode::Text, true).generate(false));
-        let mut buf = Vec::new();
-        frame.format(&mut buf).unwrap();
-
-        self.write(buf);
     }
 
     /// Send binary frame
+    #[inline]
     pub fn binary<B: Into<Binary>>(&mut self, data: B) {
-        let mut frame = Frame::message(data, OpCode::Binary, true);
+        self.write(Frame::message(data, OpCode::Binary, true).generate(false));
-        let mut buf = Vec::new();
-        frame.format(&mut buf).unwrap();
-
-        self.write(buf);
     }
 
     /// Send ping frame
+    #[inline]
     pub fn ping(&mut self, message: &str) {
-        let mut frame = Frame::message(Vec::from(message), OpCode::Ping, true);
+        self.write(Frame::message(Vec::from(message), OpCode::Ping, true).generate(false));
-        let mut buf = Vec::new();
-        frame.format(&mut buf).unwrap();
-
-        self.write(buf);
     }
 
     /// Send pong frame
+    #[inline]
     pub fn pong(&mut self, message: &str) {
-        let mut frame = Frame::message(Vec::from(message), OpCode::Pong, true);
+        self.write(Frame::message(Vec::from(message), OpCode::Pong, true).generate(false));
-        let mut buf = Vec::new();
-        frame.format(&mut buf).unwrap();
-
-        self.write(buf);
     }
 
     /// Send close frame
+    #[inline]
     pub fn close(&mut self, code: CloseCode, reason: &str) {
-        let mut frame = Frame::close(code, reason);
+        self.write(Frame::close(code, reason).generate(false));
-        let mut buf = Vec::new();
-        frame.format(&mut buf).unwrap();
-        self.write(buf);
     }
 
     /// Returns drain future
@@ -183,6 +181,7 @@ impl<A, S> WebsocketContext<A, S> where A: Actor<Context=Self> {
         !self.disconnected
     }
 
+    #[inline]
     fn add_frame(&mut self, frame: ContextFrame) {
         if self.stream.is_none() {
             self.stream = Some(SmallVec::new());

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 bytes::{BytesMut, BufMut};
-use byteorder::{ByteOrder, BigEndian};
+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<B: Into<Binary>>(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<Option<Frame>, Error> {
+    pub fn parse(buf: &mut BytesMut, server: bool) -> Result<Option<Frame>, 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 mask = if server {
-            let mut mask_bytes = [0u8; 4];
             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
+    /// Generate binary representation
-    pub fn format<W: Write>(&mut self, w: &mut W) -> Result<(), Error> {
+    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;
+        let (two, mask_size) = if genmask {
-
+            (0x80, 4)
-        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)?;
         } else {
-            two |= 127;
+            (0, 0)
-            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)?;
-        }
 
-        if self.mask.is_some() {
+        let payload_len = self.payload.len();
-            let mask = self.mask.take().unwrap();
+        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)?;
+            buf.extend_from_slice(&mask);
-            w.write_all(payload.as_ref())?;
+            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{}
 </FRAME>",
@@ -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::<String>())
@@ -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 frame = Frame::message(Vec::from("data"), OpCode::Ping, true);
-        let mut buf = Vec::new();
+        let res = frame.generate(false);
-        frame.format(&mut buf).unwrap();
 
         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 frame = Frame::message(Vec::from("data"), OpCode::Pong, true);
-        let mut buf = Vec::new();
+        let res = frame.generate(false);
-        frame.format(&mut buf).unwrap();
 
         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 frame = Frame::close(CloseCode::Normal, "data");
-        let mut buf = Vec::new();
+        let res = frame.generate(false);
-        frame.format(&mut buf).unwrap();
 
         let mut v = vec![136u8, 6u8, 3u8, 232u8];
         v.extend(b"data");
-        assert_eq!(buf, v);
+        assert_eq!(res, FrameData::Complete(v.into()));
     }
 }

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();