refactor websockets handling

examples/websocket-chat/src/main.rs

@@ -92,8 +92,7 @@ impl Handler<session::Message> for WsChatSession {
 }
 
 /// WebSocket message handler
-impl Handler<ws::Message> for WsChatSession {
-    type Result = ();
+impl StreamHandler<ws::Message, ws::WsError> for WsChatSession {
 
     fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
         println!("WEBSOCKET MESSAGE: {:?}", msg);
@@ -161,7 +160,7 @@ impl Handler<ws::Message> for WsChatSession {
             },
             ws::Message::Binary(bin) =>
                 println!("Unexpected binary"),
-            ws::Message::Close(_) | ws::Message::Error => {
+            ws::Message::Close(_) => {
                 ctx.stop();
             }
         }

examples/websocket/src/client.rs

@@ -12,7 +12,7 @@ use std::time::Duration;
 
 use actix::*;
 use futures::Future;
-use actix_web::ws::{Message, WsClientError, WsClient, WsClientWriter};
+use actix_web::ws::{Message, WsError, WsClient, WsClientWriter};
 
 
 fn main() {
@@ -93,7 +93,7 @@ impl Handler<ClientCommand> for ChatClient {
 }
 
 /// Handle server websocket messages
-impl StreamHandler<Message, WsClientError> for ChatClient {
+impl StreamHandler<Message, WsError> for ChatClient {
 
     fn handle(&mut self, msg: Message, ctx: &mut Context<Self>) {
         match msg {

examples/websocket/src/main.rs

@@ -25,8 +25,7 @@ impl Actor for MyWebSocket {
 }
 
 /// Handler for `ws::Message`
-impl Handler<ws::Message> for MyWebSocket {
-    type Result = ();
+impl StreamHandler<ws::Message, ws::WsError> for MyWebSocket {
 
     fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
         // process websocket messages
@@ -35,7 +34,7 @@ impl Handler<ws::Message> for MyWebSocket {
             ws::Message::Ping(msg) => ctx.pong(&msg),
             ws::Message::Text(text) => ctx.text(text),
             ws::Message::Binary(bin) => ctx.binary(bin),
-            ws::Message::Close(_) | ws::Message::Error => {
+            ws::Message::Close(_) => {
                 ctx.stop();
             }
             _ => (),

guide/src/qs_9.md

@@ -21,9 +21,8 @@ impl Actor for Ws {
     type Context = ws::WebsocketContext<Self>;
 }
 
-/// Define Handler for ws::Message message
-impl Handler<ws::Message> for Ws {
-    type Result=();
+/// Handler for ws::Message message
+impl StreamHandler<ws::Message, ws::WsError> for Ws {
 
     fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
         match msg {

src/error.rs

@@ -24,7 +24,7 @@ use body::Body;
 use handler::Responder;
 use httprequest::HttpRequest;
 use httpresponse::HttpResponse;
-use httpcodes::{self, HTTPBadRequest, HTTPMethodNotAllowed, HTTPExpectationFailed};
+use httpcodes::{self, HTTPExpectationFailed};
 
 /// A specialized [`Result`](https://doc.rust-lang.org/std/result/enum.Result.html)
 /// for actix web operations
@@ -341,82 +341,6 @@ impl ResponseError for ExpectError {
     }
 }
 
-/// Websocket handshake errors
-#[derive(Fail, PartialEq, Debug)]
-pub enum WsHandshakeError {
-    /// Only get method is allowed
-    #[fail(display="Method not allowed")]
-    GetMethodRequired,
-    /// Upgrade header if not set to websocket
-    #[fail(display="Websocket upgrade is expected")]
-    NoWebsocketUpgrade,
-    /// Connection header is not set to upgrade
-    #[fail(display="Connection upgrade is expected")]
-    NoConnectionUpgrade,
-    /// Websocket version header is not set
-    #[fail(display="Websocket version header is required")]
-    NoVersionHeader,
-    /// Unsupported websocket version
-    #[fail(display="Unsupported version")]
-    UnsupportedVersion,
-    /// Websocket key is not set or wrong
-    #[fail(display="Unknown websocket key")]
-    BadWebsocketKey,
-}
-
-impl ResponseError for WsHandshakeError {
-
-    fn error_response(&self) -> HttpResponse {
-        match *self {
-            WsHandshakeError::GetMethodRequired => {
-                HTTPMethodNotAllowed
-                    .build()
-                    .header(header::ALLOW, "GET")
-                    .finish()
-                    .unwrap()
-            }
-            WsHandshakeError::NoWebsocketUpgrade =>
-                HTTPBadRequest.with_reason("No WebSocket UPGRADE header found"),
-            WsHandshakeError::NoConnectionUpgrade =>
-                HTTPBadRequest.with_reason("No CONNECTION upgrade"),
-            WsHandshakeError::NoVersionHeader =>
-                HTTPBadRequest.with_reason("Websocket version header is required"),
-            WsHandshakeError::UnsupportedVersion =>
-                HTTPBadRequest.with_reason("Unsupported version"),
-            WsHandshakeError::BadWebsocketKey =>
-                HTTPBadRequest.with_reason("Handshake error"),
-        }
-    }
-}
-
-/// 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 {
@@ -769,22 +693,6 @@ mod tests {
         assert_eq!(resp.status(), StatusCode::EXPECTATION_FAILED);
     }
 
-    #[test]
-    fn test_wserror_http_response() {
-        let resp: HttpResponse = WsHandshakeError::GetMethodRequired.error_response();
-        assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
-        let resp: HttpResponse = WsHandshakeError::NoWebsocketUpgrade.error_response();
-        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
-        let resp: HttpResponse = WsHandshakeError::NoConnectionUpgrade.error_response();
-        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
-        let resp: HttpResponse = WsHandshakeError::NoVersionHeader.error_response();
-        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
-        let resp: HttpResponse = WsHandshakeError::UnsupportedVersion.error_response();
-        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
-        let resp: HttpResponse = WsHandshakeError::BadWebsocketKey.error_response();
-        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
-    }
-
     macro_rules! from {
         ($from:expr => $error:pat) => {
             match ParseError::from($from) {

src/ws/client.rs

@@ -17,14 +17,14 @@ use byteorder::{ByteOrder, NetworkEndian};
 use actix::prelude::*;
 
 use body::{Body, Binary};
-use error::{WsError, UrlParseError};
+use error::UrlParseError;
 use payload::PayloadHelper;
 
 use client::{ClientRequest, ClientRequestBuilder, ClientResponse,
              ClientConnector, SendRequest, SendRequestError,
              HttpResponseParserError};
 
-use super::Message;
+use super::{Message, WsError};
 use super::frame::Frame;
 use super::proto::{CloseCode, OpCode};
 
@@ -106,6 +106,7 @@ pub struct WsClient {
     origin: Option<HeaderValue>,
     protocols: Option<String>,
     conn: Addr<Unsync, ClientConnector>,
+    max_size: usize,
 }
 
 impl WsClient {
@@ -123,6 +124,7 @@ impl WsClient {
             http_err: None,
             origin: None,
             protocols: None,
+            max_size: 65_536,
             conn,
         };
         cl.request.uri(uri.as_ref());
@@ -158,6 +160,14 @@ impl WsClient {
         self
     }
 
+    /// Set max frame size
+    ///
+    /// By default max size is set to 64kb
+    pub fn max_frame_size(mut self, size: usize) -> Self {
+        self.max_size = size;
+        self
+    }
+
     /// Set request header
     pub fn header<K, V>(mut self, key: K, value: V) -> Self
         where HeaderName: HttpTryFrom<K>, HeaderValue: HttpTryFrom<V>
@@ -167,12 +177,12 @@ impl WsClient {
     }
 
     /// Connect to websocket server and do ws handshake
-    pub fn connect(&mut self) -> WsHandshake {
+    pub fn connect(&mut self) -> WsClientHandshake {
         if let Some(e) = self.err.take() {
-            WsHandshake::new(None, Some(e), &self.conn)
+            WsClientHandshake::error(e)
         }
         else if let Some(e) = self.http_err.take() {
-            WsHandshake::new(None, Some(e.into()), &self.conn)
+            WsClientHandshake::error(e.into())
         } else {
             // origin
             if let Some(origin) = self.origin.take() {
@@ -189,23 +199,22 @@ impl WsClient {
             }
             let request = match self.request.finish() {
                 Ok(req) => req,
-                Err(err) => return WsHandshake::new(None, Some(err.into()), &self.conn),
+                Err(err) => return WsClientHandshake::error(err.into()),
             };
 
             if request.uri().host().is_none() {
-                return WsHandshake::new(None, Some(WsClientError::InvalidUrl), &self.conn)
+                return WsClientHandshake::error(WsClientError::InvalidUrl)
             }
             if let Some(scheme) = request.uri().scheme_part() {
                 if scheme != "http" && scheme != "https" && scheme != "ws" && scheme != "wss" {
-                    return WsHandshake::new(
-                        None, Some(WsClientError::InvalidUrl), &self.conn)
+                    return WsClientHandshake::error(WsClientError::InvalidUrl)
                 }
             } else {
-                return WsHandshake::new(None, Some(WsClientError::InvalidUrl), &self.conn)
+                return WsClientHandshake::error(WsClientError::InvalidUrl)
             }
 
             // start handshake
-            WsHandshake::new(Some(request), None, &self.conn)
+            WsClientHandshake::new(request, &self.conn, self.max_size)
         }
     }
 }
@@ -216,17 +225,17 @@ struct WsInner {
     closed: bool,
 }
 
-pub struct WsHandshake {
+pub struct WsClientHandshake {
     request: Option<SendRequest>,
     tx: Option<UnboundedSender<Bytes>>,
     key: String,
     error: Option<WsClientError>,
+    max_size: usize,
 }
 
-impl WsHandshake {
-    fn new(request: Option<ClientRequest>,
-           err: Option<WsClientError>,
-           conn: &Addr<Unsync, ClientConnector>) -> WsHandshake
+impl WsClientHandshake {
+    fn new(mut request: ClientRequest,
+           conn: &Addr<Unsync, ClientConnector>, max_size: usize) -> WsClientHandshake
     {
         // Generate a random key for the `Sec-WebSocket-Key` header.
         // a base64-encoded (see Section 4 of [RFC4648]) value that,
@@ -234,34 +243,36 @@ impl WsHandshake {
         let sec_key: [u8; 16] = rand::random();
         let key = base64::encode(&sec_key);
 
-        if let Some(mut request) = request {
-            request.headers_mut().insert(
-                HeaderName::try_from("SEC-WEBSOCKET-KEY").unwrap(),
-                HeaderValue::try_from(key.as_str()).unwrap());
+        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()))));
+        let (tx, rx) = unbounded();
+        request.set_body(Body::Streaming(
+            Box::new(rx.map_err(|_| io::Error::new(
+                io::ErrorKind::Other, "disconnected").into()))));
 
-            WsHandshake {
-                key,
-                request: Some(request.with_connector(conn.clone())),
-                tx: Some(tx),
-                error: err,
-            }
-        } else {
-            WsHandshake {
-                key,
-                request: None,
-                tx: None,
-                error: err,
-            }
+        WsClientHandshake {
+            key,
+            max_size,
+            request: Some(request.with_connector(conn.clone())),
+            tx: Some(tx),
+            error: None,
+        }
+    }
+
+    fn error(err: WsClientError) -> WsClientHandshake {
+        WsClientHandshake {
+            key: String::new(),
+            request: None,
+            tx: None,
+            error: Some(err),
+            max_size: 0
         }
     }
 }
 
-impl Future for WsHandshake {
+impl Future for WsClientHandshake {
     type Item = (WsClientReader, WsClientWriter);
     type Error = WsClientError;
 
@@ -330,13 +341,15 @@ impl Future for WsHandshake {
 
         let inner = Rc::new(UnsafeCell::new(inner));
         Ok(Async::Ready(
-            (WsClientReader{inner: Rc::clone(&inner)}, WsClientWriter{inner})))
+            (WsClientReader{inner: Rc::clone(&inner), max_size: self.max_size},
+             WsClientWriter{inner})))
     }
 }
 
 
 pub struct WsClientReader {
-    inner: Rc<UnsafeCell<WsInner>>
+    inner: Rc<UnsafeCell<WsInner>>,
+    max_size: usize,
 }
 
 impl fmt::Debug for WsClientReader {
@@ -354,29 +367,36 @@ impl WsClientReader {
 
 impl Stream for WsClientReader {
     type Item = Message;
-    type Error = WsClientError;
+    type Error = WsError;
 
     fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        let max_size = self.max_size;
         let inner = self.as_mut();
         if inner.closed {
             return Ok(Async::Ready(None))
         }
 
         // read
-        match Frame::parse(&mut inner.rx, false) {
+        match Frame::parse(&mut inner.rx, false, max_size) {
             Ok(Async::Ready(Some(frame))) => {
-                // trace!("WsFrame {}", frame);
-                let (_finished, opcode, payload) = frame.unpack();
+                let (finished, opcode, payload) = frame.unpack();
+
+                // continuation is not supported
+                if !finished {
+                    inner.closed = true;
+                    return Err(WsError::NoContinuation)
+                }
 
                 match opcode {
                     OpCode::Continue => unimplemented!(),
-                    OpCode::Bad =>
-                        Ok(Async::Ready(Some(Message::Error))),
+                    OpCode::Bad => {
+                        inner.closed = true;
+                        Err(WsError::BadOpCode)
+                    },
                     OpCode::Close => {
                         inner.closed = true;
                         let code = NetworkEndian::read_uint(payload.as_ref(), 2) as u16;
-                        Ok(Async::Ready(
-                            Some(Message::Close(CloseCode::from(code)))))
+                        Ok(Async::Ready(Some(Message::Close(CloseCode::from(code)))))
                     },
                     OpCode::Ping =>
                         Ok(Async::Ready(Some(
@@ -393,17 +413,19 @@ impl Stream for WsClientReader {
                         match String::from_utf8(tmp) {
                             Ok(s) =>
                                 Ok(Async::Ready(Some(Message::Text(s)))),
-                            Err(_) =>
-                                Ok(Async::Ready(Some(Message::Error))),
+                            Err(_) => {
+                                inner.closed = true;
+                                Err(WsError::BadEncoding)
+                            }
                         }
                     }
                 }
             }
             Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
             Ok(Async::NotReady) => Ok(Async::NotReady),
-            Err(err) => {
+            Err(e) => {
                 inner.closed = true;
-                Err(err.into())
+                Err(e)
             }
         }
     }

src/ws/context.rs

@@ -18,7 +18,7 @@ use ws::frame::Frame;
 use ws::proto::{OpCode, CloseCode};
 
 
-/// Http actor execution context
+/// `WebSockets` actor execution context
 pub struct WebsocketContext<A, S=()> where A: Actor<Context=WebsocketContext<A, S>>,
 {
     inner: ContextImpl<A>,

src/ws/frame.rs

@@ -6,8 +6,10 @@ use futures::{Async, Poll, Stream};
 use rand;
 
 use body::Binary;
-use error::{WsError, PayloadError};
+use error::{PayloadError};
 use payload::PayloadHelper;
+
+use ws::WsError;
 use ws::proto::{OpCode, CloseCode};
 use ws::mask::apply_mask;
 
@@ -50,7 +52,8 @@ impl Frame {
     }
 
     /// Parse the input stream into a frame.
-    pub fn parse<S>(pl: &mut PayloadHelper<S>, server: bool) -> Poll<Option<Frame>, WsError>
+    pub fn parse<S>(pl: &mut PayloadHelper<S>, server: bool, max_size: usize)
+                    -> Poll<Option<Frame>, WsError>
         where S: Stream<Item=Bytes, Error=PayloadError>
     {
         let mut idx = 2;
@@ -99,6 +102,11 @@ impl Frame {
             len as usize
         };
 
+        // check for max allowed size
+        if length > max_size {
+            return Err(WsError::Overflow)
+        }
+
         let mask = if server {
             let buf = match pl.copy(idx + 4)? {
                 Async::Ready(Some(buf)) => buf,
@@ -267,13 +275,13 @@ mod tests {
     fn test_parse() {
         let mut buf = PayloadHelper::new(
             once(Ok(BytesMut::from(&[0b00000001u8, 0b00000001u8][..]).freeze())));
-        assert!(is_none(Frame::parse(&mut buf, false)));
+        assert!(is_none(Frame::parse(&mut buf, false, 1024)));
 
         let mut buf = BytesMut::from(&[0b00000001u8, 0b00000001u8][..]);
         buf.extend(b"1");
         let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
 
-        let frame = extract(Frame::parse(&mut buf, false));
+        let frame = extract(Frame::parse(&mut buf, false, 1024));
         println!("FRAME: {}", frame);
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
@@ -285,7 +293,7 @@ mod tests {
         let buf = BytesMut::from(&[0b00000001u8, 0b00000000u8][..]);
         let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
 
-        let frame = extract(Frame::parse(&mut buf, false));
+        let frame = extract(Frame::parse(&mut buf, false, 1024));
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
         assert!(frame.payload.is_empty());
@@ -295,14 +303,14 @@ mod tests {
     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)));
+        assert!(is_none(Frame::parse(&mut buf, false, 1024)));
 
         let mut buf = BytesMut::from(&[0b00000001u8, 126u8][..]);
         buf.extend(&[0u8, 4u8][..]);
         buf.extend(b"1234");
         let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
 
-        let frame = extract(Frame::parse(&mut buf, false));
+        let frame = extract(Frame::parse(&mut buf, false, 1024));
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
         assert_eq!(frame.payload.as_ref(), &b"1234"[..]);
@@ -312,14 +320,14 @@ mod tests {
     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)));
+        assert!(is_none(Frame::parse(&mut buf, false, 1024)));
 
         let mut buf = BytesMut::from(&[0b00000001u8, 127u8][..]);
         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 = extract(Frame::parse(&mut buf, false));
+        let frame = extract(Frame::parse(&mut buf, false, 1024));
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
         assert_eq!(frame.payload.as_ref(), &b"1234"[..]);
@@ -332,9 +340,9 @@ mod tests {
         buf.extend(b"1");
         let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
 
-        assert!(Frame::parse(&mut buf, false).is_err());
+        assert!(Frame::parse(&mut buf, false, 1024).is_err());
 
-        let frame = extract(Frame::parse(&mut buf, true));
+        let frame = extract(Frame::parse(&mut buf, true, 1024));
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
         assert_eq!(frame.payload, vec![1u8].into());
@@ -346,14 +354,28 @@ mod tests {
         buf.extend(&[1u8]);
         let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
 
-        assert!(Frame::parse(&mut buf, true).is_err());
+        assert!(Frame::parse(&mut buf, true, 1024).is_err());
 
-        let frame = extract(Frame::parse(&mut buf, false));
+        let frame = extract(Frame::parse(&mut buf, false, 1024));
         assert!(!frame.finished);
         assert_eq!(frame.opcode, OpCode::Text);
         assert_eq!(frame.payload, vec![1u8].into());
     }
 
+    #[test]
+    fn test_parse_frame_max_size() {
+        let mut buf = BytesMut::from(&[0b00000001u8, 0b00000010u8][..]);
+        buf.extend(&[1u8, 1u8]);
+        let mut buf = PayloadHelper::new(once(Ok(buf.freeze())));
+
+        assert!(Frame::parse(&mut buf, true, 1).is_err());
+
+        if let Err(WsError::Overflow) = Frame::parse(&mut buf, false, 0) {
+        } else {
+            panic!("error");
+        }
+    }
+
     #[test]
     fn test_ping_frame() {
         let frame = Frame::message(Vec::from("data"), OpCode::Ping, true, false);

src/ws/mod.rs

@@ -23,9 +23,8 @@
 //!     type Context = ws::WebsocketContext<Self>;
 //! }
 //!
-//! // Define Handler for ws::Message message
-//! impl Handler<ws::Message> for Ws {
-//!     type Result = ();
+//! // Handler for ws::Message messages
+//! impl StreamHandler<ws::Message, ws::WsError> for Ws {
 //!
 //!     fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
 //!         match msg {
@@ -48,13 +47,14 @@ use http::{Method, StatusCode, header};
 use futures::{Async, Poll, Stream};
 use byteorder::{ByteOrder, NetworkEndian};
 
-use actix::{Actor, AsyncContext, Handler};
+use actix::{Actor, AsyncContext, StreamHandler};
 
 use body::Binary;
 use payload::PayloadHelper;
-use error::{Error, WsHandshakeError, PayloadError};
+use error::{Error, PayloadError, ResponseError};
 use httprequest::HttpRequest;
 use httpresponse::{ConnectionType, HttpResponse, HttpResponseBuilder};
+use httpcodes::{HTTPBadRequest, HTTPMethodNotAllowed};
 
 mod frame;
 mod proto;
@@ -66,7 +66,8 @@ use self::frame::Frame;
 use self::proto::{hash_key, OpCode};
 pub use self::proto::CloseCode;
 pub use self::context::WebsocketContext;
-pub use self::client::{WsClient, WsClientError, WsClientReader, WsClientWriter, WsHandshake};
+pub use self::client::{WsClient, WsClientError,
+                       WsClientReader, WsClientWriter, WsClientHandshake};
 
 const SEC_WEBSOCKET_ACCEPT: &str = "SEC-WEBSOCKET-ACCEPT";
 const SEC_WEBSOCKET_KEY: &str = "SEC-WEBSOCKET-KEY";
@@ -74,6 +75,94 @@ const SEC_WEBSOCKET_VERSION: &str = "SEC-WEBSOCKET-VERSION";
 // const SEC_WEBSOCKET_PROTOCOL: &'static str = "SEC-WEBSOCKET-PROTOCOL";
 
 
+/// 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),
+    /// Bad web socket op code
+    #[fail(display="Bad web socket op code")]
+    BadOpCode,
+    /// A payload reached size limit.
+    #[fail(display="A payload reached size limit.")]
+    Overflow,
+    /// Continuation is not supproted
+    #[fail(display="Continuation is not supproted.")]
+    NoContinuation,
+    /// Bad utf-8 encoding
+    #[fail(display="Bad utf-8 encoding.")]
+    BadEncoding,
+    /// 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)
+    }
+}
+
+/// Websocket handshake errors
+#[derive(Fail, PartialEq, Debug)]
+pub enum WsHandshakeError {
+    /// Only get method is allowed
+    #[fail(display="Method not allowed")]
+    GetMethodRequired,
+    /// Upgrade header if not set to websocket
+    #[fail(display="Websocket upgrade is expected")]
+    NoWebsocketUpgrade,
+    /// Connection header is not set to upgrade
+    #[fail(display="Connection upgrade is expected")]
+    NoConnectionUpgrade,
+    /// Websocket version header is not set
+    #[fail(display="Websocket version header is required")]
+    NoVersionHeader,
+    /// Unsupported websocket version
+    #[fail(display="Unsupported version")]
+    UnsupportedVersion,
+    /// Websocket key is not set or wrong
+    #[fail(display="Unknown websocket key")]
+    BadWebsocketKey,
+}
+
+impl ResponseError for WsHandshakeError {
+
+    fn error_response(&self) -> HttpResponse {
+        match *self {
+            WsHandshakeError::GetMethodRequired => {
+                HTTPMethodNotAllowed
+                    .build()
+                    .header(header::ALLOW, "GET")
+                    .finish()
+                    .unwrap()
+            }
+            WsHandshakeError::NoWebsocketUpgrade =>
+                HTTPBadRequest.with_reason("No WebSocket UPGRADE header found"),
+            WsHandshakeError::NoConnectionUpgrade =>
+                HTTPBadRequest.with_reason("No CONNECTION upgrade"),
+            WsHandshakeError::NoVersionHeader =>
+                HTTPBadRequest.with_reason("Websocket version header is required"),
+            WsHandshakeError::UnsupportedVersion =>
+                HTTPBadRequest.with_reason("Unsupported version"),
+            WsHandshakeError::BadWebsocketKey =>
+                HTTPBadRequest.with_reason("Handshake error"),
+        }
+    }
+}
+
 /// `WebSocket` Message
 #[derive(Debug, PartialEq, Message)]
 pub enum Message {
@@ -82,19 +171,18 @@ pub enum Message {
     Ping(String),
     Pong(String),
     Close(CloseCode),
-    Error
 }
 
 /// Do websocket handshake and start actor
 pub fn start<A, S>(req: HttpRequest<S>, actor: A) -> Result<HttpResponse, Error>
-    where A: Actor<Context=WebsocketContext<A, S>> + Handler<Message>,
+    where A: Actor<Context=WebsocketContext<A, S>> + StreamHandler<Message, WsError>,
           S: 'static
 {
     let mut resp = handshake(&req)?;
     let stream = WsStream::new(req.clone());
 
     let mut ctx = WebsocketContext::new(req, actor);
-    ctx.add_message_stream(stream);
+    ctx.add_stream(stream);
 
     Ok(resp.body(ctx)?)
 }
@@ -168,33 +256,52 @@ pub fn handshake<S>(req: &HttpRequest<S>) -> Result<HttpResponseBuilder, WsHands
 pub struct WsStream<S> {
     rx: PayloadHelper<S>,
     closed: bool,
+    max_size: usize,
 }
 
 impl<S> WsStream<S> where S: Stream<Item=Bytes, Error=PayloadError> {
+    /// Create new websocket frames stream
     pub fn new(stream: S) -> WsStream<S> {
         WsStream { rx: PayloadHelper::new(stream),
-                   closed: false }
+                   closed: false,
+                   max_size: 65_536,
+        }
+    }
+
+    /// Set max frame size
+    ///
+    /// By default max size is set to 64kb
+    pub fn max_size(mut self, size: usize) -> Self {
+        self.max_size = size;
+        self
     }
 }
 
 impl<S> Stream for WsStream<S> where S: Stream<Item=Bytes, Error=PayloadError> {
     type Item = Message;
-    type Error = ();
+    type Error = WsError;
 
     fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
         if self.closed {
             return Ok(Async::Ready(None))
         }
 
-        match Frame::parse(&mut self.rx, true) {
+        match Frame::parse(&mut self.rx, true, self.max_size) {
             Ok(Async::Ready(Some(frame))) => {
-                // trace!("WsFrame {}", frame);
-                let (_finished, opcode, payload) = frame.unpack();
+                let (finished, opcode, payload) = frame.unpack();
+
+                // continuation is not supported
+                if !finished {
+                    self.closed = true;
+                    return Err(WsError::NoContinuation)
+                }
 
                 match opcode {
                     OpCode::Continue => unimplemented!(),
-                    OpCode::Bad =>
-                        Ok(Async::Ready(Some(Message::Error))),
+                    OpCode::Bad => {
+                        self.closed = true;
+                        Err(WsError::BadOpCode)
+                    }
                     OpCode::Close => {
                         self.closed = true;
                         let code = NetworkEndian::read_uint(payload.as_ref(), 2) as u16;
@@ -215,17 +322,19 @@ impl<S> Stream for WsStream<S> where S: Stream<Item=Bytes, Error=PayloadError> {
                         match String::from_utf8(tmp) {
                             Ok(s) =>
                                 Ok(Async::Ready(Some(Message::Text(s)))),
-                            Err(_) =>
-                                Ok(Async::Ready(Some(Message::Error))),
+                            Err(_) => {
+                                self.closed = true;
+                                Err(WsError::BadEncoding)
+                            }
                         }
                     }
                 }
             }
             Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
             Ok(Async::NotReady) => Ok(Async::NotReady),
-            Err(_) => {
+            Err(e) => {
                 self.closed = true;
-                Ok(Async::Ready(Some(Message::Error)))
+                Err(e)
             }
         }
     }
@@ -306,4 +415,20 @@ mod tests {
         assert_eq!(StatusCode::SWITCHING_PROTOCOLS,
                    handshake(&req).unwrap().finish().unwrap().status());
     }
+
+    #[test]
+    fn test_wserror_http_response() {
+        let resp: HttpResponse = WsHandshakeError::GetMethodRequired.error_response();
+        assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
+        let resp: HttpResponse = WsHandshakeError::NoWebsocketUpgrade.error_response();
+        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
+        let resp: HttpResponse = WsHandshakeError::NoConnectionUpgrade.error_response();
+        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
+        let resp: HttpResponse = WsHandshakeError::NoVersionHeader.error_response();
+        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
+        let resp: HttpResponse = WsHandshakeError::UnsupportedVersion.error_response();
+        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
+        let resp: HttpResponse = WsHandshakeError::BadWebsocketKey.error_response();
+        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
+    }
 }

tests/test_ws.rs

@@ -16,8 +16,7 @@ impl Actor for Ws {
     type Context = ws::WebsocketContext<Self>;
 }
 
-impl Handler<ws::Message> for Ws {
-    type Result = ();
+impl StreamHandler<ws::Message, ws::WsError> for Ws {
 
     fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
         match msg {