refactor ws frame parser

Cargo.toml

@@ -77,7 +77,6 @@ tokio-tls = { version="0.1", optional = true }
 openssl = { version="0.10", optional = true }
 tokio-openssl = { version="0.2", optional = true }
 
-backtrace="*"
 [dependencies.actix]
 version = "0.5"
 

src/client/parser.rs

@@ -151,7 +151,9 @@ impl HttpResponseParser {
             }
         }
 
-        let decoder = if let Some(len) = hdrs.get(header::CONTENT_LENGTH) {
+        let decoder = if status == StatusCode::SWITCHING_PROTOCOLS {
+            Some(Decoder::eof())
+        } else if let Some(len) = hdrs.get(header::CONTENT_LENGTH) {
             // Content-Length
             if let Ok(s) = len.to_str() {
                 if let Ok(len) = s.parse::<u64>() {
@@ -167,8 +169,6 @@ impl HttpResponseParser {
         } else if chunked(&hdrs)? {
             // Chunked encoding
             Some(Decoder::chunked())
-        } else if hdrs.contains_key(header::UPGRADE) {
-            Some(Decoder::eof())
         } else {
             None
         };

src/error.rs

@@ -236,8 +236,6 @@ pub enum PayloadError {
 
 impl From<IoError> for PayloadError {
     fn from(err: IoError) -> PayloadError {
-        use backtrace;
-        println!("IO ERROR {:?}", backtrace::Backtrace::new());
         PayloadError::Io(err)
     }
 }
@@ -391,6 +389,34 @@ impl ResponseError for WsHandshakeError {
     }
 }
 
+/// Websocket errors
+#[derive(Fail, Debug)]
+pub enum WsError {
+    /// Received an unmasked frame from client
+    #[fail(display="Received an unmasked frame from client")]
+    UnmaskedFrame,
+    /// Received a masked frame from server
+    #[fail(display="Received a masked frame from server")]
+    MaskedFrame,
+    /// Encountered invalid opcode
+    #[fail(display="Invalid opcode: {}", _0)]
+    InvalidOpcode(u8),
+    /// Invalid control frame length
+    #[fail(display="Invalid control frame length: {}", _0)]
+    InvalidLength(usize),
+    /// Payload error
+    #[fail(display="Payload error: {}", _0)]
+    Payload(#[cause] PayloadError),
+}
+
+impl ResponseError for WsError {}
+
+impl From<PayloadError> for WsError {
+    fn from(err: PayloadError) -> WsError {
+        WsError::Payload(err)
+    }
+}
+
 /// A set of errors that can occur during parsing urlencoded payloads
 #[derive(Fail, Debug)]
 pub enum UrlencodedError {

src/handler.rs

@@ -395,7 +395,6 @@ impl<S> Handler<S> for NormalizePath {
                         }
                     }
                 } else if p.ends_with('/') {
-                    println!("=== {:?}", p);
                     // try to remove trailing slash
                     let p = p.as_ref().trim_right_matches('/');
                     if router.has_route(p) {

src/lib.rs

@@ -32,7 +32,7 @@
 //! * Supported *HTTP/1.x* and *HTTP/2.0* protocols
 //! * Streaming and pipelining
 //! * Keep-alive and slow requests handling
-//! * WebSockets server/client
+//! * *WebSockets* server/client
 //! * Transparent content compression/decompression (br, gzip, deflate)
 //! * Configurable request routing
 //! * Multipart streams
@@ -44,7 +44,7 @@
     specialization, // for impl ErrorResponse for std::error::Error
 ))]
 #![cfg_attr(feature = "cargo-clippy", allow(
-    decimal_literal_representation,))]
+    decimal_literal_representation,suspicious_arithmetic_impl,))]
 
 #[macro_use]
 extern crate log;
@@ -97,8 +97,6 @@ extern crate openssl;
 #[cfg(feature="openssl")]
 extern crate tokio_openssl;
 
-extern crate backtrace;
-
 mod application;
 mod body;
 mod context;

src/multipart.rs

@@ -492,10 +492,10 @@ impl<S> InnerField<S> where S: Stream<Item=Bytes, Error=PayloadError> {
                             if &chunk[..2] == b"\r\n" && &chunk[2..4] == b"--" &&
                                 &chunk[4..] == boundary.as_bytes()
                             {
-                                payload.unread_data(chunk);
+                                payload.unread_data(chunk.freeze());
                                 Ok(Async::Ready(None))
                             } else {
-                                Ok(Async::Ready(Some(chunk)))
+                                Ok(Async::Ready(Some(chunk.freeze())))
                             }
                         }
                     }

