actix/actix-web
1
0
Fork 0
mirror of https://github.com/fafhrd91/actix-web synced 2025-06-26 15:07:42 +02:00
Code Issues Releases Wiki Activity

add wsload tool; optimize ws frame parser

Browse Source
This commit is contained in:
Nikolay Kim
2018-02-09 17:20:28 -08:00
parent 74377ef73d
commit 78da98a16d
8 changed files with 386 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
src/client/writer.rs
View File

@ -4,16 +4,17 @@ use std::fmt::Write;
use bytes::BufMut;
use futures::{Async, Poll};
use tokio_io::AsyncWrite;
// use http::header::{HeaderValue, CONNECTION, DATE};
use body::Binary;
use server::{WriterState, MAX_WRITE_BUFFER_SIZE};
use server::WriterState;
use server::shared::SharedBytes;
use client::ClientRequest;
const AVERAGE_HEADER_SIZE: usize = 30; // totally scientific
const LOW_WATERMARK: usize = 1024;
const HIGH_WATERMARK: usize = 8 * LOW_WATERMARK;
const AVERAGE_HEADER_SIZE: usize = 30;
bitflags! {
struct Flags: u8 {
@ -29,6 +30,8 @@ pub(crate) struct HttpClientWriter {
written: u64,
headers_size: u32,
buffer: SharedBytes,
low: usize,
high: usize,
}
impl HttpClientWriter {
@ -39,6 +42,8 @@ impl HttpClientWriter {
written: 0,
headers_size: 0,
buffer: buf,
low: LOW_WATERMARK,
high: HIGH_WATERMARK,
}
}
@ -50,6 +55,12 @@ impl HttpClientWriter {
self.flags.contains(Flags::KEEPALIVE) && !self.flags.contains(Flags::UPGRADE)
}
/// Set write buffer capacity
pub fn set_buffer_capacity(&mut self, low_watermark: usize, high_watermark: usize) {
self.low = low_watermark;
self.high = high_watermark;
}
fn write_to_stream<T: AsyncWrite>(&mut self, stream: &mut T) -> io::Result<WriterState> {
while !self.buffer.is_empty() {
match stream.write(self.buffer.as_ref()) {
@ -61,7 +72,7 @@ impl HttpClientWriter {
let _ = self.buffer.split_to(n);
},
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
if self.buffer.len() > MAX_WRITE_BUFFER_SIZE {
if self.buffer.len() > self.high {
return Ok(WriterState::Pause)
} else {
return Ok(WriterState::Done)
@ -117,7 +128,7 @@ impl HttpClientWriter {
self.buffer.extend_from_slice(payload.as_ref())
}
if self.buffer.len() > MAX_WRITE_BUFFER_SIZE {
if self.buffer.len() > self.high {
Ok(WriterState::Pause)
} else {
Ok(WriterState::Done)
@ -125,7 +136,7 @@ impl HttpClientWriter {
}
pub fn write_eof(&mut self) -> io::Result<WriterState> {
if self.buffer.len() > MAX_WRITE_BUFFER_SIZE {
if self.buffer.len() > self.high {
Ok(WriterState::Pause)
} else {
Ok(WriterState::Done)

1
src/lib.rs
View File

@ -51,6 +51,7 @@ extern crate log;
extern crate time;
extern crate base64;
extern crate bytes;
extern crate byteorder;
extern crate sha1;
extern crate regex;
#[macro_use]

44
src/ws/client.rs
View File

@ -92,9 +92,9 @@ impl From<HttpResponseParserError> for WsClientError {
}
}
/// WebSocket client
/// `WebSocket` client
///
/// Example of WebSocket client usage is available in
/// Example of `WebSocket` client usage is available in
/// [websocket example](
/// https://github.com/actix/actix-web/blob/master/examples/websocket/src/client.rs#L24)
pub struct WsClient {
@ -317,7 +317,7 @@ impl Future for WsHandshake {
return Err(WsClientError::InvalidChallengeResponse)
}
let inner = Rc::new(UnsafeCell::new(Inner{inner: inner}));
let inner = Rc::new(UnsafeCell::new(inner));
Ok(Async::Ready(
(WsClientReader{inner: Rc::clone(&inner)},
WsClientWriter{inner: inner})))
@ -332,12 +332,8 @@ impl Future for WsHandshake {
}
struct Inner {
inner: WsInner,
}
pub struct WsClientReader {
inner: Rc<UnsafeCell<Inner>>
inner: Rc<UnsafeCell<WsInner>>
}
impl fmt::Debug for WsClientReader {
@ -348,7 +344,7 @@ impl fmt::Debug for WsClientReader {
impl WsClientReader {
#[inline]
fn as_mut(&mut self) -> &mut Inner {
fn as_mut(&mut self) -> &mut WsInner {
unsafe{ &mut *self.inner.get() }
}
}
@ -361,10 +357,10 @@ impl Stream for WsClientReader {
let inner = self.as_mut();
let mut done = false;
match utils::read_from_io(&mut inner.inner.conn, &mut inner.inner.parser_buf) {
match utils::read_from_io(&mut inner.conn, &mut inner.parser_buf) {
Ok(Async::Ready(0)) => {
done = true;
inner.inner.closed = true;
inner.closed = true;
},
Ok(Async::Ready(_)) | Ok(Async::NotReady) => (),
Err(err) =>
@ -372,10 +368,10 @@ impl Stream for WsClientReader {
}
// write
let _ = inner.inner.writer.poll_completed(&mut inner.inner.conn, false);
let _ = inner.writer.poll_completed(&mut inner.conn, false);
// read
match Frame::parse(&mut inner.inner.parser_buf) {
match Frame::parse(&mut inner.parser_buf) {
Ok(Some(frame)) => {
// trace!("WsFrame {}", frame);
let (_finished, opcode, payload) = frame.unpack();
@ -385,8 +381,8 @@ impl Stream for WsClientReader {
OpCode::Bad =>
Ok(Async::Ready(Some(Message::Error))),
OpCode::Close => {
inner.inner.closed = true;
inner.inner.error_sent = true;
inner.closed = true;
inner.error_sent = true;
Ok(Async::Ready(Some(Message::Closed)))
},
OpCode::Ping =>
@ -413,9 +409,9 @@ impl Stream for WsClientReader {
Ok(None) => {
if done {
Ok(Async::Ready(None))
} else if inner.inner.closed {
if !inner.inner.error_sent {
inner.inner.error_sent = true;
} else if inner.closed {
if !inner.error_sent {
inner.error_sent = true;
Ok(Async::Ready(Some(Message::Closed)))
} else {
Ok(Async::Ready(None))
@ -425,8 +421,8 @@ impl Stream for WsClientReader {
}
},
Err(err) => {
inner.inner.closed = true;
inner.inner.error_sent = true;
inner.closed = true;
inner.error_sent = true;
Err(err.into())
}
}
@ -434,12 +430,12 @@ impl Stream for WsClientReader {
}
pub struct WsClientWriter {
inner: Rc<UnsafeCell<Inner>>
inner: Rc<UnsafeCell<WsInner>>
}
impl WsClientWriter {
#[inline]
fn as_mut(&mut self) -> &mut Inner {
fn as_mut(&mut self) -> &mut WsInner {
unsafe{ &mut *self.inner.get() }
}
}
@ -449,8 +445,8 @@ impl WsClientWriter {
/// Write payload
#[inline]
fn write<B: Into<Binary>>(&mut self, data: B) {
if !self.as_mut().inner.closed {
let _ = self.as_mut().inner.writer.write(&data.into());
if !self.as_mut().closed {
let _ = self.as_mut().writer.write(&data.into());
} else {
warn!("Trying to write to disconnected response");
}

7
src/ws/context.rs
View File

@ -211,11 +211,8 @@ impl<A, S> ActorHttpContext for WebsocketContext<A, S> where A: Actor<Context=Se
mem::transmute(self as &mut WebsocketContext<A, S>)
};
if self.inner.alive() {
match self.inner.poll(ctx) {
Ok(Async::NotReady) | Ok(Async::Ready(())) => (),
Err(_) => return Err(ErrorInternalServerError("error").into()),
}
if self.inner.alive() && self.inner.poll(ctx).is_err() {
return Err(ErrorInternalServerError("error").into())
}
// frames

197
src/ws/frame.rs
View File

@ -2,6 +2,7 @@ use std::{fmt, mem};
use std::io::{Write, Error, ErrorKind};
use std::iter::FromIterator;
use bytes::BytesMut;
use byteorder::{ByteOrder, BigEndian};
use body::Binary;
use ws::proto::{OpCode, CloseCode};
@ -39,7 +40,6 @@ impl Frame {
header_length += 8;
}
}
if self.mask.is_some() {
header_length += 4;
}
@ -84,136 +84,107 @@ impl Frame {
/// Parse the input stream into a frame.
pub fn parse(buf: &mut BytesMut) -> Result<Option<Frame>, Error> {
let mut idx = 2;
let mut size = buf.len();
let (frame, length) = {
let mut size = buf.len();
if size < 2 {
return Ok(None)
}
size -= 2;
let first = buf[0];
let second = buf[1];
let finished = first & 0x80 != 0;
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 mut head = [0u8; 2];
let len = u64::from(BigEndian::read_u16(&buf[idx..]));
size -= 2;
head.copy_from_slice(&buf[..2]);
trace!("Parsed headers {:?}", head);
let first = head[0];
let second = head[1];
trace!("First: {:b}", first);
trace!("Second: {:b}", second);
let finished = first & 0x80 != 0;
let rsv1 = first & 0x40 != 0;
let rsv2 = first & 0x20 != 0;
let rsv3 = first & 0x10 != 0;
let opcode = OpCode::from(first & 0x0F);
trace!("Opcode: {:?}", opcode);
let masked = second & 0x80 != 0;
trace!("Masked: {:?}", masked);
let mut header_length = 2;
let mut length = u64::from(second & 0x7F);
if length == 126 {
if size < 2 {
return Ok(None)
}
let mut length_bytes = [0u8; 2];
length_bytes.copy_from_slice(&buf[idx..idx+2]);
size -= 2;
idx += 2;
length = u64::from(unsafe{
let mut wide: u16 = mem::transmute(length_bytes);
wide = u16::from_be(wide);
wide});
header_length += 2;
} else if length == 127 {
if size < 8 {
return Ok(None)
}
let mut length_bytes = [0u8; 8];
length_bytes.copy_from_slice(&buf[idx..idx+8]);
size -= 8;
idx += 8;
unsafe { length = mem::transmute(length_bytes); }
length = u64::from_be(length);
header_length += 8;
}
trace!("Payload length: {}", length);
let mask = if masked {
let mut mask_bytes = [0u8; 4];
if size < 4 {
return Ok(None)
} else {
header_length += 4;
size -= 4;
mask_bytes.copy_from_slice(&buf[idx..idx+4]);
idx += 4;
Some(mask_bytes)
}
} else {
None
};
let length = length as usize;
if size < length {
idx += 2;
len
} else if len == 127 {
if size < 8 {
return Ok(None)
}
let len = BigEndian::read_u64(&buf[idx..]);
size -= 8;
idx += 8;
len
} else {
u64::from(len)
};
let mut data = Vec::with_capacity(length);
if length > 0 {
data.extend_from_slice(&buf[idx..idx+length]);
let mask = if masked {
let mut mask_bytes = [0u8; 4];
if size < 4 {
return Ok(None)
} else {
size -= 4;
mask_bytes.copy_from_slice(&buf[idx..idx+4]);
idx += 4;
Some(mask_bytes)
}
} else {
None
};
// Disallow bad opcode
if let OpCode::Bad = opcode {
let length = length as usize;
if size < length {
return Ok(None)
}
// get body
buf.split_to(idx);
let mut data = if length > 0 {
buf.split_to(length)
} else {
BytesMut::new()
};
// Disallow bad opcode
if let OpCode::Bad = opcode {
return Err(
Error::new(
ErrorKind::Other,
format!("Encountered invalid opcode: {}", first & 0x0F)))
}
// control frames must have length <= 125
match opcode {
OpCode::Ping | OpCode::Pong if length > 125 => {
return Err(
Error::new(
ErrorKind::Other,
format!("Encountered invalid opcode: {}", first & 0x0F)))
format!("Rejected WebSocket handshake.Received control frame with length: {}.", length)))
}
// control frames must have length <= 125
match opcode {
OpCode::Ping | OpCode::Pong if length > 125 => {
return Err(
Error::new(
ErrorKind::Other,
format!("Rejected WebSocket handshake.Received control frame with length: {}.", length)))
}
OpCode::Close if length > 125 => {
debug!("Received close frame with payload length exceeding 125. Morphing to protocol close frame.");
return Ok(Some(Frame::close(CloseCode::Protocol, "Received close frame with payload length exceeding 125.")))
}
_ => ()
OpCode::Close if length > 125 => {
debug!("Received close frame with payload length exceeding 125. Morphing to protocol close frame.");
return Ok(Some(Frame::close(CloseCode::Protocol, "Received close frame with payload length exceeding 125.")))
}
_ => ()
}
// unmask
if let Some(ref mask) = mask {
apply_mask(&mut data, mask);
}
// unmask
if let Some(ref mask) = mask {
apply_mask(&mut data, mask);
}
let frame = Frame {
finished: finished,
rsv1: rsv1,
rsv2: rsv2,
rsv3: rsv3,
opcode: opcode,
mask: mask,
payload: data.into(),
};
(frame, header_length + length)
};
buf.split_to(length);
Ok(Some(frame))
Ok(Some(Frame {
finished: finished,
rsv1: rsv1,
rsv2: rsv2,
rsv3: rsv3,
opcode: opcode,
mask: mask,
payload: data.into(),
}))
}
/// Write a frame out to a buffer