1
0
Fork 0
mirror of https://github.com/actix/actix-web.git synced 2024-06-02 13:29:24 +00:00

optimize ws frame generation

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

View file

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

View file

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

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

View file

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