From b4b5c78b510257c41912820fa39449e48f84065b Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Fri, 9 Feb 2018 20:43:14 -0800
Subject: [PATCH] optimize ws frame generation
---
src/ws/client.rs | 50 +++++-----
src/ws/context.rs | 53 ++++++-----
src/ws/frame.rs | 232 ++++++++++++++++++++++++++--------------------
src/ws/mod.rs | 2 +-
4 files changed, 179 insertions(+), 158 deletions(-)
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<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);
- 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<B: Into<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<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);
- 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<B: Into<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<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());
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<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 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<W: Write>(&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{}
</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 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();