src/payload.rs

@@ -297,9 +297,9 @@ impl Inner {
                 buf.extend_from_slice(&chunk.split_to(rem));
                 if !chunk.is_empty() {
                     self.items.push_front(chunk);
-                    return Ok(Async::Ready(buf.freeze()))
                 }
             }
+            return Ok(Async::Ready(buf.freeze()))
         }
 
         if let Some(err) = self.err.take() {
@@ -423,7 +423,7 @@ impl<S> PayloadHelper<S> where S: Stream<Item=Bytes, Error=PayloadError> {
         PayloadHelper {
             len: 0,
             items: VecDeque::new(),
-            stream: stream,
+            stream,
         }
     }
 
@@ -445,6 +445,10 @@ impl<S> PayloadHelper<S> where S: Stream<Item=Bytes, Error=PayloadError> {
         self.len
     }
 
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
     pub fn readany(&mut self) -> Poll<Option<Bytes>, PayloadError> {
         if let Some(data) = self.items.pop_front() {
             self.len -= data.len();
@@ -458,7 +462,7 @@ impl<S> PayloadHelper<S> where S: Stream<Item=Bytes, Error=PayloadError> {
         }
     }
 
-    pub fn readexactly(&mut self, size: usize) -> Poll<Option<Bytes>, PayloadError> {
+    pub fn readexactly(&mut self, size: usize) -> Poll<Option<BytesMut>, PayloadError> {
         if size <= self.len {
             let mut buf = BytesMut::with_capacity(size);
             while buf.len() < size {
@@ -468,13 +472,34 @@ impl<S> PayloadHelper<S> where S: Stream<Item=Bytes, Error=PayloadError> {
                 buf.extend_from_slice(&chunk.split_to(rem));
                 if !chunk.is_empty() {
                     self.items.push_front(chunk);
-                    return Ok(Async::Ready(Some(buf.freeze())))
+                }
+            }
+            return Ok(Async::Ready(Some(buf)))
+        }
+
+        match self.poll_stream()? {
+            Async::Ready(true) => self.readexactly(size),
+            Async::Ready(false) => Ok(Async::Ready(None)),
+            Async::NotReady => Ok(Async::NotReady),
+        }
+    }
+
+    pub fn copy(&mut self, size: usize) -> Poll<Option<BytesMut>, PayloadError> {
+        if size <= self.len {
+            let mut buf = BytesMut::with_capacity(size);
+            for chunk in &self.items {
+                if buf.len() < size {
+                    let rem = cmp::min(size - buf.len(), chunk.len());
+                    buf.extend_from_slice(&chunk[..rem]);
+                }
+                if buf.len() == size {
+                    return Ok(Async::Ready(Some(buf)))
                 }
             }
         }
 
         match self.poll_stream()? {
-            Async::Ready(true) => self.readexactly(size),
+            Async::Ready(true) => self.copy(size),
             Async::Ready(false) => Ok(Async::Ready(None)),
             Async::NotReady => Ok(Async::NotReady),
         }

src/ws/client.rs

@@ -8,24 +8,28 @@ use std::cell::UnsafeCell;
 use base64;
 use rand;
 use cookie::Cookie;
-use bytes::BytesMut;
+use bytes::{Bytes, BytesMut};
 use http::{HttpTryFrom, StatusCode, Error as HttpError};
 use http::header::{self, HeaderName, HeaderValue};
 use sha1::Sha1;
 use futures::{Async, Future, Poll, Stream};
 use futures::future::{Either, err as FutErr};
+use futures::unsync::mpsc::{unbounded, UnboundedSender};
 use tokio_core::net::TcpStream;
+use byteorder::{ByteOrder, NetworkEndian};
 
 use actix::prelude::*;
 
-use body::Binary;
+use body::{Body, Binary};
-use error::UrlParseError;
+use error::{WsError, UrlParseError};
+use payload::PayloadHelper;
 use server::shared::SharedBytes;
 
 use server::{utils, IoStream};
-use client::{ClientRequest, ClientRequestBuilder,
+use client::{ClientRequest, ClientRequestBuilder, ClientResponse,
              HttpResponseParser, HttpResponseParserError, HttpClientWriter};
-use client::{Connect, Connection, ClientConnector, ClientConnectorError};
+use client::{Connect, Connection, ClientConnector, ClientConnectorError,
+             SendRequest, SendRequestError};
 
 use super::Message;
 use super::frame::Frame;
@@ -52,7 +56,9 @@ pub enum WsClientError {
     #[fail(display="Response parsing error")]
     ResponseParseError(HttpResponseParserError),
     #[fail(display="{}", _0)]
-    Connector(ClientConnectorError),
+    SendRequest(SendRequestError),
+    #[fail(display="{}", _0)]
+    Protocol(#[cause] WsError),
     #[fail(display="{}", _0)]
     Io(io::Error),
     #[fail(display="Disconnected")]
@@ -71,9 +77,15 @@ impl From<UrlParseError> for WsClientError {
     }
 }
 
-impl From<ClientConnectorError> for WsClientError {
+impl From<SendRequestError> for WsClientError {
-    fn from(err: ClientConnectorError) -> WsClientError {
+    fn from(err: SendRequestError) -> WsClientError {
-        WsClientError::Connector(err)
+        WsClientError::SendRequest(err)
+    }
+}
+
+impl From<WsError> for WsClientError {
+    fn from(err: WsError) -> WsClientError {
+        WsClientError::Protocol(err)
     }
 }
 
@@ -206,21 +218,17 @@ impl WsClient {
 }
 
 struct WsInner {
-    conn: Connection,
+    tx: UnboundedSender<Bytes>,
-    writer: HttpClientWriter,
+    rx: PayloadHelper<ClientResponse>,
-    parser: HttpResponseParser,
-    parser_buf: BytesMut,
     closed: bool,
-    error_sent: bool,
 }
 
 pub struct WsHandshake {
     inner: Option<WsInner>,
-    request: Option<ClientRequest>,
+    request: Option<SendRequest>,
-    sent: bool,
+    tx: Option<UnboundedSender<Bytes>>,
     key: String,
     error: Option<WsClientError>,
-    stream: Option<Box<Future<Item=Result<Connection, WsClientError>, Error=WsClientError>>>,
 }
 
 impl WsHandshake {
@@ -235,31 +243,29 @@ impl WsHandshake {
         let key = base64::encode(&sec_key);
 
         if let Some(mut request) = request {
-            let stream = Box::new(
-                conn.send(Connect(request.uri().clone()))
-                    .map(|res| res.map_err(|e| e.into()))
-                    .map_err(|_| WsClientError::Disconnected));
-
             request.headers_mut().insert(
                 HeaderName::try_from("SEC-WEBSOCKET-KEY").unwrap(),
                 HeaderValue::try_from(key.as_str()).unwrap());
 
+            let (tx, rx) = unbounded();
+            request.set_body(Body::Streaming(
+                Box::new(rx.map_err(|_| io::Error::new(
+                    io::ErrorKind::Other, "disconnected").into()))));
+
             WsHandshake {
                 key: key,
                 inner: None,
-                request: Some(request),
+                request: Some(request.with_connector(conn.clone())),
-                sent: false,
+                tx: Some(tx),
                 error: err,
-                stream: Some(stream),
             }
         } else {
             WsHandshake {
                 key: key,
                 inner: None,
                 request: None,
-                sent: false,
+                tx: None,
                 error: err,
-                stream: None,
             }
         }
     }
@@ -274,94 +280,67 @@ impl Future for WsHandshake {
             return Err(err)
         }
 
-        if self.stream.is_some() {
+        let resp = match self.request.as_mut().unwrap().poll()? {
-            match self.stream.as_mut().unwrap().poll()? {
+            Async::Ready(response) => {
-                Async::Ready(result) => match result {
+                self.request.take();
-                    Ok(conn) => {
+                response
-                        let inner = WsInner {
+            },
-                            conn: conn,
+            Async::NotReady => return Ok(Async::NotReady)
-                            writer: HttpClientWriter::new(SharedBytes::default()),
+        };
-                            parser: HttpResponseParser::default(),
+
-                            parser_buf: BytesMut::new(),
+        // verify response
-                            closed: false,
+        if resp.status() != StatusCode::SWITCHING_PROTOCOLS {
-                            error_sent: false,
+            return Err(WsClientError::InvalidResponseStatus)
-                        };
+        }
-                        self.stream.take();
+        // Check for "UPGRADE" to websocket header
-                        self.inner = Some(inner);
+        let has_hdr = if let Some(hdr) = resp.headers().get(header::UPGRADE) {
-                    }
+            if let Ok(s) = hdr.to_str() {
-                    Err(err) => return Err(err),
+                s.to_lowercase().contains("websocket")
-                },
+            } else {
-                Async::NotReady => return Ok(Async::NotReady)
+                false
             }
+        } else {
+            false
+        };
+        if !has_hdr {
+            return Err(WsClientError::InvalidUpgradeHeader)
+        }
+        // Check for "CONNECTION" header
+        let has_hdr = if let Some(conn) = resp.headers().get(header::CONNECTION) {
+            if let Ok(s) = conn.to_str() {
+                s.to_lowercase().contains("upgrade")
+            } else { false }
+        } else { false };
+        if !has_hdr {
+            return Err(WsClientError::InvalidConnectionHeader)
         }
 
-        let mut inner = self.inner.take().unwrap();
+        let match_key = if let Some(key) = resp.headers().get(
-
+            HeaderName::try_from("SEC-WEBSOCKET-ACCEPT").unwrap())
-        if !self.sent {
+        {
-            self.sent = true;
+            // field is constructed by concatenating /key/
-            inner.writer.start(self.request.as_mut().unwrap())?;
+            // with the string "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" (RFC 6455)
-        }
+            const WS_GUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
-        if let Err(err) = inner.writer.poll_completed(&mut inner.conn, false) {
+            let mut sha1 = Sha1::new();
-            return Err(err.into())
+            sha1.update(self.key.as_ref());
+            sha1.update(WS_GUID);
+            key.as_bytes() == base64::encode(&sha1.digest().bytes()).as_bytes()
+        } else {
+            false
+        };
+        if !match_key {
+            return Err(WsClientError::InvalidChallengeResponse)
         }
 
-        match inner.parser.parse(&mut inner.conn, &mut inner.parser_buf) {
+        let inner = WsInner {
-            Ok(Async::Ready(resp)) => {
+            tx: self.tx.take().unwrap(),
-                // verify response
+            rx: PayloadHelper::new(resp),
-                if resp.status() != StatusCode::SWITCHING_PROTOCOLS {
+            closed: false,
-                    return Err(WsClientError::InvalidResponseStatus)
+        };
-                }
-                // Check for "UPGRADE" to websocket header
-                let has_hdr = if let Some(hdr) = resp.headers().get(header::UPGRADE) {
-                    if let Ok(s) = hdr.to_str() {
-                        s.to_lowercase().contains("websocket")
-                    } else {
-                        false
-                    }
-                } else {
-                    false
-                };
-                if !has_hdr {
-                    return Err(WsClientError::InvalidUpgradeHeader)
-                }
-                // Check for "CONNECTION" header
-                let has_hdr = if let Some(conn) = resp.headers().get(header::CONNECTION) {
-                    if let Ok(s) = conn.to_str() {
-                        s.to_lowercase().contains("upgrade")
-                    } else { false }
-                } else { false };
-                if !has_hdr {
-                    return Err(WsClientError::InvalidConnectionHeader)
-                }
 
-                let match_key = if let Some(key) = resp.headers().get(
+        let inner = Rc::new(UnsafeCell::new(inner));
-                    HeaderName::try_from("SEC-WEBSOCKET-ACCEPT").unwrap())
+        Ok(Async::Ready(
-                {
+            (WsClientReader{inner: Rc::clone(&inner)}, WsClientWriter{inner: inner})))
-                    // field is constructed by concatenating /key/
-                    // with the string "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" (RFC 6455)
-                    const WS_GUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
-                    let mut sha1 = Sha1::new();
-                    sha1.update(self.key.as_ref());
-                    sha1.update(WS_GUID);
-                    key.as_bytes() == base64::encode(&sha1.digest().bytes()).as_bytes()
-                } else {
-                    false
-                };
-                if !match_key {
-                    return Err(WsClientError::InvalidChallengeResponse)
-                }
-
-                let inner = Rc::new(UnsafeCell::new(inner));
-                Ok(Async::Ready(
-                    (WsClientReader{inner: Rc::clone(&inner)},
-                     WsClientWriter{inner: inner})))
-            },
-            Ok(Async::NotReady) => {
-                self.inner = Some(inner);
-                Ok(Async::NotReady)
-            },
-            Err(err) => Err(err.into())
-        }
     }
 }
 
@@ -389,24 +368,13 @@ impl Stream for WsClientReader {
 
     fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
         let inner = self.as_mut();
-        let mut done = false;
+        if inner.closed {
-
+            return Ok(Async::Ready(None))
-        match utils::read_from_io(&mut inner.conn, &mut inner.parser_buf) {
-            Ok(Async::Ready(0)) => {
-                done = true;
-                inner.closed = true;
-            },
-            Ok(Async::Ready(_)) | Ok(Async::NotReady) => (),
-            Err(err) =>
-                return Err(err.into())
         }
 
-        // write
-        let _ = inner.writer.poll_completed(&mut inner.conn, false);
-
         // read
-        match Frame::parse(&mut inner.parser_buf, false) {
+        match Frame::parse(&mut inner.rx, false) {
-            Ok(Some(frame)) => {
+            Ok(Async::Ready(Some(frame))) => {
                 // trace!("WsFrame {}", frame);
                 let (_finished, opcode, payload) = frame.unpack();
 
@@ -416,8 +384,9 @@ impl Stream for WsClientReader {
                         Ok(Async::Ready(Some(Message::Error))),
                     OpCode::Close => {
                         inner.closed = true;
-                        inner.error_sent = true;
+                        let code = NetworkEndian::read_uint(payload.as_ref(), 2) as u16;
-                        Ok(Async::Ready(Some(Message::Closed)))
+                        Ok(Async::Ready(
+                            Some(Message::Close(CloseCode::from(code)))))
                     },
                     OpCode::Ping =>
                         Ok(Async::Ready(Some(
@@ -440,23 +409,10 @@ impl Stream for WsClientReader {
                     }
                 }
             }
-            Ok(None) => {
+            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
-                if done {
+            Ok(Async::NotReady) => Ok(Async::NotReady),
-                    Ok(Async::Ready(None))
-                } else if inner.closed {
-                    if !inner.error_sent {
-                        inner.error_sent = true;
-                        Ok(Async::Ready(Some(Message::Closed)))
-                    } else {
-                        Ok(Async::Ready(None))
-                    }
-                } else {
-                    Ok(Async::NotReady)
-                }
-            },
             Err(err) => {
                 inner.closed = true;
-                inner.error_sent = true;
                 Err(err.into())
             }
         }
@@ -478,9 +434,9 @@ impl WsClientWriter {
 
     /// Write payload
     #[inline]
-    fn write(&mut self, data: Binary) {
+    fn write(&mut self, mut data: Binary) {
         if !self.as_mut().closed {
-            let _ = self.as_mut().writer.write(data);
+            let _ = self.as_mut().tx.unbounded_send(data.take());
         } else {
             warn!("Trying to write to disconnected response");
         }

src/ws/frame.rs

@@ -1,11 +1,13 @@
 use std::{fmt, mem};
-use std::io::{Error, ErrorKind};
 use std::iter::FromIterator;
-use bytes::{BytesMut, BufMut};
+use bytes::{Bytes, BytesMut, BufMut};
 use byteorder::{ByteOrder, BigEndian, NetworkEndian};
+use futures::{Async, Poll, Stream};
 use rand;
 
 use body::Binary;
+use error::{WsError, PayloadError};
+use payload::PayloadHelper;
 use ws::proto::{OpCode, CloseCode};
 use ws::mask::apply_mask;
 
@@ -48,14 +50,15 @@ impl Frame {
     }
 
     /// Parse the input stream into a frame.
-    pub fn parse(buf: &mut BytesMut, server: bool) -> Result<Option<Frame>, Error> {
+    pub fn parse<S>(pl: &mut PayloadHelper<S>, server: bool) -> Poll<Option<Frame>, WsError>
+        where S: Stream<Item=Bytes, Error=PayloadError>
+    {
         let mut idx = 2;
-        let mut size = buf.len();
+        let buf = match pl.copy(2)? {
-
+            Async::Ready(Some(buf)) => buf,
-        if size < 2 {
+            Async::Ready(None) => return Ok(Async::Ready(None)),
-            return Ok(None)
+            Async::NotReady => return Ok(Async::NotReady),
-        }
+        };
-        size -= 2;
         let first = buf[0];
         let second = buf[1];
         let finished = first & 0x80 != 0;
@@ -63,11 +66,9 @@ impl Frame {
         // check masking
         let masked = second & 0x80 != 0;
         if !masked && server {
-            return Err(Error::new(
+            return Err(WsError::UnmaskedFrame)
-                ErrorKind::Other, "Received an unmasked frame from client"))
         } else if masked && !server {
-            return Err(Error::new(
+            return Err(WsError::MaskedFrame)
-                ErrorKind::Other, "Received a masked frame from server"))
         }
 
         let rsv1 = first & 0x40 != 0;
@@ -77,19 +78,21 @@ impl Frame {
         let len = second & 0x7F;
 
         let length = if len == 126 {
-            if size < 2 {
+            let buf = match pl.copy(4)? {
-                return Ok(None)
+                Async::Ready(Some(buf)) => buf,
-            }
+                Async::Ready(None) => return Ok(Async::Ready(None)),
+                Async::NotReady => return Ok(Async::NotReady),
+            };
             let len = NetworkEndian::read_uint(&buf[idx..], 2) as usize;
-            size -= 2;
             idx += 2;
             len
         } else if len == 127 {
-            if size < 8 {
+            let buf = match pl.copy(10)? {
-                return Ok(None)
+                Async::Ready(Some(buf)) => buf,
-            }
+                Async::Ready(None) => return Ok(Async::Ready(None)),
+                Async::NotReady => return Ok(Async::NotReady),
+            };
             let len = NetworkEndian::read_uint(&buf[idx..], 8) as usize;
-            size -= 8;
             idx += 8;
             len
         } else {
@@ -97,50 +100,42 @@ impl Frame {
         };
 
         let mask = if server {
-            if size < 4 {
+            let buf = match pl.copy(idx + 4)? {
-                return Ok(None)
+                Async::Ready(Some(buf)) => buf,
-            } else {
+                Async::Ready(None) => return Ok(Async::Ready(None)),
-                let mut mask_bytes = [0u8; 4];
+                Async::NotReady => return Ok(Async::NotReady),
-                size -= 4;
+            };
-                mask_bytes.copy_from_slice(&buf[idx..idx+4]);
+
-                idx += 4;
+            let mut mask_bytes = [0u8; 4];
-                Some(mask_bytes)
+            mask_bytes.copy_from_slice(&buf[idx..idx+4]);
-            }
+            idx += 4;
+            Some(mask_bytes)
         } else {
             None
         };
 
-        if size < length {
+        let mut data = match pl.readexactly(idx + length)? {
-            return Ok(None)
+            Async::Ready(Some(buf)) => buf,
-        }
+            Async::Ready(None) => return Ok(Async::Ready(None)),
+            Async::NotReady => return Ok(Async::NotReady),
+        };
 
         // get body
-        buf.split_to(idx);
+        data.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(
+            return Err(WsError::InvalidOpcode(first & 0x0F))
-                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(
+                return Err(WsError::InvalidLength(length))
-                    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::default()))
+                return Ok(Async::Ready(Some(Frame::default())))
             }
             _ => ()
         }
@@ -150,14 +145,14 @@ impl Frame {
             apply_mask(&mut data, mask);
         }
 
-        Ok(Some(Frame {
+        Ok(Async::Ready(Some(Frame {
             finished: finished,
             rsv1: rsv1,
             rsv2: rsv2,
             rsv3: rsv3,
             opcode: opcode,
             payload: data.into(),
-        }))
+        })))
     }
 
     /// Generate binary representation
@@ -258,13 +253,33 @@ impl fmt::Display for Frame {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use futures::stream::once;
+
+    fn is_none(frm: Poll<Option<Frame>, WsError>) -> bool {
+        match frm {
+            Ok(Async::Ready(None)) => true,
+            _ => false,
+        }
+    }
+    
+    fn extract(frm: Poll<Option<Frame>, WsError>) -> Frame {
+        match frm {
+            Ok(Async::Ready(Some(frame))) => frame,
+            _ => panic!("error"),
+        }
+    }
 
     #[test]
     fn test_parse() {
+        let mut buf = PayloadHelper::new(
+            once(Ok(BytesMut::from(&[0b00000001u8, 0b00000001u8][..]).freeze())));
+        assert!(is_none(Frame::parse(&mut buf, false)));
+
         let mut buf = BytesMut::from(&[0b00000001u8, 0b00000001u8][..]);
-        assert!(Frame::parse(&mut buf, false).unwrap().is_none());
         buf.extend(b"1");
-        let frame = Frame::parse(&mut buf, false).unwrap().unwrap();
+        let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
+
+        let frame = extract(Frame::parse(&mut buf, false));
         println!("FRAME: {}", frame);
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
@@ -273,8 +288,10 @@ mod tests {
 
     #[test]
     fn test_parse_length0() {
-        let mut buf = BytesMut::from(&[0b00000001u8, 0b00000000u8][..]);
+        let buf = BytesMut::from(&[0b00000001u8, 0b00000000u8][..]);
-        let frame = Frame::parse(&mut buf, false).unwrap().unwrap();
+        let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
+
+        let frame = extract(Frame::parse(&mut buf, false));
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
         assert!(frame.payload.is_empty());
@@ -282,12 +299,16 @@ mod tests {
 
     #[test]
     fn test_parse_length2() {
+        let buf = BytesMut::from(&[0b00000001u8, 126u8][..]);
+        let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
+        assert!(is_none(Frame::parse(&mut buf, false)));
+
         let mut buf = BytesMut::from(&[0b00000001u8, 126u8][..]);
-        assert!(Frame::parse(&mut buf, false).unwrap().is_none());
         buf.extend(&[0u8, 4u8][..]);
         buf.extend(b"1234");
+        let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
 
-        let frame = Frame::parse(&mut buf, false).unwrap().unwrap();
+        let frame = extract(Frame::parse(&mut buf, false));
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
         assert_eq!(frame.payload.as_ref(), &b"1234"[..]);
@@ -295,12 +316,16 @@ mod tests {
 
     #[test]
     fn test_parse_length4() {
+        let buf = BytesMut::from(&[0b00000001u8, 127u8][..]);
+        let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
+        assert!(is_none(Frame::parse(&mut buf, false)));
+
         let mut buf = BytesMut::from(&[0b00000001u8, 127u8][..]);
-        assert!(Frame::parse(&mut buf, false).unwrap().is_none());
         buf.extend(&[0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 4u8][..]);
         buf.extend(b"1234");
+        let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
 
-        let frame = Frame::parse(&mut buf, false).unwrap().unwrap();
+        let frame = extract(Frame::parse(&mut buf, false));
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
         assert_eq!(frame.payload.as_ref(), &b"1234"[..]);
@@ -311,10 +336,11 @@ mod tests {
         let mut buf = BytesMut::from(&[0b00000001u8, 0b10000001u8][..]);
         buf.extend(b"0001");
         buf.extend(b"1");
+        let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
 
         assert!(Frame::parse(&mut buf, false).is_err());
 
-        let frame = Frame::parse(&mut buf, true).unwrap().unwrap();
+        let frame = extract(Frame::parse(&mut buf, true));
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
         assert_eq!(frame.payload, vec![1u8].into());
@@ -324,10 +350,11 @@ mod tests {
     fn test_parse_frame_no_mask() {
         let mut buf = BytesMut::from(&[0b00000001u8, 0b00000001u8][..]);
         buf.extend(&[1u8]);
+        let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
 
         assert!(Frame::parse(&mut buf, true).is_err());
 
-        let frame = Frame::parse(&mut buf, false).unwrap().unwrap();
+        let frame = extract(Frame::parse(&mut buf, false));
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
         assert_eq!(frame.payload, vec![1u8].into());

src/ws/mod.rs

@@ -43,15 +43,16 @@
 //! #      .finish();
 //! # }
 //! ```
-use bytes::BytesMut;
+use bytes::Bytes;
 use http::{Method, StatusCode, header};
 use futures::{Async, Poll, Stream};
+use byteorder::{ByteOrder, NetworkEndian};
 
 use actix::{Actor, AsyncContext, Handler};
 
 use body::Binary;
-use payload::Payload;
+use payload::PayloadHelper;
-use error::{Error, WsHandshakeError};
+use error::{Error, WsHandshakeError, PayloadError};
 use httprequest::HttpRequest;
 use httpresponse::{ConnectionType, HttpResponse, HttpResponseBuilder};
 
@@ -80,8 +81,7 @@ pub enum Message {
     Binary(Binary),
     Ping(String),
     Pong(String),
-    Close,
+    Close(CloseCode),
-    Closed,
     Error
 }
 
@@ -165,104 +165,67 @@ pub fn handshake<S>(req: &HttpRequest<S>) -> Result<HttpResponseBuilder, WsHands
 }
 
 /// Maps `Payload` stream into stream of `ws::Message` items
-pub struct WsStream {
+pub struct WsStream<S> {
-    rx: Payload,
+    rx: PayloadHelper<S>,
-    buf: BytesMut,
     closed: bool,
-    error_sent: bool,
 }
 
-impl WsStream {
+impl<S> WsStream<S> where S: Stream<Item=Bytes, Error=PayloadError> {
-    pub fn new(payload: Payload) -> WsStream {
+    pub fn new(stream: S) -> WsStream<S> {
-        WsStream { rx: payload,
+        WsStream { rx: PayloadHelper::new(stream),
-                   buf: BytesMut::new(),
+                   closed: false }
-                   closed: false,
-                   error_sent: false }
     }
 }
 
-impl Stream for WsStream {
+impl<S> Stream for WsStream<S> where S: Stream<Item=Bytes, Error=PayloadError> {
     type Item = Message;
     type Error = ();
 
     fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
-        let mut done = false;
+        if self.closed {
+            return Ok(Async::Ready(None))
+        }
 
-        if !self.closed {
+        match Frame::parse(&mut self.rx, true) {
-            loop {
+            Ok(Async::Ready(Some(frame))) => {
-                match self.rx.poll() {
+                // trace!("WsFrame {}", frame);
-                    Ok(Async::Ready(Some(chunk))) => {
+                let (_finished, opcode, payload) = frame.unpack();
-                        self.buf.extend_from_slice(&chunk)
+
-                    }
+                match opcode {
-                    Ok(Async::Ready(None)) => {
+                    OpCode::Continue => unimplemented!(),
-                        done = true;
+                    OpCode::Bad =>
+                        Ok(Async::Ready(Some(Message::Error))),
+                    OpCode::Close => {
                         self.closed = true;
-                        break;
+                        let code = NetworkEndian::read_uint(payload.as_ref(), 2) as u16;
-                    }
+                        Ok(Async::Ready(
-                    Ok(Async::NotReady) => break,
+                            Some(Message::Close(CloseCode::from(code)))))
-                    Err(_) => {
+                    },
-                        self.closed = true;
+                    OpCode::Ping =>
-                        break;
+                        Ok(Async::Ready(Some(
+                            Message::Ping(
+                                String::from_utf8_lossy(payload.as_ref()).into())))),
+                    OpCode::Pong =>
+                        Ok(Async::Ready(Some(
+                            Message::Pong(String::from_utf8_lossy(payload.as_ref()).into())))),
+                    OpCode::Binary =>
+                        Ok(Async::Ready(Some(Message::Binary(payload)))),
+                    OpCode::Text => {
+                        let tmp = Vec::from(payload.as_ref());
+                        match String::from_utf8(tmp) {
+                            Ok(s) =>
+                                Ok(Async::Ready(Some(Message::Text(s)))),
+                            Err(_) =>
+                                Ok(Async::Ready(Some(Message::Error))),
+                        }
                     }
                 }
             }
-        }
+            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
-
+            Ok(Async::NotReady) => Ok(Async::NotReady),
-        loop {
+            Err(_) => {
-            match Frame::parse(&mut self.buf, true) {
+                self.closed = true;
-                Ok(Some(frame)) => {
+                Ok(Async::Ready(Some(Message::Error)))
-                    // trace!("WsFrame {}", frame);
-                    let (_finished, opcode, payload) = frame.unpack();
-
-                    match opcode {
-                        OpCode::Continue => continue,
-                        OpCode::Bad =>
-                            return Ok(Async::Ready(Some(Message::Error))),
-                        OpCode::Close => {
-                            self.closed = true;
-                            self.error_sent = true;
-                            return Ok(Async::Ready(Some(Message::Closed)))
-                        },
-                        OpCode::Ping =>
-                            return Ok(Async::Ready(Some(
-                                Message::Ping(
-                                    String::from_utf8_lossy(payload.as_ref()).into())))),
-                        OpCode::Pong =>
-                            return Ok(Async::Ready(Some(
-                                Message::Pong(
-                                    String::from_utf8_lossy(payload.as_ref()).into())))),
-                        OpCode::Binary =>
-                            return Ok(Async::Ready(Some(Message::Binary(payload)))),
-                        OpCode::Text => {
-                            let tmp = Vec::from(payload.as_ref());
-                            match String::from_utf8(tmp) {
-                                Ok(s) =>
-                                    return Ok(Async::Ready(Some(Message::Text(s)))),
-                                Err(_) =>
-                                    return Ok(Async::Ready(Some(Message::Error))),
-                            }
-                        }
-                    }
-                }
-                Ok(None) => {
-                    if done {
-                        return Ok(Async::Ready(None))
-                    } else if self.closed {
-                        if !self.error_sent {
-                            self.error_sent = true;
-                            return Ok(Async::Ready(Some(Message::Closed)))
-                        } else {
-                            return Ok(Async::Ready(None))
-                        }
-                    } else {
-                        return Ok(Async::NotReady)
-                    }
-                },
-                Err(_) => {
-                    self.closed = true;
-                    self.error_sent = true;
-                    return Ok(Async::Ready(Some(Message::Error)));
-                }
             }
         }
     }

tests/test_ws.rs

@@ -24,6 +24,7 @@ impl Handler<ws::Message> for Ws {
             ws::Message::Ping(msg) => ctx.pong(&msg),
             ws::Message::Text(text) => ctx.text(text),
             ws::Message::Binary(bin) => ctx.binary(bin),
+            ws::Message::Close(reason) => ctx.close(reason, ""),
             _ => (),
         }
     }
@@ -49,5 +50,5 @@ fn test_simple() {
 
     writer.close(ws::CloseCode::Normal, "");
     let (item, _) = srv.execute(reader.into_future()).unwrap();
-    assert!(item.is_none());
+    assert_eq!(item, Some(ws::Message::Close(ws::CloseCode::Normal)));
 }