actix/actix-web
1
0
Fork 0
mirror of https://github.com/fafhrd91/actix-web synced 2025-08-20 04:35:38 +02:00
Code Issues Releases Wiki Activity

optimize ws frame generation

Browse Source
This commit is contained in:
Nikolay Kim
2018-02-09 20:43:14 -08:00
parent 78da98a16d
commit b4b5c78b51
4 changed files with 179 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files

232
src/ws/frame.rs
View File

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